|
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] |
|
if len(messages) == 0: |
|
continue |
|
|
|
prompt = [] |
|
response = [] |
|
for turn_idx, message in enumerate(messages): |
|
if turn_idx % 2 == 0: |
|
accept_tags = [dataset_attr.user_tag, dataset_attr.observation_tag] |
|
else: |
|
accept_tags = [dataset_attr.assistant_tag, dataset_attr.function_tag] |
|
|
|
if message[dataset_attr.role_tag] not in accept_tags: |
|
raise ValueError("Invalid role tag in {}.".format(messages)) |
|
|
|
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) |
|
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 "") |
|
|
|
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) |
|
|
