add repeat prompt

configs/context_numbers_10.yaml

@@ -2,6 +2,7 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
+add_repeat_prompt: false
 # eval_on_start: True
 # eval_strategy: "steps"
 # eval_steps: 500

configs/context_numbers_128.yaml

@@ -2,6 +2,7 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
+add_repeat_prompt: false
 # eval_on_start: True
 # eval_strategy: "steps"
 # eval_steps: 500

configs/context_numbers_128_only.yaml

@@ -2,6 +2,7 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
+add_repeat_prompt: false
 # eval_on_start: True
 # eval_strategy: "steps"
 # eval_steps: 500

configs/context_numbers_256.yaml

@@ -2,6 +2,7 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
+add_repeat_prompt: false
 # eval_on_start: True
 # eval_strategy: "steps"
 # eval_steps: 500

configs/context_numbers_256_only.yaml

@@ -2,6 +2,7 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
+add_repeat_prompt: false
 # eval_on_start: True
 # eval_strategy: "steps"
 # eval_steps: 500

configs/context_numbers_32.yaml

@@ -2,6 +2,7 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
+add_repeat_prompt: false
 # eval_on_start: True
 # eval_strategy: "steps"
 # eval_steps: 500

configs/context_numbers_32_only.yaml

@@ -2,6 +2,7 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
+add_repeat_prompt: false
 # eval_on_start: True
 # eval_strategy: "steps"
 # eval_steps: 500

configs/context_numbers_64.yaml

@@ -2,6 +2,7 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
+add_repeat_prompt: false
 # eval_on_start: True
 # eval_strategy: "steps"
 # eval_steps: 500

configs/context_numbers_64_only.yaml

@@ -2,6 +2,7 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
+add_repeat_prompt: false
 # eval_on_start: True
 # eval_strategy: "steps"
 # eval_steps: 500

configs/context_numbers_debug.yaml

@@ -2,6 +2,7 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
+add_repeat_prompt: false
 eval_on_start: True
 eval_strategy: "steps"
 eval_steps: 500

configs/context_numbers_easy.yaml

@@ -2,6 +2,7 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
+add_repeat_prompt: false
 eval_on_start: True
 eval_strategy: "steps"
 eval_steps: 500

configs/default.yaml

@@ -2,6 +2,7 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
+add_repeat_prompt: false
 eval_on_start: True
 eval_strategy: "steps"
 eval_steps: 500

configs/pwc.yaml

@@ -18,7 +18,7 @@ label_names: ["labels"]
 
 per_device_train_batch_size: 32
 per_device_eval_batch_size: 32
-max_val_samples_per_ds: 20
+max_val_samples_per_ds: 1000
 # optim: schedule_free_adamw
 learning_rate: 0.00001
 # lr_scheduler_type: "constant_with_warmup"

configs/pwc_and_ctx_numbers_256.yaml

@@ -2,6 +2,7 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
+add_repeat_prompt: false
 # eval_on_start: True
 # eval_strategy: "steps"
 # eval_steps: 500
@@ -18,7 +19,7 @@ label_names: ["labels"]
 
 per_device_train_batch_size: 32
 per_device_eval_batch_size: 1
-max_val_samples_per_ds: 20
+max_val_samples_per_ds: 1000
 # optim: schedule_free_adamw
 learning_rate: 0.00001
 # lr_scheduler_type: "constant_with_warmup"

hyperlora/configs.py

@@ -235,6 +235,10 @@ class CtxTrainingArguments:
         default=None,
         metadata={"help": "Wandb notes for the experiment."},
     )
+    add_repeat_prompt: bool = field(
+        default=True,
+        metadata={"help": "Whether to add repeat prompt to the dataset."},
+    )
 
 
 @dataclass

hyperlora/data_utils.py

@@ -19,17 +19,34 @@ def validate_columns(tokenized_ds):
     ), 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 add_repeat_prompt_fn(samples):
+    ctxs = samples["context"] * 2
+    responses = samples["answer"] * 2
+    prompts = samples["prompt"]
+    prompts += ["Repeat the text above."] * len(prompts)
+
+    return {"context": ctxs, "response": responses, "prompt": prompts}
+
+
 def get_tokenized_dataset(
     ds_name: str,
     split: str,
     tokenizer: PreTrainedTokenizerBase,
     tokenizer_kwargs: dict[str, Any],
     add_ctx_to_chat: bool,
+    add_repeat_prompt: bool,
 ) -> dict[str, Any]:
 
     need_ctx_ids = not add_ctx_to_chat
     ds = load_dataset(ds_name, split=split)
     ds = ds.map(get_preprocessing_fn(ds_name))
+    ds = ds.filter(filter_long_samples, batched=True)
+    if add_repeat_prompt and "context_numbers" not in ds_name:
+        ds = ds.map(add_repeat_prompt_fn, batched=True, remove_columns=ds.column_names)
     # for sft + chat_model, we need to convert the dataset to chat format
     # add "messages" field
     ds = ds.map(

hyperlora/intx_sft.py

@@ -266,6 +266,7 @@ def main(output_dir):
         tokenizer=tokenizer,
         tokenizer_kwargs=tokenizer_kwargs,
         add_ctx_to_chat=add_ctx_to_chat,
+        add_repeat_prompt=ctx_args.add_repeat_prompt,
         # get_ctx_features_fn=get_ctx_features_fn,
     )
     tokenized_ds = {}