From 35e051886a58229d1b6a7c118461ec824b0e42b2 Mon Sep 17 00:00:00 2001 From: 51616 Date: Fri, 30 May 2025 08:14:25 +0000 Subject: [PATCH] repeat intx + repeat prob --- src/ctx_to_lora/configs.py | 6 ++ src/ctx_to_lora/data/definitions.py | 106 +++++++++++++++------------- src/ctx_to_lora/data/processing.py | 39 +++++++--- src/ctx_to_lora/eval.py | 0 4 files changed, 95 insertions(+), 56 deletions(-) create mode 100644 src/ctx_to_lora/eval.py diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index 81222a6..1b7a381 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -326,6 +326,12 @@ class CtxTrainingArguments: default=False, metadata={"help": "Whether to add repeat prompt to the dataset."}, ) + repeat_prob: float = field( + default=0.0, + metadata={ + "help": "Probability that a context will be used with repeat prompts." + }, + ) add_negative_prompt: bool = field( default=False, metadata={"help": "Whether to add negative prompt to the dataset."}, diff --git a/src/ctx_to_lora/data/definitions.py b/src/ctx_to_lora/data/definitions.py index eb5e40c..a0de952 100644 --- a/src/ctx_to_lora/data/definitions.py +++ b/src/ctx_to_lora/data/definitions.py @@ -8,9 +8,39 @@ FW_QA_PATHS = [ TRANSFORMED_DATA_DIR = "data/processed_datasets" +REPEAT_PROMPTS = [ + "Repeat the text.", + "Repeat the text above.", + "Repeat the above text.", + "Repeat exactly what was written.", + "Duplicate the text.", + "Echo the text. Do not output other words.", + "Repeat the content provided.", + "Copy the content. Reply with just the content.", + "Repeat what you see above. Avoid explanation.", + "Replicate the given text precisely.", + "Repeat the entire passage.", + "Duplicate exactly the provided text.", + "Repeat the text that appears above.", + "Reproduce exactly what appears above.", + "Repeat the article.", + "Repeat verbatim the given text.", + "Output the same text as shown.", + "Output an exact copy of all of the content.", + "Provide a replica of the text.", + "Output an identical copy of the text.", + "Repeat the content without any changes.", + "Repeat the provided text.", + "Repeat the provided text above.", + "Repeat precisely what is written in the text above.", +] + # approximate length of the datasets (train split) # needed for streaming datasets DS_LEN = { + "fw_qa_3_intx_pretrain": 90_000_000, + "fw_qa_3_intx_pretrain_medium": 40_000_000, + "fw_qa_3_intx_pretrain_mini": 100_000, "fw_qa_3_mini": 100_000, "fw_qa_3_medium": 121_000_000, "fw_qa_3": 270_000_000, @@ -78,42 +108,6 @@ DS_KWARGS = { 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", @@ -133,18 +127,6 @@ DS_KWARGS = { 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", @@ -171,6 +153,34 @@ DS_KWARGS = { split="train", ), ), + "fw_qa_3_intx_pretrain_mini": dict( + train=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_intx_pretrain/000_00001.parquet"), + split="train[:100000]", + ) + ), + "fw_qa_3_intx_pretrain_medium": dict( + train=dict( + path="parquet", + data_files=glob( + "data/raw_datasets/fw_qa_intx_pretrain/00[0-5]*[!val].parquet" + ), + split="train", + ) + ), + "fw_qa_3_intx_pretrain": dict( + train=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_intx_pretrain/*[!val].parquet"), + split="train", + ), + validation=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_intx_pretrain/*val.parquet"), + split="train", + ), + ), "ctx_qa": dict( train=dict( path="parquet", diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index b516f27..bd2c046 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -16,6 +16,7 @@ from ctx_to_lora.data.definitions import ( DS_KWARGS, EVAL_INTX_TEMPLATES, IGNORE_INDEX, + REPEAT_PROMPTS, TRANSFORMED_DATA_DIR, ) from ctx_to_lora.utils import TRAINING_TASK @@ -30,6 +31,10 @@ def closed_qa_prompting(prompt: str): return template.format(input=prompt) +def get_repeat_prompt(): + return random.choice(REPEAT_PROMPTS) + + # TODO: dont use closed_qa_prompting when training on self-generated responses def get_preprocessing_fn( ds_name: str, is_eval: bool @@ -210,6 +215,17 @@ def get_preprocessing_fn( "response": sample["answer"], } + # elif "intx_pretrain" in ds_name: + # # concat ctx and QAs + # # use all as labels + # # what about prefix?? + # # neede for pre-training style lm objective + # def f(sample): + # return { + # "context": sample["context"], + # "prompt": sample["context"] + sample["qas"], + # } + if is_eval and ds_name in EVAL_INTX_TEMPLATES: prompt_template = EVAL_INTX_TEMPLATES[ds_name] @@ -273,22 +289,24 @@ def filter_long_chat(samples): return [len(chat) < 2**13 for chat in samples["chat"]] -def add_repeat_prompt_fn(samples): +def add_repeat_prompt_fn(samples: dict[str, list[str]], prob: float): unique_contexts = set() ctxs, prompts, responses = [], [], [] for ctx, prompt, response in zip( samples["context"], samples["prompt"], samples["response"] ): + if random.random() > prob: + continue # 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 + # # 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.") + prompts.append(get_repeat_prompt()) logger.debug("Adding repeat prompt...") logger.debug(f"# unique contexts: {len(unique_contexts)}") @@ -371,6 +389,7 @@ def _load_and_process_dataset( split: str, add_negative_prompt: bool, add_repeat_prompt: bool, + repeat_prob: float, streaming: bool, ds_path: str, num_proc: int, @@ -428,9 +447,9 @@ def _load_and_process_dataset( 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, + fn_kwargs={"prob": repeat_prob}, batched=True, batch_size=100_000, num_proc=num_proc, @@ -452,20 +471,24 @@ def get_tokenized_dataset( add_repeat_prompt: bool, add_negative_prompt: bool, use_kl_loss: bool, + repeat_prob: float = 0.0, # only used when add_repeat_prompt is True 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" + if add_repeat_prompt: + assert repeat_prob > 0, f"add_repeat_prompt is set but repeat_prob = 0" 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, + base_model_max_len=base_model_max_len, + ctx_model_max_len=ctx_model_max_len, add_negative_prompt=add_negative_prompt, add_repeat_prompt=add_repeat_prompt, + repeat_prob=repeat_prob, streaming=streaming, ) num_proc = None if streaming and split == "train" else 8 diff --git a/src/ctx_to_lora/eval.py b/src/ctx_to_lora/eval.py new file mode 100644 index 0000000..e69de29