diff --git a/configs/context_numbers_10.yaml b/configs/context_numbers_10.yaml index 157e51f..1ed026b 100644 --- a/configs/context_numbers_10.yaml +++ b/configs/context_numbers_10.yaml @@ -3,6 +3,7 @@ bf16: true model_name_or_path: meta-llama/Llama-3.2-1B-Instruct label_names: ["labels"] add_repeat_prompt: false +add_negative_prompt: false # eval_on_start: True # eval_strategy: "steps" # eval_steps: 500 diff --git a/configs/context_numbers_128.yaml b/configs/context_numbers_128.yaml index b57fb2f..fed257d 100644 --- a/configs/context_numbers_128.yaml +++ b/configs/context_numbers_128.yaml @@ -3,6 +3,8 @@ bf16: true model_name_or_path: meta-llama/Llama-3.2-1B-Instruct label_names: ["labels"] add_repeat_prompt: false +add_negative_prompt: false + # eval_on_start: True # eval_strategy: "steps" # eval_steps: 500 diff --git a/configs/context_numbers_128_only.yaml b/configs/context_numbers_128_only.yaml index 2cd5f1d..28a5f6f 100644 --- a/configs/context_numbers_128_only.yaml +++ b/configs/context_numbers_128_only.yaml @@ -3,6 +3,8 @@ bf16: true model_name_or_path: meta-llama/Llama-3.2-1B-Instruct label_names: ["labels"] add_repeat_prompt: false +add_negative_prompt: false + # eval_on_start: True # eval_strategy: "steps" # eval_steps: 500 diff --git a/configs/context_numbers_256.yaml b/configs/context_numbers_256.yaml index f36f4f4..ca48999 100644 --- a/configs/context_numbers_256.yaml +++ b/configs/context_numbers_256.yaml @@ -3,6 +3,8 @@ bf16: true model_name_or_path: meta-llama/Llama-3.2-1B-Instruct label_names: ["labels"] add_repeat_prompt: false +add_negative_prompt: false + # eval_on_start: True # eval_strategy: "steps" # eval_steps: 500 diff --git a/configs/context_numbers_256_only.yaml b/configs/context_numbers_256_only.yaml index ce3f84b..47ce81c 100644 --- a/configs/context_numbers_256_only.yaml +++ b/configs/context_numbers_256_only.yaml @@ -3,6 +3,8 @@ bf16: true model_name_or_path: meta-llama/Llama-3.2-1B-Instruct label_names: ["labels"] add_repeat_prompt: false +add_negative_prompt: false + # eval_on_start: True # eval_strategy: "steps" # eval_steps: 500 diff --git a/configs/context_numbers_32.yaml b/configs/context_numbers_32.yaml index 81a329e..f39138d 100644 --- a/configs/context_numbers_32.yaml +++ b/configs/context_numbers_32.yaml @@ -3,6 +3,8 @@ bf16: true model_name_or_path: meta-llama/Llama-3.2-1B-Instruct label_names: ["labels"] add_repeat_prompt: false +add_negative_prompt: false + # eval_on_start: True # eval_strategy: "steps" # eval_steps: 500 diff --git a/configs/context_numbers_32_only.yaml b/configs/context_numbers_32_only.yaml index 7bf48d8..5cc0c73 100644 --- a/configs/context_numbers_32_only.yaml +++ b/configs/context_numbers_32_only.yaml @@ -3,6 +3,8 @@ bf16: true model_name_or_path: meta-llama/Llama-3.2-1B-Instruct label_names: ["labels"] add_repeat_prompt: false +add_negative_prompt: false + # eval_on_start: True # eval_strategy: "steps" # eval_steps: 500 diff --git a/configs/context_numbers_64.yaml b/configs/context_numbers_64.yaml index ed6f48e..4c69eaf 100644 --- a/configs/context_numbers_64.yaml +++ b/configs/context_numbers_64.yaml @@ -3,6 +3,8 @@ bf16: true model_name_or_path: meta-llama/Llama-3.2-1B-Instruct label_names: ["labels"] add_repeat_prompt: false +add_negative_prompt: false + # eval_on_start: True # eval_strategy: "steps" # eval_steps: 500 diff --git a/configs/context_numbers_64_only.yaml b/configs/context_numbers_64_only.yaml index 082421d..42f6bfe 100644 --- a/configs/context_numbers_64_only.yaml +++ b/configs/context_numbers_64_only.yaml @@ -3,6 +3,8 @@ bf16: true model_name_or_path: meta-llama/Llama-3.2-1B-Instruct label_names: ["labels"] add_repeat_prompt: false +add_negative_prompt: false + # eval_on_start: True # eval_strategy: "steps" # eval_steps: 500 diff --git a/configs/context_numbers_debug.yaml b/configs/context_numbers_debug.yaml index c9d3102..9a5827d 100644 --- a/configs/context_numbers_debug.yaml +++ b/configs/context_numbers_debug.yaml @@ -3,6 +3,8 @@ bf16: true model_name_or_path: meta-llama/Llama-3.2-1B-Instruct label_names: ["labels"] add_repeat_prompt: false +add_negative_prompt: false + eval_on_start: True eval_strategy: "steps" eval_steps: 500 diff --git a/configs/context_numbers_easy.yaml b/configs/context_numbers_easy.yaml index fe392b9..316114a 100644 --- a/configs/context_numbers_easy.yaml +++ b/configs/context_numbers_easy.yaml @@ -3,6 +3,8 @@ bf16: true model_name_or_path: meta-llama/Llama-3.2-1B-Instruct label_names: ["labels"] add_repeat_prompt: false +add_negative_prompt: false + eval_on_start: True eval_strategy: "steps" eval_steps: 500 diff --git a/configs/pwc_and_ctx_numbers_256.yaml b/configs/pwc_and_ctx_numbers_256.yaml index f38788b..ea4896b 100644 --- a/configs/pwc_and_ctx_numbers_256.yaml +++ b/configs/pwc_and_ctx_numbers_256.yaml @@ -2,7 +2,6 @@ output_dir: "" # just a placeholder bf16: true model_name_or_path: meta-llama/Llama-3.2-1B-Instruct label_names: ["labels"] -add_repeat_prompt: true # eval_on_start: True # eval_strategy: "steps" # eval_steps: 500 diff --git a/hyperlora/configs.py b/hyperlora/configs.py index 57280c4..8a4953c 100644 --- a/hyperlora/configs.py +++ b/hyperlora/configs.py @@ -239,6 +239,10 @@ class CtxTrainingArguments: default=True, metadata={"help": "Whether to add repeat prompt to the dataset."}, ) + add_negative_prompt: bool = field( + default=True, + metadata={"help": "Whether to add negative prompt to the dataset."}, + ) @dataclass diff --git a/hyperlora/data_utils.py b/hyperlora/data_utils.py index c1f338c..efecfe6 100644 --- a/hyperlora/data_utils.py +++ b/hyperlora/data_utils.py @@ -1,4 +1,6 @@ import logging +import random +import numpy as np from copy import copy from typing import Any, Callable, Dict, Iterator, List, Literal, Optional, Tuple, Union @@ -27,12 +29,77 @@ def filter_long_samples(samples): def add_repeat_prompt_fn(samples): - responses = samples["answer"] + samples["context"] - ctxs = samples["context"] * 2 - prompts = samples["prompt"] - prompts += ["Repeat the text above."] * len(prompts) + 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 - return {"context": ctxs, "response": responses, "prompt": prompts} + unique_contexts.add(ctx) + ctxs.append(ctx) + responses.append(ctx) + prompts.append("Repeat the text above.") + + 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 + if any(keyword in prompt for keyword in keywords): + # Skip samples where the prompt contains any of the specified keywords + continue + + unique_contexts.add(ctx) + ctxs.append(ctx) + prompts.append(prompt) + responses.append(response) + + # 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 get_tokenized_dataset( @@ -42,6 +109,7 @@ def get_tokenized_dataset( tokenizer_kwargs: dict[str, Any], add_ctx_to_chat: bool, add_repeat_prompt: bool, + add_negative_prompt: bool, ) -> dict[str, Any]: need_ctx_ids = not add_ctx_to_chat @@ -54,8 +122,20 @@ def get_tokenized_dataset( return None ds = ds.map(get_preprocessing_fn(ds_name)) ds = ds.filter(filter_long_samples, batched=True) + if add_negative_prompt: + cols_to_remove = [ + col + for col in ds.column_names + if col not in ["context", "prompt", "response"] + ] + ds = ds.map(add_negative_prompt_fn, batched=True, remove_columns=cols_to_remove) 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) + cols_to_remove = [ + col + for col in ds.column_names + if col not in ["context", "prompt", "response"] + ] + ds = ds.map(add_repeat_prompt_fn, batched=True, remove_columns=cols_to_remove) # for sft + chat_model, we need to convert the dataset to chat format # add "messages" field ds = ds.map( diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py index ff0e6fb..61cdbec 100644 --- a/hyperlora/intx_sft.py +++ b/hyperlora/intx_sft.py @@ -262,6 +262,7 @@ def main(output_dir): tokenizer_kwargs=tokenizer_kwargs, add_ctx_to_chat=add_ctx_to_chat, add_repeat_prompt=ctx_args.add_repeat_prompt, + add_negative_prompt=ctx_args.add_negative_prompt, ) tokenized_ds = {} for split, ds_names in zip(