mirror of
https://github.com/hiyouga/LLaMA-Factory.git
synced 2025-08-04 20:52:59 +08:00
110 lines
4.5 KiB
Python
110 lines
4.5 KiB
Python
from functools import partial
|
|
from typing import TYPE_CHECKING, Any, Dict, List, Union
|
|
|
|
from .utils import Role
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from datasets import Dataset, IterableDataset
|
|
|
|
from ..hparams import DataArguments
|
|
from .parser import DatasetAttr
|
|
|
|
|
|
def convert_alpaca(examples: Dict[str, List[Any]], dataset_attr: "DatasetAttr") -> Dict[str, List[Any]]:
|
|
outputs = {"prompt": [], "response": [], "system": [], "tools": []}
|
|
for i in range(len(examples[dataset_attr.prompt])):
|
|
prompt = []
|
|
if dataset_attr.history and isinstance(examples[dataset_attr.history][i], list):
|
|
for old_prompt, old_response in examples[dataset_attr.history][i]:
|
|
prompt.append({"role": Role.USER, "content": old_prompt})
|
|
prompt.append({"role": Role.ASSISTANT, "content": old_response})
|
|
|
|
instruction = examples[dataset_attr.prompt][i]
|
|
if dataset_attr.query and examples[dataset_attr.query][i]:
|
|
instruction += "\n" + examples[dataset_attr.query][i]
|
|
prompt.append({"role": Role.USER, "content": instruction})
|
|
|
|
if dataset_attr.response and isinstance(examples[dataset_attr.response][i], list):
|
|
response = [{"role": Role.ASSISTANT, "content": content} for content in examples[dataset_attr.response][i]]
|
|
elif dataset_attr.response and isinstance(examples[dataset_attr.response][i], str):
|
|
response = [{"role": Role.ASSISTANT, "content": examples[dataset_attr.response][i]}]
|
|
else:
|
|
response = []
|
|
|
|
outputs["prompt"].append(prompt)
|
|
outputs["response"].append(response)
|
|
outputs["system"].append(examples[dataset_attr.system][i] if dataset_attr.system else "")
|
|
outputs["tools"].append("")
|
|
|
|
return outputs
|
|
|
|
|
|
def convert_sharegpt(examples: Dict[str, List[Any]], dataset_attr: "DatasetAttr") -> Dict[str, List[Any]]:
|
|
outputs = {"prompt": [], "response": [], "system": [], "tools": []}
|
|
tag_mapping = {
|
|
dataset_attr.user_tag: Role.USER,
|
|
dataset_attr.assistant_tag: Role.ASSISTANT,
|
|
dataset_attr.observation_tag: Role.OBSERVATION,
|
|
dataset_attr.function_tag: Role.FUNCTION,
|
|
}
|
|
for i, messages in enumerate(examples[dataset_attr.messages]):
|
|
messages = messages[: len(messages) // 2 * 2] # should be multiples of 2
|
|
if len(messages) == 0:
|
|
continue
|
|
|
|
n_sys = 0
|
|
prompt = []
|
|
response = []
|
|
for turn_idx, message in enumerate(messages):
|
|
accept_tags = [dataset_attr.user_tag, dataset_attr.observation_tag, dataset_attr.assistant_tag, dataset_attr.function_tag]
|
|
|
|
if message[dataset_attr.role_tag] == "system":
|
|
outputs["system"].append(message[dataset_attr.content_tag])
|
|
n_sys += 1
|
|
elif message[dataset_attr.role_tag] not in accept_tags:
|
|
print("sytem attr", dataset_attr.system)
|
|
print("accepted tags", accept_tags)
|
|
raise ValueError("Invalid role tag in {}.".format(messages))
|
|
else:
|
|
prompt.append(
|
|
{"role": tag_mapping[message[dataset_attr.role_tag]], "content": message[dataset_attr.content_tag]}
|
|
)
|
|
|
|
last_message = prompt.pop(-1)
|
|
response.append(last_message)
|
|
outputs["prompt"].append(prompt)
|
|
outputs["response"].append(response)
|
|
if n_sys == 0:
|
|
outputs["system"].append(examples[dataset_attr.system][i] if dataset_attr.system else "")
|
|
outputs["tools"].append(examples[dataset_attr.tools][i] if dataset_attr.tools else "")
|
|
assert n_sys <= 1
|
|
return outputs
|
|
|
|
|
|
def align_dataset(
|
|
dataset: Union["Dataset", "IterableDataset"], dataset_attr: "DatasetAttr", data_args: "DataArguments"
|
|
) -> Union["Dataset", "IterableDataset"]:
|
|
r"""
|
|
Aligned dataset:
|
|
prompt: [{"role": "user", "content": "..."}]
|
|
response: [{"role": "assistant", "content": "..."}]
|
|
system: "..."
|
|
tools: "..."
|
|
"""
|
|
if dataset_attr.formatting == "alpaca":
|
|
convert_func = partial(convert_alpaca, dataset_attr=dataset_attr)
|
|
else:
|
|
convert_func = partial(convert_sharegpt, dataset_attr=dataset_attr)
|
|
|
|
column_names = list(next(iter(dataset)).keys())
|
|
kwargs = {}
|
|
if not data_args.streaming:
|
|
kwargs = dict(
|
|
num_proc=data_args.preprocessing_num_workers,
|
|
load_from_cache_file=(not data_args.overwrite_cache),
|
|
desc="Converting format of dataset",
|
|
)
|
|
|
|
return dataset.map(convert_func, batched=True, remove_columns=column_names, **kwargs)
|