diff --git a/src/ctx_to_lora/data/__init__.py b/src/ctx_to_lora/data/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/ctx_to_lora/data/collator.py b/src/ctx_to_lora/data/collator.py new file mode 100644 index 0000000..2cf62d8 --- /dev/null +++ b/src/ctx_to_lora/data/collator.py @@ -0,0 +1,58 @@ +import torch +from transformers.data import DataCollatorWithFlattening + +flattener = DataCollatorWithFlattening() + + +def train_packed_collator(inp_list): + # no padding + packed_inputs = flattener(inp_list, return_tensors="pt") + if "ctx_ids" in inp_list[0]: + ctx_ids = [{"input_ids": example["ctx_ids"]} for example in inp_list] + packed_ctx = flattener(ctx_ids, return_tensors="pt") + packed_inputs["ctx_ids"] = packed_ctx["input_ids"] + packed_inputs["ctx_position_ids"] = packed_ctx["position_ids"] + return packed_inputs + + +def train_collator(inp_list, tokenizer): + # input is a list of tokenized sequences + padding_kwargs = dict( + padding=True, + padding_side="right", + pad_to_multiple_of=8, + return_tensors="pt", + ) + + ctx_ids = None + if "ctx_ids" in inp_list[0]: + # have to be manual since it has [ctx_len, features] shape + # pad to the longest ctx_len in the batch + # which can have a different length from the input_ids, attn_mask, labels + + ctx_ids = [example.pop("ctx_ids") for example in inp_list] + ctx_ids = torch.nn.utils.rnn.pad_sequence( + ctx_ids, + batch_first=True, + padding_value=0, + ) + # exotic keys won't be padded, so we need to pad them as well + ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list] + ctx_attn_mask = torch.nn.utils.rnn.pad_sequence( + ctx_attn_mask, + batch_first=True, + padding_value=0, + ) + + labels = [x.pop("labels") for x in inp_list] + padded_seq = tokenizer.pad(inp_list, **padding_kwargs) + + # hacky explicit padding since the labels are not padded by default + labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"] + labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels) + out = {**padded_seq, "labels": labels} + if ctx_ids is not None: + out["ctx_ids"] = ctx_ids + out["ctx_attn_mask"] = ctx_attn_mask + + return out diff --git a/src/ctx_to_lora/data/definitions.py b/src/ctx_to_lora/data/definitions.py new file mode 100644 index 0000000..eb5e40c --- /dev/null +++ b/src/ctx_to_lora/data/definitions.py @@ -0,0 +1,377 @@ +from glob import glob + +IGNORE_INDEX = -100 + +FW_QA_PATHS = [ + f"data/raw_datasets/fw_qa/{i:05d}.parquet" for i in [0, 1, 6, 7, 8, 10, 22, 30, 35] +] + +TRANSFORMED_DATA_DIR = "data/processed_datasets" + +# approximate length of the datasets (train split) +# needed for streaming datasets +DS_LEN = { + "fw_qa_3_mini": 100_000, + "fw_qa_3_medium": 121_000_000, + "fw_qa_3": 270_000_000, + "fw_qa_xl": 27_000_000, + "ctx_qa": 300_000, + "pwc": 240_000, + "hotpot_qa": 100_000, + "squad": 90_000, + "drop": 77_000, + "narrativeqa": 40_000, + "quoref": 11_000, + "ropes": 11_000, + "synthetic_convqa": 40_000, +} + +LONGBENCH_TASKS = [ + "longbench/narrativeqa", + "longbench/qasper", + "longbench/multifieldqa_en", + "longbench/multifieldqa_zh", + "longbench/hotpotqa", + "longbench/2wikimqa", + "longbench/musique", + "longbench/dureader", + "longbench/gov_report", + "longbench/qmsum", + "longbench/multi_news", + "longbench/vcsum", +] + +LONGBENCH_E_TASKS = [ + "longbench/qasper_e", + "longbench/multifieldqa_en_e", + "longbench/hotpotqa_e", + "longbench/2wikimqa_e", + "longbench/gov_report_e", + "longbench/multi_news_e", +] + +DS_KWARGS = { + "hotpot_qa": dict( + train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:]"), + validation=dict( + path="hotpotqa/hotpot_qa", name="fullwiki", split="train[:900]" + ), + test=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="validation"), + ), + "hotpot_qa_tiny": dict( + train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:2000]"), + validation=dict( + path="hotpotqa/hotpot_qa", name="fullwiki", split="train[:900]" + ), + ), + "pwc": dict( + train=dict(path="sggetao/PwC", split="train[900:]"), + validation=dict(path="sggetao/PwC", split="train[:900]"), + test=dict(path="sggetao/PwC", split="test"), + ), + "pwc_tiny": dict( + train=dict(path="sggetao/PwC", split="train[900:2000]"), + validation=dict(path="sggetao/PwC", split="train[:900]"), + ), + "squad": dict( + train=dict(path="rajpurkar/squad", split="train[900:]"), + validation=dict(path="rajpurkar/squad", split="train[:900]"), + test=dict(path="rajpurkar/squad", split="validation"), + ), + # "fw_qa_tiny": dict( + # train=dict( + # path="parquet", + # data_files="data/raw_datasets/fw_qa/00000.parquet", + # split="train", + # ), + # validation=dict( + # path="parquet", + # data_files="data/raw_datasets/fw_qa/00000_val.parquet", + # split="train", + # ), + # ), + # "fw_qa": dict( + # train=dict( + # path="parquet", + # data_files=FW_QA_PATHS, + # split="train", + # ), + # validation=dict( + # path="parquet", + # data_files="data/raw_datasets/fw_qa/*_val.parquet", + # split="train", + # ), + # ), + # "fw_qa_large": dict( + # train=dict( + # path="parquet", + # data_files=glob("data/raw_datasets/fw_qa_large/*res.parquet"), + # split="train", + # ), + # validation=dict( + # path="parquet", + # data_files=glob("data/raw_datasets/fw_qa_large/*val.parquet"), + # split="train", + # ), + # ), + "fw_qa_xl": dict( + train=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_xl/*[!val].parquet"), + split="train", + ), + validation=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_xl/*val.parquet"), + split="train", + ), + # there is no explcit test for this ds + # used only for sanity check + test=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_xl/*val.parquet"), + split="train", + ), + ), + # "fw_qa_2": dict( + # train=dict( + # path="parquet", + # data_files=glob("data/raw_datasets/fw_qa_2/*[!val].parquet"), + # split="train", + # ), + # validation=dict( + # path="parquet", + # data_files=glob("data/raw_datasets/fw_qa_2/*val.parquet"), + # split="train", + # ), + # ), + "fw_qa_3_mini": dict( + train=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_3/000_00001.parquet"), + split="train[:100000]", + ), + ), + "fw_qa_3_medium": dict( + train=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_3/00[0-5]*[!val].parquet"), + split="train", + ), + ), + "fw_qa_3": dict( + train=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_3/*[!val].parquet"), + split="train", + ), + validation=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_3/*val.parquet"), + split="train", + ), + ), + "ctx_qa": dict( + train=dict( + path="parquet", + data_files=glob("data/raw_datasets/ctx_qa/*res.parquet"), + split="train", + ), + validation=dict( + path="parquet", + data_files=glob("data/raw_datasets/ctx_qa/*val.parquet"), + split="train", + ), + # there is no explcit test for this ds + # used only for sanity check + test=dict( + path="parquet", + data_files=glob("data/raw_datasets/ctx_qa/*val.parquet"), + split="train", + ), + ), + "drop": dict( + train=dict( + path="ucinlp/drop", + split="train", + ), + ), + "narrativeqa": dict( + train=dict( + path="nvidia/ChatQA-Training-Data", + name="narrativeqa", + split="train", + ), + ), + "quoref": dict( + train=dict( + path="nvidia/ChatQA-Training-Data", + name="quoref", + split="train", + ), + ), + "ropes": dict( + train=dict( + path="nvidia/ChatQA-Training-Data", + name="ropes", + split="train", + ), + ), + "synthetic_convqa": dict( + train=dict( + path="nvidia/ChatQA-Training-Data", + name="synthetic_convqa", + split="train", + ), + ), + "booksum": dict( + train=dict( + path="kmfoda/booksum", + split="train", + ), + validation=dict( + path="kmfoda/booksum", + split="validation", + ), + test=dict( + path="kmfoda/booksum", + split="test", + ), + ), + "gov_report": dict( + train=dict( + path="ccdv/govreport-summarization", + split="train", + ), + validation=dict( + path="ccdv/govreport-summarization", + split="validation", + ), + test=dict( + path="ccdv/govreport-summarization", + split="test", + ), + ), + # "wikitext-2": dict( + # train=dict( + # path="EleutherAI/wikitext_document_level", + # name="wikitext-2-raw-v1", + # split="train", + # ), + # ), + # "wikitext-103": dict( + # train=dict( + # path="EleutherAI/wikitext_document_level", + # name="wikitext-103-raw-v1", + # split="train", + # ), + # ), + "gsm8k": dict( + train=dict( + path="openai/gsm8k", + name="main", + split="train[100:]", + ), + validation=dict( + path="openai/gsm8k", + name="main", + split="train[:100]", + ), + ), + "openmathintx-2": dict( + train=dict( + path="nvidia/OpenMathInstruct-2", + split="train_1M[:100000]", + ), + ), + "opencoder-edu": dict( + train=dict( + path="OpenCoder-LLM/opc-sft-stage2", + name="educational_instruct", + split="train[900:]", + ), + validation=dict( + path="OpenCoder-LLM/opc-sft-stage2", + name="educational_instruct", + split="train[:900]", + ), + ), + "triviaqa_retrieved": dict( + test=dict( + path="json", + data_files="data/eval/triviaqa/triviaqa_retrieved.jsonl", + split="train", + ), + ), + # for negative training and evaluating + "negative_nq": dict( + test=dict( + path="json", + data_files="data/raw_datasets/negative_natural_questions/validation.jsonl", + split="train", + ), + ), +} + +# LongBench kwargs +for ds_name in LONGBENCH_TASKS + LONGBENCH_E_TASKS: + DS_KWARGS[ds_name] = dict( + test=dict( + path="THUDM/LongBench", + name=ds_name.split("/")[-1], + split="test", + ) + ) + +# for training closed qa datasets, e.g., hotpot_qa, squad, etc. +CLOSED_QA_INTX_TEMPLATES = [ + "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}", + "Answer without any explanation.\n\nQuestion: {input}", + "Based on the provided text, what is the answer to the following question? Provide only the answer.\n\nQuestion: {input}", + "Extract the answer to the question from the text. Be concise. Do not explain.\n\nQuestion: {input}", + "What is the answer to this question, based on the context? Respond with the answer only.\n\nQuestion: {input}", + "Provide a direct answer to the question using the given passages. Do not give any explanation.\n\nQuestion: {input}", + "Answer the question using only information from the provided text. No extra words.\n\nQuestion: {input}", + "From the passages, answer the question. Just the answer, please.\n\nQuestion: {input}", + "Give the answer to the question. Do not include any other text.\n\nQuestion: {input}", + "The answer to the question is in the text. Find it and state it clearly. No need for explanation.\n\nQuestion: {input}", + "Concisely answer the question based on the text provided. Don't include any other words. Just the answer.\n\nQuestion: {input}", + "Read the passages and answer the question with the minimal necessary words.\n\nQuestion: {input}", + "What is the direct response to the question, according to the text? Avoid explanation.\n\nQuestion: {input}", + "Please provide only the answer to the question, derived from the text.\n\nQuestion: {input}", + "Using the provided context, answer the question. Output the answer and nothing else.\n\nQuestion: {input}", + "Identify the answer in the text and present it without elaboration.\n\nQuestion: {input}", + "Answer the following question based on the text. Your answer should be brief and to the point. No explanation.\n\nQuestion: {input}", + "Based on the information given, what is the answer to the question? Only state the answer.\n\nQuestion: {input}", + "Find the answer to the question in the provided passages and write it down. No explanations.\n\nQuestion: {input}", + "The question is: {input}. Provide the answer based on the text, and nothing more.", + "Question: {input}\nAnswer directly based on the text provided. No extra words.", + "Question: {input}\nPlease provide the answer based on the text. No explanation is needed.", +] + + +EVAL_INTX_TEMPLATES = { + "negative_nq": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}", + "triviaqa_retrieved": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}", + "hotpot_qa": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}", + "squad": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}", + "longbench/qasper": 'Answer the question as concisely as you can, using a single phrase or sentence if possible.\nIf the question cannot be answered based on the information in the article, write "unanswerable".\nIf the question is a yes/no question, answer "yes", "no", or "unanswerable". Do not provide any explanation.\n\nQuestion: {input}', + "longbench/narrativeqa": "Answer the question as concisely as you can, using a single phrase if possible. Do not provide any explanation.\n\nQuestion: {input}", + "longbench/multifieldqa_en": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}", + # "longbench/multifieldqa_zh": "阅读以下文字并用中文简短回答：\n\n现在请基于上面的文章回答下面的问题，只告诉我答案，不要输出任何其他字词。\n\n问题：{input}", + "longbench/hotpotqa": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}", + "longbench/2wikimqa": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}", + "longbench/musique": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}", + # + # "longbench/dureader": "请基于给定的文章回答下述问题\n\n请基于上述文章回答下面的问题。\n\n问题：{input}", + "longbench/gov_report": "You are given a report by a government agency. Write a one-page summary of the report.", + "longbench/qmsum": "You are given a meeting transcript and a query containing a question or instruction. Answer the query based on the above meeting transcript.\n\nQuery: {input}", + "longbench/multi_news": "You are given several news passages. Write a one-page summary of all the news.", + # "longbench/vcsum": "下面有一段会议记录，请你阅读后，写一段总结，总结会议的内容。", +} +for ds_name in LONGBENCH_E_TASKS: + EVAL_INTX_TEMPLATES[ds_name] = EVAL_INTX_TEMPLATES[ds_name[:-2]] + +for ds_name in DS_KWARGS: + if ds_name not in EVAL_INTX_TEMPLATES: + EVAL_INTX_TEMPLATES[ds_name] = "{input}" diff --git a/src/ctx_to_lora/data/pretrain.py b/src/ctx_to_lora/data/pretrain.py new file mode 100644 index 0000000..213c900 --- /dev/null +++ b/src/ctx_to_lora/data/pretrain.py @@ -0,0 +1,5 @@ +# when generating dont forget to remove dup contexts +# merge QAs from dup contexts to construct +# a doc with QAs at the end +# maybe mix with some that doesn;t have QAs +# TODO: implement diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py new file mode 100644 index 0000000..b516f27 --- /dev/null +++ b/src/ctx_to_lora/data/processing.py @@ -0,0 +1,985 @@ +import hashlib +import json +import logging +import random +from collections.abc import Callable, Iterator +from os import path +from typing import Any + +import datasets +import numpy as np +from datasets import load_dataset +from transformers import PreTrainedTokenizerBase + +from ctx_to_lora.data.definitions import ( + CLOSED_QA_INTX_TEMPLATES, + DS_KWARGS, + EVAL_INTX_TEMPLATES, + IGNORE_INDEX, + TRANSFORMED_DATA_DIR, +) +from ctx_to_lora.utils import TRAINING_TASK + +logger = logging.getLogger() + +# TODO: add continued pre-training pipeline + + +def closed_qa_prompting(prompt: str): + template = random.choice(CLOSED_QA_INTX_TEMPLATES) + return template.format(input=prompt) + + +# TODO: dont use closed_qa_prompting when training on self-generated responses +def get_preprocessing_fn( + ds_name: str, is_eval: bool +) -> Callable[[dict[str, Any]], dict[str, Any]]: + """ + Get preprocessing function for a specific dataset. + + Args: + ds_name: Name of the dataset + + Returns: + A preprocessing function that takes and returns a dictionary + """ + f = lambda x: x + + if ds_name.startswith("pwc"): + + def f(sample): + return { + "context": sample["input"], + "prompt": sample["prompt"], + "response": sample["answer"], + } + + elif ds_name.startswith("hotpot_qa"): + + def f(sample): + txt = "" + for p in sample["context"]["sentences"]: + txt += " " + "".join(p) + + q = sample["question"] + prompt = closed_qa_prompting(q) if not is_eval else q + + return { + "context": txt.strip(), + "prompt": prompt, + "response": sample["answer"], + } + + elif "squad" in ds_name: + + def f(sample): + q = sample["question"] + prompt = closed_qa_prompting(q) if not is_eval else q + return { + "context": sample["context"], + "prompt": prompt, + "response": sample["answers"]["text"][0], + } + + elif ds_name == "drop": + + def f(sample): + q = sample["question"] + prompt = closed_qa_prompting(q) if not is_eval else q + return { + "context": sample["passage"], + "prompt": prompt, + "response": ", ".join(sample["answers_spans"]["spans"]), + } + + elif ds_name in ["narrativeqa", "quoref", "ropes"]: + + def f(sample): + response = sample["answers"][0] + if isinstance(response, list): + response = response[0] + q = sample["messages"][-1]["content"] + prompt = closed_qa_prompting(q) if not is_eval else q + return { + "context": sample["document"], + "prompt": prompt, + "response": response, + } + + elif ds_name == "synthetic_convqa": + + def f(sample): + response = sample["answers"][0] + if isinstance(response, list): + response = response[0] + return { + "context": sample["document"], + "prompt": sample["messages"][-1]["content"], + "response": response, + } + + elif ds_name == "booksum": + prompt_templates = [ + "Summarization the provided text.", + "# Summary", + "### Summary", + "Summary of the text", + ] + # TODO: use these templates for training + # analysis_templates = [] + # summary_templates = [] + + def f(sample): + return { + "context": sample["chapter"], + "prompt": "Summarize the provided text.", + "response": sample["summary_text"], + } + + elif ds_name == "gov_report": + + def f(sample): + return { + "context": sample["report"], + "prompt": "Summarize the provided text.", + "response": sample["summary"], + } + + elif "wikitext" in ds_name: + + def f(sample): + return { + "context": sample["page"], + "prompt": "PLAECHOLDER", + "response": "PLAECHOLDER", + } + + elif ds_name == "openmathintx-2": + + def f(sample): + return { + "context": sample["problem"], + "prompt": sample["problem"], + "response": sample["generated_solution"], + } + + elif "gsm8k" in ds_name: + + def f(sample): + return { + "context": sample["question"], + "prompt": sample["question"], + "response": sample["answer"], + } + + elif "opencoder-edu" in ds_name: + + def f(sample): + return { + "context": sample["instruction"], + "prompt": sample["instruction"], + "response": "```python\n" + sample["code"].strip() + "\n```", + } + + elif ds_name.startswith("longbench"): + + def f(sample): + return { + "context": sample["context"], + "prompt": sample["input"], + "response": sample["answers"][0], + } + + elif ds_name == "negative_nq": + + def f(sample): + q = sample["prompt"] + prompt = closed_qa_prompting(q) if not is_eval else q + return { + "context": sample["context"], + "prompt": prompt, + "response": sample["answer"], + } + + elif ds_name == "triviaqa_retrieved": + # only used for eval + def f(sample): + return { + "context": sample["context"], + "prompt": sample["prompt"], + "response": sample["answer"], + } + + if is_eval and ds_name in EVAL_INTX_TEMPLATES: + prompt_template = EVAL_INTX_TEMPLATES[ds_name] + + def eval_intx_decorator(f): + def g(sample): + sample = f(sample) + sample["prompt"] = prompt_template.format(input=sample["prompt"]) + return sample + + return g + + f = eval_intx_decorator(f) + + return f + + +def get_ds_kwargs(ds_name: str, split: str) -> dict[str, Any]: + if (ds_name not in DS_KWARGS) or (split not in DS_KWARGS[ds_name]): + kwargs = dict(path=ds_name, split=split) + logger.warning( + f"No dataset kwargs found for '{ds_name}' with split '{split}'.\n" + f"Using default kwargs: {kwargs}" + ) + # return kwargs + else: + # return DS_KWARGS[ds_name][split] + kwargs = DS_KWARGS[ds_name][split] + + # custom logic for slicing iterable datasets + take, skip = None, None + kw_split = kwargs["split"] + if ("[" in kw_split) and kw_split.endswith("]"): + kwargs["split"], slice = kw_split.split("[") + slice = slice.strip("]") + skip = slice.split(":")[0] + if skip: + kwargs["skip"] = int(skip) + take = slice.split(":")[1] + if take: + kwargs["take"] = int(take) + return kwargs + + +def validate_columns(tokenized_ds): + cols = ["input_ids", "attention_mask", "labels"] + if "ctx_ids" in tokenized_ds.column_names: + cols += ["ctx_ids", "ctx_attn_mask"] + if "chat_ids" in tokenized_ds.column_names: + cols += ["chat_ids", "chat_attn_mask", "chat_labels"] + ref_cols = set(cols) + assert set(tokenized_ds.column_names) == ref_cols, ( + f"Columns mismatch: {set(tokenized_ds.column_names)} != {ref_cols}" + ) + + +def filter_long_samples(samples): + return [len(ctx) < 10000 for ctx in samples["context"]] + + +def filter_long_chat(samples): + return [len(chat) < 2**13 for chat in samples["chat"]] + + +def add_repeat_prompt_fn(samples): + unique_contexts = set() + ctxs, prompts, responses = [], [], [] + for ctx, prompt, response in zip( + samples["context"], samples["prompt"], samples["response"] + ): + # Only process if the context is not already in the set + if ctx in unique_contexts: + continue + # remove too long contexts + if len(ctx) > 1000: + continue + unique_contexts.add(ctx) + ctxs.append(ctx) + responses.append(ctx) + prompts.append("Repeat the text above.") + + logger.debug("Adding repeat prompt...") + logger.debug(f"# unique contexts: {len(unique_contexts)}") + + return dict( + context=ctxs + samples["context"], + response=responses + samples["response"], + prompt=prompts + samples["prompt"], + ) + + +def add_negative_prompt_fn(samples): + unique_contexts = set() + ctxs, prompts, responses = [], [], [] + keywords = [ + "repeat", + "rephrase", + "summarize", + "rewrite", + "title", + "keyword", + "continuation", + ] + + for ctx, prompt, response in zip( + samples["context"], samples["prompt"], samples["response"] + ): + if ctx in unique_contexts: + continue + unique_contexts.add(ctx) + if len(ctx) > 3000: + continue + if any(keyword in prompt for keyword in keywords): + # Skip samples where the prompt contains any of the specified keywords + continue + + ctxs.append(ctx) + prompts.append(prompt) + responses.append(response) + + logger.debug("Adding negative prompt...") + logger.debug(f"# unique contexts: {len(unique_contexts)}") + + # remove one last sample if the number of samples is odd + if len(ctxs) % 2 != 0: + ctxs.pop() + prompts.pop() + responses.pop() + + # to make sure that the negative prompt/response is not the same as the original + indices = list(np.random.permutation(len(ctxs))) + list( + np.random.permutation(len(ctxs)) + ) + neg_ctxs, neg_prompts, neg_responses = [], [], [] + for idx in range(0, len(indices), 2): + i = indices[idx] + j = indices[idx + 1] + neg_ctxs.append(ctxs[i]) + neg_prompts.append(ctxs[j] + "\n\n" + prompts[j]) + neg_responses.append(responses[j]) + + return dict( + context=neg_ctxs + samples["context"], + prompt=neg_prompts + samples["prompt"], + response=neg_responses + samples["response"], + ) + + +def filter_none(samples): + out = [True] * len(samples["context"]) + for k in samples: + for i, sample in enumerate(samples[k]): + if not bool(sample): + out[i] = False + return out + + +def _load_and_process_dataset( + ds_name: str, + split: str, + add_negative_prompt: bool, + add_repeat_prompt: bool, + streaming: bool, + ds_path: str, + num_proc: int, +): + try: + ds_kwargs = get_ds_kwargs(ds_name, split) + skip = ds_kwargs.pop("skip", None) + take = ds_kwargs.pop("take", None) + # only load as stream if the dataset is local (parquet) + load_as_stream = ( + ("fw_qa" in ds_name) + and ("tiny" not in ds_name) + and ("mini" not in ds_name) + and streaming + and ds_kwargs["path"] == "parquet" + ) + ds = load_dataset( + **ds_kwargs, + trust_remote_code=True, + streaming=load_as_stream, + ) + if split == "train" and streaming and (not load_as_stream): + # if the dataset is hosted on HF, we load it first then convert to iterable + ds = ds.to_iterable_dataset(num_shards=128) + if skip is not None: + ds = ds.skip(skip) + if take is not None: + ds = ds.take(take) + except ValueError as e: + logger.warning( + f"Failed to load dataset {ds_name} with split {split}. Error: {e}\nSkipping..." + ) + return None + cols_to_remove = [ + col for col in ds.column_names if col not in ["context", "prompt", "response"] + ] + is_eval = split != "train" + ds = ds.map( + get_preprocessing_fn(ds_name, is_eval), + remove_columns=cols_to_remove, + num_proc=num_proc, + ) + # ds = ds.remove_columns(cols_to_remove) + ds = ds.filter(filter_none, batched=True, num_proc=num_proc) + + if split == "train": + # TODO: drop samples or split into multiple contexts + ds = ds.filter(filter_long_samples, batched=True, num_proc=num_proc) + if add_negative_prompt: + # TODO: make explicit dataset for this + ds = ds.map( + add_negative_prompt_fn, + batched=True, + batch_size=100_000, + num_proc=num_proc, + ) + if add_repeat_prompt and "context_numbers" not in ds_name: + # TODO: convert to AE/LM + ds = ds.map( + add_repeat_prompt_fn, + batched=True, + batch_size=100_000, + num_proc=num_proc, + ) + ds.save_to_disk(ds_path, num_proc=num_proc) + return ds + + +def get_tokenized_dataset( + ds_name: str, + split: str, + base_model_max_len: int, + tokenizer: PreTrainedTokenizerBase, + tokenizer_kwargs: dict[str, Any], + ctx_model_max_len: int, + ctx_tokenizer: PreTrainedTokenizerBase, + ctx_tokenizer_kwargs: dict[str, Any], + add_ctx_to_chat: bool, + add_repeat_prompt: bool, + add_negative_prompt: bool, + use_kl_loss: bool, + set_format: str | None = None, + streaming: bool = False, +) -> dict[str, Any]: + # TODO: update hash function to include max_len + # TODO: max_len could be included in the kwargs + assert not use_kl_loss, "KL loss is deprecated" + logger.debug(f"Loading dataset {ds_name} with split {split}...") + need_ctx_ids = not add_ctx_to_chat and bool(ctx_model_max_len) + + load_and_process_kwargs = dict( + ds_name=ds_name, + split=split, + add_negative_prompt=add_negative_prompt, + add_repeat_prompt=add_repeat_prompt, + streaming=streaming, + ) + num_proc = None if streaming and split == "train" else 8 + + ds_hash = hashlib.sha256(json.dumps(load_and_process_kwargs).encode()).hexdigest() + ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}" + + if path.exists(ds_path) and split == "train": + # load the cached ds + logger.info(f"Loaded processed dataset from {ds_path}") + else: + logger.info(f"Loading dataset {ds_name} with split {split}...") + ds = _load_and_process_dataset( + **load_and_process_kwargs, + ds_path=ds_path, + num_proc=num_proc, + ) + if split == "train": + # force loading the cached version for newly created ds + ds = datasets.load_from_disk(ds_path) + + tokenized_ds = construct_and_tokenize_ctx_qa( + base_model_max_len, + tokenizer, + tokenizer_kwargs, + ctx_model_max_len, + ctx_tokenizer, + ctx_tokenizer_kwargs, + add_ctx_to_chat, + use_kl_loss, + need_ctx_ids, + ds, + split, + set_format, + num_proc, + ) + return tokenized_ds + + +def construct_and_tokenize_ctx_qa( + base_model_max_len, + tokenizer, + tokenizer_kwargs, + ctx_model_max_len, + ctx_tokenizer, + ctx_tokenizer_kwargs, + add_ctx_to_chat, + use_kl_loss, + need_ctx_ids, + ds, + split, + set_format=None, + num_proc=None, +): + is_eval = split != "train" + # TODO: update hash function to include max_len + # TODO: max_len could be included in the kwargs + kwargs = dict( + tokenizer=repr(tokenizer), + tokenizer_kwargs=json.dumps(tokenizer_kwargs), + ctx_tokenizer=repr(ctx_tokenizer), + ctx_tokenizer_kwargs=json.dumps(ctx_tokenizer_kwargs), + add_ctx_to_chat=add_ctx_to_chat, + use_kl_loss=use_kl_loss, + need_ctx_ids=need_ctx_ids, + ds=ds._fingerprint, + # TODO: add split (after moving to continued pre-training pipeline) + set_format=set_format, + ) + kwargs_str = json.dumps(kwargs) + logger.debug(f"Tokenizing dataset with kwargs: {kwargs_str}") + ds_hash = hashlib.sha256(kwargs_str.encode()).hexdigest() + ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}" + if path.exists(ds_path) and split == "train": + # load the cached ds + logger.info(f"Loaded tokenized dataset from {ds_path}") + ds = datasets.load_from_disk(ds_path) + return ds + # for sft + chat_model, we need to convert the dataset to chat format + # add "messages" field + ds = ds.map( + convert_ctx_prompt_response_to_messages, + fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat}, + num_proc=num_proc, + ) + # add "chat" field + ds = ds.map( + get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer), + batched=True, + batch_size=100_000, + num_proc=num_proc, + ) + # tokenize the chat + mask the assistant inputs + # TODO: drop samples or split into multiple contexts + if split == "train": + ds = ds.filter( + filter_long_chat, + batched=True, + batch_size=100_000, + num_proc=num_proc, + ) + + # add "input_ids", "attention_mask", "labels" + tokenized_ds = ds.map( + tokenize_chat_messages, + fn_kwargs={ + "tokenizer": tokenizer, + "mask_assistant_inputs": True, + "tokenizer_kwargs": tokenizer_kwargs, + }, + num_proc=num_proc, + ) + if split == "train": + # drop? + # base model cant handle these samples naturally + # should skip + ... + else: + tokenized_ds = tokenized_ds.map( + truncate_middle_if_too_long, + fn_kwargs={ + "max_length": base_model_max_len, + "columns": ["input_ids", "attention_mask"], + "max_new_tokens": 256, + }, + batched=True, + batch_size=100_000, + num_proc=num_proc, + ) + tokenized_ds = tokenized_ds.map( + add_length_info, + fn_kwargs={"columns": ["input_ids"]}, + batched=True, + batch_size=100_000, + num_proc=num_proc, + ) + + # for use_kl_loss, we need "chat_ids" and "chat_attn_mask" + if use_kl_loss: + raise NotImplementedError("KL loss deprecated") + tokenized_ds = tokenized_ds.map( + convert_ctx_prompt_response_to_messages, + fn_kwargs={"add_ctx_to_chat": True}, + remove_columns=["messages"], + ) + tokenized_ds = tokenized_ds.map( + get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer), + remove_columns=["chat"], + ) + tokenized_ds = tokenized_ds.map( + tokenize_chat_messages, + fn_kwargs={ + "tokenizer": tokenizer, + "mask_assistant_inputs": True, + "tokenizer_kwargs": tokenizer_kwargs, + "for_kl_loss": True, + }, + ) + + if need_ctx_ids: + # tokenize the ctx_text to get ctx_ids and ctx_attn_mask + tokenized_ds = tokenized_ds.map( + tokenize_ctx_text, + fn_kwargs={"tokenizer": ctx_tokenizer}, + batched=True, + batch_size=100_000, + num_proc=num_proc, + ) + if split == "train": + # TODO: do something if ctx length is longer than the ctx model length + # e.g., drop or split for multi-lora training + ... + + else: + # truncate in the middle + tokenized_ds = tokenized_ds.map( + truncate_middle_if_too_long, + fn_kwargs={ + "max_length": ctx_model_max_len, + "columns": ["ctx_ids", "ctx_attn_mask"], + # no need to add new_tokens here (used only for the ctx encoder) + "max_new_tokens": 0, + }, + batched=True, + batch_size=100_000, + num_proc=num_proc, + ) + tokenized_ds = tokenized_ds.map( + add_length_info, + fn_kwargs={"columns": ["ctx_ids"]}, + batched=True, + batch_size=100_000, + num_proc=num_proc, + ) + tokenized_ds = tokenized_ds.remove_columns( + ["messages", "chat", "context", "prompt", "response"] + ) + if set_format: + tokenized_ds.set_format(type=set_format) + # if isinstance(tokenized_ds, IterableDataset): + # # the columns are unknown when using streaming dataset + # tokenized_ds = tokenized_ds._resolve_features() + # validate_columns(tokenized_ds) + if split == "train": + tokenized_ds.save_to_disk(ds_path, num_proc=num_proc) + del tokenized_ds + return datasets.load_from_disk(ds_path) + return tokenized_ds + + +def get_sft_prompt_formatting_fn( + sft_mode: TRAINING_TASK, + tokenizer: PreTrainedTokenizerBase, +) -> Callable[[dict[str, Any]], dict[str, Any]]: + """ + Get a function that formats examples for supervised fine-tuning. + + Args: + sft_mode: The training task type + tokenizer: The tokenizer to use for chat template application + + Returns: + A function that takes a training example and returns formatted data + + Raises: + NotImplementedError: If sft_mode is not COMPLETION or tokenizer has no chat template + """ + if sft_mode != TRAINING_TASK.COMPLETION: + raise NotImplementedError( + f"Training task {sft_mode} not supported. " + "Only completion is supported for now." + ) + + if tokenizer.chat_template is None: + raise NotImplementedError( + "Only chat models + SFT are supported. " + "Training with pre-training data is not supported yet." + ) + # TODO: add support for causal_lm training (for pre-training data) + # TODO: add support for recon training (for pre-training data) + # TODO: add support for non-chat models + + # def f(example): + # output_texts = ( + # dict(text=[]) if sft_mode == "causal_lm" else dict(prompt=[], response=[]) + # ) + # df = pd.DataFrame(dict(example)) + # for i, inp_txt in df.iterrows(): + # if sft_mode == "causal_lm": + # text = metadata["text_template"].format(**inp_txt) + # output_texts["text"].append(text) + # elif sft_mode == "completion": + # prompt = metadata["user_prompt_template"].format(**inp_txt) + # output_texts["prompt"].append(prompt) + # output_texts["response"].append(str(inp_txt[metadata["response_field"]])) + # return output_texts + + def f_intx(example): + chat_text = tokenizer.apply_chat_template( + example["messages"], tokenize=False, add_generation_prompt=False + ) + return dict(chat=chat_text) + + # return f if not apply_chat_template_fn is not None else f_intx + return f_intx + + +def convert_ctx_prompt_response_to_messages( + example: dict[str, Any], + add_ctx_to_chat: bool = True, +) -> dict[str, Any]: + """ + Convert context/prompt/response format to chat messages format. + + Args: + example: Dictionary containing 'prompt' and 'response' keys + add_ctx_to_chat: Whether to prepend context to the user message + + Returns: + Dictionary with added 'messages' key containing chat format + + Raises: + ValueError: If 'prompt' or 'response' keys are missing + """ + if "prompt" not in example or "response" not in example: + raise ValueError(f"'prompt' and 'response' are required. Got: {example}") + + system_msg = "" + if "system_message" in example: + system_msg = example["system_message"].strip() + + user_msg = example["prompt"].strip() + if add_ctx_to_chat: + user_msg = example["context"].strip() + "\n\n" + user_msg.strip() + + messages = [ + {"role": "system", "content": system_msg.strip()}, + {"role": "user", "content": user_msg.strip()}, + {"role": "assistant", "content": example["response"].strip()}, + ] + + return dict(messages=messages) + + +# adapted from https://github.com/huggingface/trl/issues/632#issuecomment-1972630547 +def get_assistant_start_end_indices( + messages: list[dict[str, str]], + tokenizer: PreTrainedTokenizerBase, + conversation_ids: list[int], +) -> list[tuple[int, int]]: + """ + Get the start and end indices of assistant messages in conversation text using tokenized messages. + + Args: + messages: List of message dictionaries with 'role' and 'content' keys + tokenizer: Tokenizer used for tokenization + conversation_ids: Tokenized conversation + + Returns: + List of (start, end) index tuples for assistant messages in the tokenized conversation + """ + assistant_ranges = [] + current_idx = 0 + + for idx, message in enumerate(messages): + # if message["role"] == "assistant": + # Tokenize the assistant message content + tokenized_message = tokenizer.encode( + message["content"], add_special_tokens=False + ) + # Find the start index of the assistant message in the tokenized conversation + for i in range(current_idx, len(conversation_ids)): + if conversation_ids[i : i + len(tokenized_message)] == tokenized_message: + start_idx = i + end_idx = i + len(tokenized_message) + current_idx = end_idx + if message["role"] == "assistant": + if idx == (len(messages) - 1): + # include everything after the last assistant message + assistant_ranges.append((start_idx, len(conversation_ids))) + else: + assistant_ranges.append((start_idx, end_idx)) + break + + return assistant_ranges + + +def get_masked_labels( + conversation_ids: dict[str, list[Any]], + assistant_ranges: list[tuple[int, int]], +) -> Iterator[int]: + """ + Generate masked labels for conversation, masking non-assistant tokens. + NOTE: This will also includes extra tokens between assistant and user messages + in multi-turn conversations. + E.g., {assistant_msg} <|eot_id|><|start_header_id|>user<|end_header_id|> {user_msg} + for Llama 3 models. + + Args: + conversation_ids: Dictionary with tokenized conversation info + assistant_ranges: List of (start, end) indices for assistant messages + + Yields: + Token ID or IGNORE_INDEX for each position + """ + for id_, (id_s, id_e) in list( + zip(conversation_ids["input_ids"], conversation_ids["offset_mapping"]) + ): + if any(id_s >= s and id_e <= e for s, e in assistant_ranges): + yield id_ + else: + yield IGNORE_INDEX + + +def tokenize_chat_messages( + example: dict[str, Any], + tokenizer: PreTrainedTokenizerBase, + mask_assistant_inputs: bool = True, + for_kl_loss: bool = False, + tokenizer_kwargs: dict[str, Any] | None = None, +) -> dict[str, list[int]]: + """ + Tokenize chat messages and optionally mask non-assistant tokens. + + Args: + example: Dictionary containing 'chat' and 'messages' keys + tokenizer: Tokenizer to use + mask_assistant_inputs: Whether to mask non-assistant tokens + for_kl_loss: change the column names to "chat_ids" and "chat_attn_mask" + tokenizer_kwargs: Additional arguments to pass to tokenizer + + Returns: + Dictionary containing tokenized inputs and labels + """ + # should be used only with chat models + text = example["chat"] + messages = example["messages"] + n_response = len([m for m in messages if m["role"] == "assistant"]) + if n_response != 1: + raise ValueError(f"Expected 1 assistant response. Got {n_response}.") + conversation_ids = tokenizer( + text, + # return_offsets_mapping=mask_assistant_inputs, + add_special_tokens=False, + truncation=False, + **(tokenizer_kwargs or {}), + ) + if tokenizer_kwargs and ( + len(conversation_ids["input_ids"]) >= tokenizer_kwargs["max_length"] + ): + raise ValueError( + f"Conversation length {len(conversation_ids['input_ids'])} exceeds max length {tokenizer_kwargs['max_length']}" + ) + + if mask_assistant_inputs: + assistant_ranges = get_assistant_start_end_indices( + messages, tokenizer, conversation_ids["input_ids"] + ) + # labels = get_masked_labels(conversation_ids, assistant_ranges) + labels = [ + id_ if any(s <= i < e for s, e in assistant_ranges) else IGNORE_INDEX + for i, id_ in enumerate(conversation_ids["input_ids"]) + ] + conversation_ids["labels"] = labels + # del conversation_ids["offset_mapping"] + else: + conversation_ids["labels"] = conversation_ids["input_ids"] + if for_kl_loss: + conversation_ids["chat_ids"] = conversation_ids.pop("input_ids") + conversation_ids["chat_attn_mask"] = conversation_ids.pop("attention_mask") + conversation_ids["chat_labels"] = conversation_ids.pop("labels") + return conversation_ids + + +def add_length_info(samples: dict[str, any], columns: list[str]) -> dict[str, any]: + for col in columns: + samples[f"{col}_len"] = [len(t) for t in samples[col]] + return samples + + +def truncate_middle_if_too_long( + samples: dict[str, any], + max_length: int, + columns: list[str], + max_new_tokens: int = 256, +) -> dict[str, any]: + """ + Truncate the middle of a list of tokens to fit within a maximum length. + + Args: + tokens: List of token IDs + max_length: Maximum length for the truncated tokens + + Returns: + List of truncated token IDs + """ + max_new_tokens_half = max_new_tokens // 2 + # leave max_new_tokens for generation + half = max_length // 2 - max_new_tokens_half + for col in columns: + # if len(samples[col]) <= max_length: + # return samples[col] + samples[col] = [ + t[:half] + t[-half:] if len(t) > max_length else t for t in samples[col] + ] + return samples + + +def tokenize_ctx_text( + samples: dict[str, Any], + tokenizer: PreTrainedTokenizerBase, +) -> dict[str, Any]: + if tokenizer.chat_template: + txt = tokenizer.apply_chat_template( + [ + [{"role": "user", "content": text.strip()}] + for text in samples["context"] + ], + tokenize=False, + add_generation_prompt=False, + ) + # special tokens are already added by the chat template + tokenized_text = tokenizer(txt, truncation=False, add_special_tokens=False) + else: + txt = samples["context"] + tokenized_text = tokenizer(txt, truncation=False) + + ctx_ids = tokenized_text["input_ids"] + ctx_attn_mask = tokenized_text["attention_mask"] + return dict(ctx_ids=ctx_ids, ctx_attn_mask=ctx_attn_mask) + + +if __name__ == "__main__": + from transformers import AutoTokenizer + + # model_name = "meta-llama/Llama-3.2-1B-Instruct" + # model_name = "meta-llama/Llama-2-7b-chat-hf" + model_name = "google/gemma-2-2b-it" + messages = [ + {"role": "user", "content": "Hello!"}, + {"role": "assistant", "content": "Hello!"}, + # {"role": "user", "content": "Not too bad"}, + # {"role": "assistant", "content": "Cooooooool"}, + ] + tokenizer = AutoTokenizer.from_pretrained(model_name) + chat = tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=False, add_special_tokens=False + ) + print(chat) + model_inputs = tokenize_chat_messages( + {"chat": chat, "messages": messages}, + tokenizer, + # for_kl_loss=True, + ) + + print(tokenizer(chat, add_special_tokens=False)) + print(model_inputs)