From 28eb85a2f39f1683bd27a3dd348b64d4ea0c7a6e Mon Sep 17 00:00:00 2001 From: 51616 Date: Mon, 30 Jun 2025 16:10:52 +0900 Subject: [PATCH] tiny configs + fix dataset gen (remove repeated response + qa template for compact ds) + add ropes and drop eval ds --- configs/qa_short_ctx_no_fw_qa.yaml | 58 +++++++++++++++++ configs/qa_short_ctx_self_gen.yaml | 2 +- configs/qa_short_ctx_self_gen_lv0_tiny.yaml | 52 +++++++++++++++ configs/qa_short_ctx_self_gen_lv1_tiny.yaml | 52 +++++++++++++++ configs/qa_short_ctx_self_gen_lv2_tiny.yaml | 52 +++++++++++++++ configs/qa_short_ctx_self_gen_lv3_tiny.yaml | 52 +++++++++++++++ ..._short_ctx_self_gen_lv3_tiny_temp_0.3.yaml | 52 +++++++++++++++ ..._short_ctx_self_gen_lv3_tiny_temp_0.5.yaml | 52 +++++++++++++++ ..._short_ctx_self_gen_lv3_tiny_temp_0.7.yaml | 52 +++++++++++++++ ..._short_ctx_self_gen_lv3_tiny_temp_1.0.yaml | 52 +++++++++++++++ data/build_drop_compact.py | 4 +- data/build_pwc_compact.py | 4 +- data/build_ropes_compact.py | 4 +- data/build_squad_compact.py | 4 +- data/generate_fw_edu_qa_v2.py | 30 ++++++++- data/generate_fw_edu_qa_v2_repeat.py | 31 ++++++++- data/self_generate_qa.py | 65 +++++++++++++++---- src/ctx_to_lora/data/definitions.py | 43 ++++++++---- src/ctx_to_lora/data/processing.py | 1 + 19 files changed, 626 insertions(+), 36 deletions(-) create mode 100644 configs/qa_short_ctx_no_fw_qa.yaml create mode 100644 configs/qa_short_ctx_self_gen_lv0_tiny.yaml create mode 100644 configs/qa_short_ctx_self_gen_lv1_tiny.yaml create mode 100644 configs/qa_short_ctx_self_gen_lv2_tiny.yaml create mode 100644 configs/qa_short_ctx_self_gen_lv3_tiny.yaml create mode 100644 configs/qa_short_ctx_self_gen_lv3_tiny_temp_0.3.yaml create mode 100644 configs/qa_short_ctx_self_gen_lv3_tiny_temp_0.5.yaml create mode 100644 configs/qa_short_ctx_self_gen_lv3_tiny_temp_0.7.yaml create mode 100644 configs/qa_short_ctx_self_gen_lv3_tiny_temp_1.0.yaml diff --git a/configs/qa_short_ctx_no_fw_qa.yaml b/configs/qa_short_ctx_no_fw_qa.yaml new file mode 100644 index 0000000..e18a3b2 --- /dev/null +++ b/configs/qa_short_ctx_no_fw_qa.yaml @@ -0,0 +1,58 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: google/gemma-3-1b-it +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00004 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 + +warmup_steps: 100 + +dataloader_prefetch_factor: 16 +dataloader_num_workers: 8 + +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# short ctx data +train_ds_names: + - squad_compact + - pwc_compact + - drop_compact + - ropes_compact + +# the contexts are kinda long-ish +# - narrativeqa +# - quoref +# - synthetic_convqa + +val_ds_names: + - fw_qa_v2_2k_len_level_0 + - squad + - drop + - ropes + - pwc + diff --git a/configs/qa_short_ctx_self_gen.yaml b/configs/qa_short_ctx_self_gen.yaml index 5c8a6f0..15e49c5 100644 --- a/configs/qa_short_ctx_self_gen.yaml +++ b/configs/qa_short_ctx_self_gen.yaml @@ -39,7 +39,7 @@ target_modules: # data train_ds_names: - - self_gen/google/gemma-2-2b-it/fw_qa_v2_2k_len_level_3 + - self_gen/google/gemma-2-2b-it/fw_qa_v2/min_0_to_2000 - self_gen/google/gemma-2-2b-it/squad_compact - self_gen/google/gemma-2-2b-it/pwc_compact - self_gen/google/gemma-2-2b-it/ropes_compact diff --git a/configs/qa_short_ctx_self_gen_lv0_tiny.yaml b/configs/qa_short_ctx_self_gen_lv0_tiny.yaml new file mode 100644 index 0000000..3bb1aeb --- /dev/null +++ b/configs/qa_short_ctx_self_gen_lv0_tiny.yaml @@ -0,0 +1,52 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: google/gemma-3-1b-it +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00004 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 + +warmup_steps: 100 + +dataloader_prefetch_factor: 16 +dataloader_num_workers: 8 + +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: + - self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_0_tiny + - self_gen/google/gemma-2-2b-it_temp_0.0/squad_compact + - self_gen/google/gemma-2-2b-it_temp_0.0/pwc_compact + - self_gen/google/gemma-2-2b-it_temp_0.0/ropes_compact + - self_gen/google/gemma-2-2b-it_temp_0.0/drop_compact + +val_ds_names: + - squad + - pwc + - drop + - ropes diff --git a/configs/qa_short_ctx_self_gen_lv1_tiny.yaml b/configs/qa_short_ctx_self_gen_lv1_tiny.yaml new file mode 100644 index 0000000..b49e699 --- /dev/null +++ b/configs/qa_short_ctx_self_gen_lv1_tiny.yaml @@ -0,0 +1,52 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: google/gemma-3-1b-it +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00004 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 + +warmup_steps: 100 + +dataloader_prefetch_factor: 16 +dataloader_num_workers: 8 + +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: + - self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_1_tiny + - self_gen/google/gemma-2-2b-it_temp_0.0/squad_compact + - self_gen/google/gemma-2-2b-it_temp_0.0/pwc_compact + - self_gen/google/gemma-2-2b-it_temp_0.0/ropes_compact + - self_gen/google/gemma-2-2b-it_temp_0.0/drop_compact + +val_ds_names: + - squad + - pwc + - drop + - ropes diff --git a/configs/qa_short_ctx_self_gen_lv2_tiny.yaml b/configs/qa_short_ctx_self_gen_lv2_tiny.yaml new file mode 100644 index 0000000..78467c0 --- /dev/null +++ b/configs/qa_short_ctx_self_gen_lv2_tiny.yaml @@ -0,0 +1,52 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: google/gemma-3-1b-it +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00004 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 + +warmup_steps: 100 + +dataloader_prefetch_factor: 16 +dataloader_num_workers: 8 + +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: + - self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_2_tiny + - self_gen/google/gemma-2-2b-it_temp_0.0/squad_compact + - self_gen/google/gemma-2-2b-it_temp_0.0/pwc_compact + - self_gen/google/gemma-2-2b-it_temp_0.0/ropes_compact + - self_gen/google/gemma-2-2b-it_temp_0.0/drop_compact + +val_ds_names: + - squad + - pwc + - drop + - ropes diff --git a/configs/qa_short_ctx_self_gen_lv3_tiny.yaml b/configs/qa_short_ctx_self_gen_lv3_tiny.yaml new file mode 100644 index 0000000..0b4e48a --- /dev/null +++ b/configs/qa_short_ctx_self_gen_lv3_tiny.yaml @@ -0,0 +1,52 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: google/gemma-3-1b-it +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00004 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 + +warmup_steps: 100 + +dataloader_prefetch_factor: 16 +dataloader_num_workers: 8 + +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: + - self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_3_tiny + - self_gen/google/gemma-2-2b-it_temp_0.0/squad_compact + - self_gen/google/gemma-2-2b-it_temp_0.0/pwc_compact + - self_gen/google/gemma-2-2b-it_temp_0.0/ropes_compact + - self_gen/google/gemma-2-2b-it_temp_0.0/drop_compact + +val_ds_names: + - squad + - pwc + - drop + - ropes diff --git a/configs/qa_short_ctx_self_gen_lv3_tiny_temp_0.3.yaml b/configs/qa_short_ctx_self_gen_lv3_tiny_temp_0.3.yaml new file mode 100644 index 0000000..e60d8ee --- /dev/null +++ b/configs/qa_short_ctx_self_gen_lv3_tiny_temp_0.3.yaml @@ -0,0 +1,52 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: google/gemma-3-1b-it +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00004 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 + +warmup_steps: 100 + +dataloader_prefetch_factor: 16 +dataloader_num_workers: 8 + +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: + - self_gen/google/gemma-2-2b-it_temp_0.3/fw_qa_v2_2k_len_level_3_tiny + - self_gen/google/gemma-2-2b-it_temp_0.3/squad_compact + - self_gen/google/gemma-2-2b-it_temp_0.3/pwc_compact + - self_gen/google/gemma-2-2b-it_temp_0.3/ropes_compact + - self_gen/google/gemma-2-2b-it_temp_0.3/drop_compact + +val_ds_names: + - squad + - pwc + - drop + - ropes diff --git a/configs/qa_short_ctx_self_gen_lv3_tiny_temp_0.5.yaml b/configs/qa_short_ctx_self_gen_lv3_tiny_temp_0.5.yaml new file mode 100644 index 0000000..43e069c --- /dev/null +++ b/configs/qa_short_ctx_self_gen_lv3_tiny_temp_0.5.yaml @@ -0,0 +1,52 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: google/gemma-3-1b-it +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00004 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 + +warmup_steps: 100 + +dataloader_prefetch_factor: 16 +dataloader_num_workers: 8 + +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: + - self_gen/google/gemma-2-2b-it_temp_0.5/fw_qa_v2_2k_len_level_3_tiny + - self_gen/google/gemma-2-2b-it_temp_0.5/squad_compact + - self_gen/google/gemma-2-2b-it_temp_0.5/pwc_compact + - self_gen/google/gemma-2-2b-it_temp_0.5/ropes_compact + - self_gen/google/gemma-2-2b-it_temp_0.5/drop_compact + +val_ds_names: + - squad + - pwc + - drop + - ropes diff --git a/configs/qa_short_ctx_self_gen_lv3_tiny_temp_0.7.yaml b/configs/qa_short_ctx_self_gen_lv3_tiny_temp_0.7.yaml new file mode 100644 index 0000000..eecf57d --- /dev/null +++ b/configs/qa_short_ctx_self_gen_lv3_tiny_temp_0.7.yaml @@ -0,0 +1,52 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: google/gemma-3-1b-it +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00004 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 + +warmup_steps: 100 + +dataloader_prefetch_factor: 16 +dataloader_num_workers: 8 + +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: + - self_gen/google/gemma-2-2b-it_temp_0.7/fw_qa_v2_2k_len_level_3_tiny + - self_gen/google/gemma-2-2b-it_temp_0.7/squad_compact + - self_gen/google/gemma-2-2b-it_temp_0.7/pwc_compact + - self_gen/google/gemma-2-2b-it_temp_0.7/ropes_compact + - self_gen/google/gemma-2-2b-it_temp_0.7/drop_compact + +val_ds_names: + - squad + - pwc + - drop + - ropes diff --git a/configs/qa_short_ctx_self_gen_lv3_tiny_temp_1.0.yaml b/configs/qa_short_ctx_self_gen_lv3_tiny_temp_1.0.yaml new file mode 100644 index 0000000..88bfc19 --- /dev/null +++ b/configs/qa_short_ctx_self_gen_lv3_tiny_temp_1.0.yaml @@ -0,0 +1,52 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: google/gemma-3-1b-it +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00004 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.01 + +warmup_steps: 100 + +dataloader_prefetch_factor: 16 +dataloader_num_workers: 8 + +# LoRA +lora_r: 8 +lora_dropout: 0.0 +target_modules: + - down_proj + +# data +train_ds_names: + - self_gen/google/gemma-2-2b-it_temp_1.0/fw_qa_v2_2k_len_level_3_tiny + - self_gen/google/gemma-2-2b-it_temp_1.0/squad_compact + - self_gen/google/gemma-2-2b-it_temp_1.0/pwc_compact + - self_gen/google/gemma-2-2b-it_temp_1.0/ropes_compact + - self_gen/google/gemma-2-2b-it_temp_1.0/drop_compact + +val_ds_names: + - squad + - pwc + - drop + - ropes diff --git a/data/build_drop_compact.py b/data/build_drop_compact.py index 261e20a..d126c71 100644 --- a/data/build_drop_compact.py +++ b/data/build_drop_compact.py @@ -3,6 +3,8 @@ import gc from datasets import Dataset, load_dataset from tqdm import tqdm +from ctx_to_lora.data.processing import closed_qa_prompting + if __name__ == "__main__": ds_name = "ucinlp/drop" @@ -14,7 +16,7 @@ if __name__ == "__main__": ctx = sample["passage"] if ctx not in ctx_qa_dict: ctx_qa_dict[ctx] = {"prompts": [], "responses": []} - question = sample["question"] + question = closed_qa_prompting(sample["question"]) answer = ", ".join(set(sample["answers_spans"]["spans"])) ctx_qa_dict[ctx]["prompts"].append(question) ctx_qa_dict[ctx]["responses"].append(answer) diff --git a/data/build_pwc_compact.py b/data/build_pwc_compact.py index 75c3882..63e4f72 100644 --- a/data/build_pwc_compact.py +++ b/data/build_pwc_compact.py @@ -3,6 +3,8 @@ import gc from datasets import Dataset, load_dataset from tqdm import tqdm +from ctx_to_lora.data.processing import closed_qa_prompting + if __name__ == "__main__": ds_name = "sggetao/PwC" @@ -14,7 +16,7 @@ if __name__ == "__main__": ctx = sample["input"] if ctx not in ctx_qa_dict: ctx_qa_dict[ctx] = {"prompts": [], "responses": []} - question = sample["prompt"] + question = closed_qa_prompting(sample["prompt"]) answer = sample["answer"] ctx_qa_dict[ctx]["prompts"].append(question) ctx_qa_dict[ctx]["responses"].append(answer) diff --git a/data/build_ropes_compact.py b/data/build_ropes_compact.py index 8ea13a9..1d82db1 100644 --- a/data/build_ropes_compact.py +++ b/data/build_ropes_compact.py @@ -3,6 +3,8 @@ import gc from datasets import Dataset, load_dataset from tqdm import tqdm +from ctx_to_lora.data.processing import closed_qa_prompting + if __name__ == "__main__": ds_name = "allenai/ropes" @@ -16,7 +18,7 @@ if __name__ == "__main__": bg_txt = sample["background"] situation_txt = sample["situation"] ctx = ctx_template.format(background=bg_txt, situation=situation_txt) - q = sample["question"] + q = closed_qa_prompting(sample["question"]) if ctx not in ctx_qa_dict: ctx_qa_dict[ctx] = {"prompts": [], "responses": []} ctx_qa_dict[ctx]["prompts"].append(q) diff --git a/data/build_squad_compact.py b/data/build_squad_compact.py index 6cfa984..9594f52 100644 --- a/data/build_squad_compact.py +++ b/data/build_squad_compact.py @@ -3,6 +3,8 @@ import gc from datasets import Dataset, load_dataset from tqdm import tqdm +from ctx_to_lora.data.processing import closed_qa_prompting + if __name__ == "__main__": ds_name = "rajpurkar/squad" @@ -14,7 +16,7 @@ if __name__ == "__main__": ctx = sample["context"] if ctx not in ctx_qa_dict: ctx_qa_dict[ctx] = {"prompts": [], "responses": []} - question = sample["question"] + question = closed_qa_prompting(sample["question"]) answer = sample["answers"]["text"][0] ctx_qa_dict[ctx]["prompts"].append(question) ctx_qa_dict[ctx]["responses"].append(answer) diff --git a/data/generate_fw_edu_qa_v2.py b/data/generate_fw_edu_qa_v2.py index a53c288..83d1404 100644 --- a/data/generate_fw_edu_qa_v2.py +++ b/data/generate_fw_edu_qa_v2.py @@ -7,6 +7,10 @@ import pandas as pd from datasets import Dataset, load_dataset from vllm import LLM, SamplingParams +STOP_STRINGS = { + "google/gemma-2-2b-it": ["", ""], +} + SYSTEM_TEMPLATE = ( "You are a creative and helpful assistant.\n" "You are tasked with generating questions for reading comprehension tests.\n" @@ -57,6 +61,14 @@ def get_prompt(context, n_qa_pairs): return prompt +def check_should_skip(txt: str, vllm_model: str) -> bool: + """Check if the response should be skipped based on stop strings.""" + for stop in STOP_STRINGS[vllm_model]: + if stop in txt[-len(stop) :]: + return (txt.split(stop)[0], False) # Found a valid stop string + return (txt, True) # No valid stop string found, skip this response + + def postprocess_qa_pairs(res_txt: str): """ Postprocesses the QA pairs from the response text. @@ -83,10 +95,21 @@ def postprocess_qa_pairs(res_txt: str): out_q = [] out_a = [] + n_skips = 0 if (len(questions) > 0) and (len(answers) > 0): - for i in range(min(len(questions), len(answers))): + n_gen_pairs = min(len(questions), len(answers)) + has_left_over = n_gen_pairs < len(questions) or n_gen_pairs < len(answers) + for i in range(n_gen_pairs): + response = answers[i].strip() + if (not has_left_over) and (i == n_gen_pairs - 1): + response, skip = check_should_skip(response, vllm_model) + if skip: + print(f"Skipping due to missing stop string") + n_skips += 1 + continue out_q.append(questions[i].strip()) out_a.append(answers[i].strip()) + print(f"Skipped {n_skips} responses due to missing stop strings") return out_q, out_a @@ -193,7 +216,10 @@ if __name__ == "__main__": messages, sampling_params=SamplingParams( max_tokens=2048, - temperature=0.7, + temperature=0.0, + # needed for checking if stop tokens are present + skip_special_tokens=False, + include_stop_str_in_output=True, ), ) samples = [] diff --git a/data/generate_fw_edu_qa_v2_repeat.py b/data/generate_fw_edu_qa_v2_repeat.py index 53634c5..76e13dd 100644 --- a/data/generate_fw_edu_qa_v2_repeat.py +++ b/data/generate_fw_edu_qa_v2_repeat.py @@ -7,6 +7,10 @@ from glob import glob from datasets import load_dataset from vllm import LLM, SamplingParams +STOP_STRINGS = { + "google/gemma-2-2b-it": ["", ""], +} + SYSTEM_TEMPLATE = ( "You are a creative and helpful assistant.\n" "You are tasked with generating questions for reading comprehension tests.\n" @@ -64,6 +68,14 @@ def get_prompt(context, example_qa_pairs, n_qa_pairs): return prompt +def check_should_skip(txt: str, vllm_model: str) -> bool: + """Check if the response should be skipped based on stop strings.""" + for stop in STOP_STRINGS[vllm_model]: + if stop in txt[-len(stop) :]: + return (txt.split(stop)[0], False) # Found a valid stop string + return (txt, True) # No valid stop string found, skip this response + + def postprocess_qa_pairs(res_txt: str): """ Postprocesses the QA pairs from the response text. @@ -90,10 +102,21 @@ def postprocess_qa_pairs(res_txt: str): out_q = [] out_a = [] + n_skips = 0 if (len(questions) > 0) and (len(answers) > 0): - for i in range(min(len(questions), len(answers))): + n_gen_pairs = min(len(questions), len(answers)) + has_left_over = n_gen_pairs < len(questions) or n_gen_pairs < len(answers) + for i in range(n_gen_pairs): + response = answers[i].strip() + if (not has_left_over) and (i == n_gen_pairs - 1): + response, skip = check_should_skip(response, vllm_model) + if skip: + print(f"Skipping due to missing stop string") + n_skips += 1 + continue out_q.append(questions[i].strip()) out_a.append(answers[i].strip()) + print(f"Skipped {n_skips} responses due to missing stop strings") return out_q, out_a @@ -217,8 +240,10 @@ if __name__ == "__main__": completions = llm.chat( messages, sampling_params=SamplingParams( - max_tokens=2048, - temperature=0.7, + temperature=0.0, + # needed for checking if stop tokens are present + skip_special_tokens=False, + include_stop_str_in_output=True, ), ) samples = [] diff --git a/data/self_generate_qa.py b/data/self_generate_qa.py index 5fa93a4..9446eb6 100644 --- a/data/self_generate_qa.py +++ b/data/self_generate_qa.py @@ -1,5 +1,6 @@ import argparse import os +import random from glob import glob import pandas as pd @@ -15,6 +16,10 @@ from ctx_to_lora.data.processing import ( ) from ctx_to_lora.utils import clear_gpu +STOP_STRINGS = { + "google/gemma-2-2b-it": ["", ""], +} + SYSTEM_TEMPLATE = ( "### SYSTEM INSTRUCTIONS ###\n" "You are a creative and helpful assistant.\n" @@ -42,6 +47,14 @@ def load_config(config_path: str) -> dict: return config +def check_should_skip(txt: str, vllm_model: str) -> bool: + """Check if the response should be skipped based on stop strings.""" + for stop in STOP_STRINGS[vllm_model]: + if stop in txt[-len(stop) :]: + return (txt.split(stop)[0], False) # Found a valid stop string + return (txt, True) # No valid stop string found, skip this response + + def get_dataset_configs( ds_names: list[str] | None, config: dict | None, @@ -147,6 +160,10 @@ def self_generate( print(f"Loading dataset: {ds_name} with split: {split}") ds = load_and_process_dataset(**kwargs, streaming=False, num_proc=8) + + temp = 0.0 + if "_temp_" in ds_name: + temp = float(ds_name.split("_temp_")[-1].split("/")[0]) if args.debug: ds = ds.take(10) @@ -163,19 +180,38 @@ def self_generate( completions = llm.chat( messages, sampling_params=SamplingParams( - max_tokens=2048, - temperature=1.0, # TODO: lower the temp + max_tokens=1024, + temperature=temp, + # needed for checking if stop tokens are present + skip_special_tokens=False, + include_stop_str_in_output=True, ), ) self_gen_data = {ctx: {"prompts": [], "responses": []} for ctx in ctxs} c = 0 + n_skips = 0 for ctx, q_list in zip(ctxs, questions): for i, q in enumerate(q_list): + response = completions[c + i].outputs[0].text + response, skip = check_should_skip(response, args.vllm_model) + # skip = True + # for stop in STOP_STRINGS[args.vllm_model]: + # if stop == response[-len(stop) :]: + # # Check if response ends with stop string + # response = response.split(stop)[0] + # skip = False + # break + if skip: + print(f"Skipping due to missing stop string") + n_skips += 1 + continue + self_gen_data[ctx]["prompts"].append(q) - self_gen_data[ctx]["responses"].append(completions[c + i].outputs[0].text) + self_gen_data[ctx]["responses"].append(response) c += i + 1 + print(f"Skipped {n_skips} responses due to missing stop strings") samples = [ { "context": ctx, @@ -193,20 +229,21 @@ def self_generate( print("=" * 80) print(f"Generated {len(samples)} samples") - # random.shuffle(samples) + random.shuffle(samples) # Save results - df = pd.DataFrame(samples) - ds_out = Dataset.from_pandas(df) - fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}/{ds_name}/{split}/ds{shard_name}" - os.makedirs(os.path.dirname(fpath), exist_ok=True) - fpath = f"{fpath}.parquet" - ds_out.to_parquet(fpath) - print(f"Saved to {fpath}") + if not args.debug: + df = pd.DataFrame(samples) + ds_out = Dataset.from_pandas(df) + fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}/{ds_name}/{split}/ds{shard_name}" + os.makedirs(os.path.dirname(fpath), exist_ok=True) + fpath = f"{fpath}.parquet" + ds_out.to_parquet(fpath) + print(f"Saved to {fpath}") - # Cleanup - del samples, df, ds_out, completions, messages, ctxs, questions - clear_gpu() + # Cleanup + del samples, df, ds_out, completions, messages, ctxs, questions + clear_gpu() def parse_args() -> argparse.Namespace: diff --git a/src/ctx_to_lora/data/definitions.py b/src/ctx_to_lora/data/definitions.py index 6291818..96d9341 100644 --- a/src/ctx_to_lora/data/definitions.py +++ b/src/ctx_to_lora/data/definitions.py @@ -149,8 +149,10 @@ DS_KWARGS = { "fw_qa_v2_2k_len_level_0_tiny": dict( train=dict( path="parquet", - data_files="data/raw_datasets/fw_qa_v2/min_0_to_2000/000_00000_level_0.parquet", - split="train", + data_files=glob( + "data/raw_datasets/fw_qa_v2/min_0_to_2000/013*level_0.parquet" + ), + split="train[:400000]", ), ), "fw_qa_v2_2k_len_level_0": dict( @@ -181,6 +183,15 @@ DS_KWARGS = { # split="train", # ), ), + "fw_qa_v2_2k_len_level_1_tiny": dict( + train=dict( + path="parquet", + data_files=glob( + "data/raw_datasets/fw_qa_v2/min_0_to_2000/013*level_1.parquet" + ), + split="train[:200000]", + ), + ), "fw_qa_v2_2k_len_level_2": dict( train=dict( path="parquet", @@ -195,11 +206,13 @@ DS_KWARGS = { # split="train", # ), ), - "fw_qa_v2_2k_len_level_3_tiny": dict( + "fw_qa_v2_2k_len_level_2_tiny": dict( train=dict( path="parquet", - data_files="data/raw_datasets/fw_qa_v2/min_0_to_2000/000_00000_level_3.parquet", - split="train", + data_files=glob( + "data/raw_datasets/fw_qa_v2/min_0_to_2000/013*level_2.parquet" + ), + split="train[:100000]", ), ), "fw_qa_v2_2k_len_level_3": dict( @@ -216,13 +229,15 @@ DS_KWARGS = { # 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_v2_2k_len_level_3_tiny": dict( + train=dict( + path="parquet", + data_files=glob( + "data/raw_datasets/fw_qa_v2/min_0_to_2000/013*level_3.parquet" + ), + split="train[:50000]", + ), + ), # "fw_qa_3_small": dict( # train=dict( # path="parquet", @@ -491,6 +506,8 @@ CLOSED_QA_DATASETS = { "squad", "triviaqa_retrieved", "negative_nq", + "ropes", + "drop", } MULTI_ANSWER_DATASETS = { @@ -543,6 +560,8 @@ CLOSED_QA_INTX_TEMPLATES = [ EVAL_INTX_TEMPLATES = { + "ropes": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}", + "drop": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}", "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}", diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index 64d9794..556655f 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -185,6 +185,7 @@ def get_preprocessing_fn( situation_txt = sample["situation"] ctx = ctx_template.format(background=bg_txt, situation=situation_txt) q = sample["question"] + q = closed_qa_prompting(q) if not is_eval else q return {"context": ctx, "prompt": q, "response": response} elif ds_name in ["narrativeqa", "quoref"]: # , "ropes"]: