diff --git a/belle_sft.py b/belle_sft.py index 4d6a521..8818f4a 100644 --- a/belle_sft.py +++ b/belle_sft.py @@ -1,10 +1,39 @@ from datasets import load_dataset from utils import process_dataset + +def build_prompt(query, history) -> str: + ret_prompt = "" + if len(history) > 0: + for his_query, his_response in history: + ret_prompt += f"<|user|> {his_query} <|system|> {his_response}\n" + if query is not None: + ret_prompt += f"<|user|> {query} <|system|> " + return ret_prompt + + +def process_func(input_dict: dict): + conversations = input_dict["conversations"] + n = len(conversations) // 2 + examples = [] + + for i in range(n): + user_msg = conversations[2*i]["value"] + assistant_msg = conversations[2*i+1]["value"] + examples.append((user_msg, assistant_msg)) + + content = { + "text": build_prompt(None, examples) + } + print(content) + return content + + if __name__ == "__main__": dataset = load_dataset("BelleGroup/train_3.5M_CN") - # process_dataset( - # dataset_dict=dataset, - # output_subdir="belle_sft", - # max_chunk_size=5, - # ) \ No newline at end of file + process_dataset( + dataset_dict=dataset, + output_subdir="belle_sft", + chunk_size=5, + process_func=process_func, + ) \ No newline at end of file diff --git a/utils.py b/utils.py index 8dc740d..7330a26 100644 --- a/utils.py +++ b/utils.py @@ -1,4 +1,4 @@ -from typing import List, Callable +from typing import List, Callable, Union from datasets import DatasetDict from tokenizer import BpeTokenizer from tqdm import tqdm @@ -64,7 +64,7 @@ def process_dataset( chunk_size: int = 1000000, split_name: str = "train", column_name: str = "text", - process_func: Callable[[dict], dict] = None, + process_func: Callable[[Union[dict, List[dict]]], dict] = None, normalization_func=comprehensive_normalization, ): train_dataset = dataset_dict[split_name]