self-gen w/ closed_qa format + perceiver scripts + no strip response

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

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

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": ["<eos>", "<end_of_turn>"],
 }
 
+# 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)

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 \

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 \

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},
             ]
         )