diff --git a/pipeline/strategies/chatml.py b/pipeline/strategies/chatml.py index cd8411c..b4f1493 100644 --- a/pipeline/strategies/chatml.py +++ b/pipeline/strategies/chatml.py @@ -14,17 +14,18 @@ class ChatMLStrategy(PromptStrategy): def __init__( self, tokenizer: AutoTokenizer, - user_start: str = "<|im▁start|>user\n", - user_end: str = "<|im▁end|>\n", - assistant_start: str = "<|im▁start|>assistant\n", - assistant_end: str = "<|im▁end|>\n", + user_start: str = "<|im▁start|>user", + user_end: str = "<|im▁end|>", + assistant_start: str = "<|im▁start|>assistant", + assistant_end: str = "<|im▁end|>", ): super().__init__(tokenizer) + nl_id = tokenizer.encode("a\nb", add_special_tokens=False)[1] - self._user_start_ids = self._encode_format(user_start) - self._user_end_ids = self._encode_format(user_end) - self._assistant_start_ids = self._encode_format(assistant_start) - self._assistant_end_ids = self._encode_format(assistant_end) + self._user_start_ids = self._encode_format(user_start) + [nl_id] + self._user_end_ids = self._encode_format(user_end) + [nl_id] + self._assistant_start_ids = self._encode_format(assistant_start) + [nl_id] + self._assistant_end_ids = self._encode_format(assistant_end) + [nl_id] @property def name(self) -> str: diff --git a/scripts/supervised_finetuning/sft_metamathqa.py b/scripts/supervised_finetuning/sft_metamathqa.py new file mode 100644 index 0000000..a26ff59 --- /dev/null +++ b/scripts/supervised_finetuning/sft_metamathqa.py @@ -0,0 +1,17 @@ +from datasets import load_dataset +from pipeline import export_dataset + + +def process_func(sample: dict) -> dict: + return {"query": sample["query"], "response": sample["response"]} + + +if __name__ == "__main__": + dataset = load_dataset("meta-math/MetaMathQA", split="train") + export_dataset( + dataset=dataset, + output_dir="./dataset", + output_prefix="MetaMathQA", + process_func=process_func, + chunk_size=1_000_000, + )