mirror of
https://github.com/hiyouga/LlamaFactory.git
synced 2026-02-01 20:23:37 +00:00
[v1] add accelerator (#9607)
This commit is contained in:
Yaowei Zheng
@@ -13,12 +13,13 @@
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
from typing import Callable, TypedDict
|
||||
from typing import Any, Literal, TypedDict
|
||||
|
||||
from typing_extensions import NotRequired, Required
|
||||
from typing_extensions import NotRequired
|
||||
|
||||
from ....extras import logging
|
||||
from ...extras.types import DPOSample, Sample, SFTSample
|
||||
from ...utils import logging
|
||||
from ...utils.plugin import BasePlugin
|
||||
from ...utils.types import DPOSample, Sample, SFTSample
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
@@ -26,35 +27,48 @@ logger = logging.get_logger(__name__)
|
||||
|
||||
class AlpacaSample(TypedDict, total=False):
|
||||
system: NotRequired[str]
|
||||
instruction: NotRequired[str]
|
||||
instruction: str
|
||||
input: NotRequired[str]
|
||||
output: NotRequired[str]
|
||||
output: str
|
||||
|
||||
|
||||
ShareGPTMessage = TypedDict(
|
||||
"ShareGPTMessage",
|
||||
{
|
||||
"from": Required[str], # Role of the message sender (e.g., "human", "gpt", "system")
|
||||
"value": Required[str], # Content of the message
|
||||
},
|
||||
SharegptMessage = TypedDict(
|
||||
"SharegptMessage", {"from": Literal["human", "gpt", "system", "function_call", "observation"], "value": str}
|
||||
)
|
||||
|
||||
|
||||
class ShareGPTSample(TypedDict, total=False):
|
||||
"""Type definition for raw ShareGPT sample."""
|
||||
class SharegptSample(TypedDict, total=False):
|
||||
conversations: list[SharegptMessage]
|
||||
tools: NotRequired[str]
|
||||
|
||||
conversations: Required[list[ShareGPTMessage]]
|
||||
|
||||
class OpenaiMessage(TypedDict, total=False):
|
||||
role: Literal["user", "assistant", "tool"]
|
||||
content: str
|
||||
|
||||
|
||||
class OpenaiSample(TypedDict, total=False):
|
||||
messages: list[OpenaiMessage]
|
||||
|
||||
|
||||
class PairSample(TypedDict, total=False):
|
||||
prompt: NotRequired[str]
|
||||
chosen: NotRequired[list[dict]]
|
||||
rejected: NotRequired[list[dict]]
|
||||
chosen: list[OpenaiMessage]
|
||||
rejected: list[OpenaiMessage]
|
||||
|
||||
|
||||
class DataConverterPlugin(BasePlugin):
|
||||
"""Plugin for data converters."""
|
||||
|
||||
def __call__(self, raw_sample: dict[str, Any]) -> Sample:
|
||||
return super().__call__(raw_sample)
|
||||
|
||||
|
||||
@DataConverterPlugin("alpaca").register
|
||||
def alpaca_converter(raw_sample: AlpacaSample) -> SFTSample:
|
||||
"""Convert Alpaca sample to SFT sample.
|
||||
|
||||
See raw example at: https://huggingface.co/datasets/llamafactory/alpaca_gpt4_en
|
||||
|
||||
Args:
|
||||
raw_sample (AlpacaSample): Alpaca sample.
|
||||
|
||||
@@ -67,20 +81,6 @@ def alpaca_converter(raw_sample: AlpacaSample) -> SFTSample:
|
||||
{"role": "system", "content": [{"type": "text", "value": raw_sample["system"]}], "loss_weight": 0.0}
|
||||
)
|
||||
|
||||
if "history" in raw_sample:
|
||||
for idx, item in enumerate(raw_sample["history"]):
|
||||
if len(item) != 2:
|
||||
logger.warning_rank0(
|
||||
f"Warning: History item at index {idx} has invalid length (expected 2, got {len(item)}). Skipping."
|
||||
)
|
||||
continue
|
||||
|
||||
old_prompt, old_response = item
|
||||
messages.append({"role": "user", "content": [{"type": "text", "value": old_prompt}], "loss_weight": 0.0})
|
||||
messages.append(
|
||||
{"role": "assistant", "content": [{"type": "text", "value": old_response}], "loss_weight": 1.0}
|
||||
)
|
||||
|
||||
if "instruction" in raw_sample or "input" in raw_sample:
|
||||
messages.append(
|
||||
{
|
||||
@@ -100,149 +100,85 @@ def alpaca_converter(raw_sample: AlpacaSample) -> SFTSample:
|
||||
return {"messages": messages}
|
||||
|
||||
|
||||
def sharegpt_converter(raw_sample: ShareGPTSample) -> SFTSample:
|
||||
"""Converts a raw ShareGPT sample into a formatted SFT (Supervised Fine-Tuning) sample.
|
||||
@DataConverterPlugin("sharegpt").register
|
||||
def sharegpt_converter(raw_sample: SharegptSample) -> SFTSample:
|
||||
"""Convert ShareGPT sample to SFT sample.
|
||||
|
||||
Retains only SFT-relevant scenarios and removes parity checks.
|
||||
See raw example at: https://huggingface.co/datasets/llamafactory/glaive_toolcall_en
|
||||
|
||||
Args:
|
||||
raw_sample (ShareGPTSample): A raw sample in ShareGPT format.
|
||||
raw_sample (SharegptSample): ShareGPT sample.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the formatted 'messages' list for SFT training.
|
||||
Returns an empty list if the input data is invalid.
|
||||
SFTSample: SFT sample.
|
||||
"""
|
||||
tag_mapping = {
|
||||
"system": "system",
|
||||
"human": "user",
|
||||
"gpt": "assistant",
|
||||
"observation": "observation",
|
||||
"function_call": "function",
|
||||
"observation": "tool",
|
||||
"function_call": "assistant",
|
||||
}
|
||||
messages = raw_sample.get("conversations", [])
|
||||
aligned_messages = []
|
||||
system_content = ""
|
||||
messages = []
|
||||
tools = raw_sample.get("tools", "")
|
||||
|
||||
# Extract system message if present (typically the first message)
|
||||
if messages and messages[0]["from"] == "system":
|
||||
system_content = messages[0]["value"]
|
||||
messages = messages[1:]
|
||||
for message in raw_sample.get("conversations", []):
|
||||
tag = message["from"]
|
||||
if tag not in tag_mapping:
|
||||
logger.warning_rank0(f"Unsupported role tag {tag} in message: {message}")
|
||||
elif tag == "function_call":
|
||||
messages.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "tool_calls", "value": message["value"]}],
|
||||
"loss_weight": 1.0,
|
||||
}
|
||||
)
|
||||
else:
|
||||
messages.append(
|
||||
{
|
||||
"role": tag_mapping[tag],
|
||||
"content": [{"type": "text", "value": message["value"]}],
|
||||
"loss_weight": 1.0 if tag == "gpt" else 0.0,
|
||||
}
|
||||
)
|
||||
|
||||
if system_content:
|
||||
aligned_messages.append(
|
||||
{"role": "system", "content": [{"type": "text", "value": system_content}], "loss_weight": 0.0}
|
||||
)
|
||||
if tools:
|
||||
if messages and messages[0]["role"] == "system":
|
||||
messages[0]["content"].append({"type": "tools", "value": tools})
|
||||
else:
|
||||
messages.insert(0, {"role": "system", "content": [{"type": "tools", "value": tools}], "loss_weight": 0.0})
|
||||
|
||||
has_invalid_role = False
|
||||
for message in messages:
|
||||
sender = message["from"]
|
||||
# validate sender is in supported tags
|
||||
if sender not in tag_mapping:
|
||||
logger.warning_rank0(f"Unsupported role tag '{sender}' in message: {message}")
|
||||
has_invalid_role = True
|
||||
break
|
||||
|
||||
aligned_messages.append(
|
||||
{
|
||||
"role": tag_mapping[sender],
|
||||
"content": [{"type": "text", "value": message["value"]}],
|
||||
"loss_weight": 0.0 if sender in ("human", "observation") else 1.0,
|
||||
}
|
||||
)
|
||||
|
||||
if has_invalid_role:
|
||||
logger.warning_rank0("Skipping invalid example due to unsupported role tags.")
|
||||
return {"messages": []}
|
||||
|
||||
return {"messages": aligned_messages}
|
||||
return {"messages": messages}
|
||||
|
||||
|
||||
@DataConverterPlugin("pair").register
|
||||
def pair_converter(raw_sample: PairSample) -> DPOSample:
|
||||
"""Convert Pair sample to standard DPO sample.
|
||||
"""Convert Pair sample to DPO sample.
|
||||
|
||||
See raw example at: https://huggingface.co/datasets/HuggingFaceH4/orca_dpo_pairs
|
||||
|
||||
Args:
|
||||
raw_sample (PairSample): pair sample with prompt, chosen, rejected fields.
|
||||
see raw example at: https://huggingface.co/datasets/HuggingFaceH4/orca_dpo_pairs
|
||||
raw_sample (PairSample): pair sample with chosen, rejected fields.
|
||||
|
||||
Returns:
|
||||
DPOSample: DPO sample with chosen_messages and rejected_messages.
|
||||
see the standard DPO sample at: https://huggingface.co/datasets/frozenleaves/v1-dpo-demo/raw/main/v1-dpo-demo.jsonl
|
||||
"""
|
||||
chosen_messages = []
|
||||
assert "chosen" in raw_sample, "chosen field is required in pair sample."
|
||||
assert "rejected" in raw_sample, "rejected field is required in pair sample."
|
||||
assert isinstance(raw_sample["chosen"], list) and isinstance(raw_sample["rejected"], list), (
|
||||
"chosen and rejected field should be a list[dict], or you may need to implement your custom converter."
|
||||
)
|
||||
|
||||
if "chosen" in raw_sample:
|
||||
value = raw_sample.get("chosen", "")
|
||||
for item in value:
|
||||
if item.get("role", "") == "system":
|
||||
chosen_messages.append(
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "value": item.get("content", "")}],
|
||||
"loss_weight": 0.0,
|
||||
}
|
||||
)
|
||||
if item.get("role", "") == "user":
|
||||
chosen_messages.append(
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "value": item.get("content", "")}],
|
||||
"loss_weight": 0.0,
|
||||
}
|
||||
)
|
||||
if item.get("role", "") == "assistant":
|
||||
chosen_messages.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "value": item.get("content", "")}],
|
||||
"loss_weight": 1.0,
|
||||
}
|
||||
)
|
||||
def process_message(raw_messages: list[OpenaiMessage]):
|
||||
messages = []
|
||||
for message in raw_messages:
|
||||
messages.append(
|
||||
{
|
||||
"role": message["role"],
|
||||
"content": [{"type": "text", "value": message["content"]}],
|
||||
"loss_weight": 1.0 if message["role"] == "assistant" else 0.0,
|
||||
}
|
||||
)
|
||||
|
||||
rejected_messages = []
|
||||
if "rejected" in raw_sample:
|
||||
value = raw_sample.get("rejected", "")
|
||||
for item in value:
|
||||
if item.get("role", "") == "system":
|
||||
rejected_messages.append(
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "value": item.get("content", "")}],
|
||||
"loss_weight": 0.0,
|
||||
}
|
||||
)
|
||||
if item.get("role", "") == "user":
|
||||
rejected_messages.append(
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "value": item.get("content", "")}],
|
||||
"loss_weight": 0.0,
|
||||
}
|
||||
)
|
||||
if item.get("role", "") == "assistant":
|
||||
rejected_messages.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "value": item.get("content", "")}],
|
||||
"loss_weight": 1.0,
|
||||
}
|
||||
)
|
||||
return messages
|
||||
|
||||
chosen_messages = process_message(raw_sample.get("chosen", []))
|
||||
rejected_messages = process_message(raw_sample.get("rejected", []))
|
||||
|
||||
return {"chosen_messages": chosen_messages, "rejected_messages": rejected_messages}
|
||||
|
||||
|
||||
CONVERTERS = {
|
||||
"alpaca": alpaca_converter,
|
||||
"pair": pair_converter,
|
||||
"sharegpt": sharegpt_converter,
|
||||
}
|
||||
|
||||
|
||||
def get_converter(converter_name: str) -> Callable[[dict], Sample]:
|
||||
if converter_name not in CONVERTERS:
|
||||
raise ValueError(f"Converter {converter_name} not found.")
|
||||
|
||||
return CONVERTERS[converter_name]
|
||||
|
||||
Reference in New Issue
Block a user