From c9e958f9636bb70c158f9724eecdb3bc7742bf97 Mon Sep 17 00:00:00 2001 From: 51616 Date: Tue, 8 Jul 2025 14:21:43 +0000 Subject: [PATCH] self-gen w/ closed_qa format + perceiver scripts + no strip response --- ...ort_ctx_self_gen_lv1_closed_qa_1_tiny.yaml | 52 +++++++++++++++++++ ...ort_ctx_self_gen_lv3_closed_qa_1_tiny.yaml | 52 +++++++++++++++++++ data/self_generate_qa.py | 47 +++++++++++++---- .../short_ctx/gemma_qa_short_ctx_perceiver.sh | 2 +- .../gemma_qa_short_ctx_perceiver_l2l.sh | 2 +- src/ctx_to_lora/data/processing.py | 2 +- 6 files changed, 145 insertions(+), 12 deletions(-) create mode 100644 configs/tiny_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_tiny.yaml create mode 100644 configs/tiny_exp/qa_short_ctx_self_gen_lv3_closed_qa_1_tiny.yaml diff --git a/configs/tiny_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_tiny.yaml b/configs/tiny_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_tiny.yaml new file mode 100644 index 0000000..5b729fe --- /dev/null +++ b/configs/tiny_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_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_closed_qa_prob_1.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/tiny_exp/qa_short_ctx_self_gen_lv3_closed_qa_1_tiny.yaml b/configs/tiny_exp/qa_short_ctx_self_gen_lv3_closed_qa_1_tiny.yaml new file mode 100644 index 0000000..456d0cb --- /dev/null +++ b/configs/tiny_exp/qa_short_ctx_self_gen_lv3_closed_qa_1_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_closed_qa_prob_1.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/data/self_generate_qa.py b/data/self_generate_qa.py index e7ecd2c..08e0e59 100644 --- a/data/self_generate_qa.py +++ b/data/self_generate_qa.py @@ -1,6 +1,7 @@ import argparse import os import random +import re from glob import glob import pandas as pd @@ -8,7 +9,11 @@ import yaml from datasets import Dataset, load_dataset from vllm import LLM, SamplingParams -from ctx_to_lora.data.definitions import RAW_DATA_DIR, SELF_GEN_DATA_DIR +from ctx_to_lora.data.definitions import ( + CLOSED_QA_INTX_TEMPLATES, + RAW_DATA_DIR, + SELF_GEN_DATA_DIR, +) from ctx_to_lora.data.processing import ( filter_none, get_preprocessing_fn, @@ -20,11 +25,17 @@ STOP_STRINGS = { "google/gemma-2-2b-it": ["", ""], } +# TODO: !!!! +# see https://huggingface.co/datasets/YuxinJiang/LTE_train_data/viewer/default/train?row=97&views%5B%5D=train&sql_row=2 +# based on https://arxiv.org/pdf/2402.11905 SYSTEM_TEMPLATE = ( - "### SYSTEM INSTRUCTIONS ###\n" - "You are a creative and helpful assistant.\n" + "### SYSTEM INSTRUCTION ###\n" + "You are an honest and helpful assistant.\n" + "You must use the information provided in the context for responding to the question.\n" + "If the context does not contain enough information to answer the question, you must say so.\n" + "If that is the case, you can answer based on your knowledge, but you must clearly state that you are doing so.\n" "**DO NOT** hallucinate or make up information.\n" - "### END OF SYSTEM INSTRUCTIONS ###" + "### END OF SYSTEM INSTRUCTION ###" ) PROMPT_TEMPLATE = "### Context ###\n{context}\n\n\n### Question ###\n{question}" @@ -40,6 +51,12 @@ def get_prompt(context: str, q: str) -> str: return PROMPT_TEMPLATE.format(context=context, question=q) +def add_closed_qa_prompt(q: str, closed_qa_prob: float = 0.1) -> str: + if random.random() <= closed_qa_prob: + q = random.choice(CLOSED_QA_INTX_TEMPLATES).format(input=q) + return q + + def load_config(config_path: str) -> dict: """Load dataset names from YAML config file.""" with open(config_path) as f: @@ -141,17 +158,25 @@ def self_generate( shard_name = "" temp = 0.0 - if "_temp_" in ds_name: - temp = float(ds_name.split("_temp_")[-1].split("/")[0]) + closed_qa_prob = 0.0 + if ds_name is not None and "_temp_" in ds_name: + temp_match = re.search(r"_temp_([\d.]+)", ds_name) + if temp_match: + temp = float(temp_match.group(1)) + if ds_name is not None and "_closed_qa_prob_" in ds_name: + prob_match = re.search(r"_closed_qa_prob_([\d.]+)", ds_name) + if prob_match: + closed_qa_prob = float(prob_match.group(1)) print(f"Processing dataset: {ds_name}, split: {split}") print(f"Using temperature: {temp}") + print(f"Using closed QA prompt probability: {closed_qa_prob}") if parquet_file: print(f"Loading dataset from parquet file: {parquet_file}") split = "train" ds_name = "/".join(parquet_file.split(RAW_DATA_DIR)[-1].split("/")[:-1]) - print(ds_name) + shard_name = "_" + os.path.basename(parquet_file).replace(".parquet", "") ds = load_dataset(path="parquet", data_files=[parquet_file], split="train") processing_fn = get_preprocessing_fn(ds_name, is_eval=False, is_pretrain=False) @@ -171,6 +196,7 @@ def self_generate( ) ds = load_and_process_dataset(**kwargs, streaming=False, num_proc=8) + print(f"Loaded dataset: {ds_name} with split: {split}") if args.debug: ds = ds.take(10) @@ -178,7 +204,10 @@ def self_generate( ds = ds.filter(filter_none, batched=False, num_proc=8) ctxs = [sample["context"] for sample in ds] - questions = [sample["prompts"] for sample in ds] + questions = [ + [add_closed_qa_prompt(q, closed_qa_prob) for q in sample["prompts"] if q] + for sample in ds + ] print(f"Loaded {len(ctxs)} contexts and {len(questions)} questions") messages = create_messages(ctxs, questions, args.vllm_model, SYSTEM_TEMPLATE) @@ -243,7 +272,7 @@ def self_generate( 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}" + fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}/{ds_name}/{split}/ds{shard_name}" os.makedirs(os.path.dirname(fpath), exist_ok=True) fpath = f"{fpath}.parquet" ds_out.to_parquet(fpath) diff --git a/scripts/short_ctx/gemma_qa_short_ctx_perceiver.sh b/scripts/short_ctx/gemma_qa_short_ctx_perceiver.sh index dc9ac29..3143954 100644 --- a/scripts/short_ctx/gemma_qa_short_ctx_perceiver.sh +++ b/scripts/short_ctx/gemma_qa_short_ctx_perceiver.sh @@ -19,7 +19,7 @@ uv run accelerate launch --main_process_port $port \ --per_device_eval_batch_size=64 \ --target_modules=down_proj \ --num_blocks=8 \ ---num_self_attn_per_block=6 \ +--num_self_attn_per_block=3 \ --n_latent_queries=208 \ --num_pre_head_layers=1 \ --lora_r=8 \ diff --git a/scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l.sh b/scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l.sh index c6d229a..9d9f5c4 100644 --- a/scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l.sh +++ b/scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l.sh @@ -19,7 +19,7 @@ uv run accelerate launch --main_process_port $port \ --per_device_eval_batch_size=32 \ --target_modules=down_proj \ --num_blocks=8 \ ---num_self_attn_per_block=6 \ +--num_self_attn_per_block=3 \ --ctx_encoder_type=per_layer_activations \ --n_latent_queries=8 \ --num_pre_head_layers=1 \ diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index b0bbd9e..e42c328 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -994,7 +994,7 @@ def convert_ctx_prompt_response_to_messages( [ {"role": "system", "content": system_msg.strip()}, {"role": "user", "content": user_msg.strip()}, - {"role": "assistant", "content": response.strip()}, + {"role": "assistant", "content": response}, ] )