apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-23 17:01:04 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

iclr cleanup

Browse source
This commit is contained in:
51616 2025-09-30 14:53:21 +00:00
parent 8745db6e11
commit b6679ba755
171 changed files with 219 additions and 248151 deletions
Download patch file Download diff file Expand all files Collapse all files

237
README.md
View file

@ -7,13 +7,6 @@
---
[Project doc](https://docs.google.com/document/d/1RCQDzlVU7YGoTwR84gLQfhxTv0RFnW6srQlqfCC2bvQ/edit?usp=sharing)
<!-- ### Finetuning the base model with LoRA adaptor
```bash
WANDB_MODE=disabled python hyperlora/train.py configs/default.yaml --model_name_or_path=meta-llama/Llama-3.2-1B-Instruct --num_train_epochs=50 --per_device_train_batch_size=128 --per_device_eval_batch_size=128 --exp_setup=lora --train_ds_name=data/raw_datasets/context_numbers_small
``` -->
## 🚀 API Usage [WIP]
```python
from ctx_to_lora.modeling import ModulatedPretrainedModel
@ -28,151 +21,8 @@ outputs = model.generate(ctx_ids, input_ids)
print(model.decode(outputs))
```
## 🏋️ Training
### 🔢 HyperLoRA w/ context_numbers_10
```bash
WANDB_MODE=disabled run uv run train.py configs/context_numbers_10.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=1 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208
```
KL loss
```bash
WANDB_MODE=disabled run uv run train.py configs/context_numbers_10_self_gen.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=2 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=100 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=2048 --max_packed_ctx_len=2048 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 --use_kl_loss=True
```
### Synthetic
ctx numbers
```bash
WANDB_MODE=disabled run uv run train.py configs/gemma-3-1b-it/toy_exp/ctx_numbers_64_128.yaml --model_name_or_path=google/gemma-3-1b-it --num_train_epochs=3 --per_device_train_batch_size=-1 --gradient_accumulation_steps=2 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=100 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.0 --use_sequence_packing=True --max_packed_inp_len=2048 --max_packed_ctx_len=2048 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=per_layer_activations --n_latent_queries=8 --use_kl_loss=False --eval_on_start=True
```
ctx magic num
```bash
WANDB_PROJECT=ctx-magic-num srun --partition=aiscilow --gpus=1 --unbuffered uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=32 --per_device_eval_batch_size=16 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=100 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=2048 --max_packed_ctx_len=2048 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=512 --min_ctx_chunk_len=25 --num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' --max_val_samples_per_ds=100 --seed=1
```
<!-- ### FWQA-v2 Level-0 Tiny
```bash
WANDB_MODE=disabled run uv run train.py configs/fw_qa_v2_level_0_tiny.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=64 --gradient_accumulation_steps=8 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_self_attends_per_block=8 --num_latent_factor=1 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --add_negative_prompt=False --add_repeat_prompt=False --use_sequence_packing=True --max_packed_inp_len=24000 --max_packed_ctx_len=48000 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --logging_steps=50
``` -->
### Squad only
```bash
WANDB_PROJECT=ctx-squad-test run uv run train.py configs/squad.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=-1 --gradient_accumulation_steps=16 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --ctx_encoder_type=per_layer_activations --n_latent_queries=8 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=1 --use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=4096 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --logging_steps=10 --max_ctx_chunk_len=-1
```
### Synthetic data generation
```bash
# old number repeat
uv run data/generate_fav_num.py
# new num repeat
uv run data/generate_ctx_numbers.py
# ctx magic num (NIAH)
uv run data/generate_ctx_magic_number.py --tokenizer-name google/gemma-2-2b-it
```
Self-gen for the number toy dataset
```bash
# old numbers
run uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names context_numbers_2_10 --split train
run uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names context_numbers_2_10 --split validation
# new ctx numbers
run uv run data/self_generate_qa.py --vllm_model google/gemma-3-1b-it --ds_names ctx_numbers_64_128 --split train --remove_qa_template --max_new_tokens 150
```
### SQuAD
```bash
# for some reason download directly through `load_dataset` does not work
HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download --repo-type dataset rajpurkar/squad --local-dir data/raw_datasets/squad
```
### Self-gen data upload/download
```bash
# create a bucket (needed only once)
# gsutil mb -l EU gs://ctx-to-lora/
# uploading from login node to gcp bucket
gsutil -m rsync -r data/raw_datasets/self_gen gs://ctx-to-lora/data/raw_datasets/self_gen
# downloading from gcp bucket to login node
mkdir -p data/raw_datasets/self_gen
gsutil -m rsync -r gs://ctx-to-lora/data/raw_datasets/self_gen data/raw_datasets/self_gen
# download only self-gen eval
mkdir -p data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/
gsutil -m cp -r gs://ctx-to-lora/data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/
```
### Upload/download checkpoints
```bash
# upload to bucket
gsutil -m rsync -r train_outputs gs://ctx-to-lora/train_outputs
# download from bucket
gsutil -m rsync -r gs://ctx-to-lora/train_outputs train_outputs
```
<!-- ### Data Generation (v2)
***Download a subset of generated Fineweb-QA directly***
```bash
uv run data/download_generated_fineweb_qa.py
uv run data/download_self_gen_qa.py
``` -->
Loading self-generated data
```python
from ctx_to_lora.model_loading import get_tokenizer
from ctx_to_lora.data.processing import load_and_process_dataset, get_tokenized_dataset
base_model_name = "google/gemma-2-2b-it"
# currently available self_gen dataset names
# ["drop_compact", "pwc_compact", "ropes_compact", "squad_compact", "fw_qa_v2/min_0_to_2000"]
ds_name = "fw_qa_v2/min_0_to_2000"
ds = load_and_process_dataset(f"self_gen/{base_model_name}/{ds_name}",
split="train",
add_negative_prompt=False,
add_repeat_prompt=False,
repeat_prob=0,
is_pretrain=False,
streaming=False,
num_proc=8
)
# or load the tokenized version
base_model_max_len = 2**13
ctx_model_max_len = 2**13
# load via a custom function because of custom chat_template
tokenizer = get_tokenizer(base_model_name)
ctx_tokenizer = get_tokenizer(base_model_name)
ds = get_tokenized_dataset(
ds_name,
split="train",
base_model_max_len=base_model_max_len,
tokenizer=tokenizer,
tokenizer_kwargs={},
ctx_model_max_len=ctx_model_max_len,
ctx_tokenizer=ctx_tokenizer,
ctx_tokenizer_kwargs={},
add_ctx_to_chat=False,
add_repeat_prompt=False,
add_negative_prompt=False,
use_kl_loss=False,
)
```
***Generate data from scratch***
0. download fineweb_edu to `data/raw_datasets/fineweb_edu
```bash
# [WIP] there are other datasets where we used openAI to gen QA pairs
# 0. download fineweb_edu to `data/raw_datasets/fineweb_edu
uv run data/download_fineweb_edu.py
```
@ -182,9 +32,6 @@ uv run data/download_fineweb_edu.py
# run from 000 to 0013
run uv run data/generate_fw_edu_qa_v2.py --shard_pattern "000_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048;
run uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
run uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_1*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
run uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_2*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it
run uv run python data/generate_fw_edu_qa_v3.py --shard_pattern '*' --question_weight 3 --use_case_weight 1 --creative_weight 1 --generic_weight 1
```
2. Self-generated response QA data (depends on step 0 and 1)
@ -200,50 +47,11 @@ uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/qa_short_ctx_self_gen_no_fw_qa.yaml
```
<!--
### GSM8k LoRA
```bash
run python train.py configs/gsm8k.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=10 --per_device_train_batch_size=16 --gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=lora --target_modules=up_proj,down_proj --eval_steps=5000 --save_steps=5000 --learning_rate=1e-4 --neftune_noise_alpha=5 --load_best_model_at_end=True --metric_for_best_model=gsm8k_loss
```
### Multitask LoRA
```bash
run python train.py configs/math_and_code.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=4 --gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=lora --target_modules=up_proj,down_proj --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --neftune_noise_alpha=5 --load_best_model_at_end=True --metric_for_best_model=gsm8k_loss
```
### HyperLoRA finetune on GSM8K
```bash
run python train.py configs/gsm8k.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=16 --gradient_accumulation_steps=1 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=1 --num_self_attends_per_block=16 --self_attention_widening_factor=1 --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --neftune_noise_alpha=5 --use_light_weight_lora=True --light_weight_latent_size=512 --load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False --add_repeat_prompt=False --ctx_encoder_model_name_or_path=meta-llama/Llama-3.2-11B-Vision-Instruct
``` -->
### Continue from a checkpoint
```bash
run python train.py configs/...yaml ... --from_pretrained_checkpoint=train_outputs/runs/May09_16-25-35_slurm0-a3nodeset-4_59459_ea85a571/checkpoint-10000/pytorch_model.bin --resume_from_checkpoint=train_outputs/runs/May09_16-25-35_slurm0-a3nodeset-4_59459_ea85a571/checkpoint-10000
```
### Evaluation
ctx magic num
See [eval scripts](scripts/eval/).
```bash
# base model
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets ctx_magic_number_32_1024 ctx_magic_number_1024_2048 ctx_magic_number_2048_3072 ctx_magic_number_3072_4096 ctx_magic_number_4096_5120 ctx_magic_number_5120_6144 ctx_magic_number_6144_7168 ctx_magic_number_7168_8192 ctx_magic_number_8192_9216 ctx_magic_number_9216_10240 ctx_magic_number_10240_11264 ctx_magic_number_11264_12288 ctx_magic_number_12288_13312 ctx_magic_number_13312_14336 ctx_magic_number_14336_15360 ctx_magic_number_15360_16384 --split test --eval_batch_size_gen=16
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 4
# hypernet w/ 1024 max chunk size
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/Aug26_05-46-31_slurm0-aiscinodeset-1_81810_f78c9c91/checkpoint-1000/pytorch_model.bin --datasets ctx_magic_number_32_1024 ctx_magic_number_1024_2048 ctx_magic_number_2048_3072 ctx_magic_number_3072_4096 ctx_magic_number_4096_5120 ctx_magic_number_5120_6144 ctx_magic_number_6144_7168 ctx_magic_number_7168_8192 ctx_magic_number_8192_9216 ctx_magic_number_9216_10240 ctx_magic_number_10240_11264 ctx_magic_number_11264_12288 ctx_magic_number_12288_13312 ctx_magic_number_13312_14336 ctx_magic_number_14336_15360 ctx_magic_number_15360_16384 --max_ctx_chunk_len=1024 --split test
```
LongBench
```bash
# generative
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/.../pytorch_model.bin --datasets negative_nq triviaqa_retrieved squad longbench_e --split test
# hypernet checkpoint
WANDB_MODE=disabled run uv run run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad ropes drop longbench/gov_report_e longbench/multifieldqa_en_e longbench/2wikimqa_e --split test --max_ctx_chunk_len -1 --eval_batch_size_gen 8
# squad only
WANDB_MODE=disabled run uv run run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad --split test
# chunking
WANDB_MODE=disabled run uv run run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad --split validation --max_ctx_chunk_len 100 --max_val_samples_per_ds 10
# base model
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --eval_batch_size 2
@ -251,45 +59,4 @@ WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it
# base model w/o context
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --remove_context
# # benchmark
# cd LongBench/LongBench
# run python pred_ctx_to_lora.py --checkpoint_path ../../train_outputs/runs/Mar16_12-38-01_slurm0-a3nodeset-12_54818_32426662/checkpoint-136782/pytorch_model.bin
# run python eval_ctx_to_lora.py --model_name Mar16_12-38-01_slurm0-a3nodeset-12_54818_32426662/checkpoint-136782 --checkpoint_path ../../train_outputs/runs/Mar16_12-38-01_slurm0-a3nodeset-12_54818_32426662/checkpoint-136782/pytorch_model.bin
```
### LLM-comparator
```bash
# install nvm
curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.3/install.sh | bash
nvm install 16
nvm use 16
git clone https://github.com/PAIR-code/llm-comparator.git
cd llm-comparator
npm install
npm run build
# running llm-comparator webui
npm run serve
# in another terminal
# run http server for file fetching
# cd back to root folder first
# taken from https://stackoverflow.com/a/79135787
alias srv='echo -e "from sys import argv as a\nfrom http.server import HTTPServer as H, SimpleHTTPRequestHandler as HH, test as t\nclass C(HH):\n def end_headers (self):\n self.send_header(a[2],a[3])\n HH.end_headers(self)\nt(C,H,port=int(a[1]))" | /usr/bin/env python3 -- - 8001 "Access-Control-Allow-Origin" "*"'
srv
# copy-paste the relative path from the root
# e.g., http://localhost:8001/train_outputs/runs/May08_13-56-31_slurm0-a3nodeset-5_59383_906acb28/eval-results-105000/comparator.json
```
## :warning: Known Issues
For the ICLR version, we use two `<bos>` tokens for `gemma-2-2b-it`.
This is because of a bug from `vllm` which we didn't realize until very close to the deadline (7 Sep 2025) (see https://github.com/vllm-project/vllm/pull/16081).
This bug causes the self-generated data to have two `<bos>` tokens (only the `input_ids` not `ctx_ids`)
NOTE: previously we were using `vllm==0.8.5` which does not have the double `<bos>` problem.
However it has another bug that randomly causes some detokenization error (see https://github.com/vllm-project/vllm/issues/17448)
Therefore we, not knowing about the double `<bos>` bug, recently downgraded `vllm` to `0.8.4`.

24
chat_templates/google/gemma-2-2b-it.jinja.bak(w_trim)
View file

@ -1,24 +0,0 @@
{{- bos_token }}
{%- if messages[0]['role'] == 'system' %}
{%- set system_message = messages[0]['content'] %}
{%- set loop_messages = messages[1:] %}
{%- else %}
{%- set loop_messages = messages %}
{%- endif %}
{% for message in loop_messages %}
{% if (message['role'] == 'assistant') %}
{% set role = 'model' %}
{% else %}
{% set role = message['role'] %}
{% endif %}
{%- if message['role'] == 'user' and loop.first and system_message is defined %}
{{ '<start_of_turn>' + role + '\n' + system_message + '\n\n' + message['content'] | trim + '<end_of_turn>\n' }}
{%- elif message['role'] == 'assistant' %}
{{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] | trim + '<end_of_turn>\n' }}{% endgeneration %}
{%- else %}
{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}
{%- endif %}
{% endfor %}
{% if add_generation_prompt %}
{{'<start_of_turn>model\n'}}
{% endif %}

24
chat_templates/google/gemma-2-9b-it.jinja
View file

@ -1,24 +0,0 @@
{{- bos_token }}
{%- if messages[0]['role'] == 'system' %}
{%- set system_message = messages[0]['content'] %}
{%- set loop_messages = messages[1:] %}
{%- else %}
{%- set loop_messages = messages %}
{%- endif %}
{% for message in loop_messages %}
{% if (message['role'] == 'assistant') %}
{% set role = 'model' %}
{% else %}
{% set role = message['role'] %}
{% endif %}
{%- if message['role'] == 'user' and loop.first and system_message is defined %}
{{ '<start_of_turn>' + role + '\n' + system_message + '\n\n' + message['content'] | trim + '<end_of_turn>\n' }}
{%- elif message['role'] == 'assistant' %}
{{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] + '<end_of_turn>' }}{% endgeneration %}{{ '\n' }}
{%- else %}
{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}
{%- endif %}
{% endfor %}
{% if add_generation_prompt %}
{{'<start_of_turn>model\n'}}
{% endif %}

30
chat_templates/google/gemma-3-1b-it.jinja
View file

@ -1,30 +0,0 @@
{{ bos_token }}
{%- if messages[0]['role'] == 'system' -%}
{%- if messages[0]['content'] is string -%}
{%- set first_user_prefix = messages[0]['content'] + '\n\n' -%}
{%- else -%}
{%- set first_user_prefix = messages[0]['content'][0]['text'] + '\n\n' -%}
{%- endif -%}
{%- set loop_messages = messages[1:] -%}
{%- else -%}
{%- set first_user_prefix = "" -%}
{%- set loop_messages = messages -%}
{%- endif -%}
{%- for message in loop_messages -%}
{%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}
{{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }}
{%- endif -%}
{%- if (message['role'] == 'assistant') -%}
{%- set role = "model" -%}
{%- else -%}
{%- set role = message['role'] -%}
{%- endif -%}
{%- if (message['role'] == 'user') -%}
{{ '<start_of_turn>' + role + '\n' + (first_user_prefix if loop.first else "") + message['content'] | trim + '<end_of_turn>\n' }}
{%- elif (message['role'] == 'assistant') -%}
{{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] + '<end_of_turn>' }}{% endgeneration %}{{ '\n' }}
{%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt -%}
{{'<start_of_turn>model\n'}}
{%- endif -%}

30
chat_templates/google/gemma-3-4b-it.jinja
View file

@ -1,30 +0,0 @@
{{ bos_token }}
{%- if messages[0]['role'] == 'system' -%}
{%- if messages[0]['content'] is string -%}
{%- set first_user_prefix = messages[0]['content'] + '\n\n' -%}
{%- else -%}
{%- set first_user_prefix = messages[0]['content'][0]['text'] + '\n\n' -%}
{%- endif -%}
{%- set loop_messages = messages[1:] -%}
{%- else -%}
{%- set first_user_prefix = "" -%}
{%- set loop_messages = messages -%}
{%- endif -%}
{%- for message in loop_messages -%}
{%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}
{{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }}
{%- endif -%}
{%- if (message['role'] == 'assistant') -%}
{%- set role = "model" -%}
{%- else -%}
{%- set role = message['role'] -%}
{%- endif -%}
{%- if (message['role'] == 'user') -%}
{{ '<start_of_turn>' + role + '\n' + (first_user_prefix if loop.first else "") + message['content'] | trim + '<end_of_turn>\n' }}
{%- elif (message['role'] == 'assistant') -%}
{{ '<start_of_turn>' + role + '\n' }}{% generation %}{{ message['content'] + '<end_of_turn>' }}{% endgeneration %}{{ '\n' }}
{%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt -%}
{{'<start_of_turn>model\n'}}
{%- endif -%}

57
configs/context_numbers_10.yaml
View file

@ -1,57 +0,0 @@
output_dir: "" # just a placeholder
bf16: true
model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
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: 64
per_device_eval_batch_size: 128
max_new_tokens: 64
gen_per_device_eval_batch_size: 128
max_val_samples_per_ds: 500
# optim: schedule_free_adamw
learning_rate: 0.0001
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
# LoRA
lora_r: 16
lora_dropout: 0.0
target_modules:
- down_proj
# data
train_ds_names:
- data/raw_datasets/context_numbers_2
- data/raw_datasets/context_numbers_3
- data/raw_datasets/context_numbers_4
- data/raw_datasets/context_numbers_5
- data/raw_datasets/context_numbers_6
- data/raw_datasets/context_numbers_7
- data/raw_datasets/context_numbers_8
- data/raw_datasets/context_numbers_9
- data/raw_datasets/context_numbers_10
val_ds_names:
- data/raw_datasets/context_numbers_2
- data/raw_datasets/context_numbers_3
- data/raw_datasets/context_numbers_4
- data/raw_datasets/context_numbers_5
- data/raw_datasets/context_numbers_6
- data/raw_datasets/context_numbers_7
- data/raw_datasets/context_numbers_8
- data/raw_datasets/context_numbers_9
- data/raw_datasets/context_numbers_10

12
configs/context_numbers_10_self_gen.yaml
View file

@ -1,12 +0,0 @@
# 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_0.0/context_numbers_2_10
val_ds_names:
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/context_numbers_2_10

46
configs/fw_qa_v2_level_0.yaml
View file

@ -1,46 +0,0 @@
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:
- fw_qa_v2_2k_len_level_0
val_ds_names:
- fw_qa_v2_2k_len_level_0
- squad

46
configs/fw_qa_v2_level_1.yaml
View file

@ -1,46 +0,0 @@
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:
- fw_qa_v2_2k_len_level_1
val_ds_names:
- fw_qa_v2_2k_len_level_0
- squad

46
configs/fw_qa_v2_level_2.yaml
View file

@ -1,46 +0,0 @@
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:
- fw_qa_v2_2k_len_level_2
val_ds_names:
- fw_qa_v2_2k_len_level_0
- squad

46
configs/fw_qa_v2_level_3.yaml
View file

@ -1,46 +0,0 @@
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:
- fw_qa_v2_2k_len_level_3
val_ds_names:
- fw_qa_v2_2k_len_level_0
- squad

60
configs/large_exp/fw_qa_v3.yaml
View file

@ -1,60 +0,0 @@
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
use_kl_loss: true
ctx_encoder_type: per_layer_activations
n_latent_queries: 8
num_blocks: 8
num_self_attn_per_block: 0
# gradient_accumulation_steps: 11
# max_packed_inp_len: 6144
# max_packed_ctx_len: 6144
gradient_accumulation_steps: 8
max_packed_inp_len: 8192
max_packed_ctx_len: 8192
# data
train_ds_names:
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_3*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.0_summ_prob_0.0_struct_prob_0.1/smol-rewrite
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.0_summ_prob_0.0_struct_prob_0.1/smol-summarize
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.0_summ_prob_1.0_struct_prob_1.0/bookqa
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.0_summ_prob_1.0_struct_prob_1.0/gov_report_qa
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/openhermes
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/smol-magpie-ultra
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/systemchats
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.2_summ_prob_0.0_struct_prob_0.0/ctx-q-gsm
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.1_cot_prob_0.0_summ_prob_1.0_struct_prob_1.0/synthetic_convqa
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.1_cot_prob_0.0_summ_prob_1.0_struct_prob_1.0/tatqa
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.8_cot_prob_0.0_summ_prob_0.0_struct_prob_0.0/quoref
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.8_cot_prob_0.0_summ_prob_1.0_struct_prob_1.0/hotpot_qa
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.8_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/drop
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.8_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/pwc
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.8_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/ropes
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.8_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/squad
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.01_cot_prob_0.1_summ_prob_0.05_struct_prob_0.1/fw_qa_v3/min_0_to_8000
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.01_cot_prob_0.1_summ_prob_0.05_struct_prob_0.1/fw_qa_v3/min_0_to_10000
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.01_cot_prob_0.1_summ_prob_0.05_struct_prob_0.1/fw_qa_v3/min_0_to_15000
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.01_cot_prob_0.1_summ_prob_0.05_struct_prob_0.1/fw_qa_v3/min_0_to_20000
val_ds_names:
- squad
- pwc
- drop
- ropes
- hotpot_qa
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet

68
configs/large_exp/fw_qa_v3_only_0_to_20k.yaml
View file

@ -1,68 +0,0 @@
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
use_kl_loss: true
ctx_encoder_type: per_layer_activations
n_latent_queries: 8
num_blocks: 8
num_self_attn_per_block: 0
# gradient_accumulation_steps: 11
# max_packed_inp_len: 6144
# max_packed_ctx_len: 6144
gradient_accumulation_steps: 8
max_packed_inp_len: 8192
max_packed_ctx_len: 8192
# data
train_ds_names:
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_3*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.0_summ_prob_0.0_struct_prob_0.1/smol-rewrite
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.0_summ_prob_0.0_struct_prob_0.1/smol-summarize
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.0_summ_prob_1.0_struct_prob_1.0/bookqa
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.0_summ_prob_1.0_struct_prob_1.0/gov_report_qa
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/openhermes
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/smol-magpie-ultra
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/systemchats
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.2_summ_prob_0.0_struct_prob_0.0/ctx-q-gsm
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.1_cot_prob_0.0_summ_prob_1.0_struct_prob_1.0/synthetic_convqa
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.1_cot_prob_0.0_summ_prob_1.0_struct_prob_1.0/tatqa
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.8_cot_prob_0.0_summ_prob_0.0_struct_prob_0.0/quoref
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.8_cot_prob_0.0_summ_prob_1.0_struct_prob_1.0/hotpot_qa
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.8_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/drop
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.8_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/pwc
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.8_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/ropes
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.8_cot_prob_0.1_summ_prob_0.0_struct_prob_0.0/squad
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.01_cot_prob_0.1_summ_prob_0.05_struct_prob_0.1/fw_qa_v3/min_0_to_8000
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.01_cot_prob_0.1_summ_prob_0.05_struct_prob_0.1/fw_qa_v3/min_0_to_10000
# - data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.01_cot_prob_0.1_summ_prob_0.05_struct_prob_0.1/fw_qa_v3/min_0_to_15000
- data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.01_cot_prob_0.1_summ_prob_0.05_struct_prob_0.1/fw_qa_v3/min_0_to_20000
val_ds_names:
- squad
- pwc
- drop
- ropes
- hotpot_qa
- booksum
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_qa/validation/*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet
# - negative_nq

0
configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_l2l.yaml → configs/main_exp/self_gen_lv1_closed_qa_1_l2l.yaml
View file

42
configs/medium_exp/qa_short_ctx_self_gen_lv1_closed_1_and_lv3_l2l.yaml
View file

@ -1,42 +0,0 @@
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
use_kl_loss: true
ctx_encoder_type: per_layer_activations
n_latent_queries: 8
num_blocks: 9
num_self_attn_per_block: 0
gradient_accumulation_steps: 11
max_packed_inp_len: 6144
max_packed_ctx_len: 6144
# data
train_ds_names:
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_3*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes
# - hotpot_qa
# - negative_nq
# - booksum
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_qa/validation/*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet

45
configs/pwc_tiny.yaml
View file

@ -1,45 +0,0 @@
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:
- pwc_tiny
val_ds_names:
- pwc

59
configs/qa_short_ctx.yaml
View file

@ -1,59 +0,0 @@
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:
- fw_qa_v2_2k_len_level_3
- squad
- pwc
- drop
- ropes
# the contexts are kinda long-ish
# - narrativeqa
# - quoref
# - synthetic_convqa
val_ds_names:
- fw_qa_v2_2k_len_level_0
- squad
- drop
- ropes
- pwc

59
configs/qa_short_ctx_compact.yaml
View file

@ -1,59 +0,0 @@
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:
- fw_qa_v2_2k_len_level_3
- 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

52
configs/qa_short_ctx_self_gen.yaml
View file

@ -1,52 +0,0 @@
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/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
- self_gen/google/gemma-2-2b-it/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

51
configs/qa_short_ctx_self_gen_no_fw_qa.yaml
View file

@ -1,51 +0,0 @@
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_0.0/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

45
configs/qa_short_ctx_self_gen_squad.yaml
View file

@ -1,45 +0,0 @@
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/squad_compact
val_ds_names:
- squad

38
configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1.yaml
View file

@ -1,38 +0,0 @@
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
use_kl_loss: true
n_latent_queries: 208
num_blocks: 9
num_self_attn_per_block: 0
gradient_accumulation_steps: 8
max_packed_inp_len: 8192
max_packed_ctx_len: 8192
# data
train_ds_names:
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes
- hotpot_qa
- negative_nq
- booksum
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_qa/validation/*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet

39
configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_and_lv3.yaml
View file

@ -1,39 +0,0 @@
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
use_kl_loss: true
n_latent_queries: 208
num_blocks: 9
num_self_attn_per_block: 0
gradient_accumulation_steps: 16
max_packed_inp_len: 4096
max_packed_ctx_len: 4096
# data
train_ds_names:
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_3*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes
- hotpot_qa
- negative_nq
- booksum
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_qa/validation/*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet

37
configs/small_exp/qa_short_ctx_self_gen_lv1_closed_qa_1_next_token_pred.yaml
View file

@ -1,37 +0,0 @@
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
use_kl_loss: false
n_latent_queries: 208
num_blocks: 9
num_self_attn_per_block: 0
gradient_accumulation_steps: 8
max_packed_inp_len: 8192
max_packed_ctx_len: 8192
# data
train_ds_names:
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes
- hotpot_qa
- negative_nq
- booksum
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
# - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_qa/validation/*.parquet
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet

45
configs/squad.yaml
View file

@ -1,45 +0,0 @@
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:
- squad
val_ds_names:
- squad

45
configs/squad_compact.yaml
View file

@ -1,45 +0,0 @@
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:
- squad_compact
val_ds_names:
- squad

47
configs/squad_fw_qa_v2_level_3.yaml
View file

@ -1,47 +0,0 @@
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:
- squad
- fw_qa_v2_2k_len_level_3
val_ds_names:
- fw_qa_v2_2k_len_level_0
- squad

47
configs/squad_hotpot_qa.yaml
View file

@ -1,47 +0,0 @@
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:
- hotpot_qa
- squad
val_ds_names:
- hotpot_qa
- squad

18
configs/toy_exp/ctx_magic_number_32_256_self_gen.yaml
View file

@ -1,18 +0,0 @@
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
use_kl_loss: true
# data
train_ds_names:
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.0_summ_prob_0.0_struct_prob_0.0/ctx_magic_number_32_128
- self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0_cot_prob_0.0_summ_prob_0.0_struct_prob_0.0/ctx_magic_number_128_256
val_ds_names:
- ctx_magic_number_32_128
- ctx_magic_number_128_256

51
data/create_test_data.py
View file

@ -1,51 +0,0 @@
import argparse
import pandas as pd
def make_test_samples():
return [
{
"context": "Alice was beginning to get very tired of sitting by her sister on the bank.",
"prompts_level_0": [
"Who is sitting by Alice?",
"How is Alice feeling at the start?"
],
"responses_level_0": [
"Her sister is sitting by her.",
"She is very tired."
],
},
{
"context": "<think>debug info</think>Then she saw a White Rabbit with pink eyes run close by her.",
"prompts_level_0": [
"What did Alice see run by?",
"What color were its eyes?"
],
"responses_level_0": [
"She saw a White Rabbit.",
"Its eyes were pink."
],
},
{
"context": "x" * 200,
"prompts_level_0": ["How long is this context?"],
"responses_level_0": ["It is two hundred characters long."],
},
]
def main(output_path: str):
samples = make_test_samples()
df = pd.DataFrame(samples)
df.to_parquet(output_path, index=False)
print(f"Wrote {len(df)} test rows to {output_path}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--output",
"-o",
type=str,
default="test_shard.parquet",
help="Where to write the test Parquet file"
)
args = parser.parse_args()
main(args.output)

20
data/download_facts_grounding.py
View file

@ -1,20 +0,0 @@
# Install dependencies as needed:
# pip install kagglehub[hf-datasets]
import kagglehub
from kagglehub import KaggleDatasetAdapter
# Set the path to the file you'd like to load
file_path = "data/raw_datasets/FACTS/ds.parquet"
# Load the latest version
hf_dataset = kagglehub.load_dataset(
KaggleDatasetAdapter.HUGGING_FACE,
"deepmind/FACTS-grounding-examples",
file_path,
# Provide any additional arguments like
# sql_query, hf_kwargs, or pandas_kwargs. See
# the documenation for more information:
# https://github.com/Kaggle/kagglehub/blob/main/README.md#kaggledatasetadapterhugging_face
)
print("Hugging Face Dataset:", hf_dataset)

9
data/download_generated_fineweb_qa.py
View file

@ -1,9 +0,0 @@
from huggingface_hub import snapshot_download
if __name__ == "__main__":
fw_dir = "./data/raw_datasets/fw_qa_v2/"
snapshot_download(
"SakanaAI/fineweb_qa",
repo_type="dataset",
local_dir=fw_dir,
)

9
data/download_self_gen_qa.py
View file

@ -1,9 +0,0 @@
from huggingface_hub import snapshot_download
if __name__ == "__main__":
fw_dir = "./data/raw_datasets/self_gen/"
snapshot_download(
"SakanaAI/self_gen_qa",
repo_type="dataset",
local_dir=fw_dir,
)

9
data/download_self_gen_qa_logprobs.py
View file

@ -1,9 +0,0 @@
from huggingface_hub import snapshot_download
if __name__ == "__main__":
fw_dir = "./data/raw_datasets/self_gen/"
snapshot_download(
"Rujikorn/self_gen_qa_logprobs",
repo_type="dataset",
local_dir=fw_dir,
)

220535
data/eval/facts/examples.csv
View file

File diff suppressed because one or more lines are too long

11313
data/eval/triviaqa/triviaqa_retrieved.jsonl
View file

File diff suppressed because one or more lines are too long

243
data/generate_ctx_kv.py
View file

@ -1,243 +0,0 @@
import argparse
import json
import os
import random
import uuid
from collections.abc import Iterable
import wonderwords
# -----------------------------
# Simple, editable config knobs
# -----------------------------
# You can change these defaults or pass CLI flags: --key-type / --value-type
DEFAULT_KEY_TYPE = "uuid" # choices: "uuid" | "digits" | "words"
DEFAULT_VALUE_TYPE = "uuid" # choices: "uuid" | "digits" | "words"
PAIR_SEP = "\n" # between pairs
KV_SEP = " : " # between key and value
TOKENS_PER_PAIR = 10 # rough heuristic used for bucketing by context length
BASE_SAMPLES_PER_BIN = 128_000
RNG_SEED = 42
nouns = wonderwords.random_word._get_words_from_text_file("nounlist.txt")
adjs = wonderwords.random_word._get_words_from_text_file("adjectivelist.txt")
# verbs = wonderwords.random_word._get_words_from_text_file("verblist.txt")
words = [f"{adj}-{noun}" for adj in adjs for noun in nouns]
words = sorted(list(set(words)))
def save_jsonl(data: list[dict], filepath: str) -> None:
"""Save data to a JSONL file."""
parent_dir = os.path.dirname(filepath)
if parent_dir:
os.makedirs(parent_dir, exist_ok=True)
with open(filepath, "w") as f:
for entry in data:
json.dump(entry, f)
f.write("\n")
def _gen_uuid8() -> str:
return uuid.uuid4().hex[:8]
def _gen_digits6() -> str:
return f"{random.randint(0, 999_999):06d}"
def _make_generator(kind: str):
kind = kind.lower()
if kind == "uuid":
return _gen_uuid8
if kind == "digits":
return _gen_digits6
if kind == "words":
return _gen_word
raise ValueError(f"Unknown kind '{kind}', expected 'uuid', 'digits', or 'words'")
def _gen_digits4() -> str:
return f"{random.randint(0, 9_999):04d}"
def _gen_word() -> str:
return random.choice(words)
def _make_value_generator(kind: str):
"""Like _make_generator, but when kind is 'digits' use 4 digits for values."""
kind = kind.lower()
if kind == "uuid":
return _gen_uuid8
if kind == "digits":
return _gen_digits4
if kind == "words":
return _gen_word
raise ValueError(f"Unknown kind '{kind}', expected 'uuid', 'digits', or 'words'")
def _unique(seq: Iterable[str]) -> list[str]:
s = set()
out = []
for x in seq:
if x not in s:
s.add(x)
out.append(x)
return out
def render_context(pairs: list[tuple[str, str]]) -> str:
"""Render key-value pairs as a single context string.
Example: "k1:v1 | k2:v2 | k3:v3"
"""
return PAIR_SEP.join([f"{k}{KV_SEP}{v}" for k, v in pairs])
def generate_kv_dataset(n: int, k: int, key_type: str, value_type: str):
"""Generate n examples. Each example has k key-value pairs in the context.
- context: "k1:v1 | k2:v2 | ..."
- prompt: "What is the value of `{k}`?"
- response: "v" (exactly the value corresponding to the chosen k)
"""
key_gen = _make_generator(key_type)
val_gen = _make_value_generator(value_type)
dataset = []
for _ in range(n):
# Generate unique keys to avoid ambiguity
keys: list[str] = []
while len(keys) < k:
keys = _unique([key_gen() for _ in range(k * 2)]) # oversample, then dedupe
keys = keys[:k]
values = [val_gen() for _ in range(k)]
pairs = list(zip(keys, values))
# Choose a random key from the context to ask about
q_key = random.choice(keys)
q_val = values[keys.index(q_key)]
entry = {
"context": f"key{KV_SEP}value\n" + render_context(pairs),
"prompt": f'What is the value of key "{q_key}"? Reply with only the value.',
"response": q_val,
}
dataset.append(entry)
# Split into train/val/test matching the same structure as generate_ctx_numbers.py
total_size = len(dataset)
train_size = int(0.98 * total_size)
val_size = int(0.01 * total_size)
train_data = dataset[:train_size]
val_data = dataset[train_size : train_size + val_size]
test_data = dataset[train_size + val_size :]
return train_data, val_data, test_data
def main():
parser = argparse.ArgumentParser(
description="Generate key-value context QA data with the same structure as generate_ctx_numbers.py",
)
parser.add_argument(
"--key-type",
choices=["uuid", "digits", "words"],
default=DEFAULT_KEY_TYPE,
help="Type of keys to generate: 8-char uuid, 6-digit numbers, or words",
)
parser.add_argument(
"--value-type",
choices=["uuid", "digits", "words"],
default=DEFAULT_VALUE_TYPE,
help="Type of values to generate: 8-char uuid, 4-digit numbers, or words",
)
parser.add_argument("--seed", type=int, default=RNG_SEED, help="Random seed")
parser.add_argument(
"--base-samples-per-bin",
type=int,
default=BASE_SAMPLES_PER_BIN,
help="Baseline number of samples per token bin (scaled by bin width)",
)
parser.add_argument(
"--out-prefix",
type=str,
default="data/raw_datasets/ctx_kv",
help="Output directory prefix (bin range will be appended)",
)
parser.add_argument(
"--tokens-per-pair",
type=int,
default=TOKENS_PER_PAIR,
help="Heuristic tokens per key:value pair for bucketing",
)
parser.add_argument(
"--only-first-n-bins",
type=int,
default=None,
help="For quick tests: only generate the first N token bins",
)
parser.add_argument(
"--dry-run",
action="store_true",
help="Print a small sample and exit without writing files",
)
args = parser.parse_args()
random.seed(args.seed)
# Same token bins as generate_ctx_numbers.py
tok_bins = [(64, 128), (128, 256), (256, 512)] + [
(512 + 256 * i, 512 + 256 * (i + 1)) for i in range(14)
]
if args.only_first_n_bins is not None:
tok_bins = tok_bins[: args.only_first_n_bins]
# Map token bins to pair-length bins using a simple heuristic
len_bins = [
(lo // args.tokens_per_pair, hi // args.tokens_per_pair)
for (lo, hi) in tok_bins
]
# Optional dry-run: build a tiny sample to preview format and exit
if args.dry_run:
train, val, test = generate_kv_dataset(
n=3,
k=max(1, len_bins[0][0] or 1),
key_type=args.key_type,
value_type=args.value_type,
)
print("Sample entry:")
print(json.dumps(train[0], indent=2))
return
for len_bin, tok_bin in zip(len_bins, tok_bins):
bin_size = max(1, len_bin[1] - len_bin[0]) # avoid div-by-zero for tiny bins
save_dir = f"{args.out_prefix}_{tok_bin[0]}_{tok_bin[1]}"
train_data, val_data, test_data = [], [], []
# Iterate over different context lengths (number of pairs)
for k in range(max(1, len_bin[0]), max(1, len_bin[1])):
# scale like the numbers script
per_k = max(1, args.base_samples_per_bin // bin_size)
train, val, test = generate_kv_dataset(
n=per_k, k=k, key_type=args.key_type, value_type=args.value_type
)
train_data += train
val_data += val
test_data += test
os.makedirs(save_dir, exist_ok=True)
save_jsonl(train_data, f"{save_dir}/train.jsonl")
save_jsonl(val_data, f"{save_dir}/val.jsonl")
save_jsonl(test_data, f"{save_dir}/test.jsonl")
print(f"Dataset generated and saved at {save_dir}")
if __name__ == "__main__":
main()

95
data/generate_ctx_numbers.py
View file

@ -1,95 +0,0 @@
import json
import os
import random
def save_jsonl(data: list[dict], filepath: str) -> None:
"""Save data to a JSONL file."""
parent_dir = os.path.dirname(filepath)
if parent_dir: # Only create directories if there's a parent path
os.makedirs(parent_dir, exist_ok=True)
with open(filepath, "w") as f:
for entry in data:
json.dump(entry, f)
f.write("\n")
def get_random_combinations(numbers: list[int], n: int, k: int) -> list[list[int]]:
"""Get n random combinations of k numbers from the list."""
# using itertools.combinations hangs with large numbers
# so we explicitly generate the n indices
indices = [random.sample(range(len(numbers)), k=k) for _ in range(n)]
return [[numbers[i] for i in ind] for ind in indices]
def generate_number_dataset(max_num: int = 1000, n: int = 12000, k: int = 3):
"""
Generate a dataset of numbers with corresponding query and answer,
split into train/val/test sets.
Args:
max_num: Maximum number in the range (exclusive)
k: Number of elements in each combination
"""
# Generate list of all numbers and shuffle them
numbers = list(range(max_num))
random.shuffle(numbers)
# Create dataset entries
dataset = []
query = f"Repeat the information above exactly. Do not output anything else."
# Generate all unique combinations of k numbers
combinations = get_random_combinations(numbers, n, k)
# random.shuffle(combinations)
for combination in combinations:
num_list = [f"{i + 1}. {num}" for i, num in enumerate(combination)]
entry = {
"context": "\n".join(num_list),
"prompt": query,
"response": "\n".join(num_list),
}
dataset.append(entry)
# Calculate split sizes
total_size = len(dataset)
train_size = int(0.98 * total_size)
val_size = int(0.01 * total_size)
# Split dataset
train_data = dataset[:train_size]
val_data = dataset[train_size : train_size + val_size]
test_data = dataset[train_size + val_size :]
return train_data, val_data, test_data
if __name__ == "__main__":
# Set random seed for reproducibility
random.seed(42)
tok_bins = [(64, 128), (128, 256), (256, 512)] + [
(512 + 256 * i, 512 + 256 * (i + 1)) for i in range(14)
]
# roughly 9 tokens per number
tok_per_num = [9 if bin[0] >= 1024 else 8 for bin in tok_bins]
len_bins = [
(bin[0] // tok, bin[1] // tok) for bin, tok in zip(tok_bins, tok_per_num)
]
for len_bin, tok_bin in zip(len_bins, tok_bins):
bin_size = len_bin[1] - len_bin[0]
save_dir = f"data/raw_datasets/ctx_numbers_{tok_bin[0]}_{tok_bin[1]}"
train_data, val_data, test_data = [], [], []
for k in range(*len_bin):
train, val, test = generate_number_dataset(n=128_000 // bin_size, k=k)
train_data += train
val_data += val
test_data += test
# Save splits to separate files
os.makedirs(save_dir, exist_ok=True)
save_jsonl(train_data, f"{save_dir}/train.jsonl")
save_jsonl(val_data, f"{save_dir}/val.jsonl")
save_jsonl(test_data, f"{save_dir}/test.jsonl")
print(f"Dataset generated and saved at {save_dir}")

104
data/generate_fav.py
View file

@ -1,104 +0,0 @@
import json
import os
import random
def save_jsonl(data: list[dict], filepath: str) -> None:
"""Save data to a JSONL file."""
parent_dir = os.path.dirname(filepath)
if parent_dir: # Only create directories if there's a parent path
os.makedirs(parent_dir, exist_ok=True)
with open(filepath, "w") as f:
for entry in data:
json.dump(entry, f)
f.write("\n")
def get_random_combinations(numbers: list[int], n: int, k: int) -> list[list[int]]:
"""Get n random combinations of k numbers from the list."""
# using itertools.combinations hangs with large numbers
# so we explicitly generate the n indices
indices = [random.sample(range(len(numbers)), k=k) for _ in range(n)]
return [[numbers[i] for i in ind] for ind in indices]
def generate_number_dataset(
max_num: int = 1000, n: int = 12000, k: int = 3, save_dir: str = None
):
"""
Generate a dataset of numbers with corresponding query and answer,
split into train/val/test sets.
Args:
max_num: Maximum number in the range (exclusive)
k: Number of elements in each combination
"""
# Generate list of all numbers and shuffle them
numbers = list(range(max_num))
random.shuffle(numbers)
# Create dataset entries
dataset = []
query = f"Repeat the text above."
# Generate all unique combinations of k numbers
combinations = get_random_combinations(numbers, n, k)
# random.shuffle(combinations)
for combination in combinations:
entry = {
"context": ", ".join(map(str, combination)),
"prompt": query,
"response": ", ".join(map(str, combination)),
}
dataset.append(entry)
# Calculate split sizes
total_size = len(dataset)
train_size = int(0.98 * total_size)
val_size = int(0.01 * total_size)
# Split dataset
train_data = dataset[:train_size]
val_data = dataset[train_size : train_size + val_size]
test_data = dataset[train_size + val_size :]
save_dir = "" if save_dir is None else save_dir
# Save splits to separate files
save_jsonl(train_data, f"{save_dir}/train.jsonl")
save_jsonl(val_data, f"{save_dir}/val.jsonl")
save_jsonl(test_data, f"{save_dir}/test.jsonl")
if __name__ == "__main__":
# Set random seed for reproducibility
random.seed(42)
for k in range(2, 11):
save_dir = f"data/raw_datasets/context_numbers_{k}"
os.makedirs(save_dir, exist_ok=True)
generate_number_dataset(n=12_000, k=k, save_dir=save_dir)
print(f"Dataset generated and saved at {save_dir}")
# # Generate dataset
# for k in range(2, 257):
# save_dir = f"data/raw_datasets/context_numbers_{k}"
# os.makedirs(save_dir, exist_ok=True)
# generate_number_dataset(n=12_000, k=k, save_dir=save_dir)
# print(f"Dataset generated and saved at {save_dir}")
# # for k in range(144, 257, 16):
# # save_dir = f"context_numbers_{k}"
# # os.makedirs(save_dir, exist_ok=True)
# # generate_number_dataset(n=100_000, k=k, save_dir=save_dir)
# # print(f"Dataset generated and saved at {save_dir}")
# for k in [512]:
# save_dir = f"data/raw_datasets/context_numbers_{k}"
# os.makedirs(save_dir, exist_ok=True)
# generate_number_dataset(n=100_000, k=k, save_dir=save_dir)
# print(f"Dataset generated and saved at {save_dir}")

83
data/generate_fav_num_big.py
View file

@ -1,83 +0,0 @@
import json
import os
import random
def save_jsonl(data: list[dict], filepath: str) -> None:
"""Save data to a JSONL file."""
parent_dir = os.path.dirname(filepath)
if parent_dir: # Only create directories if there's a parent path
os.makedirs(parent_dir, exist_ok=True)
with open(filepath, "w") as f:
for entry in data:
json.dump(entry, f)
f.write("\n")
def get_random_combinations(numbers: list[int], n: int, k: int) -> list[list[int]]:
"""Get n random combinations of k numbers from the list."""
# using itertools.combinations hangs with large numbers
# so we explicitly generate the n indices
indices = [random.sample(range(len(numbers)), k=k) for _ in range(n)]
return [[numbers[i] for i in ind] for ind in indices]
def generate_number_dataset(
max_num: int = 1000, n: int = 12000, k: int = 3, save_dir: str = None
):
"""
Generate a dataset of numbers with corresponding query and answer,
split into train/val/test sets.
Args:
max_num: Maximum number in the range (exclusive)
k: Number of elements in each combination
"""
# Generate list of all numbers and shuffle them
numbers = list(range(max_num))
random.shuffle(numbers)
# Create dataset entries
dataset = []
query = f"Repeat the text above."
# Generate all unique combinations of k numbers
combinations = get_random_combinations(numbers, n, k)
# random.shuffle(combinations)
for combination in combinations:
entry = {
"context": ", ".join(map(str, combination)),
"prompt": query,
"response": ", ".join(map(str, combination)),
}
dataset.append(entry)
# Calculate split sizes
total_size = len(dataset)
train_size = int(0.98 * total_size)
val_size = int(0.01 * total_size)
# Split dataset
train_data = dataset[:train_size]
val_data = dataset[train_size : train_size + val_size]
test_data = dataset[train_size + val_size :]
save_dir = "" if save_dir is None else save_dir
# Save splits to separate files
save_jsonl(train_data, f"{save_dir}/train.jsonl")
save_jsonl(val_data, f"{save_dir}/val.jsonl")
save_jsonl(test_data, f"{save_dir}/test.jsonl")
if __name__ == "__main__":
# Set random seed for reproducibility
random.seed(42)
# Generate dataset
for k in [32, 64, 128, 256]:
save_dir = f"context_numbers_{k}_big"
os.makedirs(save_dir, exist_ok=True)
generate_number_dataset(n=240_000, k=k, save_dir=save_dir)
print(f"Dataset generated and saved at {save_dir}")

157
data/generate_fw_edu_augment_vllm.py
View file

@ -1,157 +0,0 @@
import os
import random
import sys
from glob import glob
import pandas as pd
from datasets import Dataset, load_dataset
from vllm import LLM, SamplingParams
from ctx_to_lora.utils import clear_gpu
MODEL_CTX_LEN = {
"google/gemma-2-27b-it": 8192,
"google/gemma-2-2b-it": 8192,
"google/gemma-2-9b-it": 8192,
# not actually 16k but it's faster
"mistralai/Mistral-Small-3.1-24B-Instruct-2503": 2**14, # 11264,
}
# based on New News: system-2 finetuning (https://arxiv.org/pdf/2505.01812)
SYSTEM_TEMPLATE = (
"You are an careful and creative paraphraser. Paraphrase the given content without leaving out any important information. "
"You will be given a context and you must generate a new variation of the provided context.\n"
"You only output the paraphrase itself.\n"
"**DO NOT** halucinate or make up information."
)
# based on On the generalization of language models from in-context learning and finetuning: a controlled study
# https://arxiv.org/pdf/2505.00661
PROMPT_TEMPLATE = (
"### Context ###\n"
"{context}"
"\n\n\n"
"### Instructions ###\n"
"- Please generate rephrasings that can be inferred from each sentence. Even a simple sentence has some alternative phrasings that are logically equivalent.\n"
"- Reversing a sentence is also valid but make sure that the reversed sentence is still logically equivalent.\n"
"- You can shuffle the order of sentences as long as the content is coherent and the meaning is preserved.\n"
"- You can combine multiple sentences into one, or split a long sentence into multiple shorter ones.\n"
"- You can combine reversal and paraphasing, e.g., reversing a sentence and then paraphrasing it.\n"
"- You can also add new sentences that can be logically inferred from the context. "
"Please only use logic and language to draw your conclusions, regardless of whether the entities in question actually exist.\n"
"- Try to keep the length to be similar to the original context\n"
"- Please accurately maintain facts, such as names and numbers, mentioned in the context."
"\n\nOnly output a new variation of the context. Don't include any other words."
)
def get_prompt(context):
prompt = PROMPT_TEMPLATE.format(context=context)
return prompt
def length_filter(samples):
return [len(text) < 10_000 for text in samples["text"]]
if __name__ == "__main__":
vllm_model = os.environ.get("vllm_model")
print(f"Using model: {vllm_model}")
llm_kwargs = dict(
model=vllm_model,
dtype="bfloat16",
enable_prefix_caching=True,
enable_chunked_prefill=True,
max_model_len=MODEL_CTX_LEN[vllm_model],
)
if vllm_model == "mistralai/Mistral-Small-3.1-24B-Instruct-2503":
mm_kwargs = {"image": 0}
llm_kwargs["tokenizer_mode"] = "mistral"
llm_kwargs["config_format"] = "mistral"
llm_kwargs["load_format"] = "mistral"
llm_kwargs["limit_mm_per_prompt"] = mm_kwargs
SYSTEM_TEMPLATE = (
"You are Mistral Small 3.1, a Large Language Model (LLM) created by Mistral AI, a French startup headquartered in Paris.\n"
"You power an AI assistant called Le Chat.\n"
"Your knowledge base was last updated on 2023-10-01.\n\n\n"
) + SYSTEM_TEMPLATE
llm = LLM(**llm_kwargs)
tokenizer = llm.get_tokenizer()
shard_pattern = sys.argv[1]
paths = glob(
f"./data/raw_datasets/fineweb_edu/sample/100BT/{shard_pattern}.parquet"
)
for path in paths:
ds = load_dataset(
"parquet",
data_files=path,
split="train",
# streaming=True,
)
ds = ds.filter(length_filter, batched=True)
ctxs = [sample["text"] for sample in iter(ds)]
if "gemma" in vllm_model:
messages = [
[
{
"role": "user",
"content": SYSTEM_TEMPLATE + "\n\n\n" + get_prompt(ctx),
},
]
for ctx in ctxs
]
else:
messages = [
[
{"role": "system", "content": SYSTEM_TEMPLATE},
{"role": "user", "content": get_prompt(ctx)},
]
for ctx in ctxs
]
print(f"Generating from {len(messages)} contexts")
completions = llm.chat(
messages,
sampling_params=SamplingParams(
max_tokens=2048,
temperature=(
0.2
if vllm_model == "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
else 0.7
),
# frequency_penalty=0.1,
),
)
samples = []
for ctx, completion in zip(ctxs, completions):
samples.append(
{
"context": ctx,
"variation": completion.outputs[0].text.strip(),
}
)
for sample in samples:
print(f"{sample['context']=}")
print(f"{sample['variation']=}")
print("=" * 60)
print(f"Generated {len(samples)} samples")
random.shuffle(samples)
df = pd.DataFrame(samples)
ds = Dataset.from_pandas(df)
val_ds = ds.take(10)
ds = ds.skip(10)
shard_name = path.split("/")[-1].split(".")[0]
ds.to_parquet(f"data/raw_datasets/fw_qa_aug_pretrain/{shard_name}.parquet")
val_ds.to_parquet(
f"data/raw_datasets/fw_qa_aug_pretrain/{shard_name}_val.parquet"
)
print(f"Saved to data/raw_datasets/fw_qa_aug_pretrain/{shard_name}.parquet")
print(f"Saved to data/raw_datasets/fw_qa_aug_pretrain/{shard_name}_val.parquet")
del ds, val_ds, df, samples, completions, messages, ctxs
clear_gpu()

2
data/generate_fw_edu_qa_v2.py
View file

@ -157,7 +157,7 @@ if __name__ == "__main__":
parser.add_argument(
"--max_length",
type=int,
default=10_000,
default=2000,
help="Maximum length of the context to consider for generation",
)
parser.add_argument(

317
data/generate_fw_edu_qa_v3.py
View file

@ -1,317 +0,0 @@
import argparse
import os
import random
import re
from glob import glob
import pandas as pd
from datasets import Dataset, load_dataset
from q_generation_prompts import ( # get_structuring_seed_prompt,; get_summarization_seed_prompt,
PROMPT_TEMPLATE,
SYSTEM_PROMPT,
get_creative_seed_prompt,
get_generic_seed_prompt,
get_question_seed_prompt,
get_use_case_seed_prompt,
)
from vllm import LLM, SamplingParams
from ctx_to_lora.data.processing import load_and_process_dataset
STOP_STRINGS = {
"google/gemma-3-12b-it": ["<eos>", "<end_of_turn>"],
"google/gemma-3-4b-it": ["<eos>", "<end_of_turn>"],
}
def build_messages(
context: str,
question_weight: float,
use_case_weight: float,
creative_weight: float,
generic_weight: float,
) -> list[dict]:
# Create list of instruction functions with their weights
instruction_options = [
("question", get_question_seed_prompt, question_weight),
("use_case", get_use_case_seed_prompt, use_case_weight),
("creative", get_creative_seed_prompt, creative_weight),
("generic", get_generic_seed_prompt, generic_weight),
]
q_types, functions, weights = zip(*instruction_options)
chosen_idx = random.choices(range(len(instruction_options)), weights=weights)[0]
q_type = q_types[chosen_idx]
function = functions[chosen_idx]
custom_instruction = function()
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{
"role": "user",
"content": PROMPT_TEMPLATE.format(
context=context, instruction=custom_instruction
),
},
]
return messages, q_type
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_questions(res_txt: str):
"""
Postprocesses the questions from the response text.
Args:
res_txt: The response text.
Returns:
A list containing the questions.
"""
# capture everything after each "Question {number}:" until the next question or end
res_txt = res_txt.split("</think>")[-1]
q_pattern = r"Message:(.*?)(?=Message \d+|$)" # thanks chatgpt
questions = re.findall(q_pattern, res_txt, flags=re.S)
out_q = []
n_skips = 0
if len(questions) > 0:
for i, question in enumerate(questions):
question = question.strip()
if not question:
print(f"Skipping empty question at index {i}")
continue
# Check if this is the last question and handle stop strings
if i == len(questions) - 1:
question, skip = check_should_skip(question, vllm_model)
if skip:
print(f"Skipping due to missing stop string")
n_skips += 1
continue
out_q.append(question.strip())
print(f"Skipped {n_skips} responses due to missing stop strings")
return out_q
def length_filter(sample, min_len, max_len):
return min_len <= len(sample["text"]) <= max_len
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Generate QA pairs from FineWeb Edu dataset"
)
parser.add_argument(
"--vllm_model",
type=str,
default=os.environ.get("vllm_model", "google/gemma-3-4b-it"),
help="VLLM model to use for generation",
)
parser.add_argument(
"--shard_pattern",
type=str,
default=None,
help="Pattern to match shard files (e.g., '000_0000*')",
)
parser.add_argument(
"--ds_names", # list of names
type=str,
nargs="+",
default=None,
help="dataset names",
)
parser.add_argument(
"--split",
type=str,
default=None,
help="Dataset split to use (e.g., 'train', 'validation')",
)
# parser.add_argument(
# "--n_qa_pairs",
# type=int,
# required=True,
# help="Number of question-answer pairs to generate per context",
# )
parser.add_argument(
"--min_length",
type=int,
default=0,
help="Minimum length of the context to consider for generation",
)
parser.add_argument(
"--max_length",
type=int,
default=20_000,
help="Maximum length of the context to consider for generation",
)
parser.add_argument(
"--max_model_length",
type=int,
default=2**13,
help="Maximum length of the model input (context + prompt + response) in tokens",
)
parser.add_argument(
"--debug",
action="store_true",
help="Debug mode - process only first 100 samples",
)
parser.add_argument(
"--question_weight",
type=float,
default=0,
help="Weight for question seed prompts",
)
parser.add_argument(
"--use_case_weight",
type=float,
default=0,
help="Weight for use case seed prompts",
)
parser.add_argument(
"--creative_weight",
type=float,
default=0,
help="Weight for creative seed prompts",
)
parser.add_argument(
"--generic_weight",
type=float,
default=0,
help="Weight for generic seed prompts",
)
args = parser.parse_args()
assert bool(args.shard_pattern) ^ bool(args.ds_names)
if args.ds_names:
assert bool(args.split)
weights = [
args.question_weight,
args.use_case_weight,
args.creative_weight,
args.generic_weight,
]
total_weight = sum(weights)
assert all(w >= 0 for w in weights), "All weights must be non-negative"
assert total_weight > 0, "At least one weight must be greater than zero"
vllm_model = args.vllm_model
print(f"Using model: {vllm_model}")
llm_kwargs = dict(
model=vllm_model,
dtype="bfloat16",
enable_prefix_caching=True,
enable_chunked_prefill=True,
max_model_len=args.max_model_length,
limit_mm_per_prompt={"image": 0},
)
llm = LLM(**llm_kwargs)
tokenizer = llm.get_tokenizer()
shard_pattern = args.shard_pattern
# n_qa_pairs = args.n_qa_pairs
if args.shard_pattern is not None:
paths = glob(
f"./data/raw_datasets/fineweb_edu/sample/100BT/{shard_pattern}.parquet"
)
elif args.ds_names:
paths = args.ds_names
else:
raise ValueError("Either shard_pattern or ds_names must be provided.")
split = "train[:100]" if args.debug else args.split or "train"
for path in paths:
if args.shard_pattern is not None:
ds = load_dataset(
"parquet",
data_files=path,
split=split,
)
elif args.ds_names:
ds = load_and_process_dataset(path, args.split, num_proc=16)
ds = ds.rename_column("context", "text")
print(f"Loaded {len(ds)} samples from {path}")
ds = ds.filter(
length_filter,
fn_kwargs={"min_len": args.min_length, "max_len": args.max_length},
num_proc=8,
)
messages = []
q_types = []
for ctx in ds["text"]:
message, q_type = build_messages(
ctx,
args.question_weight,
args.use_case_weight,
args.creative_weight,
args.generic_weight,
)
messages.append(message)
q_types.append(q_type)
print(f"Generating from {len(messages)} contexts")
completions = llm.chat(
messages,
sampling_params=SamplingParams(
max_tokens=256,
temperature=0.0,
# needed for checking if stop tokens are present
skip_special_tokens=False,
include_stop_str_in_output=True,
),
)
samples = []
for ctx, completion, q_type in zip(ds["text"], completions, q_types):
questions = postprocess_questions(completion.outputs[0].text)
samples.append(
{
"context": ctx,
"prompts": questions,
"type": q_type,
}
)
if args.debug:
print(f"{ctx=}")
print(f"{completion.outputs[0].text=}")
print(f"{q_type=}")
for q in questions:
print(f"{q=}")
print()
print("=" * 80)
print(f"Generated {len(samples)} samples")
df = pd.DataFrame(samples)
ds = Dataset.from_pandas(df)
val_ds = ds.take(10)
ds = ds.skip(10)
shard_name = path.split("/")[-1].split(".")[0]
shard_name += "_level_0"
if args.debug:
shard_name += "_debug"
ds.to_parquet(
f"data/raw_datasets/fw_qa_v3/min_{args.min_length}_to_{args.max_length}/{shard_name}.parquet"
)
val_ds.to_parquet(
f"data/raw_datasets/fw_qa_v3/min_{args.min_length}_to_{args.max_length}/{shard_name}_val.parquet"
)
print(
f"Saved to data/raw_datasets/fw_qa_v3/min_{args.min_length}_to_{args.max_length}/{shard_name}.parquet"
)
print(
f"Saved to data/raw_datasets/fw_qa_v3/min_{args.min_length}_to_{args.max_length}/{shard_name}_val.parquet"
)

235
data/generate_fw_edu_qa_vllm.py
View file

@ -1,235 +0,0 @@
import os
import random
import re
import sys
from glob import glob
import pandas as pd
from datasets import Dataset, load_dataset
from vllm import LLM, SamplingParams
SYSTEM_TEMPLATE = (
"You are a creative and helpful assistant.\n"
"You are tasked with generating questions for reading comprehension tests.\n"
"You will be given a context and you need to generate questions and corresponding answers from the given context.\n"
"The questions should be highly specific to the information provided in the context, not general questions that suits any context.\n"
"**DO NOT** halucinate or make up information."
)
# based on Make Your LLM Fully Utilize the Context (https://arxiv.org/pdf/2404.16811)
PROMPT_TEMPLATE = (
"### Instructions ###\n"
"Generate questions and corresponding answers from the given context. The questions should be highly specific to the "
"information provided in the context, not general questions that suits any context.\n\n\n"
"### Context ###\n"
"{context}"
"\n\n\n"
"### Additional Instructions ###\n"
"Rules to follow when generate the questions:\n"
"1. The questions must be specific to the given context and fully answerable from information present in given context.\n"
"2. Make sure the questions are clear and unambiguous.\n"
"3. Phrases like 'based on the provided context', 'according to the context', etc, are **NOT ALLOWED** to appear in "
"the questions.\n"
"4. The questions should not overlap. They should be diverse, covering many aspects of the context.\n"
"5. Do not give away too much information in the questions. For example, ask 'Who is X' instead of 'Who is X that did Y' when Y is clear from the context.\n"
"6. Ignore the text formatting of the context, e.g., bold, italic, underline, etc.\n"
"7. Ignore typos, spacing and grammatical errors in the context.\n\n"
"Rules to follow when generate the answers:\n"
"1. The answers must use the information provided in the context.\n"
"2. Do not just copy words from the context. Answer the question in your own words.\n"
"3. The answers should be detailed and comprehensive. Please include additional specific details from the context when possible.\n\n"
"Response with {n_qa_pairs} question-answer pairs.\n"
"Always use proper grammar and punctuation.\n"
"Try to use different question forms and styles.\n"
"Use simple words and make sure that the answers are clear and comprehensive.\n\n"
"The question-answer pairs should be in the following format:\n"
"Question 1: {{question_1}}\n"
"Answer 1: {{answer_1}}\n"
"Question 2: {{question_2}}\n"
"Answer 2: {{answer_2}}\n"
"..."
)
def get_prompt(context, n_qa_pairs):
prompt = PROMPT_TEMPLATE.format(context=context, n_qa_pairs=n_qa_pairs)
return prompt
def postprocess_qa_pairs(res_txt: str):
"""
Postprocesses the QA pairs from the response text.
Args:
res_txt: The response text.
n_qa_pairs: The number of QA pairs.
Returns:
A tuple of two lists, the first containing the questions and the second containing the answers.
"""
# capture everything after each "Question {number}:" until "Answer"
q_pattern = r"Question \d+:(.*?)(?=Answer|$)" # thanks chatgpt
questions = re.findall(q_pattern, res_txt, flags=re.S)
a_pattern = r"Answer \d+:(.*?)(?=Question|$)" # thanks chatgpt
answers = re.findall(a_pattern, res_txt, flags=re.S)
if len(questions) != len(answers):
print(f"Warning---number of questions and answers do not match")
print(f"Number of questions: {len(questions)}")
print(f"Number of answers: {len(answers)}")
out_q = []
out_a = []
if (len(questions) > 0) and (len(answers) > 0):
for i in range(min(len(questions), len(answers))):
out_q.append(questions[i].strip())
out_a.append(answers[i].strip())
return out_q, out_a
def length_filter(samples):
return [len(text) < 10_000 for text in samples["text"]]
if __name__ == "__main__":
# api_key = os.environ.get("vllm_api_key")
# vllm_port = os.environ.get("vllm_port")
# vllm_ip = os.environ.get("vllm_ip") # "172.16.0.62"
vllm_model = os.environ.get("vllm_model") # "google/gemma-2-27b-it"
print(f"Using model: {vllm_model}")
# print(f"Using API key: {api_key}")
# print(f"Using VLLM IP: {vllm_ip}")
# print(f"Using VLLM port: {vllm_port}")
# client = OpenAI(
# base_url=f"http://{vllm_ip}:{vllm_port}/v1",
# api_key=api_key,
# )
mm_kwargs = {"image": 0} # if "gemma-3" in vllm_model else {}
llm_kwargs = dict(
model=vllm_model,
dtype="bfloat16",
# tokenizer_mode="mistral",
# config_format="mistral",
# load_format="mistral",
enable_prefix_caching=True,
max_model_len=2**14, # 11264,
limit_mm_per_prompt=mm_kwargs,
# enable_chunked_prefill=True,
)
if vllm_model == "mistralai/Mistral-Small-3.1-24B-Instruct-2503":
llm_kwargs["tokenizer_mode"] = "mistral"
llm_kwargs["config_format"] = "mistral"
llm_kwargs["load_format"] = "mistral"
SYSTEM_TEMPLATE = (
"You are Mistral Small 3.1, a Large Language Model (LLM) created by Mistral AI, a French startup headquartered in Paris.\n"
"You power an AI assistant called Le Chat.\n"
"Your knowledge base was last updated on 2023-10-01.\n\n\n"
) + SYSTEM_TEMPLATE
llm = LLM(**llm_kwargs)
tokenizer = llm.get_tokenizer()
shard_pattern = sys.argv[1]
n_qa_pairs = int(sys.argv[2])
paths = glob(
f"./data/raw_datasets/fineweb_edu/sample/100BT/{shard_pattern}.parquet"
)
for path in paths:
ds = load_dataset(
"parquet",
data_files=path,
split="train",
streaming=True,
)
ds = ds.filter(length_filter, batched=True)
ctxs = [sample["text"] for sample in iter(ds)]
messages = [
[
{"role": "system", "content": SYSTEM_TEMPLATE},
{"role": "user", "content": get_prompt(ctx, n_qa_pairs)},
]
for ctx in ctxs
]
print(f"Generating from {len(messages)} contexts")
completions = llm.chat(
messages,
sampling_params=SamplingParams(
max_tokens=2048,
temperature=(
0.1
if vllm_model == "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
else 0.7
),
# frequency_penalty=0.1,
),
)
samples = []
for ctx, completion in zip(ctxs, completions):
questions, answers = postprocess_qa_pairs(completion.outputs[0].text)
for q, a in zip(questions, answers):
samples.append(
{
"context": ctx,
"prompt": q,
"response": a,
}
)
print(f"Generated {len(samples)} samples")
random.shuffle(samples)
df = pd.DataFrame(samples)
ds = Dataset.from_pandas(df)
val_ds = ds.take(10)
ds = ds.skip(10)
shard_name = path.split("/")[-1].split(".")[0]
ds.to_parquet(f"data/raw_datasets/fw_qa_3/{shard_name}.parquet")
val_ds.to_parquet(f"data/raw_datasets/fw_qa_3/{shard_name}_val.parquet")
print(f"Saved to data/raw_datasets/fw_qa_3/{shard_name}.parquet")
print(f"Saved to data/raw_datasets/fw_qa_3/{shard_name}_val.parquet")
# debug
# ds = load_dataset(
# "parquet",
# data_files=paths[0],
# split="train[:10]",
# )
# ctxs = [sample["text"] for sample in iter(ds)]
# messages = [
# [
# {"role": "system", "content": SYSTEM_TEMPLATE},
# {"role": "user", "content": get_prompt(ctx, n_qa_pairs)},
# ]
# for ctx in ctxs
# ]
#
# print(f"Generating from {len(messages)} contexts")
# completions = llm.chat(
# messages,
# sampling_params=SamplingParams(
# max_tokens=2048,
# temperature=0.1 if vllm_model=="mistralai/Mistral-Small-3.1-24B-Instruct-2503" else 0.7,
# # frequency_penalty=0.1,
# ),
# )
# samples = []
# for ctx, completion in zip(ctxs, completions):
# questions, answers = postprocess_qa_pairs(completion.outputs[0].text)
# for q, a in zip(questions, answers):
# samples.append(
# {
# "context": ctx,
# "prompt": q,
# "response": a,
# }
# )
# for sample in samples:
# print(f"{sample['context']=}")
# print(f"{sample['prompt']=}")
# print(f"{sample['response']=}")
# print()

64
data/generate_negative_natural_questions.py
View file

@ -1,64 +0,0 @@
import os
from datasets import load_dataset
def process_sample(sample):
"""Process a single sample and return the processed data or None if invalid."""
ctx = " ".join(
[
token
for is_html, token in zip(
sample["document"]["tokens"]["is_html"],
sample["document"]["tokens"]["token"],
)
if not is_html
]
).strip()
if not ctx:
return {"prompt": None, "context": None, "answer": None}
answers = []
for answer in sample["annotations"]["short_answers"]:
answers += answer["text"]
if not answers:
return {"prompt": None, "context": None, "answer": None}
answer = answers[0].strip()
if not answer:
return {"prompt": None, "context": None, "answer": None}
prompt = sample["question"]["text"].capitalize().strip() + " ?"
return {"prompt": prompt, "context": ctx, "answer": answer.capitalize()}
def shift_contexts(examples):
"""Shift contexts by one position to create negative examples."""
contexts = examples["context"]
shifted_contexts = contexts[1:] + contexts[:1]
examples["context"] = shifted_contexts
return examples
if __name__ == "__main__":
ds = load_dataset("google-research-datasets/natural_questions", split="validation")
ds = ds.shuffle(seed=42)
# for split in ["validation"]: # , "train"
# Process samples in parallel using map
processed_ds = ds.map(process_sample, remove_columns=ds.column_names, num_proc=16)
# Filter out None values
processed_ds = processed_ds.filter(lambda x: x["context"] is not None)
# Shift contexts to create negative examples
processed_ds = processed_ds.map(
shift_contexts, batched=True, batch_size=len(processed_ds)
)
print(processed_ds)
save_dir = "data/raw_datasets/negative_natural_questions"
os.makedirs(save_dir, exist_ok=True)
processed_ds.to_json(f"{save_dir}/validation.jsonl")

37
data/generate_pretrain_from_fw_qa.py
View file

@ -1,37 +0,0 @@
import gc
from collections import defaultdict
from glob import glob
from datasets import Dataset, load_dataset
from tqdm import tqdm
QA_TEMPLATE = "\n\nQuestion: {question}\nAnswer: {answer}"
if __name__ == "__main__":
root_data_dir = "./data/raw_datasets/fw_qa_3"
files = sorted(glob(f"{root_data_dir}/*.parquet"))
for file in files:
ctx_qa_dict = defaultdict(str)
ds = load_dataset("parquet", data_files=file, split="train")
print(f"Loading dataset from {file}")
print(f"Original size: {len(ds)}")
for i, sample in tqdm(enumerate(ds)):
ctx = sample["context"]
question = sample["prompt"]
answer = sample["response"]
ctx_qa_dict[ctx] += QA_TEMPLATE.format(question=question, answer=answer)
print(f"Unique contexts: {len(ctx_qa_dict)}")
sampled_data = ctx_qa_dict[ctx]
print(f"Sampled context-qa pairs: {ctx}{''.join(sampled_data)}")
# convert ctx_qa_dict to a list of dictionaries
samples = [
{"context": ctx, "qas": qa_pairs} for ctx, qa_pairs in ctx_qa_dict.items()
]
# save to a new dataset
ds = Dataset.from_list(samples)
save_path = f"./data/raw_datasets/fw_qa_intx_pretrain/{file.split('/')[-1]}"
print(f"Saving dataset to {save_path}")
ds.to_parquet(save_path)
print("=" * 80)
del ds, samples, ctx_qa_dict
gc.collect()

143
data/prepare_distillation_dataset.py
View file

@ -1,143 +0,0 @@
import argparse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from ctx_to_lora.data.processing import load_and_process_dataset
def add_logits(row, model, tok, device, args):
ctx = row["context"]
qs = row["prompts"]
ans = row["responses"]
responses_masks = []
all_logits_values = []
all_logits_positions = []
for question, answer in zip(qs, ans):
messages = [
{
"role": "user",
"content": f"Answer the question based on the following information:\n{ctx}.\n\n{question}",
},
{"role": "assistant", "content": answer},
]
all_tokens = tok.apply_chat_template(
messages,
tokenize=True,
add_special_tokens=False,
truncation=False,
padding=False,
add_generation_prompt=False,
return_dict=True,
)
input = torch.tensor([all_tokens["input_ids"]], device=device)
with torch.no_grad():
logits = model(input).logits
top_logits, top_logits_positions = torch.topk(
logits, args.max_logits_to_store, dim=-1
)
# Convert logits to specified precision
dtype_map = {
"float32": torch.float32,
"bfloat16": torch.bfloat16,
"float16": torch.float16,
}
target_dtype = dtype_map[args.logit_precision]
top_logits = top_logits.to(target_dtype)
tokens_to_mask = tok.apply_chat_template(
[messages[0]], # Just the user message
tokenize=True,
add_special_tokens=False,
truncation=False,
padding=False,
add_generation_prompt=True,
return_dict=True,
)
number_of_tokens_to_mask = len(tokens_to_mask["input_ids"])
# Create mask: False for prompt tokens, True for response tokens
responses_mask = torch.zeros(len(all_tokens["input_ids"]), dtype=torch.bool)
responses_mask[number_of_tokens_to_mask:] = True
responses_masks.append(responses_mask)
all_logits_values.append(top_logits.cpu())
all_logits_positions.append(top_logits_positions.cpu())
row["labels_logits_values"] = all_logits_values
row["labels_logits_positions"] = all_logits_positions
row["responses_mask"] = responses_masks
return row
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--ds_name", required=True, help="name of the dataset")
ap.add_argument("--split", required=True, help="split of the dataset")
ap.add_argument(
"--max_logits_to_store",
type=int,
default=100,
help="max number of logits kept per completion",
)
ap.add_argument(
"--limit_samples",
type=int,
default=None,
help="process only the first n dataset rows",
)
ap.add_argument(
"--logit_precision",
default="bfloat16",
choices=["float32", "bfloat16", "float16"],
help="dtype for saved logits",
)
ap.add_argument(
"--model_name", required=True, help="model name to use for generating logits"
)
ap.add_argument(
"--output_path", required=True, help="path to save the processed dataset"
)
args = ap.parse_args()
# Load model and tokenizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AutoModelForCausalLM.from_pretrained(args.model_name)
model = model.to(device)
model.eval()
tok = AutoTokenizer.from_pretrained(args.model_name)
# Load dataset
ds = load_and_process_dataset(
args.ds_name,
args.split,
add_negative_prompt=False,
add_repeat_prompt=False,
repeat_prob=0,
is_pretrain=False,
streaming=False,
num_proc=8,
)
# Limit samples if specified
if args.limit_samples:
ds = ds.select(range(min(args.limit_samples, len(ds))))
# Process dataset with logits
ds = ds.map(
lambda row: add_logits(row, model, tok, device, args),
batched=False,
desc="Adding logits to dataset",
)
# Save the processed dataset
ds.save_to_disk(args.output_path)
print(f"Dataset saved to {args.output_path}")
if __name__ == "__main__":
main()

114
data/q_generation_prompts.py
View file

@ -1,114 +0,0 @@
import random
SYSTEM_PROMPT = (
"You are a creative and helpful assistant. "
"You will be given a context and you need to generate questions from the given context. "
"**DO NOT** hallucinate or make up information."
)
PROMPT_TEMPLATE = (
"### Context ###\n{context}\n\n"
"### Instruction ###\n{instruction}\n\n"
"### Rules ###\n"
"Phrases like 'based on the provided context', 'according to the context', etc., must not to appear in your response.\n\n"
# "2. The questions should not overlap. They should be diverse, covering many aspects of the context.\n"
# "3. Do not give away too much information in the questions. For example, ask 'Who is X?' instead of 'Who is X that did Y?' when Y is clear from the context.\n"
# "4. Ignore the text formatting of the context, e.g., bold, italic, underline, etc.\n"
# "5. Ignore typos, spacing, and grammatical errors in the context.\n"
# "6. Always use proper grammar and punctuation.\n"
# "7. Try to use different question forms and styles.\n"
"### Output Format ###\n"
"The question/instruction/task/request should be in the following format:\n\n"
"Message: {{question}}\n\n"
"Use this output template regardless of the type of the output."
)
# taken from https://github.com/HazyResearch/cartridges/blob/2ac563d79c2f3367a9e780a7bb0b3cf3039a8d50/cartridges/data/resources.py#L195
# def get_structuring_seed_prompt() -> str:
# DATA_FORMATS = [
# "JSON",
# "YAML",
# "TOML",
# "INI",
# "XML",
# "plain text",
# ]
# data_format = random.choice(DATA_FORMATS)
# EXAMPLES = [
# dedent(f"""
# Can you structure the information in the context in the following format: {data_format}? Be sure to include precise information like any dates, times, names, and numerical values.
# """).strip(),
# ]
# example = random.choice(EXAMPLES)
# return dedent(f"""
# Please generate a single chat message instructing an LLM to structure the information in {data_format}. The message can follow the following template, filling in details from the context:
# '{example}'
# """).strip()
# def get_summarization_seed_prompt() -> str:
# prompts = [
# dedent("""
# Please generate a single chat message instructing an LLM to summarize part of the context.
# Make sure the instruction is very explicit about the section of the context that you want to summarize.
# Include details (ids, names, titles, dates, etc.) that make it clear what you are asking about.
# """).strip(),
# dedent("""
# Please generate a single chat message instructing an LLM to summarize a section.
# Make sure the instruction is explicit about the section that should be summarized and the document it is from.
# """).strip(),
# ]
# prompt = random.choice(prompts)
# return prompt
def get_question_seed_prompt() -> str:
prompts = [
(
"Generate a question for an LLM that will test its knowledge of the information in the context above. "
"In your question be sure to include details (ids, names, titles, dates, etc.) that make it clear what you are asking about. "
"Output only a single question. Do NOT include any other text or explanation other than the question."
),
(
"Generate a message for an LLM that will test its knowledge of the information in the context above. "
"Be sure to include details (ids, names, titles, dates, etc.) in the question so that it can be answered without access to the context (i.e. closed-book setting). "
"Output only a single question. Do NOT include any other text or explanation other than the question."
),
(
"You are helping to quiz a user about the information in the context. "
"Please generate a question about the subsection of the context above. "
"Be sure to include details (ids, names, titles, dates, etc.) in the question to make it clear what you are asking about. "
"Answer only with the question, do not include any other text."
),
]
prompt = random.choice(prompts)
return prompt
def get_use_case_seed_prompt() -> str:
prompt = (
"Your primary goal is to think about practical, real-world tasks or applications that someone could achieve using the knowledge contained within the provided context. "
"Consider how a user might want to apply this information in a real-world scenario, not just recall it. Put yourself into someone else's shoes and think how you might use the information. "
"After considering potential use cases, your task will be to generate an instruction or task that reflects one of these downstream applications. "
"This instruction or task should be something a user, who has access to this context, might ask when trying to accomplish their specific goal. "
"Output only a single instruction or task. Do NOT include any other text or explanation."
)
return prompt
def get_creative_seed_prompt() -> str:
prompt = (
"You are having a creative open-ended conversation inspired by the information in the context. "
"Please generate an open question for your conversation partner to start off the discussion. "
"Answer only with the question, do not include any other text."
)
return prompt
def get_generic_seed_prompt() -> str:
return "Please generate a single chat message to begin a conversation about the information in the context. Make an open-ended request or provide a task."

124
data/self_generate_qa.py
View file

@ -11,11 +11,8 @@ from vllm import LLM, SamplingParams
from ctx_to_lora.data.definitions import (
CLOSED_QA_INTX_TEMPLATES,
COT_TEMPLATES,
RAW_DATA_DIR,
SELF_GEN_DATA_DIR,
STRUCTURING_PROMPTS,
SUMMARIZATION_PROMPTS,
)
from ctx_to_lora.data.processing import (
filter_none,
@ -67,29 +64,15 @@ def truncate_middle_if_too_long(
return input_ids
# TODO: self-gen for magic ctx num
def get_prompt(context: str, q: str, remove_qa_template: bool) -> str:
prompt = QA_PROMPT_TEMPLATE if not remove_qa_template else PROMPT_TEMPLATE
return prompt.format(context=context, question=q)
def augment_prompt(sample, closed_qa_prob, cot_prob):
z = random.random()
should_add_closed_qa = ("type" not in sample) or (sample["type"] == "question")
if should_add_closed_qa and (z < closed_qa_prob):
sample["prompts"] = add_closed_qa_prompt(sample["prompts"])
elif closed_qa_prob <= z < (closed_qa_prob + cot_prob):
sample["prompts"] = add_cot_prompt(sample["prompts"])
return sample
def add_cot_prompt(prompts: list[str]) -> str:
return [random.choice(COT_TEMPLATES).format(input=q) for q in prompts]
def add_closed_qa_prompt(prompts: list[str]) -> str:
return [random.choice(CLOSED_QA_INTX_TEMPLATES).format(input=q) for q in prompts]
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:
@ -99,49 +82,6 @@ def load_config(config_path: str) -> dict:
return config
def add_paraphrasing_prompt(samples, summ_prob: float, struct_prob: float):
"""Augment batch with summarization / structuring prompts given separate probabilities.
The original samples are extended (in-place) with newly created paraphrased variants.
"""
if summ_prob == 0 and struct_prob == 0:
return samples
n_samples = len(samples["ctx_ids"]) if "ctx_ids" in samples else 0
if n_samples == 0:
return samples
n_summ = int(n_samples * summ_prob)
n_struct = int(n_samples * struct_prob)
all_indices = list(range(n_samples))
random.shuffle(all_indices)
summ_indices = all_indices[:n_summ]
struct_indices = all_indices[n_summ : n_summ + n_struct]
def _add(kind_indices, prompts_source, prompt_type):
out = dict()
if not kind_indices:
return out
for k in samples:
if k == "prompts":
sampled_prompts = random.choices(prompts_source, k=len(kind_indices))
out[k] = [[p] for p in sampled_prompts]
elif k == "type":
out[k] = [prompt_type] * len(kind_indices)
else:
out[k] = [samples[k][i] for i in kind_indices]
return out
summ_samples = _add(summ_indices, SUMMARIZATION_PROMPTS, "summarization")
struct_samples = _add(struct_indices, STRUCTURING_PROMPTS, "structuring")
# Extend original samples
for k in samples:
if summ_indices:
samples[k] += summ_samples[k]
if struct_indices:
samples[k] += struct_samples[k]
return samples
def get_dataset_configs(
ds_names: list[str] | None,
config: dict | None,
@ -334,46 +274,18 @@ def self_generate(
keep_in_memory=True,
)
ds = ds.map(
augment_prompt,
fn_kwargs={"closed_qa_prob": closed_qa_prob, "cot_prob": cot_prob},
keep_in_memory=True,
num_proc=16,
)
ds = ds.map(
add_paraphrasing_prompt,
batched=True,
fn_kwargs={
"summ_prob": summ_prob,
"struct_prob": struct_prob,
},
batch_size=50_000,
keep_in_memory=True,
)
ctxs = [sample["context"] for sample in ds]
questions = [
[add_closed_qa_prompt(q, closed_qa_prob) for q in sample["prompts"] if q]
for sample in ds
]
questions = [q_list for q_list in ds["prompts"] if len(q_list) > 0]
# for sample in ds:
# if ("type" not in sample) or (
# "type" in sample and sample["type"] == "question"
# ):
# # for generated qa, apply to only "question" type
# questions.append(
# [
# add_closed_qa_prompt(q, closed_qa_prob)
# for q in sample["prompts"]
# if q
# ]
# )
# else:
# questions.append([q for q in sample["prompts"] if q])
print(f"Loaded {len(ctxs)} contexts and {len(questions)} questions")
k = 16
fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}_cot_prob_{cot_prob}_summ_prob_{summ_prob}_struct_prob_{struct_prob}/{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}"
chunk_size = 1_000
for chunk_idx, start in enumerate(range(0, len(ctxs), chunk_size)):
@ -670,24 +582,6 @@ def parse_args() -> argparse.Namespace:
default=0.0,
help="Probability of using closed QA prompt template (default: 0.0)",
)
parser.add_argument(
"--cot_prob",
type=float,
default=0.0,
help="Probability of using chain-of-thought prompt template (default: 0.0)",
)
parser.add_argument(
"--summ_prob",
type=float,
default=0.0,
help="Probability of adding a summarization variant (default: 0.0)",
)
parser.add_argument(
"--struct_prob",
type=float,
default=0.0,
help="Probability of adding a structuring variant (default: 0.0)",
)
parser.add_argument(
"--do_truncate",
action="store_true",

126
data/store_top_logits.py
View file

@ -1,126 +0,0 @@
import argparse
import io
import zipfile
from pathlib import Path
from tqdm import tqdm
import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, PreTrainedTokenizer
# e.g., python data/store_top_logits.py --data "temp/test_data.parquet" --model "Qwen/Qwen3-0.6B"
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--data', required=True,
help='path to a parquet shard of the dataset')
ap.add_argument('--model', required=True,
help='huggingface model name or checkpoint')
ap.add_argument('--max_logits_to_store', type=int, default=100,
help='max number of logits kept per completion')
ap.add_argument('--limit_samples', type=int, default=None,
help='process only the first n dataset rows')
ap.add_argument('--logit_precision', default='bfloat16',
choices=['float32', 'bfloat16', 'float16'],
help='dtype for saved logits')
ap.add_argument('--index_precision', default='int32',
choices=['int64', 'int32', 'int16'],
help='dtype for saved indices')
args = ap.parse_args()
float_map = {
'float32': torch.float32,
'bfloat16': torch.bfloat16,
'float16': torch.float16,
}
int_map = {
'int64': torch.int64,
'int32': torch.int32,
'int16': torch.int16,
}
l_dtype = float_map[args.logit_precision]
i_dtype = int_map[args.index_precision]
ds = load_dataset('parquet', data_files=args.data)['train']
if args.limit_samples and args.limit_samples > 0:
ds = ds.select(range(min(args.limit_samples, len(ds))))
tok: PreTrainedTokenizer = AutoTokenizer.from_pretrained(
args.model, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(args.model)
model.eval()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
base = Path(args.data).with_suffix('')
model_dir = args.model.replace('/', '__')
out_dir = base / model_dir
out_dir.mkdir(parents=True, exist_ok=True)
if args.limit_samples:
zip_name = f'limited_logits_{args.limit_samples}.zip'
else:
zip_name = 'all_logits.zip'
zip_path = out_dir / zip_name
samp_id = 0
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
for row in tqdm(ds):
ctx = row['context']
qs = row['prompts_level_0']
ans = row['responses_level_0']
q_list = []
a_list = []
vals_list = []
sel_pos_list = []
for q, a in zip(qs, ans):
messages = [
{'role': 'user', 'content': ctx + '\n' + q},
{'role': 'assistant', 'content': a},
]
tpl = tok.apply_chat_template(
messages,
tokenize=True,
add_special_token=False,
truncation=False,
add_generation_prompt=False,
return_dict=True,
)
tokens_to_mask = tok.apply_chat_template(
messages[:-1],
tokenize=True,
add_special_token=False,
truncation=False,
add_generation_prompt=True,
return_dict=True,
)
number_of_tokens_to_mask = len(tokens_to_mask['input_ids'])
assert tokens_to_mask['input_ids'] == tpl['input_ids'][:number_of_tokens_to_mask]
inp = torch.tensor([tpl['input_ids']], device=device)
with torch.no_grad():
logits = model(inp).logits
logits_for_answer = logits[0, number_of_tokens_to_mask - 1:, :]
assert logits_for_answer.shape[0] > 0
assert logits_for_answer.shape[0] + number_of_tokens_to_mask -1 == inp.size(1)
# just in case, logits to store is super large
k = min(args.max_logits_to_store, logits_for_answer.shape[-1])
vals, idxs = torch.topk(logits_for_answer, k, dim=-1)
q_list.append(q)
a_list.append(a)
vals_list.append(vals.cpu())
sel_pos_list.append(idxs.cpu())
record = {
'context': ctx,
'question': q_list,
'answer': a_list,
'top_logits': vals_list,
'top_positions_positions': sel_pos_list,
}
buf = io.BytesIO()
torch.save(record, buf)
samp_id += 1
zf.writestr(f'{samp_id:07d}.pt', buf.getvalue())
if __name__ == '__main__':
main()

15
eval_scripts/chunked_eval.sh
View file

@ -1,15 +0,0 @@
WANDB_MODE=disabled run uv run python run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad ropes drop longbench/squad_e longbench/gov_report_e test_longbench/multifieldqa_en_e --split test --max_ctx_chunk_len -1 --lora_aggregation mean
WANDB_MODE=disabled run uv run python run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad ropes drop longbench/squad_e longbench/gov_report_e test_longbench/multifieldqa_en_e --split test --max_ctx_chunk_len -1 --lora_aggregation sum
WANDB_MODE=disabled run uv run python run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad ropes drop longbench/squad_e longbench/gov_report_e test_longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 128 --lora_aggregation mean
WANDB_MODE=disabled run uv run python run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad ropes drop longbench/squad_e longbench/gov_report_e test_longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 128 --lora_aggregation sum
WANDB_MODE=disabled run uv run python run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad ropes drop longbench/squad_e longbench/gov_report_e test_longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 256 --lora_aggregation mean
WANDB_MODE=disabled run uv run python run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad ropes drop longbench/squad_e longbench/gov_report_e test_longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 256 --lora_aggregation sum
WANDB_MODE=disabled run uv run python run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad ropes drop longbench/squad_e longbench/gov_report_e test_longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 512 --lora_aggregation mean
WANDB_MODE=disabled run uv run python run_eval.py --checkpoint_path train_outputs/runs/Aug02_07-51-08_slurm0-a3nodeset-9_76501_7fdab5ea/checkpoint-50000/pytorch_model.bin --datasets squad ropes drop longbench/squad_e longbench/gov_report_e test_longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 512 --lora_aggregation sum

341
gcp_bucket_watcher.py
View file

@ -1,341 +0,0 @@
"""Minimal GCS prefix sync (one-way): download (default) or upload.
Remote -> Local (download):
python gcp_bucket_watcher.py --bucket ctx-to-lora --prefix train_outputs --dest train_outputs
Local -> Remote (upload):
python gcp_bucket_watcher.py --bucket ctx-to-lora --prefix train_outputs --dest train_outputs --upload
Mirrors only new / changed files (size or mtime difference). No deletions.
Filtering: --include / --exclude accept comma separated globs. Auth uses
Application Default Credentials or --credentials JSON.
Example
# download
python gcp_bucket_watcher.py --bucket ctx-to-lora --prefix train_outputs --dest train_outputs
# upload
python gcp_bucket_watcher.py --bucket ctx-to-lora --prefix train_outputs --dest train_outputs --upload
"""
from __future__ import annotations
import argparse
import fnmatch
import json
import os
import sys
import time
from collections.abc import Sequence
from dataclasses import dataclass
from datetime import datetime
try:
from google.cloud import storage # type: ignore
from google.oauth2 import service_account # type: ignore
except Exception as e: # pragma: no cover - import guard
print("google-cloud-storage not installed. Add to dependencies.", file=sys.stderr)
raise
@dataclass
class ObjectMeta: # Unified meta for local + remote
name: str # relative path (no leading prefix for remote)
updated: float # mtime (remote object updated timestamp or local mtime)
size: int # bytes
def parse_time_string(time_str: str) -> float:
"""Parse human-readable time string to timestamp.
Supports formats:
- ISO format: 2024-01-15T10:30:00 or 2024-01-15 10:30:00
- Date only: 2024-01-15 (assumes 00:00:00)
- Unix timestamp: 1705315800
"""
# Try parsing as unix timestamp first
try:
return float(time_str)
except ValueError:
pass
# Try various datetime formats
formats = [
"%Y-%m-%dT%H:%M:%S", # 2024-01-15T10:30:00
"%Y-%m-%d %H:%M:%S", # 2024-01-15 10:30:00
"%Y-%m-%dT%H:%M", # 2024-01-15T10:30
"%Y-%m-%d %H:%M", # 2024-01-15 10:30
"%Y-%m-%d", # 2024-01-15
]
for fmt in formats:
try:
dt = datetime.strptime(time_str, fmt)
return dt.timestamp()
except ValueError:
continue
raise ValueError(f"Unable to parse time string: {time_str}")
def parse_globs(spec: str | None) -> list[str]:
if not spec:
return []
return [g.strip() for g in spec.split(",") if g.strip()]
def should_skip(
rel_path: str,
includes: Sequence[str],
excludes: Sequence[str],
meta: ObjectMeta | None = None,
from_time: float | None = None,
) -> bool:
base = os.path.basename(rel_path)
if excludes and any(fnmatch.fnmatch(base, p) for p in excludes):
return True
if includes and not any(fnmatch.fnmatch(base, p) for p in includes):
return True
if from_time is not None and meta is not None and meta.updated < from_time:
return True
return False
def list_remote(
client: storage.Client, bucket: str, prefix: str
) -> dict[str, ObjectMeta]:
out: dict[str, ObjectMeta] = {}
p = prefix.rstrip("/")
for b in client.list_blobs(bucket_or_name=bucket, prefix=prefix):
if b.name.endswith("/"):
continue
rel = b.name[len(p) :].lstrip("/") if p else b.name
out[rel] = ObjectMeta(
rel, (b.updated.timestamp() if b.updated else 0.0), b.size or 0
)
return out
def list_local(root: str) -> dict[str, ObjectMeta]:
out: dict[str, ObjectMeta] = {}
for base, _dirs, files in os.walk(root):
for f in files:
path = os.path.join(base, f)
try:
st = os.stat(path)
except OSError:
continue
rel = os.path.relpath(path, root).replace("\\", "/")
out[rel] = ObjectMeta(rel, st.st_mtime, st.st_size)
return out
def needs_transfer(src: ObjectMeta, dst: ObjectMeta | None, tolerance: float) -> bool:
if dst is None:
return True
if src.size != dst.size:
return True
if abs(src.updated - dst.updated) > tolerance:
return True
return False
# --- Minimal persistent state to honor local deletions ---
def _state_path(dest: str) -> str:
return os.path.join(dest, ".gcs_sync_state.json")
def _load_state(dest: str) -> tuple[set[str], set[str]]:
path = _state_path(dest)
try:
with open(path, encoding="utf-8") as f:
data = json.load(f)
return set(data.get("downloaded", [])), set(data.get("tombstones", []))
except FileNotFoundError:
return set(), set()
except Exception:
# Corrupt/invalid state; ignore to keep going.
return set(), set()
def _save_state(dest: str, downloaded: set[str], tombstones: set[str]) -> None:
path = _state_path(dest)
tmp = path + ".part"
data = {
"downloaded": sorted(downloaded),
"tombstones": sorted(tombstones),
}
with open(tmp, "w", encoding="utf-8") as f:
json.dump(data, f)
os.replace(tmp, path)
def download_file(
client: storage.Client, bucket: str, prefix: str, meta: ObjectMeta, dest: str
) -> None:
blob_name = f"{prefix.rstrip('/')}/{meta.name}" if prefix else meta.name
blob = client.bucket(bucket).blob(blob_name)
local_path = os.path.join(dest, meta.name)
os.makedirs(os.path.dirname(local_path), exist_ok=True)
tmp = local_path + ".part"
blob.download_to_filename(tmp)
os.utime(tmp, (time.time(), meta.updated))
os.replace(tmp, local_path)
def upload_file(
client: storage.Client, bucket: str, prefix: str, meta: ObjectMeta, src_root: str
) -> None:
blob_name = f"{prefix.rstrip('/')}/{meta.name}" if prefix else meta.name
blob = client.bucket(bucket).blob(blob_name)
blob.upload_from_filename(os.path.join(src_root, meta.name))
def sync_cycle(args: argparse.Namespace, client: storage.Client) -> int:
includes, excludes = parse_globs(args.include), parse_globs(args.exclude)
remote = list_remote(client, args.bucket, args.prefix)
local = list_local(args.dest)
changed = 0
if args.upload: # local -> remote (original tolerance 2s)
for rel, meta in local.items():
if should_skip(rel, includes, excludes, meta, args.from_time):
continue
remote_meta = remote.get(rel)
# For uploads we consider a file up-to-date if sizes match and the remote
# object timestamp is >= local mtime (within tolerance). Using absolute
# delta caused perpetual re-uploads because GCS sets its own server time.
if (
remote_meta
and meta.size == remote_meta.size
and remote_meta.updated >= meta.updated - 2.0
):
continue
if needs_transfer(meta, remote_meta, tolerance=2.0):
print(f"[uploading ] {rel} ({meta.size} bytes)")
try:
upload_file(client, args.bucket, args.prefix, meta, args.dest)
except FileNotFoundError:
# File disappeared between scan and upload attempt; skip gracefully.
print(f"[skipped missing] {rel}")
continue
print(f"[uploaded ✓] {rel} ({meta.size} bytes)")
changed += 1
if changed == 0:
print("No new or updated local files to upload.")
else:
print("All uploads complete.")
else: # download (original tolerance 1s)
downloaded_set, tombstones = _load_state(args.dest)
state_dirty = False
for rel, meta in remote.items():
if should_skip(rel, includes, excludes, meta, args.from_time):
continue
# Respect user-deleted files (tombstones)
if rel in tombstones:
continue
# If we had downloaded this before but it's now missing locally,
# treat as user-deleted and tombstone it to avoid re-download.
if rel in downloaded_set and rel not in local:
tombstones.add(rel)
state_dirty = True
continue
local_meta = local.get(rel)
# Skip if local file is newer than remote (avoid overwriting newer local files)
if local_meta and local_meta.updated > meta.updated + 1.0: # 1s tolerance
if rel not in downloaded_set:
downloaded_set.add(rel)
state_dirty = True
continue
if needs_transfer(meta, local_meta, tolerance=1.0):
print(f"[downloading] {rel} ({meta.size} bytes)")
download_file(client, args.bucket, args.prefix, meta, args.dest)
print(f"[downloaded ✓] {rel} ({meta.size} bytes)")
changed += 1
if rel not in downloaded_set:
downloaded_set.add(rel)
state_dirty = True
if rel in tombstones:
tombstones.discard(rel)
state_dirty = True
else:
# File exists locally and is in sync; ensure it's marked as downloaded
# so future local deletions are respected without re-download.
if local_meta is not None and rel not in downloaded_set:
downloaded_set.add(rel)
state_dirty = True
if changed == 0:
print("No new or updated objects.")
else:
print("All downloads complete.")
if state_dirty:
_save_state(args.dest, downloaded_set, tombstones)
return changed
def build_client(args: argparse.Namespace) -> storage.Client:
if args.credentials:
creds = service_account.Credentials.from_service_account_file(args.credentials)
return storage.Client(credentials=creds, project=creds.project_id)
return storage.Client() # ADC
def parse_args(argv: Sequence[str]) -> argparse.Namespace:
p = argparse.ArgumentParser(description="Watch and mirror a GCS prefix locally.")
p.add_argument("--bucket", required=True, help="GCS bucket name (without gs://)")
p.add_argument(
"--prefix",
default="",
help="Prefix inside the bucket to mirror (e.g. train_outputs)",
)
p.add_argument(
"--dest", default="train_outputs", help="Local destination directory"
)
p.add_argument(
"--interval", type=int, default=60, help="Seconds between scans (default 60)"
)
p.add_argument("--once", action="store_true", help="Run a single scan and exit")
p.add_argument(
"--upload",
action="store_true",
help="Reverse direction: upload local files to bucket (default is download).",
)
p.add_argument("--include", help="Comma separated glob(s) to include (default all)")
p.add_argument("--exclude", help="Comma separated glob(s) to exclude")
p.add_argument(
"--from-time",
help="Only sync files modified after this time (e.g. '2024-01-15', '2024-01-15T10:30:00', or unix timestamp)",
)
p.add_argument("--credentials", help="Path to service account JSON (optional)")
return p.parse_args(argv)
def main(argv: Sequence[str] | None = None) -> int:
args = parse_args(argv or sys.argv[1:])
# Parse from_time if provided
if args.from_time:
try:
args.from_time = parse_time_string(args.from_time)
except ValueError as e:
print(f"Error parsing --from-time: {e}", file=sys.stderr)
return 1
client = build_client(args)
os.makedirs(args.dest, exist_ok=True)
try:
while True:
print(
time.strftime("%Y-%m-%d %H:%M:%S"),
"Scanning (upload)" if args.upload else "Scanning (download)",
)
sync_cycle(args, client)
if args.once:
break
time.sleep(args.interval)
except KeyboardInterrupt:
print("Interrupted.")
return 0
if __name__ == "__main__": # pragma: no cover
raise SystemExit(main())

0
icae_v2/__init__.py
View file

6
icae_v2/dev_v2.jsonl
View file

@ -1,6 +0,0 @@
{"context": "French senior civil servant arrested on suspicion of spying for North Korea\n\nNovember 27, 2018 by Joseph Fitsanakis\n\nA senior civil servant in the upper house of the French parliament has been arrested on suspicion of spying for North Korea, according to prosecutors. The news of the suspected spy\u2019s arrest was first reported on Monday by Quotidien, a daily politics and culture show on the Monaco-based television channel TMC. The show cited \u201ca judicial source in Paris\u201d and said that France\u2019s domestic security and counterintelligence agency, the General Directorate for Internal Security (DGSI), was in charge of the espionage case.\n\nThe senior administrator has been identified as Benoit Quennedey, a civil servant who liaises between the French Senate and the Department of Architecture and Heritage, which operates under France\u2019s Ministry of Culture. Quennedey was reportedly detained on Sunday morning and his office in the French Senate was raided by DGSI officers on the same day. Quotidien said that he was arrested on suspicion of \u201ccollecting and delivering to a foreign power information likely to subvert core national interests\u201d. The report did not provide specific information about the type of information that Quennedey is believed to have passed to North Korea. It did state, however, that a counterintelligence investigation into his activities began in March of this year.\n\nQuennedey is believed to be the president of the Franco-Korean Friendship Association, the French branch of a Spanish-based organization that lobbies in favor of international support for North Korea. Korea Friendship Association branches exist in over 30 countries and are believed to be officially sanctioned by Pyongyang. They operate as something akin to the pre-World War II Comintern (Communist International), a Moscow-sanctioned international pressure group that advocated in favor of Soviet-style communism around the world. French media reported on Monday that Quennedey traveled extensively to the Korean Peninsula in the past decade and has written a French-language book on North Korea. News reports said that the French President Emmanuel Macron had been made aware of Quennedey\u2019s arrest. The senior civil servant faces up to 30 years in prison if found guilty of espionage.\n\n\u25ba Author: Joseph Fitsanakis | Date: 27 November 2018 | Permalink\n\n", "prompt": "Identify the person arrested on suspicion of spying for North Korea.", "response": "Benoit Quennedey"}
{"context": "French senior civil servant arrested on suspicion of spying for North Korea\n\nNovember 27, 2018 by Joseph Fitsanakis\n\nA senior civil servant in the upper house of the French parliament has been arrested on suspicion of spying for North Korea, according to prosecutors. The news of the suspected spy\u2019s arrest was first reported on Monday by Quotidien, a daily politics and culture show on the Monaco-based television channel TMC. The show cited \u201ca judicial source in Paris\u201d and said that France\u2019s domestic security and counterintelligence agency, the General Directorate for Internal Security (DGSI), was in charge of the espionage case.\n\nThe senior administrator has been identified as Benoit Quennedey, a civil servant who liaises between the French Senate and the Department of Architecture and Heritage, which operates under France\u2019s Ministry of Culture. Quennedey was reportedly detained on Sunday morning and his office in the French Senate was raided by DGSI officers on the same day. Quotidien said that he was arrested on suspicion of \u201ccollecting and delivering to a foreign power information likely to subvert core national interests\u201d. The report did not provide specific information about the type of information that Quennedey is believed to have passed to North Korea. It did state, however, that a counterintelligence investigation into his activities began in March of this year.\n\nQuennedey is believed to be the president of the Franco-Korean Friendship Association, the French branch of a Spanish-based organization that lobbies in favor of international support for North Korea. Korea Friendship Association branches exist in over 30 countries and are believed to be officially sanctioned by Pyongyang. They operate as something akin to the pre-World War II Comintern (Communist International), a Moscow-sanctioned international pressure group that advocated in favor of Soviet-style communism around the world. French media reported on Monday that Quennedey traveled extensively to the Korean Peninsula in the past decade and has written a French-language book on North Korea. News reports said that the French President Emmanuel Macron had been made aware of Quennedey\u2019s arrest. The senior civil servant faces up to 30 years in prison if found guilty of espionage.\n\n\u25ba Author: Joseph Fitsanakis | Date: 27 November 2018 | Permalink\n\n", "prompt": "List the actions taken against Benoit Quennedey after his arrest.", "response": "His office in the French Senate was raided by DGSI officers."}
{"context": "French senior civil servant arrested on suspicion of spying for North Korea\n\nNovember 27, 2018 by Joseph Fitsanakis\n\nA senior civil servant in the upper house of the French parliament has been arrested on suspicion of spying for North Korea, according to prosecutors. The news of the suspected spy\u2019s arrest was first reported on Monday by Quotidien, a daily politics and culture show on the Monaco-based television channel TMC. The show cited \u201ca judicial source in Paris\u201d and said that France\u2019s domestic security and counterintelligence agency, the General Directorate for Internal Security (DGSI), was in charge of the espionage case.\n\nThe senior administrator has been identified as Benoit Quennedey, a civil servant who liaises between the French Senate and the Department of Architecture and Heritage, which operates under France\u2019s Ministry of Culture. Quennedey was reportedly detained on Sunday morning and his office in the French Senate was raided by DGSI officers on the same day. Quotidien said that he was arrested on suspicion of \u201ccollecting and delivering to a foreign power information likely to subvert core national interests\u201d. The report did not provide specific information about the type of information that Quennedey is believed to have passed to North Korea. It did state, however, that a counterintelligence investigation into his activities began in March of this year.\n\nQuennedey is believed to be the president of the Franco-Korean Friendship Association, the French branch of a Spanish-based organization that lobbies in favor of international support for North Korea. Korea Friendship Association branches exist in over 30 countries and are believed to be officially sanctioned by Pyongyang. They operate as something akin to the pre-World War II Comintern (Communist International), a Moscow-sanctioned international pressure group that advocated in favor of Soviet-style communism around the world. French media reported on Monday that Quennedey traveled extensively to the Korean Peninsula in the past decade and has written a French-language book on North Korea. News reports said that the French President Emmanuel Macron had been made aware of Quennedey\u2019s arrest. The senior civil servant faces up to 30 years in prison if found guilty of espionage.\n\n\u25ba Author: Joseph Fitsanakis | Date: 27 November 2018 | Permalink\n\n", "prompt": "What is the role of Benoit Quennedey in the French government?", "response": "He is a senior civil servant who liaises between the French Senate and the Department of Architecture and Heritage."}
{"context": "French senior civil servant arrested on suspicion of spying for North Korea\n\nNovember 27, 2018 by Joseph Fitsanakis\n\nA senior civil servant in the upper house of the French parliament has been arrested on suspicion of spying for North Korea, according to prosecutors. The news of the suspected spy\u2019s arrest was first reported on Monday by Quotidien, a daily politics and culture show on the Monaco-based television channel TMC. The show cited \u201ca judicial source in Paris\u201d and said that France\u2019s domestic security and counterintelligence agency, the General Directorate for Internal Security (DGSI), was in charge of the espionage case.\n\nThe senior administrator has been identified as Benoit Quennedey, a civil servant who liaises between the French Senate and the Department of Architecture and Heritage, which operates under France\u2019s Ministry of Culture. Quennedey was reportedly detained on Sunday morning and his office in the French Senate was raided by DGSI officers on the same day. Quotidien said that he was arrested on suspicion of \u201ccollecting and delivering to a foreign power information likely to subvert core national interests\u201d. The report did not provide specific information about the type of information that Quennedey is believed to have passed to North Korea. It did state, however, that a counterintelligence investigation into his activities began in March of this year.\n\nQuennedey is believed to be the president of the Franco-Korean Friendship Association, the French branch of a Spanish-based organization that lobbies in favor of international support for North Korea. Korea Friendship Association branches exist in over 30 countries and are believed to be officially sanctioned by Pyongyang. They operate as something akin to the pre-World War II Comintern (Communist International), a Moscow-sanctioned international pressure group that advocated in favor of Soviet-style communism around the world. French media reported on Monday that Quennedey traveled extensively to the Korean Peninsula in the past decade and has written a French-language book on North Korea. News reports said that the French President Emmanuel Macron had been made aware of Quennedey\u2019s arrest. The senior civil servant faces up to 30 years in prison if found guilty of espionage.\n\n\u25ba Author: Joseph Fitsanakis | Date: 27 November 2018 | Permalink\n\n", "prompt": "What are the charges against Benoit Quennedey?", "response": "He is suspected of collecting and delivering to a foreign power information likely to subvert core national interests."}
{"context": "French senior civil servant arrested on suspicion of spying for North Korea\n\nNovember 27, 2018 by Joseph Fitsanakis\n\nA senior civil servant in the upper house of the French parliament has been arrested on suspicion of spying for North Korea, according to prosecutors. The news of the suspected spy\u2019s arrest was first reported on Monday by Quotidien, a daily politics and culture show on the Monaco-based television channel TMC. The show cited \u201ca judicial source in Paris\u201d and said that France\u2019s domestic security and counterintelligence agency, the General Directorate for Internal Security (DGSI), was in charge of the espionage case.\n\nThe senior administrator has been identified as Benoit Quennedey, a civil servant who liaises between the French Senate and the Department of Architecture and Heritage, which operates under France\u2019s Ministry of Culture. Quennedey was reportedly detained on Sunday morning and his office in the French Senate was raided by DGSI officers on the same day. Quotidien said that he was arrested on suspicion of \u201ccollecting and delivering to a foreign power information likely to subvert core national interests\u201d. The report did not provide specific information about the type of information that Quennedey is believed to have passed to North Korea. It did state, however, that a counterintelligence investigation into his activities began in March of this year.\n\nQuennedey is believed to be the president of the Franco-Korean Friendship Association, the French branch of a Spanish-based organization that lobbies in favor of international support for North Korea. Korea Friendship Association branches exist in over 30 countries and are believed to be officially sanctioned by Pyongyang. They operate as something akin to the pre-World War II Comintern (Communist International), a Moscow-sanctioned international pressure group that advocated in favor of Soviet-style communism around the world. French media reported on Monday that Quennedey traveled extensively to the Korean Peninsula in the past decade and has written a French-language book on North Korea. News reports said that the French President Emmanuel Macron had been made aware of Quennedey\u2019s arrest. The senior civil servant faces up to 30 years in prison if found guilty of espionage.\n\n\u25ba Author: Joseph Fitsanakis | Date: 27 November 2018 | Permalink\n\n", "prompt": "Mention the organization Benoit Quennedey is believed to be the president of.", "response": "Franco-Korean Friendship Association"}
{"context": "French senior civil servant arrested on suspicion of spying for North Korea\n\nNovember 27, 2018 by Joseph Fitsanakis\n\nA senior civil servant in the upper house of the French parliament has been arrested on suspicion of spying for North Korea, according to prosecutors. The news of the suspected spy\u2019s arrest was first reported on Monday by Quotidien, a daily politics and culture show on the Monaco-based television channel TMC. The show cited \u201ca judicial source in Paris\u201d and said that France\u2019s domestic security and counterintelligence agency, the General Directorate for Internal Security (DGSI), was in charge of the espionage case.\n\nThe senior administrator has been identified as Benoit Quennedey, a civil servant who liaises between the French Senate and the Department of Architecture and Heritage, which operates under France\u2019s Ministry of Culture. Quennedey was reportedly detained on Sunday morning and his office in the French Senate was raided by DGSI officers on the same day. Quotidien said that he was arrested on suspicion of \u201ccollecting and delivering to a foreign power information likely to subvert core national interests\u201d. The report did not provide specific information about the type of information that Quennedey is believed to have passed to North Korea. It did state, however, that a counterintelligence investigation into his activities began in March of this year.\n\nQuennedey is believed to be the president of the Franco-Korean Friendship Association, the French branch of a Spanish-based organization that lobbies in favor of international support for North Korea. Korea Friendship Association branches exist in over 30 countries and are believed to be officially sanctioned by Pyongyang. They operate as something akin to the pre-World War II Comintern (Communist International), a Moscow-sanctioned international pressure group that advocated in favor of Soviet-style communism around the world. French media reported on Monday that Quennedey traveled extensively to the Korean Peninsula in the past decade and has written a French-language book on North Korea. News reports said that the French President Emmanuel Macron had been made aware of Quennedey\u2019s arrest. The senior civil servant faces up to 30 years in prison if found guilty of espionage.\n\n\u25ba Author: Joseph Fitsanakis | Date: 27 November 2018 | Permalink\n\n", "prompt": "When did the counterintelligence investigation into Quennedey's activities begin?", "response": "In March of this year."}

137
icae_v2/fine_tuned_inference.py
View file

@ -1,137 +0,0 @@
import json
import torch
from modeling_icae_multi_span import (
ICAE,
DataArguments,
ModelArguments,
TrainingArguments,
)
from peft import LoraConfig
from safetensors.torch import load_file
from tqdm import tqdm
from transformers import HfArgumentParser
# Set the computation device
device = "cuda"
# Parse model, data, and training arguments
parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
# Define Lora configuration
lora_config = LoraConfig(
r=512,
lora_alpha=32,
lora_dropout=model_args.lora_dropout,
bias="none",
task_type="CAUSAL_LM",
)
# Initialize model and send it to CUDA device
model = ICAE(model_args, training_args, lora_config)
# Load the fine-tuned checkpoint
print(f"Loading trained checkpoint from {training_args.output_dir}")
state_dict = load_file(training_args.output_dir)
model.load_state_dict(
state_dict, strict=False
) # only load lora and memory token embeddings
model = model.to(device)
# Read the data file
file_path = "./dev_v2.jsonl"
lines = None
with open(file_path) as f:
lines = f.readlines()
# Prepare the model for evaluation
max_out_length = 512
model.eval()
with torch.no_grad():
with open("ft_inference.out", "w") as f:
for line in tqdm(lines):
# Tokenize input text
data = json.loads(line)
tokenized_input = model.tokenizer(
data["input"],
truncation=True,
max_length=5120,
padding=False,
return_attention_mask=False,
)
tokenized_prompt = model.tokenizer(
data["prompt"],
truncation=False,
padding=False,
return_attention_mask=False,
add_special_tokens=False,
)
# Generate compressed outputs
input_ids = torch.LongTensor([tokenized_input["input_ids"]]).to(device)
memory_slots = model._compress(input_ids)
# decoder input has 3 parts: prefix, memory slots and suffix
# the following code is for Mistral tokenizer for example: 733, 16289, 28793 are for the Mistral instruction tempmlate
prompt_left_ids = torch.LongTensor([[1, 733, 16289, 28793]]).to(device)
prompt_right_ids = (
[model.ft_token_id]
+ tokenized_prompt["input_ids"]
+ [733, 28748, 16289, 28793]
)
prompt_right_ids = torch.LongTensor([prompt_right_ids]).to(device)
prompt_left_embs = model.tokens_to_embeddings(prompt_left_ids)
prompt_right_embs = model.tokens_to_embeddings(prompt_right_ids)
memory_slots = memory_slots.to(prompt_right_embs)
# Concatenate and clone input embeddings
decoder_input_embeddings = torch.cat(
(prompt_left_embs, memory_slots.unsqueeze(0), prompt_right_embs), dim=1
)
output = decoder_input_embeddings.clone()
generate_text = []
past_key_values = None
# Generate text output
for i in range(max_out_length):
with (
model.icae.disable_adapter()
): # no independent decoder; use self.icae
out = model.icae(
inputs_embeds=output,
past_key_values=past_key_values,
use_cache=True,
)
# out = decoder(inputs_embeds=output, past_key_values=past_key_values, use_cache=True)
logit = out.logits[:, -1, : model.vocab_size - 1]
past_key_values = out.past_key_values
next_token_id = torch.argmax(logit, dim=-1)
# print(next_token_id)
if next_token_id.item() == 2: # eos
break
output = (
model.icae.get_base_model()
.model.embed_tokens(next_token_id)
.unsqueeze(1)
.to(device)
)
generate_text.append(next_token_id.item())
generated_text = model.tokenizer.decode(generate_text)
# Structure output data
output_ = {
"input": data["input"],
"prompt": data["prompt"],
"output": generated_text,
"answer": data["answer"],
}
f.write(json.dumps(output_) + "\n")

16
icae_v2/fine_tuned_inference_script.sh
View file

@ -1,16 +0,0 @@
#!/bin/bash
# MODEL="mistralai/Mistral-7B-v0.1"
BASE_MODEL="mistralai/Mistral-7B-Instruct-v0.2"
# MODEL="meta-llama/Llama-2-7b-hf"
# MODEL="meta-llama/Llama-2-7b-chat-hf"
MODEL_NAME="${MODEL//\//-}"
maxlen=5120
mem=128
r=512
mean_compression_rate=4
ICAE_MODEL_PATH=$1 # ICAE model to use; wget "https://huggingface.co/sggetao/icae/resolve/main/mistral_7b_ft_icae.safetensors"
python fine_tuned_inference.py --mean_compression_rate $mean_compression_rate --model_max_length $maxlen --fixed_mem_size $mem --lora_r $r --output_dir $ICAE_MODEL_PATH --model_name_or_path $BASE_MODEL --bf16 --train False

118
icae_v2/instruction_finetune.py
View file

@ -1,118 +0,0 @@
from datasets import load_dataset
from modeling_icae_multi_span import (
ICAE,
DataArguments,
ModelArguments,
TrainingArguments,
)
from peft import LoraConfig
from training_utils import (
instruct_ft_tokenize_function,
train_model,
)
from transformers import HfArgumentParser
def compute_metrics(eval_pred) -> dict:
"""
Custom metrics function for the trainer
Args:
eval_pred: tuple of predictions and labels
Returns:
dictionary containing metric names (str) and values (Any)
"""
# preds, labels = eval_preds
# # predictions is generated tokens for Seq2SeqTrainer
# # decode preds and labels
# labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
# decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
# decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
# compute per token accuracy
predictions, labels = eval_pred.predictions, eval_pred.label_ids
# predictions is logits for Trainer
preds = predictions.argmax(-1)
acc = (preds == labels).mean()
return {"per_token_acc": acc}
def main():
parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
print(model_args)
print(data_args)
training_args.eval_on_start = True
training_args.eval_strategy = "steps"
training_args.eval_steps = 500
training_args.save_strategy = "no"
# training_args.save_steps = 500
training_args.logging_strategy = "steps"
training_args.logging_steps = 100
# seq2seq args for generation evaluation
# training_args.predict_with_generate = True
# training_args.generation_max_length = 100
training_args.gradient_checkpointing_kwargs = {
"use_reentrant": False
} # manually add this argument in the code
lora_config = LoraConfig(
r=model_args.lora_r,
lora_alpha=32,
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
# check model_args.mem_size and min_tokens_for_lm
assert (training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)) == 0, (
"training_args.fixed_mem_size must be a power of 2"
)
memory_size = training_args.fixed_mem_size
train_file = "../data/raw_datasets/context_numbers/train.jsonl"
eval_file = "../data/raw_datasets/context_numbers/val.jsonl"
print("Loading dataset...")
dataset = load_dataset("json", data_files={"train": train_file, "eval": eval_file})
train_dataset = dataset["train"]
eval_dataset = dataset["eval"]
model = ICAE(model_args, training_args, lora_config).to("cuda")
MEM_TOKENS = list(range(model.vocab_size, model.vocab_size + memory_size))
tokenized_train_ds = train_dataset.map(
instruct_ft_tokenize_function,
batched=True,
fn_kwargs={"model": model, "mem": MEM_TOKENS},
)
tokenized_eval_ds = {"train": train_dataset.select(range(100)), "val": eval_dataset}
for split in tokenized_eval_ds:
tokenized_eval_ds[split] = tokenized_eval_ds[split].map(
instruct_ft_tokenize_function,
batched=True,
fn_kwargs={"model": model, "mem": MEM_TOKENS},
)
train_model(
model,
tokenized_train_ds,
tokenized_eval_ds,
training_args,
compute_metrics=compute_metrics,
)
if __name__ == "__main__":
main()

334
icae_v2/modeling_icae_multi_span.py
View file

@ -1,334 +0,0 @@
# ICAE that supports multi span concat
import math
from dataclasses import dataclass, field
import torch
import torch.nn as nn
import transformers
from peft import get_peft_model
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, AutoTokenizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@dataclass
class ModelArguments:
model_name_or_path: str = field(default="mistralai/Mistral-7B-v0.1")
lora_r: int = field(default=128, metadata={"help": "lora rank"})
lora_dropout: float = field(default=0.05, metadata={"help": "lora dropout"})
train: bool = field(
default=True,
metadata={
"help": "if true, the model ckpt will be initialized for training; else, it's for inference"
},
)
@dataclass
class DataArguments:
data_path: str = field(
default=None, metadata={"help": "Path to the training data."}
)
debug_data: bool = field(
default=False,
metadata={
"help": "Enable debug dataset to quickly verify the training process"
},
)
@dataclass
class TrainingArguments(transformers.TrainingArguments):
cache_dir: str | None = field(default=None)
optim: str = field(default="adamw_torch")
model_max_length: int = field(
default=28000,
metadata={
"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
},
)
fixed_mem_size: int = field(
default=128,
metadata={"help": "Enalbing the fixed mem size."},
)
mean_compression_rate: int = field(
default=4,
metadata={"help": "Mean compression rate; default=4"},
)
min_tokens_for_lm: int = field(
default=64,
metadata={"help": "Minimum tokens for lm objective learning"},
)
leave_tokens_for_lm: int = field(
default=8,
metadata={"help": "Leave some tokens without loss for lm objective"},
)
lm_ratio: float = field(
default=0.0,
metadata={"help": "Ratio for LM training."},
)
add_special_token_for_lm: bool = field(
default=False,
metadata={
"help": "Add a special token for the prompt of language modeling; default: False"
},
)
restore_from: str = field(
default="",
metadata={
"help": "The checkpoint that should be restored from for fine-tuning"
},
)
def print_trainable_parameters(model):
trainable_parameters = 0
all_param = 0
for _, param in model.named_parameters():
all_param += param.numel()
if param.requires_grad:
trainable_parameters += param.numel()
print(
f"trainable params: {trainable_parameters} || all params: {all_param} || trainable%: {100 * trainable_parameters / all_param}"
)
# for name, param in model.named_parameters():
# if param.requires_grad:
# print(name, param.shape)
def freeze_model(model):
for _, param in model.named_parameters():
param.requires_grad = False
class ICAE(torch.nn.Module):
def __init__(self, model_args, training_args, lora_config):
super().__init__()
self.model_args = model_args
self.training_args = training_args
self.model_name = model_args.model_name_or_path
self.icae = AutoModelForCausalLM.from_pretrained(
self.model_name,
torch_dtype=torch.float16
if training_args.bf16 is False
else torch.bfloat16,
use_flash_attention_2=True,
resume_download=True,
)
self.training = self.model_args.train
if self.training: # indepedent model for gradient checkpointing
self.decoder = AutoModelForCausalLM.from_pretrained(
self.model_name,
torch_dtype=torch.float16
if training_args.bf16 is False
else torch.bfloat16,
use_flash_attention_2=True,
resume_download=True,
)
self.vocab_size = self.icae.config.vocab_size + 1 # [PAD] token
self.pad_token_id = self.vocab_size - 1
self.mean_compression_rate = training_args.mean_compression_rate
# tunable
self.mem_size = self.training_args.fixed_mem_size
self.vocab_size_with_mem = (
self.vocab_size + self.mem_size
) # so, the mem tokens are in the range [self.vocab_size, self.vocab_size + self.mem_size)
# special tokens in addition to mem and length tokens
self.ae_token_id = self.vocab_size_with_mem + 0
self.lm_token_id = self.vocab_size_with_mem + 1
self.ft_token_id = self.vocab_size_with_mem + 2
self.icae.resize_token_embeddings(self.vocab_size_with_mem + 3)
# special tokens for Llama-2/Mistral tokenizer
self.bos_id = 1
self.eos_id = 2
self.dim = self.icae.config.hidden_size
self.icae = get_peft_model(self.icae, lora_config)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.memory_token_embed = nn.Embedding(
self.mem_size + 3, self.dim, padding_idx=None
)
self.loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, use_fast=False)
self.append_sequence = torch.arange(
self.vocab_size,
self.vocab_size + self.mem_size,
dtype=torch.long,
device=device,
).unsqueeze(0) # mem tokens
if self.training:
self.init()
def init(self):
print("Freezing the decoder...")
freeze_model(self.decoder)
self.decoder.eval()
print_trainable_parameters(self)
if (
self.training_args.restore_from is not None
and self.training_args.restore_from != ""
):
print(
f"Loading from the pretrained checkpoint: {self.training_args.restore_from}..."
)
state_dict = load_file(self.training_args.restore_from)
self.load_state_dict(state_dict)
print(f"Finished loading from {self.training_args.restore_from}")
print("Enabling gradient checkpointing...")
# self.icae.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
self.decoder.gradient_checkpointing_enable(
gradient_checkpointing_kwargs={"use_reentrant": False}
)
def compute_num_segments(self, total_length):
assert total_length > 0
num_segments = math.ceil(
total_length / (self.mem_size * self.mean_compression_rate)
)
return num_segments
def forward(
self,
input_ids: torch.LongTensor = None,
prompt_answer_ids: torch.LongTensor = None,
labels: torch.LongTensor | None = None,
):
# encoder part
batch_size = input_ids.size(0)
total_length = input_ids.size(1)
num_segments = self.compute_num_segments(total_length)
segment_length = math.ceil(total_length / num_segments)
prompt_answer_embs = self.icae.get_base_model().model.embed_tokens(
prompt_answer_ids
)
max_compressed_length = num_segments * self.mem_size
compress_outputs = torch.zeros((max_compressed_length, self.dim)).to(
prompt_answer_embs
)
for segment_idx in range(num_segments):
start_idx = segment_idx * segment_length
end_idx = min((segment_idx + 1) * segment_length, total_length)
segment_input_ids = input_ids[:, start_idx:end_idx]
segment_input_ids = torch.cat(
[segment_input_ids, self.append_sequence], dim=1
)
mem_flag = segment_input_ids >= self.vocab_size
segment_input_embedding = self.icae.get_base_model().model.embed_tokens(
segment_input_ids
)
segment_input_embedding[mem_flag] = self.memory_token_embed(
segment_input_ids[mem_flag] - self.vocab_size
).to(segment_input_embedding)
# compress the current segment
segment_compress_outputs = self.icae(
inputs_embeds=segment_input_embedding, output_hidden_states=True
)
segment_compress_outputs = segment_compress_outputs.hidden_states[-1]
# collect memory tokens
compress_outputs[
segment_idx * self.mem_size : self.mem_size * (segment_idx + 1)
] = segment_compress_outputs[mem_flag]
del segment_input_ids, segment_input_embedding
torch.cuda.empty_cache()
# decoder part
decoder_mem_flag = (prompt_answer_ids >= self.vocab_size) & (
prompt_answer_ids < self.vocab_size + self.mem_size
) # only mem tokens
prompt_answer_embs[decoder_mem_flag] = compress_outputs # replace memory slots
special_prompt = prompt_answer_ids >= self.vocab_size_with_mem
prompt_answer_embs[special_prompt] = self.memory_token_embed(
prompt_answer_ids[special_prompt] - self.vocab_size
).to(
prompt_answer_embs
) # replace special token's embedding from self.memory_token_embed
if self.training: # has an independent se.f.decoder
decoder_outputs = self.decoder(
inputs_embeds=prompt_answer_embs, output_hidden_states=True
)
else:
with self.icae.disable_adapter(): # no independent decoder; use self.icae
decoder_outputs = self.icae(
inputs_embeds=prompt_answer_embs, output_hidden_states=True
)
logits = decoder_outputs.logits
effective_logits = logits[:, :-1, :].reshape(-1, logits.size(-1))
target_ids = labels[:, 1:].reshape(-1)
loss = self.loss_fct(effective_logits, target_ids)
return {"loss": loss, "logits": logits}
def tokens_to_embeddings(
self, token_ids
): # input_tokens can be either normal tokens and special tokens
embeddings = self.icae.get_base_model().model.embed_tokens(token_ids)
special_flags = token_ids >= self.vocab_size
embeddings[special_flags] = self.memory_token_embed(
token_ids[special_flags] - self.vocab_size
).to(
embeddings
) # replace special token's embedding from self.memory_token_embed
return embeddings
def _compress(
self, input_ids: torch.LongTensor = None
): # for inference; compress a fixed length of input into memory slots
batch_size = input_ids.size(0)
total_length = input_ids.size(1)
num_segments = self.compute_num_segments(total_length)
segment_length = math.ceil(total_length / num_segments)
max_compressed_length = num_segments * self.mem_size
compress_outputs = torch.zeros((max_compressed_length, self.dim))
for segment_idx in range(num_segments):
start_idx = segment_idx * segment_length
end_idx = min((segment_idx + 1) * segment_length, total_length)
segment_input_ids = input_ids[:, start_idx:end_idx]
segment_input_ids = torch.cat(
[segment_input_ids, self.append_sequence], dim=1
)
mem_flag = segment_input_ids >= self.vocab_size
segment_input_embedding = self.icae.get_base_model().model.embed_tokens(
segment_input_ids
)
segment_input_embedding[mem_flag] = self.memory_token_embed(
segment_input_ids[mem_flag] - self.vocab_size
).to(segment_input_embedding)
# compress the current segment
segment_compress_outputs = self.icae(
inputs_embeds=segment_input_embedding, output_hidden_states=True
)
segment_compress_outputs = segment_compress_outputs.hidden_states[-1]
# collect memory tokens
compress_outputs[
segment_idx * self.mem_size : self.mem_size * (segment_idx + 1)
] = segment_compress_outputs[mem_flag]
del segment_input_ids, segment_input_embedding
torch.cuda.empty_cache()
return compress_outputs

82
icae_v2/pretrain.py
View file

@ -1,82 +0,0 @@
import transformers
from datasets import load_dataset
from modeling_icae_multi_span import (
ICAE,
DataArguments,
ModelArguments,
TrainingArguments,
)
from peft import LoraConfig
from training_utils import (
DataCollatorForDynamicPadding,
pretrain_tokenize_function,
train_model,
)
def main():
parser = transformers.HfArgumentParser(
(ModelArguments, DataArguments, TrainingArguments)
)
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
print(model_args)
print(data_args)
training_args.gradient_checkpointing_kwargs = {
"use_reentrant": False
} # manually add this argument in the code
lora_config = LoraConfig(
r=model_args.lora_r,
lora_alpha=32,
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
# check model_args.mem_size and min_tokens_for_lm
assert (training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)) == 0, (
"training_args.fixed_mem_size must be a power of 2"
)
assert training_args.leave_tokens_for_lm <= training_args.min_tokens_for_lm, (
"leave_tokens_for_lm should be fewer than min_tokens_for_lm"
)
memory_size = training_args.fixed_mem_size
train_file = "/path/to/train/file"
eval_file = "/path/to/dev/file"
print("Loading dataset...")
dataset = load_dataset(
"json", data_files={"train": train_file, "eval": eval_file}, streaming=True
) # streaming can be removed if the dataset is not very large.
train_dataset = dataset["train"]
eval_dataset = dataset["eval"]
model = ICAE(model_args, training_args, lora_config)
MEM_TOKENS = list(range(model.vocab_size, model.vocab_size + memory_size))
train_dataset = train_dataset.map(
pretrain_tokenize_function,
batched=True,
batch_size=64,
fn_kwargs={
"model": model,
"mem": MEM_TOKENS,
"lm_ratio": training_args.lm_ratio,
},
)
eval_dataset = eval_dataset.map(
pretrain_tokenize_function,
batched=True,
fn_kwargs={"model": model, "mem": MEM_TOKENS},
) # don't add lm in the dev set.
data_collator = DataCollatorForDynamicPadding(model.pad_token_id)
train_model(model, train_dataset, eval_dataset, training_args, data_collator)
main()

111
icae_v2/pretrained_inference.py
View file

@ -1,111 +0,0 @@
import json
import torch
from modeling_icae_multi_span import (
ICAE,
DataArguments,
ModelArguments,
TrainingArguments,
)
from peft import LoraConfig
from safetensors.torch import load_file
from tqdm import tqdm
from transformers import HfArgumentParser
# Set the computation device
device = "cuda"
# Parse model, data, and training arguments
parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
# Define Lora configuration
lora_config = LoraConfig(
r=512,
lora_alpha=32,
lora_dropout=model_args.lora_dropout,
bias="none",
task_type="CAUSAL_LM",
)
# Initialize model and send it to CUDA device
model = ICAE(model_args, training_args, lora_config)
# Load the fine-tuned checkpoint
print(f"Loading trained checkpoint from {training_args.output_dir}")
state_dict = load_file(training_args.output_dir)
model.load_state_dict(
state_dict, strict=False
) # only load lora and memory token embeddings
model = model.to(device)
lines = [
"I don't have a favorite condiment as I don't consume food or condiments. However, I can tell you that many people enjoy condiments like ketchup, mayonnaise, mustard, soy sauce, hot sauce, and ranch dressing, among others. The favorite condiment can vary greatly from person to person, depending on their taste preferences and cultural influences."
]
# Prepare the model for evaluation
model.eval()
with torch.no_grad():
with open("tmp.out", "w") as f:
for line in tqdm(lines):
# Tokenize input text
tokenized_text = model.tokenizer(
line,
truncation=True,
max_length=5120,
padding=False,
return_attention_mask=False,
)
# Generate compressed outputs
input_ids = torch.LongTensor([tokenized_text["input_ids"]]).to(device)
memory_slots = model._compress(input_ids)
# prompt_output = model.tokenizer(data['prompt'], add_special_tokens=False, padding=False)
prompt_ids = torch.LongTensor([[model.ae_token_id]]).to(device)
prompt_answer_embs = model.tokens_to_embeddings(prompt_ids)
memory_slots = memory_slots.to(prompt_answer_embs)
# Concatenate and clone input embeddings
decoder_input_embeddings = torch.cat(
(memory_slots.unsqueeze(0), prompt_answer_embs), dim=1
)
output = decoder_input_embeddings.clone()
generate_text = []
past_key_values = None
# Generate text output
for i in range(512):
with (
model.icae.disable_adapter()
): # no independent decoder; use self.icae
out = model.icae(
inputs_embeds=output,
past_key_values=past_key_values,
use_cache=True,
)
logit = out.logits[:, -1, : model.vocab_size - 1]
past_key_values = out.past_key_values
next_token_id = torch.argmax(logit, dim=-1)
# print(next_token_id)
if next_token_id.item() == 2: # eos
break
output = (
model.icae.get_base_model()
.model.embed_tokens(next_token_id)
.unsqueeze(1)
.to(device)
)
generate_text.append(next_token_id.item())
generated_text = model.tokenizer.decode(generate_text)
# Structure output data
output_ = {"output": generated_text}
f.write(json.dumps(output_) + "\n")

16
icae_v2/pretrained_inference_script.sh
View file

@ -1,16 +0,0 @@
#!/bin/bash
# MODEL="mistralai/Mistral-7B-v0.1"
BASE_MODEL="mistralai/Mistral-7B-Instruct-v0.2"
# MODEL="meta-llama/Llama-2-7b-hf"
# MODEL="meta-llama/Llama-2-7b-chat-hf"
MODEL_NAME="${MODEL//\//-}"
maxlen=5120
mem=128
r=512
mean_compression_rate=4
ICAE_MODEL_PATH=$1 # ICAE model to use; wget "https://huggingface.co/sggetao/icae/resolve/main/mistral_7b_pretrained_icae.safetensors"
python pretrained_inference.py --mean_compression_rate $mean_compression_rate --model_max_length $maxlen --fixed_mem_size $mem --lora_r $r --output_dir $ICAE_MODEL_PATH --model_name_or_path $BASE_MODEL --bf16 --train False

256
icae_v2/training_utils.py
View file

@ -1,256 +0,0 @@
import os
import random
import torch
from transformers import Trainer
from transformers.trainer_utils import get_last_checkpoint
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def train_model(
model,
train_dataset,
eval_dataset,
training_args,
data_collator=None,
compute_metrics=None,
):
last_checkpoint = None
if (
os.path.isdir(training_args.output_dir)
and not training_args.overwrite_output_dir
):
last_checkpoint = get_last_checkpoint(training_args.output_dir)
if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
raise ValueError(
f"Output directory ({training_args.output_dir})"
" already exists and is not empty. "
"Use --overwrite_output_dir to overcome."
)
elif (
last_checkpoint is not None and training_args.resume_from_checkpoint is None
):
print(
f"Checkpoint detected, resuming training at {last_checkpoint}. "
"To avoid this behavior, change "
"the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
)
if (
max(
training_args.per_device_train_batch_size,
training_args.per_device_eval_batch_size,
)
== 1
):
data_collator = None
# print training_args at local_rank 0
local_rank = int(os.getenv("LOCAL_RANK", "0"))
if local_rank == 0:
print(training_args)
# Seq2SeqTrainer is actually just the same as Trainer
# (although it uses a different data collator, i.e., explicit prompt/answer separation)
# it just allows `predict_with_generate`
# allowing us to compute metrics on the generated outputs
# no clue why they call this seq2seq...
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
data_collator=data_collator,
compute_metrics=compute_metrics,
)
checkpoint = None
if training_args.resume_from_checkpoint is not None:
checkpoint = training_args.resume_from_checkpoint
elif last_checkpoint is not None:
checkpoint = last_checkpoint
print(f"Loaded from the checkpoint: {checkpoint}")
train_result = trainer.train(resume_from_checkpoint=checkpoint)
trainer.save_model()
trainer.log_metrics("train", train_result.metrics)
metrics = trainer.evaluate()
trainer.log_metrics("eval", metrics)
trainer.save_metrics("eval", metrics)
def text_extraction(input_ids, length, lm_ratio=0.0):
input_len = len(input_ids)
assert input_len >= 1, f"Error: invalid input length ({input_len})"
# ae
if random.random() >= lm_ratio:
if input_len <= length: # if shorter, keep the complete text
return input_ids, []
else:
last_start = input_len - length
random_start = random.randint(0, last_start)
return input_ids[random_start : random_start + length], []
# lm
if input_len <= length:
r = random.randint(0, input_len - 1)
return input_ids[: r + 1], input_ids[r + 1 :]
else:
last_start = input_len - length
random_start = random.randint(0, last_start)
return (
input_ids[random_start : random_start + length],
input_ids[random_start + length :],
)
def pretrain_tokenize_function(examples, model, mem, lm_ratio=0.0):
text_output = model.tokenizer(
examples["text"], truncation=False, padding=False, return_attention_mask=False
)
text_output["prompt_answer_ids"] = []
text_output["labels"] = []
max_len = model.training_args.model_max_length # heuristic
for idx in range(len(text_output["input_ids"])):
ae = True
a, b = text_extraction(
text_output["input_ids"][idx], max_len, lm_ratio=lm_ratio
)
length_a = len(a)
num_segments = model.compute_num_segments(length_a)
total_mem_length = num_segments * model.mem_size
if (
len(b) > model.training_args.min_tokens_for_lm
): # avoid too few tokens for lm, which is a waste of computing
ae = False
b = b[:max_len]
text_output["input_ids"][idx] = a
# decoder part: note that in v2, we add mem_tokens to the prompt_ids
# for easy implementation; which is different from v1 implementation
# where mem tokens are not in the prompt_ids
if ae: # autoencoding objective
prompt_ids = [mem[0]] * total_mem_length + [model.ae_token_id]
answer_ids = a + [model.eos_id] # if ae, eos token
else: # lm objective
prompt_ids = [mem[0]] * total_mem_length
if model.training_args.add_special_token_for_lm:
prompt_ids += [model.lm_token_id]
answer_ids = b # if lm, no eos token
text_output["prompt_answer_ids"].append(prompt_ids + answer_ids)
if ae:
labels = [-100] * len(prompt_ids) + answer_ids
else:
labels = (
[-100] * len(prompt_ids)
+ [-100] * model.training_args.leave_tokens_for_lm
+ answer_ids[model.training_args.leave_tokens_for_lm :]
) # no loss for leave_tokens_for_lm
text_output["labels"].append(labels)
assert len(text_output["prompt_answer_ids"][-1]) == len(labels)
return text_output
def instruct_ft_tokenize_function(examples, model, mem):
text_output = model.tokenizer(
examples["context"],
max_length=5120,
truncation=True,
padding=False,
return_attention_mask=False,
add_special_tokens=False,
)
prompt_output = model.tokenizer(
examples["prompt"],
truncation=False,
padding=False,
return_attention_mask=False,
add_special_tokens=False,
)
label_output = model.tokenizer(
examples["answer"],
truncation=False,
padding=False,
return_attention_mask=False,
add_special_tokens=False,
)
text_output["prompt_answer_ids"] = []
text_output["labels"] = []
max_len = model.training_args.model_max_length # heuristic
for idx in range(len(text_output["input_ids"])):
length = len(text_output["input_ids"][idx])
num_segments = model.compute_num_segments(length)
total_mem_length = num_segments * model.mem_size
prompt_ids = (
[mem[0]] * total_mem_length
+ [model.ft_token_id]
+ prompt_output["input_ids"][idx]
)
prompt_ids = (
[1, 733, 16289, 28793] + prompt_ids + [733, 28748, 16289, 28793]
) # special formats for prompt in Mistral
answer_ids = label_output["input_ids"][idx] + [model.eos_id]
text_output["prompt_answer_ids"].append(prompt_ids + answer_ids)
labels = [-100] * len(prompt_ids) + answer_ids
text_output["labels"].append(labels)
assert len(text_output["prompt_answer_ids"][-1]) == len(labels)
return text_output
class DataCollatorForDynamicPadding:
def __init__(self, pad_token_id, pad_to_multiple_of=None):
self.pad_token_id = pad_token_id
self.pad_to_multiple_of = pad_to_multiple_of
def __call__(self, examples):
input_ids = [
torch.tensor(example["input_ids"], dtype=torch.long) for example in examples
]
labels = [
torch.tensor(example["labels"], dtype=torch.long) for example in examples
]
prompt_answer_ids = [
torch.tensor(example["prompt_answer_ids"], dtype=torch.long)
for example in examples
]
input_ids = self.dynamic_padding(input_ids, fill_value=self.pad_token_id)
prompt_answer_ids = self.dynamic_padding(
prompt_answer_ids, fill_value=self.pad_token_id
)
labels = self.dynamic_padding(labels)
batch = {
"input_ids": input_ids,
"labels": labels,
"prompt_answer_ids": prompt_answer_ids,
}
return batch
def dynamic_padding(self, sequences, fill_value=-100):
max_length = max(len(x) for x in sequences)
if self.pad_to_multiple_of:
max_length = (
(max_length - 1) // self.pad_to_multiple_of + 1
) * self.pad_to_multiple_of
padded_sequences = torch.full(
(len(sequences), max_length), fill_value, dtype=torch.long
)
for i, seq in enumerate(sequences):
padded_sequences[i, : len(seq)] = seq
return padded_sequences

36
install.sh
View file

@ -1,33 +1,27 @@
# uv pip install -e .
# uv pip install flash-attn==2.6.3 --no-build-isolation
# # huggingface-cli login
# # lm-evaluation-harness
# # git clone https://github.com/EleutherAI/lm-evaluation-harness.git
# # cd lm-evaluation-harness
# # git checkout f724be699e8adf7ca8004ea0e519dfac83a06f18
# # pip install -e .
module load cuda/12.4
module load cudnn/9.1.0
module load hpcx/v2.21
# the installation script assumes the following modules (slurm)
# module load cuda/12.4
# module load cudnn/9.1.0
# module load hpcx/v2.21
uv self update
uv venv --python 3.10 --seed
# uv pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu124
# uv pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --torch-backend=auto
uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --torch-backend=cu124
# uv pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --torch-backend=cu128
uv sync
uv pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
uv pip install flashinfer-python==0.2.2 -i https://flashinfer.ai/whl/cu124/torch2.6
# uv pip install flash-attn==2.6.3 --no-build-isolation
# download squad dataset
HF_HUB_ENABLE_HF_TRANSFER=1 uv run huggingface-cli download --repo-type dataset rajpurkar/squad --local-dir data/raw_datasets/squad
uv run data/build_drop_compact.py
uv run data/build_pwc_compact.py
uv run data/build_ropes_compact.py
uv run data/build_squad_compact.py
# [Optional]
# needed for gated models
uv run huggingface-cli login
# uv run huggingface-cli login
# needed for logging with wandb
wandb login
# wandb login
# dev
uv run pre-commit install
# uv run pre-commit install

BIN
mock_results.pdf
View file

Binary file not shown.

BIN
mock_results.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 150 KiB

BIN
niah_memory_latency_vs_context_length.pdf
View file

Binary file not shown.

BIN
niah_memory_latency_vs_context_length.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 150 KiB

BIN
perf_vs_latency_plot.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 174 KiB

1
pyproject.toml
View file

@ -38,6 +38,7 @@ dependencies = [
"google-cloud-storage>=3.2.0",
"wonderwords>=2.2.0",
"tokenizers==0.21.0",
"llmlingua>=0.2.2",
]
[build-system]

BIN
qa_performance_vs_context_length.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 247 KiB

20
run_eval.py
View file

@ -1,8 +1,6 @@
import logging
from ctx_to_lora.eval_utils import (
run_eval,
)
from ctx_to_lora.eval_utils import run_eval
logger = logging.getLogger()
@ -156,7 +154,6 @@ if __name__ == "__main__":
)
cli_args = vars(parser.parse_args())
# setup_logging(output_dir, debug=os.getenv("DEBUG", False))
if cli_args["model_name_or_path"]:
assert cli_args["max_ctx_chunk_len"] <= 0, (
@ -165,23 +162,8 @@ if __name__ == "__main__":
eval_batch_size_gen = cli_args.pop("eval_batch_size_gen")
eval_batch_size = cli_args.pop("eval_batch_size")
# run_eval(
# **cli_args,
# # cli_args.checkpoint_path,
# # cli_args.model_name_or_path,
# # cli_args.eval_batch_size,
# # args,
# # split=cli_args.split,
# eval_batch_size=eval_batch_size,
# generative=False,
# )
run_eval(
**cli_args,
# cli_args.checkpoint_path,
# cli_args.model_name_or_path,
# cli_args.eval_batch_size_gen,
# args,
# split=cli_args.split,
eval_batch_size=eval_batch_size_gen,
generative=True,
)

30
scripts/large_exp/gemma_qa_exp_4_gpus.sh
View file

@ -1,30 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# Default arguments - these can be overridden by passing arguments to this script
default_args=(
"--model_name_or_path=google/gemma-2-2b-it"
# "--max_steps=50_000"
"--num_train_epochs=1"
"--target_modules=down_proj"
"--lora_r=8"
"--eval_strategy=no"
"--max_qas_len=2048"
"--max_qas_per_sample=1"
"--per_rank_gen=True"
"--per_layer_processing=True"
"--gen_lora_l1_reg_coef=0.1"
)
# Pass all script arguments to the training command
# $1 comes first, then defaults, then remaining arguments (which can override defaults)
uv run accelerate launch --config_file accelerate_config.yaml --main_process_port $port \
--num_processes=4 --gpu_ids all train.py $1 "${default_args[@]}" "${@:2}"

30
scripts/large_exp/gemma_qa_exp_8_gpus.sh
View file

@ -1,30 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=8
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# Default arguments - these can be overridden by passing arguments to this script
default_args=(
"--model_name_or_path=google/gemma-2-2b-it"
# "--max_steps=50_000"
"--num_train_epochs=1"
"--target_modules=down_proj"
"--lora_r=8"
"--eval_strategy=no"
"--max_qas_len=2048"
"--max_qas_per_sample=1"
"--per_rank_gen=True"
"--per_layer_processing=True"
"--gen_lora_l1_reg_coef=0.1"
)
# Pass all script arguments to the training command
# $1 comes first, then defaults, then remaining arguments (which can override defaults)
uv run accelerate launch --config_file accelerate_config.yaml --main_process_port $port \
--num_processes=8 --gpu_ids all train.py $1 "${default_args[@]}" "${@:2}"

8
scripts/main_exp/eval/base_model.sh Normal file
View file

@ -0,0 +1,8 @@
# no truncation
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1
# w/ truncation
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --truncate_if_too_long_inp
# no context
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --remove_context

6
scripts/main_exp/eval/cd.sh Normal file
View file

@ -0,0 +1,6 @@
# qa
uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=4
# longbench
uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=1

6
scripts/main_exp/eval/cd_oracle.sh Normal file
View file

@ -0,0 +1,6 @@
# qa
uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp
# longbench
uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp

13
scripts/main_exp/eval/d2l.sh Normal file
View file

@ -0,0 +1,13 @@
# main results
# batched
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1
# iterative
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1 --use_iterative_mode
# query internalization
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad --split test --eval_batch_size_gen=1 --flip_ctx_inp
# replaced squad context
WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad_assistant_ctx_no_passage --split test
WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad_negative_no_passage --split test

5
scripts/main_exp/eval/llmlingua.sh Normal file
View file

@ -0,0 +1,5 @@
for dataset in squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e; do
for rate in 0.9 0.8 0.6 0.4 0.2 0.1; do
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets "$dataset" --split test --eval_batch_size_gen=1 --use_llmlingua --llmlingua_compression_rate "$rate" --truncate_if_too_long_ctx
done
done

4
scripts/main_exp/eval/t2l.sh Normal file
View file

@ -0,0 +1,4 @@
# download t2l checkpoint
uv run huggingface-cli download SakanaAI/text-to-lora --local-dir . --include "trained_t2l/gemma_2b_t2l"
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen=1 --use_t2l

20
scripts/main_exp/gen_data.sh Normal file
View file

@ -0,0 +1,20 @@
# download fineweb_edu to `data/raw_datasets/fineweb_edu
uv run data/download_fineweb_edu.py
# generate qa data
# run from 000 to 013
for shard_id in $(seq -f "%03g" 0 1); do
uv run data/generate_fw_edu_qa_v2.py --shard_pattern "${shard_id}_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048;
uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/${shard_id}*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
# self-generated response QA data
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern "data/raw_datasets/fw_qa_v2/min_0_to_2000/${shard_id}*_level_1*" --closed_qa_prob 1.0
done
# val split
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_v2/min_0_to_2000/*_level_0_val.parquet'
# self-gen data for other ds
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names squad_compact ropes_compact drop_compact --closed_qa_prob 1.0
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names pwc_compact --closed_qa_prob 0.0

28
scripts/main_exp/train.sh Normal file
View file

@ -0,0 +1,28 @@
#!/bin/bash
port=29051
default_args=(
"--model_name_or_path=google/gemma-2-2b-it"
"--max_steps=50_000"
"--target_modules=down_proj"
"--lora_r=8"
"--eval_strategy=no"
"--max_qas_len=2048"
"--max_qas_per_sample=1"
"--per_rank_gen=True"
"--per_layer_processing=True"
"--gen_lora_l1_reg_coef=0.1"
)
uv run accelerate launch --config_file accelerate_config.yaml --main_process_port $port \
--num_processes=8 --gpu_ids all train.py \
configs/main_exp/self_gen_lv1_closed_qa_1_l2l.yaml \
"${default_args[@]}" \
--model_name_or_path=google/gemma-2-2b-it \
--target_modules=down_proj --lora_r=8 \
--eval_strategy=no --max_qas_len=2048 --max_qas_per_sample=1 \
--per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 \
--max_steps=80000 --gradient_accumulation_steps=8 --max_packed_inp_len=4096 \
--max_packed_ctx_len=4096 --use_per_ctx_average_loss=True --use_kl_loss=True \
--quantize_ctx_encoder=True

1
scripts/niah/0-gen_data.sh Normal file
View file

@ -0,0 +1 @@
uv run data/generate_ctx_magic_number.py

22
scripts/niah/1-train.sh Normal file
View file

@ -0,0 +1,22 @@
WANDB_MODE=disabled uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml \
--model_name_or_path=google/gemma-2-2b-it \
--num_train_epochs=1 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=32 \
--per_device_eval_batch_size=16 \
--exp_setup=hyper_lora \
--aggregator_type=perceiver --target_modules=down_proj \
--num_blocks=8 --num_self_attn_per_block=0 \
--num_pre_head_layers=1 --lora_r=8 \
--eval_steps=100 --logging_steps=10 \
--save_steps=1000 --learning_rate=4e-5 \
--lora_dropout=0.0 --neftune_noise_alpha=0 \
--per_rank_gen=True --per_layer_processing=True \
--gen_lora_l1_reg_coef=1 --use_sequence_packing=True \
--max_packed_inp_len=2048 --max_packed_ctx_len=2048 \
--dataloader_num_workers=0 --dataloader_prefetch_factor=None \
--ctx_encoder_type=early_exit --n_latent_queries=208 \
--eval_on_start=True --lora_r=8 \
--max_ctx_chunk_len=512 --min_ctx_chunk_len=25 \
--num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' \
--max_val_samples_per_ds=100 --seed=1 --use_per_ctx_average_loss=True

1
scripts/niah/2-eval.sh Normal file
View file

@ -0,0 +1 @@
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path CHECKPOINT_PATH --datasets ctx_magic_number_32_1024 ctx_magic_number_1024_2048 ctx_magic_number_2048_3072 ctx_magic_number_3072_4096 ctx_magic_number_4096_5120 ctx_magic_number_5120_6144 ctx_magic_number_6144_7168 ctx_magic_number_7168_8192 ctx_magic_number_8192_9216 ctx_magic_number_9216_10240 ctx_magic_number_10240_11264 ctx_magic_number_11264_12288 ctx_magic_number_12288_13312 ctx_magic_number_13312_14336 ctx_magic_number_14336_15360 ctx_magic_number_15360_16384 ctx_magic_number_16384_20480 ctx_magic_number_20480_24576 ctx_magic_number_24576_28672 ctx_magic_number_28672_32768 ctx_magic_number_32768_40960 ctx_magic_number_40960_49152 ctx_magic_number_49152_57344 ctx_magic_number_57344_65536 ctx_magic_number_65536_73728 ctx_magic_number_73728_81920 ctx_magic_number_81920_90112 ctx_magic_number_90112_98304 ctx_magic_number_98304_106496 ctx_magic_number_106496_114688 ctx_magic_number_114688_122880 ctx_magic_number_122880_131072 --max_ctx_chunk_len=1024 --split test --eval_batch_size_gen=4

10
scripts/sbatch_1_gpu.sh
View file

@ -1,10 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=1
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
"$@"

29
scripts/short_ctx/gemma_qa_short_ctx_exp_default.sh
View file

@ -1,29 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=4
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# Default arguments - these can be overridden by passing arguments to this script
default_args=(
"--model_name_or_path=google/gemma-2-2b-it"
"--max_steps=50_000"
"--target_modules=down_proj"
"--lora_r=8"
"--eval_strategy=no"
"--max_qas_len=2048"
"--max_qas_per_sample=1"
"--per_rank_gen=True"
"--per_layer_processing=True"
"--gen_lora_l1_reg_coef=0.1"
)
# Pass all script arguments to the training command
# $1 comes first, then defaults, then remaining arguments (which can override defaults)
uv run accelerate launch --config_file accelerate_config.yaml --main_process_port $port \
--num_processes=4 --gpu_ids all train.py $1 "${default_args[@]}" "${@:2}"

29
scripts/short_ctx/gemma_qa_short_ctx_exp_default_8_gpus.sh
View file

@ -1,29 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=sakura-gpu
#SBATCH --gpus=8
#SBATCH --output=slurm_logs/%x-%j.out
#SBATCH --error=slurm_logs/%x-%j.out
port=$((10000 + ($SLURM_JOBID % 50000)))
echo "Using port: $port"
# Default arguments - these can be overridden by passing arguments to this script
default_args=(
"--model_name_or_path=google/gemma-2-2b-it"
"--max_steps=50_000"
"--target_modules=down_proj"
"--lora_r=8"
"--eval_strategy=no"
"--max_qas_len=2048"
"--max_qas_per_sample=1"
"--per_rank_gen=True"
"--per_layer_processing=True"
"--gen_lora_l1_reg_coef=0.1"
)
# Pass all script arguments to the training command
# $1 comes first, then defaults, then remaining arguments (which can override defaults)
uv run accelerate launch --config_file accelerate_config.yaml --main_process_port $port \
--num_processes=8 --gpu_ids all train.py $1 "${default_args[@]}" "${@:2}"

48
scripts/submit_tiny_exp_jobs.sh
View file

@ -1,48 +0,0 @@
#!/bin/bash
# Script to submit SLURM jobs for all configs in configs/tiny_exp/
# with random sleep intervals between submissions
CONFIG_DIR="configs/tiny_exp"
SCRIPT_PATH="scripts/short_ctx/gemma_qa_short_ctx_exp_default.sh"
# Check if config directory exists
if [ ! -d "$CONFIG_DIR" ]; then
echo "Error: Config directory $CONFIG_DIR does not exist"
exit 1
fi
# Check if script exists
if [ ! -f "$SCRIPT_PATH" ]; then
echo "Error: Script $SCRIPT_PATH does not exist"
exit 1
fi
# Get all yaml config files
configs=($(find "$CONFIG_DIR" -name "*.yaml" -type f))
if [ ${#configs[@]} -eq 0 ]; then
echo "No .yaml config files found in $CONFIG_DIR"
exit 1
fi
echo "Found ${#configs[@]} config files:"
for config in "${configs[@]}"; do
echo " - $config"
done
echo ""
echo "Starting job submissions..."
# Submit jobs with random sleep intervals
for config in "${configs[@]}"; do
echo "Submitting job for config: $config"
sbatch "$SCRIPT_PATH" "$config"
# Random sleep between 3-30 seconds
sleep_time=$((3 + RANDOM % 28))
echo "Waiting $sleep_time seconds before next submission..."
sleep $sleep_time
done
echo "All jobs submitted!"

17
scripts/toy_exp.sh
View file

@ -1,17 +0,0 @@
# data gen
uv run data/generate_ctx_magic_num.py
# train
# WANDB_PROJECT=ctx-magic-num srun --partition=aiscilow --gpus=1 --unbuffered uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=32 --per_device_eval_batch_size=16 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=100 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=2048 --max_packed_ctx_len=2048 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=-1 --max_val_samples_per_ds=100 --seed=1
# WANDB_PROJECT=ctx-magic-num srun --partition=aiscilow --gpus=1 --unbuffered uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=32 --per_device_eval_batch_size=16 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=100 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=2048 --max_packed_ctx_len=2048 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=512 --min_ctx_chunk_len=25 --num_chunk_probs='{"1":"0.5", "2":"0.5"}' --max_val_samples_per_ds=100 --seed=1
# WANDB_PROJECT=ctx-magic-num srun --partition=aiscilow --gpus=1 --unbuffered uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=32 --per_device_eval_batch_size=16 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=100 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=2048 --max_packed_ctx_len=2048 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=512 --min_ctx_chunk_len=25 --num_chunk_probs='{"1":"0.5", "2":"0.25", "3":"0.125", "4":"0.125"}' --max_val_samples_per_ds=100 --seed=1
WANDB_PROJECT=ctx-magic-num run uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=16 --per_device_eval_batch_size=16 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=100 --logging_steps=10 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=1 --use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=512 --min_ctx_chunk_len=25 --num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' --max_val_samples_per_ds=100 --seed=1 --use_per_ctx_average_loss=True
# self-distill version
# ce
# WANDB_PROJECT=ctx-magic-num run uv run train.py configs/toy_exp/ctx_magic_number_32_256_self_gen.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=16 --per_device_eval_batch_size=16 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=100 --logging_steps=10 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=1 --use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=per_layer_activations --n_latent_queries=8 --use_kl_loss=False --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=-1 --max_val_samples_per_ds=100 --seed=1 --notes="1 l1norm + ce loss + 1 chunk" --use_per_ctx_average_loss=False
# kl
# WANDB_PROJECT=ctx-magic-num run uv run train.py configs/toy_exp/ctx_magic_number_32_256_self_gen.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=-1 --gradient_accumulation_steps=16 --per_device_eval_batch_size=16 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=100 --logging_steps=10 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=0 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=1 --use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=per_layer_activations --n_latent_queries=8 --use_kl_loss=True --eval_on_start=True --lora_r=8 --max_ctx_chunk_len=-1 --max_val_samples_per_ds=100 --seed=1 --notes="1 l1norm + kl loss + 1 chunk" --use_per_ctx_average_loss=False
# eval
WANDB_MODE=disabled srun --partition=aiscilow --gpus=1 --unbuffered uv run run_eval.py --checkpoint_path CHECKPOINT_PATH --datasets ctx_magic_number_32_1024 ctx_magic_number_1024_2048 ctx_magic_number_2048_3072 ctx_magic_number_3072_4096 ctx_magic_number_4096_5120 ctx_magic_number_5120_6144 ctx_magic_number_6144_7168 ctx_magic_number_7168_8192 ctx_magic_number_8192_9216 ctx_magic_number_9216_10240 ctx_magic_number_10240_11264 ctx_magic_number_11264_12288 ctx_magic_number_12288_13312 ctx_magic_number_13312_14336 ctx_magic_number_14336_15360 ctx_magic_number_15360_16384 ctx_magic_number_16384_20480 ctx_magic_number_20480_24576 ctx_magic_number_24576_28672 ctx_magic_number_28672_32768 ctx_magic_number_32768_40960 ctx_magic_number_40960_49152 ctx_magic_number_49152_57344 ctx_magic_number_57344_65536 ctx_magic_number_65536_73728 ctx_magic_number_73728_81920 ctx_magic_number_81920_90112 ctx_magic_number_90112_98304 ctx_magic_number_98304_106496 ctx_magic_number_106496_114688 ctx_magic_number_114688_122880 ctx_magic_number_122880_131072 --max_ctx_chunk_len=1024 --split test --eval_batch_size_gen=4

BIN
sigmoid_fit.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 34 KiB

0
slurm_logs/.gitkeep
View file

45
src/ctx_to_lora/configs.py
View file

@ -85,8 +85,6 @@ class ArgumentParser(HfArgumentParser):
raise ValueError(
f"Duplicate argument provided: {arg}, may cause unexpected behavior"
)
# else:
# raise ValueError(f"Argument provided not found in dataclass: {arg}")
obj = data_class(**inputs)
outputs.append(obj)
@ -126,9 +124,7 @@ class ArgumentParser(HfArgumentParser):
class ExperimentSetup(str, Enum):
# LORA = "lora"
HYPERLORA = "hyper_lora"
# FULL_FINETUNE = "full_finetune"
@dataclass
@ -207,7 +203,6 @@ class TrainingArguments(TrainingArguments):
default=0.01,
metadata={"help": "Weight decay for the optimizer."},
)
# TODO: is paged_adam_w faster? (bigger batch size -> fewer grad accum steps)
optim: str = field(
default="adamw_torch_fused",
metadata={"help": "Optimizer."},
@ -252,10 +247,6 @@ class TrainingArguments(TrainingArguments):
default=None,
metadata={"help": "Metric for best model."},
)
# greater_is_better: bool = field(
# default=False,
# metadata={"help": "Whether the metric is better when it is greater."},
# )
load_best_model_at_end: bool = field(
default=False,
metadata={"help": "Whether to load the best model at the end of training."},
@ -318,10 +309,6 @@ class ModelArguments:
default=True,
metadata={"help": "Whether to use flash attention."},
)
# use_peft: bool = field(
# default=False,
# metadata={"help": ("Whether to use PEFT or not for training.")},
# )
@dataclass
@ -338,10 +325,6 @@ class LoRAArguments:
default=None,
metadata={"help": ("LoRA target modules.")},
)
# modules_to_save: Optional[list[str]] = field(
# default=None,
# metadata={"help": ("Modules to save.")},
# )
@dataclass
@ -464,10 +447,10 @@ class DataArguments:
default=1000,
metadata={"help": "Maximum number of validation samples per dataset."},
)
# max_test_samples_per_ds: int | None = field(
# default=1000,
# metadata={"help": "Maximum number of test samples per dataset."},
# )
max_test_samples_per_ds: int | None = field(
default=500,
metadata={"help": "Maximum number of test samples per dataset."},
)
@dataclass
@ -555,17 +538,6 @@ class AggregatorArguments:
default="mean",
metadata={"help": "Pooling type for HyperLoRA."},
)
# feature_size: int
# num_layers: int
# num_modules: int
# output_size: int
# perceiver
# attention_probs_dropout_prob: float = field(
# default=0.0,
# metadata={"help": "Attention dropout probability for Perceiver."},
# )
# not used (keep herer for backward compatibility)
num_latent_factor: int = field(
default=8,
metadata={"help": "Number of latent factors for Perceiver."},
@ -574,10 +546,6 @@ class AggregatorArguments:
default=208, # 26 * 8
metadata={"help": "Number of latent queries of Perceiver."},
)
# num_blocks: int = field(
# default=8,
# metadata={"help": "Number of blocks for Perceiver."},
# )
num_blocks: int = field(
default=8,
@ -592,11 +560,6 @@ class AggregatorArguments:
metadata={"help": "Whether to share weights across blocks for Perceiver."},
)
# decoder_depth: int = field(
# default=1,
# metadata={"help": "Decoder depth for Perceiver."},
# )
# needed for loading model from checkpoint
# see https://github.com/huggingface/transformers/pull/34632

8
src/ctx_to_lora/data/collator.py
View file

@ -43,9 +43,6 @@ def flatten_if_not_packed(inp_list):
if "ctx_ids" in sample:
# HACK: assumes 1 ctx chunk here
# TODO: fix this after we have chunk ctx training
# print(sample["ctx_ids"])
# breakpoint()
# assert all(len(ctx_ids) == 1 for ctx_ids in sample["ctx_ids"]), (
# "ctx_ids can only have one chunk for eval. "
# "Please implement chunked ctx forward pass to handle this."
@ -68,14 +65,12 @@ def flatten_if_not_packed(inp_list):
def eval_collator(inp_list, tokenizer):
# only used for teacher-forcing eval
# input is a list of tokenized sequences
# TODO: handle list input (used to be torch.Tensor)
padding_kwargs = dict(padding=True, padding_side="right", return_tensors="pt")
has_ctx_ids = "ctx_ids" in inp_list[0]
if has_ctx_ids:
# pad to the longest ctx_len in the batch
# which can have a different length from the input_ids, attn_mask, labels
# TODO: handle chunked ctx_ids
ctx_ids = [example.pop("ctx_ids") for example in inp_list]
ctx_attn_mask = [torch.ones_like(x) for x in ctx_ids]
ctx_ids = torch.nn.utils.rnn.pad_sequence(
@ -147,7 +142,4 @@ def generation_collator(inp_list, tokenizer):
out["ctx_attn_mask"] = ctx_attn_mask
out["n_ctx_chunks"] = torch.tensor(n_chunks, dtype=torch.int32)
# print(f"ctx_ids: {ctx_ids}")
# print(f"n_chunks: {n_chunks}")
# breakpoint()
return out

680
src/ctx_to_lora/data/definitions.py
View file

@ -1,167 +1,31 @@
from glob import glob
IGNORE_INDEX = -100
FW_QA_PATHS = [
f"data/raw_datasets/fw_qa/{i:05d}.parquet" for i in [0, 1, 6, 7, 8, 10, 22, 30, 35]
]
TRANSFORMED_DATA_DIR = "data/processed_datasets"
RAW_DATA_DIR = "data/raw_datasets/"
SELF_GEN_DATA_DIR = f"{RAW_DATA_DIR}/self_gen/"
SUMMARIZATION_PROMPTS = [
"Summarize the article."
"Provide a concise summary capturing the main idea and key detail.",
"Write a one-paragraph executive summary of the passage. Avoid fluff.",
"Create a bullet list of the core points from the text. No intro line.",
"Summarize the content for a general audience with no domain knowledge.",
"Produce a headline that best represents the passage.",
"Give a headline and a subheadline separated by a newline summarizing the article.",
"Write a tweet-length summary (max 240 characters, no hashtags).",
"Condense the passage into a plain-language abstract.",
"Explain the passage as if speaking to a 10-year-old.",
"Summarize the information presented.",
"Extract and list the three most critical facts. No extra wording.",
"Provide a neutral summary (no opinion words) in one paragraph.",
"Summarize emphasizing quantitative data only. Ignore narrative.",
"Provide a facts-only summary: list raw data points, no interpretation.",
"Summarize by extracting all key figures and associated units only.",
"Provide a brief TL;DR starting with 'TL;DR:'.",
"Summarize in exactly 25 words. If needed, rephrase to fit.",
"Produce a question-and-answer style summary with 3 Q&A pairs capturing essence.",
"Provide a summary that only uses the 600 most common English words.",
"Create a concise abstract in academic tone.",
"Write a marketing-style summary appealing to a broad audience.",
"Summarize in the form of a newspaper lead (who, what, when, where, why).",
"Write a summary that preserves all numbers and proper nouns verbatim.",
"Produce a satirical-style summary (remain respectful).",
"Summarize with a formal tone using no contractions.",
"Summarize with an informal conversational tone",
"Summarize as if reporting to a technical expert (dense, precise).",
"Summarize as if explaining to a non-technical stakeholder (clear, simplified).",
"Summarize including exactly one direct quote (short) from the text if present.",
"Summarize into a structured JSON with keys: topic, key_points (array), tone. Do not add commentary.",
"Write a poetic-style 2-line summary (no rhyme required).",
"Summarize in exactly 3 bullet points.",
"Summarize in question form: pose 2-3 questions that capture essence.",
"Summarize as a press release opening paragraph.",
"Summarize focusing on benefits to end users.",
"Summarize avoiding any jargon; replace technical terms with common words.",
"Produce a summary in exactly 2 haiku-style lines (not strict syllables).",
"Summarize as if drafting the opening line of an encyclopedia entry.",
]
COT_TEMPLATES = [
"{input} Think out loud before giving an answer.",
"{input} Think step by step to ensure a correct final answer.",
"{input} List any assumptions, reason through the steps, then provide the final answer.",
"{input} Explain your reasoning clearly first.",
"{input} Reason carefully: identify what is known, what is needed, derive intermediate conclusions, then answer.",
"{input} Enumerate relevant facts, derive implications, eliminate alternatives, then state the answer.",
"{input} Think through possible interpretations, choose the most plausible, reason it out, then answer.",
"{input} Consider edge cases or exceptions while reasoning step by step before finalizing your answer.",
"{input} List data given, infer what is missing, derive needed values logically, then answer.",
"{input} Use logical deduction step by step; avoid jumping to the conclusion. Conclude with the answer.",
"{input} Reason transparently: gather, analyze, synthesize, conclude. Final answer after the reasoning.",
]
STRUCTURING_PROMPTS = [
f"Can you structure the information in the {format} format? Be sure to include precise information like any dates, times, names, and numerical values."
for format in ["JSON", "YAML", "TOML", "bullet"]
]
STRUCTURING_PROMPTS += [
"Can you structure the information in plain text that is easy to read? Be sure to include precise information like any dates, times, names, and numerical values."
]
# for chunking
CTX_AFFIXES = {
"google/gemma-3-1b-it": {
"prefix": [2, 105, 2364, 109], # <bos><start_of_turn>user\n\n\n
"suffix": [106, 107, 105, 4368, 107], # <end_of_turn>\n<start_of_turn>model\n
},
"google/gemma-2-2b-it": {
"prefix": [2, 106, 1645, 110], # <bos><start_of_turn>user\n\n\n
"suffix": [107, 108, 106, 2516, 108], # <end_of_turn>\n<start_of_turn>model\n
},
}
# approximate length of the datasets (train split)
# needed for streaming datasets
DS_LEN = {
"fw_qa_3_intx_pretrain": 90_000_000,
"fw_qa_3_intx_pretrain_medium": 40_000_000,
"fw_qa_3_intx_pretrain_mini": 100_000,
"fw_qa_3_mini": 100_000,
"fw_qa_3_medium": 121_000_000,
"fw_qa_3": 270_000_000,
"fw_qa_xl": 27_000_000,
"ctx_qa": 300_000,
"pwc": 240_000,
"hotpot_qa": 100_000,
"squad": 90_000,
"drop": 77_000,
"narrativeqa": 40_000,
"quoref": 11_000,
"ropes": 11_000,
"synthetic_convqa": 40_000,
}
LONGBENCH_TASKS = [
"longbench/narrativeqa",
"longbench/qasper",
"longbench/multifieldqa_en",
"longbench/multifieldqa_zh",
"longbench/hotpotqa",
"longbench/2wikimqa",
"longbench/musique",
"longbench/dureader",
"longbench/gov_report",
"longbench/qmsum",
"longbench/multi_news",
"longbench/vcsum",
]
LONGBENCH_E_TASKS = [
"longbench/qasper_e",
"longbench/multifieldqa_en_e",
"longbench/hotpotqa_e",
"longbench/2wikimqa_e",
"longbench/gov_report_e",
"longbench/multi_news_e",
]
DS_KWARGS = {
"context_numbers_2_10": dict(
train=dict(
path="json",
data_files=glob("data/raw_datasets/context_numbers_[2-9]/train.jsonl")
+ ["data/raw_datasets/context_numbers_10/train.jsonl"],
split="train",
),
validation=dict(
path="json",
data_files=glob("data/raw_datasets/context_numbers_[2-9]/val.jsonl")
+ ["data/raw_datasets/context_numbers_10/val.jsonl"],
split="train",
),
),
"hotpot_qa": dict(
train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train"),
validation=dict(
path="hotpotqa/hotpot_qa", name="fullwiki", split="validation[:900]"
),
test=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="validation"),
),
"hotpot_qa_tiny": dict(
train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:2000]"),
validation=dict(
path="hotpotqa/hotpot_qa", name="fullwiki", split="train[:900]"
),
),
"pwc": dict(
train=dict(path="sggetao/PwC", split="train"),
validation=dict(path="sggetao/PwC", split="test[:900]"),
@ -190,9 +54,6 @@ DS_KWARGS = {
train=dict(path="data/raw_datasets/squad", split="train"),
validation=dict(path="data/raw_datasets/squad", split="validation[:1000]"),
test=dict(path="data/raw_datasets/squad", split="validation"),
# train=dict(path="rajpurkar/squad", split="train[900:]"),
# validation=dict(path="rajpurkar/squad", split="train[:900]"),
# test=dict(path="rajpurkar/squad", split="validation"),
),
"squad_compact": dict(
train=dict(
@ -214,242 +75,6 @@ DS_KWARGS = {
"squad_assistant_ctx_no_passage": dict(
test=dict(path="data/raw_datasets/squad", split="validation"),
),
"fw_qa_xl": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_xl/*[!val].parquet"),
split="train",
),
validation=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_xl/*val.parquet"),
split="train",
),
# there is no explcit test for this ds
# used only for sanity check
test=dict(
path="parquet",
data_files=glob("data/raw_datasets/fw_qa_xl/*val.parquet"),
split="train",
),
),
"fw_qa_v2_2k_len_level_0_tiny": dict(
train=dict(
path="parquet",
data_files=glob(
"data/raw_datasets/fw_qa_v2/min_0_to_2000/000*level_0.parquet"
),
split="train",
),
),
"fw_qa_v2_2k_len_level_0_small": dict(
train=dict(
path="parquet",
data_files=glob(
"data/raw_datasets/fw_qa_v2/min_0_to_2000/00[02468]*level_0.parquet"
),
split="train",
),
),
"fw_qa_v2_2k_len_level_0": dict(
train=dict(
path="parquet",
data_files="data/raw_datasets/fw_qa_v2/min_0_to_2000/*level_0.parquet",
split="train",
),
validation=dict(
path="parquet",
data_files=glob(
"data/raw_datasets/fw_qa_v2/min_0_to_2000/*level_0_val.parquet"
),
split="train",
),
),
"fw_qa_v2_2k_len_level_1": dict(
train=dict(
path="parquet",
data_files="data/raw_datasets/fw_qa_v2/min_0_to_2000/*level_1.parquet",
split="train",
),
# validation=dict(
# path="parquet",
# data_files=glob(
# "data/raw_datasets/fw_qa_v2/min_0_to_2000/*level_1_val.parquet"
# ),
# 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/000*level_1.parquet"
),
split="train",
),
),
"fw_qa_v2_2k_len_level_1_small": dict(
train=dict(
path="parquet",
data_files=glob(
"data/raw_datasets/fw_qa_v2/min_0_to_2000/00[02468]*level_1.parquet"
),
split="train",
),
),
"fw_qa_v2_2k_len_level_2": dict(
train=dict(
path="parquet",
data_files="data/raw_datasets/fw_qa_v2/min_0_to_2000/*level_2.parquet",
split="train",
),
# validation=dict(
# path="parquet",
# data_files=glob(
# "data/raw_datasets/fw_qa_v2/min_0_to_2000/*level_2_val.parquet"
# ),
# split="train",
# ),
),
"fw_qa_v2_2k_len_level_2_tiny": dict(
train=dict(
path="parquet",
data_files=glob(
"data/raw_datasets/fw_qa_v2/min_0_to_2000/000*level_2.parquet"
),
split="train",
),
),
"fw_qa_v2_2k_len_level_2_small": dict(
train=dict(
path="parquet",
data_files=glob(
"data/raw_datasets/fw_qa_v2/min_0_to_2000/00[02468]*level_2.parquet"
),
split="train",
),
),
"fw_qa_v2_2k_len_level_3": dict(
train=dict(
path="parquet",
data_files="data/raw_datasets/fw_qa_v2/min_0_to_2000/*level_3.parquet",
split="train",
),
# validation=dict(
# path="parquet",
# data_files=glob(
# "data/raw_datasets/fw_qa_v2/min_0_to_2000/*level_3_val.parquet"
# ),
# split="train",
# ),
),
"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/000*level_3.parquet"
),
split="train",
),
),
"fw_qa_v2_2k_len_level_3_small": dict(
train=dict(
path="parquet",
data_files=glob(
"data/raw_datasets/fw_qa_v2/min_0_to_2000/00[02468]*level_3.parquet"
),
split="train",
),
),
# "fw_qa_3_small": dict(
# train=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_3/000*[!val].parquet"),
# split="train",
# ),
# ),
# "fw_qa_3_medium": dict(
# train=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_3/00[0-5]*[!val].parquet"),
# split="train",
# ),
# ),
# "fw_qa_3": dict(
# train=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_3/*[!val].parquet"),
# split="train",
# ),
# validation=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_3/*val.parquet"),
# split="train",
# ),
# ),
# "fw_qa_3_mini_pretrain": dict(
# train=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_intx_pretrain/000_00004.parquet"),
# split="train[:100000]",
# )
# ),
# "fw_qa_3_small_pretrain": dict(
# train=dict(
# path="parquet",
# data_files=glob(
# "data/raw_datasets/fw_qa_intx_pretrain/000_*[!val].parquet"
# ),
# split="train",
# )
# ),
# "fw_qa_3_small_aug_pretrain": dict(
# train=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_aug_pretrain/000_*[!val].parquet"),
# split="train",
# )
# ),
# "fw_qa_3_medium_pretrain": dict(
# train=dict(
# path="parquet",
# data_files=glob(
# "data/raw_datasets/fw_qa_intx_pretrain/00[0-5]*[!val].parquet"
# ),
# split="train",
# )
# ),
# "fw_qa_3_pretrain": dict(
# train=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_intx_pretrain/*[!val].parquet"),
# split="train",
# ),
# validation=dict(
# path="parquet",
# data_files=glob("data/raw_datasets/fw_qa_intx_pretrain/*val.parquet"),
# split="train",
# ),
# ),
"ctx_qa": dict(
train=dict(
path="parquet",
data_files=glob("data/raw_datasets/ctx_qa/*res.parquet"),
split="train",
),
validation=dict(
path="parquet",
data_files=glob("data/raw_datasets/ctx_qa/*val.parquet"),
split="train",
),
# there is no explcit test for this ds
# used only for sanity check
test=dict(
path="parquet",
data_files=glob("data/raw_datasets/ctx_qa/*val.parquet"),
split="train",
),
),
"drop": dict(
train=dict(
path="ucinlp/drop",
@ -472,27 +97,6 @@ DS_KWARGS = {
split="train",
)
),
"narrativeqa": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="narrativeqa",
split="train",
),
),
"quoref": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="quoref",
split="train",
),
),
# "ropes": dict(
# train=dict(
# path="nvidia/ChatQA-Training-Data",
# name="ropes",
# split="train",
# ),
# ),
"ropes": dict(
train=dict(
path="allenai/ropes",
@ -511,256 +115,6 @@ DS_KWARGS = {
split="train",
)
),
"synthetic_convqa": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="synthetic_convqa",
split="train",
),
),
"tatqa": dict(
train=dict(
path="nvidia/ChatQA-Training-Data",
name="tatqa",
split="train",
)
),
"booksum": dict(
train=dict(
path="kmfoda/booksum",
split="train",
),
validation=dict(
path="kmfoda/booksum",
split="validation",
),
test=dict(
path="kmfoda/booksum",
split="test",
),
),
"bookqa": dict(
train=dict(
path="parquet",
data_files="data/raw_datasets/fw_qa_v3/min_0_to_20000/booksum_level_0.parquet",
split="train",
)
),
"gov_report": dict(
train=dict(
path="ccdv/govreport-summarization",
split="train",
),
validation=dict(
path="ccdv/govreport-summarization",
split="validation",
),
test=dict(
path="ccdv/govreport-summarization",
split="test",
),
),
"gov_report_qa": dict(
train=dict(
path="parquet",
data_files="data/raw_datasets/fw_qa_v3/min_0_to_20000/gov_report_level_0.parquet",
split="train",
)
),
# "wikitext-2": dict(
# train=dict(
# path="EleutherAI/wikitext_document_level",
# name="wikitext-2-raw-v1",
# split="train",
# ),
# ),
# "wikitext-103": dict(
# train=dict(
# path="EleutherAI/wikitext_document_level",
# name="wikitext-103-raw-v1",
# split="train",
# ),
# ),
"gsm8k_assistant_ctx": dict(
train=dict(
path="openai/gsm8k",
name="main",
split="train[100:]",
),
validation=dict(
path="openai/gsm8k",
name="main",
split="train[:100]",
),
test=dict(
path="openai/gsm8k",
name="main",
split="test",
),
),
"gsm8k_negative": dict(
test=dict(
path="openai/gsm8k",
name="main",
split="test",
),
),
"gsm8k": dict(
train=dict(
path="openai/gsm8k",
name="main",
split="train[100:]",
),
validation=dict(
path="openai/gsm8k",
name="main",
split="train[:100]",
),
test=dict(
path="openai/gsm8k",
name="main",
split="test",
),
),
"gsm8k_fewshot": dict(
train=dict(
path="openai/gsm8k",
name="main",
split="train[100:]",
),
validation=dict(
path="openai/gsm8k",
name="main",
split="train[:100]",
),
test=dict(
path="openai/gsm8k",
name="main",
split="test",
),
),
"metaicl": dict(
validation=dict(
path="SakanaAI/metaicl",
split="train[:100]",
),
test=dict(
path="SakanaAI/metaicl",
split="train",
),
),
"openmathintx-2": dict(
train=dict(
path="nvidia/OpenMathInstruct-2",
split="train_1M[:100000]",
),
),
"opencoder-edu": dict(
train=dict(
path="OpenCoder-LLM/opc-sft-stage2",
name="educational_instruct",
split="train[900:]",
),
validation=dict(
path="OpenCoder-LLM/opc-sft-stage2",
name="educational_instruct",
split="train[:900]",
),
),
"triviaqa_retrieved": dict(
test=dict(
path="json",
data_files="data/eval/triviaqa/triviaqa_retrieved.jsonl",
split="train",
),
),
# for negative training and evaluating
"negative_nq": dict(
train=dict(
path="json",
data_files="data/raw_datasets/negative_natural_questions/train.jsonl",
split="train",
),
validation=dict(
path="json",
data_files="data/raw_datasets/negative_natural_questions/validation.jsonl",
split="train[:1000]",
),
test=dict(
path="json",
data_files="data/raw_datasets/negative_natural_questions/validation.jsonl",
split="train",
),
),
"facts": dict(
validation=dict(
path="csv",
data_files="data/eval/facts/examples.csv",
split="train[:400]",
),
test=dict(
path="csv",
data_files="data/eval/facts/examples.csv",
split="train",
),
),
# sft data
"openhermes": dict(
train=dict(path="HuggingFaceTB/smoltalk", name="openhermes-100k", split="train")
),
"smol-magpie-ultra": dict(
train=dict(
path="HuggingFaceTB/smoltalk", name="smol-magpie-ultra", split="train"
)
),
"smol-rewrite": dict(
train=dict(path="HuggingFaceTB/smoltalk", name="smol-rewrite", split="train")
),
"smol-summarize": dict(
train=dict(path="HuggingFaceTB/smoltalk", name="smol-summarize", split="train")
),
"systemchats": dict(
train=dict(path="HuggingFaceTB/smoltalk", name="systemchats-30k", split="train")
),
"ctx-q-gsm": dict(train=dict(path="letta-ai/stateful-gsm-symbolic", split="train")),
"automathtext_web": dict(
dict(
train=dict(
path="math-ai/AutoMathText", name="web-0.50-to-1.00", split="train"
)
)
),
"automathtext_code_python": dict(
dict(
train=dict(
path="math-ai/AutoMathText",
name="code-python-0.50-to-1.00",
split="train",
)
)
),
"automathtext_code_notebook": dict(
dict(
train=dict(
path="math-ai/AutoMathText",
name="code-jupyter-notebook-0.50-to-1.00",
split="train",
)
)
),
"automathtext_arxiv": dict(
dict(
train=dict(
path="math-ai/AutoMathText", name="arxiv-0.50-to-1.00", split="train"
)
)
),
"codeparrot": dict(
dict(train=dict(path="transformersbook/codeparrot", split="train[:2000000]"))
),
"programming_books": dict(
dict(train=dict(path="open-phi/programming_books_llama", split="train"))
),
}
# add ctx_numbers
@ -804,31 +158,23 @@ for ds_name in LONGBENCH_TASKS + LONGBENCH_E_TASKS:
CLOSED_QA_DATASETS = {
"longbench/narrativeqa",
"longbench/qasper",
"longbench/multifieldqa_en",
"longbench/hotpotqa",
"longbench/2wikimqa",
"longbench/musique",
"hotpot_qa",
"squad",
"squad_negative",
"squad_assistant_ctx",
"squad_negative_no_passage",
"squad_assistant_ctx_no_passage",
"triviaqa_retrieved",
"negative_nq",
"ropes",
"drop",
}
MULTI_ANSWER_DATASETS = {
"longbench/narrativeqa",
"longbench/qasper",
"longbench/multifieldqa_en",
"longbench/hotpotqa",
"longbench/2wikimqa",
"longbench/musique",
"squad",
"squad_negative",
"squad_assistant_ctx",
@ -847,9 +193,6 @@ for ds_name in list(MULTI_ANSWER_DATASETS):
if ds_name.startswith("longbench/"):
MULTI_ANSWER_DATASETS.add(f"{ds_name}_e")
GSM8K_DATASETS = {"gsm8k", "gsm8k_fewshot", "gsm8k_assistant_ctx", "gsm8k_negative"}
MULTI_CHOICE_DATASETS = {"metaicl"}
# for training closed qa datasets, e.g., hotpot_qa, squad, etc.
CLOSED_QA_INTX_TEMPLATES = [
"Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}",
@ -888,33 +231,10 @@ EVAL_INTX_TEMPLATES = {
"squad_assistant_ctx": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
"squad_negative_no_passage": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
"squad_assistant_ctx_no_passage": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
# retrieved noisy ctx?
"triviaqa_retrieved": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
# doc w/ distractors qa
"hotpot_qa": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
# unrelated ctx-qa
"negative_nq": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
# summary
"booksum": "Write a summary for the story. Make it a *one-paragraph monologue*. Don't use bullet point. Do not output anything else beyond the summary.",
"gov_report": "Write a summary for the report. Make it a *one-paragraph monologue*. Don't use bullet point. Do not output anything else beyond the summary.",
# longbench
"longbench/qasper": 'Answer the question as concisely as you can, using a single phrase or sentence if possible.\nIf the question cannot be answered based on the information in the article, write "unanswerable".\nIf the question is a yes/no question, answer "yes", "no", or "unanswerable". Do not provide any explanation.\n\nQuestion: {input}',
"longbench/narrativeqa": "Answer the question as concisely as you can, using a single phrase if possible. Do not provide any explanation.\n\nQuestion: {input}",
"longbench/multifieldqa_en": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
# "longbench/multifieldqa_zh": "阅读以下文字并用中文简短回答：\n\n现在请基于上面的文章回答下面的问题，只告诉我答案，不要输出任何其他字词。\n\n问题：{input}",
"longbench/hotpotqa": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
"longbench/2wikimqa": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
"longbench/musique": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
#
# "longbench/dureader": "请基于给定的文章回答下述问题\n\n请基于上述文章回答下面的问题。\n\n问题：{input}",
"longbench/gov_report": "You are given a report by a government agency. Write a one-page summary of the report.",
"longbench/qmsum": "You are given a meeting transcript and a query containing a question or instruction. Answer the query based on the above meeting transcript.\n\nQuery: {input}",
"longbench/multi_news": "You are given several news passages. Write a one-page summary of all the news.",
# "longbench/vcsum": "下面有一段会议记录，请你阅读后，写一段总结，总结会议的内容。",
}
for ds_name in LONGBENCH_E_TASKS:
EVAL_INTX_TEMPLATES[ds_name] = EVAL_INTX_TEMPLATES[ds_name[:-2]]
# for ds_name in DS_KWARGS:
# if ds_name not in EVAL_INTX_TEMPLATES:
# EVAL_INTX_TEMPLATES[ds_name] = "{input}"

85
src/ctx_to_lora/data/packing.py
View file

@ -7,7 +7,7 @@ import pprint
import numpy as np
from ctx_to_lora.utils import check_is_iterable, concat_list, setup_logging
from ctx_to_lora.utils import check_is_iterable, concat_list
logger = logging.getLogger()
@ -198,7 +198,6 @@ def pack_batch(
)
# Pack groups
# n_samples = len(boundaries) - 1
packed_batch = {
"ctx_ids": [],
"ctx_position_ids": [],
@ -218,8 +217,6 @@ def pack_batch(
for idx_pair in idx_pairs:
start_idx, end_idx = idx_pair[0], idx_pair[1]
# start_idx = start_indices[i]
# end_idx = start_indices[i + 1]
group_items = {
"ctx_ids": batch["ctx_ids"][start_idx:end_idx],
"input_ids": batch["input_ids"][start_idx:end_idx],
@ -250,16 +247,11 @@ def pack_batch(
ctx_efficiency = len(packed_item["ctx_ids"]) / max_packed_ctx_len
ctx_packing_efficiency_ratios.append(ctx_efficiency)
# # Calculate total original tokens
# total_original_inp_tokens = sum(inp_lens)
# total_original_ctx_tokens = sum(ctx_lens)
# total_original_tokens = total_original_inp_tokens + total_original_ctx_tokens
# # Calculate length statistics
# Calculate length statistics
packed_inp_lens_arr = np.array([len(x) for x in packed_batch["input_ids"]])
packed_ctx_lens_arr = np.array([len(x) for x in packed_batch["ctx_ids"]])
# # Log performance statistics
# Log performance statistics
avg_inp_packing_efficiency = (
np.mean(packing_efficiency_ratios) if packing_efficiency_ratios else 0
)
@ -296,74 +288,3 @@ def pack_batch(
logging.debug(f"Packing stats:\n{pprint.pformat(packing_stats, indent=2)}")
return packed_batch
if __name__ == "__main__":
from datasets import disable_caching
from ctx_to_lora.data.processing import get_tokenized_dataset
from ctx_to_lora.model_loading import get_model_and_tokenizer
setup_logging("tmp/packing_debug.log", debug=True)
logger.info("Starting packing script...")
model, tokenizer = get_model_and_tokenizer(
"google/gemma-2-2b-it",
train=True,
requires_grad=False,
)
base_model_max_len = model.base_model.config.max_position_embeddings
tokenizer_kwargs = {"max_length": base_model_max_len} # not used
ctx_tokenizer_kwargs = {"max_length": base_model_max_len} # not used for now
ds = get_tokenized_dataset(
ds_name="self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_2_tiny",
split="train",
max_qas_len=2048,
base_model_max_len=model.base_model.config.max_position_embeddings,
tokenizer=tokenizer,
tokenizer_kwargs={},
ctx_model_max_len=model.base_model.config.max_position_embeddings,
ctx_tokenizer=tokenizer,
ctx_tokenizer_kwargs={},
add_ctx_to_chat=False,
add_repeat_prompt=False,
repeat_prob=0.0,
add_negative_prompt=False,
use_kl_loss=False,
)
# ds.set_format("torch")
print(ds)
disable_caching()
packed_ds = ds.map(
pack_batch,
fn_kwargs={
"max_packed_inp_len": 2**13, # 8k
"max_packed_ctx_len": 2**14, # 16k
},
batched=True,
batch_size=100_000,
remove_columns=ds.column_names,
num_proc=4,
)
print(packed_ds)
# packed_ds = PackedDataset(
# dataset=ds,
# tokenizer=tokenizer,
# max_input_length=base_model_max_len,
# max_packed_length=base_model_max_len * 2,
# # max_packed_size=3, # maximum number of data points being packed
# )
# packed_ds.stat()
# packed_ds = Dataset.from_list(packed_ds)
# packed_ds.set_format("torch")
# print(packed_ds)
# print(ds[0])
# print(packed_ds[0])
orig_seq = ds[0]["ctx_ids"]
packed_seq = packed_ds[0]["ctx_ids"][: len(orig_seq)]
print(orig_seq, packed_seq)
print(orig_seq == packed_seq)

324
src/ctx_to_lora/data/preprocessing_fn.py
View file

@ -4,7 +4,6 @@ from collections.abc import Callable
from typing import Any
from ctx_to_lora.data.definitions import CLOSED_QA_INTX_TEMPLATES, EVAL_INTX_TEMPLATES
from ctx_to_lora.utils import concat_list
logger = logging.getLogger()
@ -43,48 +42,6 @@ def get_preprocessing_fn(
# already processed data, do nothing
return f
if "fw_qa_v2" in ds_name:
def f(sample):
# get questions/answers from all levels in the ds
q_cols = [col for col in sample.keys() if col.startswith("prompts_level")]
r_cols = [col for col in sample.keys() if col.startswith("responses_level")]
questions = concat_list([sample[col] for col in q_cols])
responses = concat_list([sample[col] for col in r_cols])
min_len = min(len(questions), len(responses))
if min_len == 0:
return {
"context": None,
"prompts": None,
"responses": None,
}
return {
"context": sample["context"],
"prompts": questions[:min_len],
"responses": responses[:min_len],
}
elif ds_name.startswith("ctx_qa"):
def f(sample):
ctx = sample["context"]
# remove unintended prompt
ctx = ctx.split("\n...\n\n### Code ###\n")[-1]
sample["context"] = ctx
return sample
elif ds_name == "facts":
def f(sample):
prompt = sample["system_instruction"] + sample["user_request"]
return {
"context": sample["context_document"],
"prompt": prompt,
"response": "PLACEHOLDER. PLEASE GENERATE SELF-GEN DATA.",
}
elif ds_name.startswith("longbench"):
def f(sample):
@ -94,24 +51,6 @@ def get_preprocessing_fn(
"response": sample["answers"][0],
}
elif ds_name == "negative_nq":
def f(sample):
return {
"context": sample["context"],
"prompt": sample["prompt"],
"response": sample["answer"],
}
elif ds_name == "triviaqa_retrieved":
# only used for eval
def f(sample):
return {
"context": sample["context"],
"prompt": sample["prompt"],
"response": sample["answer"],
}
elif ds_name == "pwc" or ds_name == "pwc_tiny":
# original pwc
def f(sample):
@ -121,22 +60,6 @@ def get_preprocessing_fn(
"response": sample["answer"],
}
elif ds_name.startswith("hotpot_qa"):
def f(sample):
txt = ""
for p in sample["context"]["sentences"]:
txt += " " + "".join(p)
q = sample["question"]
prompt = closed_qa_prompting(q) if not is_eval else q
return {
"context": txt.strip(),
"prompt": prompt,
"response": sample["answer"],
}
elif ds_name == "squad":
# original squad
def f(sample):
@ -211,253 +134,6 @@ def get_preprocessing_fn(
q = closed_qa_prompting(q) if not is_eval else q
return {"context": ctx, "prompt": q, "response": response}
elif ds_name in ["narrativeqa", "quoref", "tatqa"]: # , "ropes"]:
def f(sample):
response = sample["answers"][0]
if isinstance(response, list):
response = response[0]
q = sample["messages"][-1]["content"]
return {
"context": sample["document"],
"prompt": q,
"response": response,
}
elif ds_name == "synthetic_convqa":
def f(sample):
response = sample["answers"][0]
if isinstance(response, list):
response = response[0]
return {
"context": sample["document"],
"prompt": sample["messages"][-1]["content"],
"response": response,
}
elif ds_name == "booksum":
# TODO: sample from SUMM_PROMPTS once improved data gen branch is merged
prompt_templates = [
"Summarization the provided text.",
"# Summary",
"### Summary",
"Summary of the text",
]
# TODO: use these templates for training
# analysis_templates = []
# summary_templates = []
def f(sample):
return {
"context": sample["chapter"],
"prompt": "Summarize the provided text.",
"response": sample["summary_text"],
}
elif ds_name == "gov_report":
def f(sample):
return {
"context": sample["report"],
"prompt": "Summarize the provided text.",
"response": sample["summary"],
}
elif "wikitext" in ds_name:
def f(sample):
return {
"context": sample["page"],
"prompt": "PLAECHOLDER",
"response": "PLAECHOLDER",
}
elif ds_name == "openmathintx-2":
def f(sample):
return {
"context": sample["problem"],
"prompt": sample["problem"],
"response": sample["generated_solution"],
}
elif ds_name == "gsm8k":
instruction_prompt = (
"Provide the reasoning steps before giving the final integer answer."
)
def f(sample):
question = sample["question"]
last_dot_index = question.rfind(".")
context = question[: last_dot_index + 1]
prompt = question[last_dot_index + 1 :].lstrip()
return {
"prompt": prompt + "\n" + instruction_prompt,
"context": context,
"response": sample["answer"],
}
# return {
# "context": sample["question"],
# "prompt": sample["question"] + "\n" + instruction_prompt,
# "response": sample["answer"],
# }
elif ds_name == "gsm8k_negative":
with open("data/gutenburg_sample.txt") as f:
gutenburg_sample = f.read()
def f(sample):
return {
"prompt": sample["question"],
"context": gutenburg_sample,
"response": sample["answer"],
}
elif ds_name == "gsm8k_assistant_ctx":
def f(sample):
return {
"prompt": sample["question"],
"context": "You are a useful AI assistant.",
"response": sample["answer"],
}
# return {
# "context": sample["question"],
# "prompt": sample["question"] + "\n" + instruction_prompt,
# "response": sample["answer"],
# }
elif ds_name == "gsm8k_fewshot":
N_FEW_SHOT_EXAMPLES = 5
few_shot_examples = [
{
"question": "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?",
"target": "Natalia sold 48/2 = 24 clips in May.\nNatalia sold 48+24 = 72 clips altogether in April and May.\n#### 72",
},
{
"question": "Weng earns $12 an hour for babysitting. Yesterday, she just did 50 minutes of babysitting. How much did she earn?",
"target": "Weng earns 12/60 = 0.2 per minute.\nWorking 50 minutes, she earned 0.2 x 50 = $10.\n#### 10",
},
{
"question": "Betty is saving money for a new wallet which costs $100. Betty has only half of the money she needs. Her parents decided to give her $15 for that purpose, and her grandparents twice as much as her parents. How much more money does Betty need to buy the wallet?",
"target": "In the beginning, Betty has only 100 / 2 = 50.\nBetty's grandparents gave her 15 * 2 = $30.\nThis means, Betty needs 100 - 50 - 30 - 15 = 5 more.\n#### 5",
},
{
"question": "Julie is reading a 120-page book. Yesterday, she was able to read 12 pages and today, she read twice as many pages as yesterday. If she wants to read half of the remaining pages tomorrow, how many pages should she read?",
"target": "Maila read 12 x 2 = 24 pages today.\nSo she was able to read a total of 12 + 24 = 36 pages since yesterday.\nThere are 120 - 36 = 84 pages left to be read.\nSince she wants to read half of the remaining pages tomorrow, then she should read 84/2 = 42 pages.\n#### 42",
},
{
"question": "James writes a 3-page letter to 2 different friends twice a week. How many pages does he write a year?",
"target": "He writes each friend 3*2=6 pages a week\nSo he writes 6*2=12 pages every week\nThat means he writes 12*52=624 pages a year\n#### 624",
},
]
few_shot_examples_txt = ""
for i, example in enumerate(few_shot_examples[:N_FEW_SHOT_EXAMPLES]):
few_shot_examples_txt += f"# Example {i + 1}:\nQuestion: {example['question']}\nSolution: {example['target']}\n\n"
def f(sample):
instruction_prompt = (
"Provide the reasoning steps before giving the final integer answer."
)
question = sample["question"]
# Find the last sentence as prompt (removing leading space)
return {
"prompt": question + "\n" + instruction_prompt,
"context": few_shot_examples_txt,
"response": sample["answer"],
}
elif ds_name == "metaicl":
def f(sample):
prompt = f"You are a classifier. Reply must be of the following format: 'Answer: $LETTER' (without quotes) where LETTER is one of {sample['options']['options']}.\nInput:\n{sample['input']}"
context = "Examples\n"
for example in sample["examples"]:
context += f"Input: {example['input']}\nAnswer: {example['output']}\n\n"
return {
"context": context,
"prompt": prompt,
"response": sample["output"],
}
elif "opencoder-edu" in ds_name:
def f(sample):
return {
"context": sample["instruction"],
"prompt": sample["instruction"],
"response": "```python\n" + sample["code"].strip() + "\n```",
}
elif "openhermes" == ds_name:
# system prompt is from a set of predefined prompts → query
# user prompt is the content → ctx
def f(sample):
messages = sample["messages"]
if messages[0]["role"] != "system":
return {"context": None, "prompt": None, "response": None}
ctx = messages[1]["content"] # user input
q = messages[0]["content"] # system intx
response = messages[2]["content"]
return {"context": ctx, "prompt": q, "response": response}
elif "smol-magpie-ultra" == ds_name or "systemchats" == ds_name:
# ctx → all history up to the last turn
# q → last user input
def f(sample):
chat_history = chat_to_str(sample["messages"][:-2])
q = sample["messages"][-2]["content"]
response = sample["messages"][-1]["content"]
return {"context": chat_history, "prompt": q, "response": response}
elif "smol-rewrite" == ds_name or "smol-summarize" == ds_name:
# user content is a paragraph of text → ctx
# system prompt is a paraphrase/summarize instruction → q
def f(sample):
ctx = sample["messages"][1]["content"]
q = sample["messages"][0]["content"]
response = sample["messages"][2]["content"]
return {"context": ctx, "prompt": q, "response": response}
elif "ctx-q-gsm" == ds_name:
def f(sample):
ctx = sample["stateful_gsm_symbolic_context"]
q = sample["stateful_gsm_symbolic_question"]
response = sample["answer"]
return {"context": ctx, "prompt": q, "response": response}
elif ds_name.startswith("automathtext"):
def f(sample):
return {
"context": sample["text"],
"prompt": "PLACEHOLDER. PLEASE GENERATE QUESTIONS FOR THIS DATASET.",
"response": "PLACEHOLDER PLEASE GENERATE QUESTIONS FOR THIS DATASET.",
}
elif "codeparrot" == ds_name:
def f(sample):
return {
"context": sample["content"],
"prompt": "PLACEHOLDER. PLEASE GENERATE QUESTIONS FOR THIS DATASET.",
"response": "PLACEHOLDER. PLEASE GENERATE QUESTIONS FOR THIS DATASET.",
}
elif "programming_books" == ds_name:
def f(sample):
return {
"context": sample["markdown"],
"prompt": "PLACEHOLDER. PLEASE GENERATE QUESTIONS FOR THIS DATASET.",
"response": "PLACEHOLDER. PLEASE GENERATE QUESTIONS FOR THIS DATASET.",
}
if is_eval and (ds_name in EVAL_INTX_TEMPLATES):
prompt_template = EVAL_INTX_TEMPLATES[ds_name]

5
src/ctx_to_lora/data/pretrain.py
View file

@ -1,5 +0,0 @@
# when generating dont forget to remove dup contexts
# merge QAs from dup contexts to construct
# a doc with QAs at the end
# maybe mix with some that doesn;t have QAs
# TODO: implement

38
src/ctx_to_lora/data/processing.py
View file

@ -228,10 +228,7 @@ def get_tokenized_dataset(
f"`max_qas_len` should be <= {base_model_max_len=}, got {max_qas_len=}"
)
logger.info(f"Loading dataset {ds_name} with split {split}...")
# TODO: fix this to allow using both ctx_ids and add_ctx_ids
need_ctx_ids = (
ctx_model_max_len is not None
) # not add_ctx_to_chat and bool(ctx_model_max_len)
need_ctx_ids = ctx_model_max_len is not None
load_and_process_kwargs = dict(
ds_name=ds_name,
@ -368,15 +365,9 @@ def construct_and_tokenize_ctx_qa(
# add `input_ids`, `attention_mask`, `labels`
os.environ["TOKENIZERS_PARALLELISM"] = "true"
logging.debug("Tokenizing inputs")
# HACK: check magic num niah dataset
prepend_bos = False # see known issues in readme
# sample = ds[0]
# inp_txt = sample["prompts"][0]
# if inp_txt == "What is the special magic number? Reply with only the number.":
# logger.info("Detected Niah dataset, not prepending BOS")
# prepend_bos = False
tokenized_ds = ds.map(
get_sft_prompt_formatting_fn(tokenizer, prepend_bos),
get_sft_prompt_formatting_fn(tokenizer),
batched=True,
batch_size=100_000,
)
@ -402,15 +393,6 @@ def construct_and_tokenize_ctx_qa(
remove_columns=["context"],
)
# # TODO: this can be removed once we implement ctx chunking for training
# if is_train:
# # drop
# tokenized_ds = tokenized_ds.filter(
# len_filter,
# fn_kwargs={"max_length": ctx_model_max_len, "keys": ["ctx_ids"]},
# num_proc=16,
# )
# ctx chunking
# Ideally we want to chunk the raw text directly
# however, since the contexts are tokenized during self-gen
@ -519,7 +501,6 @@ def get_labels_from_input_ids(sample: dict[str, Any]) -> dict[str, Any]:
def get_sft_prompt_formatting_fn(
tokenizer: PreTrainedTokenizerBase,
prepend_bos: bool,
) -> Callable[[dict[str, Any]], dict[str, Any]]:
"""
Get a function that formats examples for supervised fine-tuning.
@ -559,14 +540,6 @@ def get_sft_prompt_formatting_fn(
return_dict=True,
)
if prepend_bos:
# HACK: add bos at the beginning (see `Known Issues` in README.md)
for i, (tok_ids, assistant_masks) in enumerate(
zip(tokens["input_ids"], tokens["assistant_masks"])
):
tokens["input_ids"][i] = [tokenizer.bos_token_id] + tok_ids
tokens["assistant_masks"][i] = [0] + assistant_masks
labels = []
for tok_ids, masks in zip(tokens["input_ids"], tokens["assistant_masks"]):
o = [id_ if mask else IGNORE_INDEX for id_, mask in zip(tok_ids, masks)]
@ -765,7 +738,6 @@ def split_too_long_qas(
max_qas_per_sample < 0 or len(input_ids[i]) <= max_qas_per_sample
):
# No need to split - add entire sample
# logger.debug(f"Sample {i} is within limits, adding as is.")
for k in samples:
out[k].append(samples[k][i])
continue
@ -815,9 +787,6 @@ def split_too_long_qas(
new_target_indices.append(target_indices)
else:
# Current batch is full, save it and start new batch
# logger.debug(
# f"sample {i}: adding batch with {len(new_input_ids)} sequences"
# )
if new_input_ids: # Only add non-empty batches
add_batch(
new_input_ids,
@ -999,7 +968,6 @@ def pack(
train_ds = interleave_datasets(
list(ds_dict.values()),
probabilities=get_ds_prob(train_ds_lens, total_samples),
# probabilities=[l / total_samples for l in train_ds_lens],
seed=seed,
stopping_strategy="all_exhausted",
)

114
src/ctx_to_lora/data/q_generation_prompts.py
View file

@ -1,114 +0,0 @@
import random
SYSTEM_PROMPT = (
"You are a creative and helpful assistant. "
"You will be given a context and you need to generate questions from the given context. "
"**DO NOT** hallucinate or make up information."
)
PROMPT_TEMPLATE = (
"### Context ###\n{context}\n\n"
"### Instruction ###\n{instruction}\n\n"
"### Rules ###\n"
"Phrases like 'based on the provided context', 'according to the context', etc., must not to appear in your response.\n\n"
# "2. The questions should not overlap. They should be diverse, covering many aspects of the context.\n"
# "3. Do not give away too much information in the questions. For example, ask 'Who is X?' instead of 'Who is X that did Y?' when Y is clear from the context.\n"
# "4. Ignore the text formatting of the context, e.g., bold, italic, underline, etc.\n"
# "5. Ignore typos, spacing, and grammatical errors in the context.\n"
# "6. Always use proper grammar and punctuation.\n"
# "7. Try to use different question forms and styles.\n"
"### Output Format ###\n"
"The question/instruction/task/request should be in the following format:\n\n"
"Message: {{question}}\n\n"
"Use this output template regardless of the type of the output."
)
# taken from https://github.com/HazyResearch/cartridges/blob/2ac563d79c2f3367a9e780a7bb0b3cf3039a8d50/cartridges/data/resources.py#L195
# def get_structuring_seed_prompt() -> str:
# DATA_FORMATS = [
# "JSON",
# "YAML",
# "TOML",
# "INI",
# "XML",
# "plain text",
# ]
# data_format = random.choice(DATA_FORMATS)
# EXAMPLES = [
# dedent(f"""
# Can you structure the information in the context in the following format: {data_format}? Be sure to include precise information like any dates, times, names, and numerical values.
# """).strip(),
# ]
# example = random.choice(EXAMPLES)
# return dedent(f"""
# Please generate a single chat message instructing an LLM to structure the information in {data_format}. The message can follow the following template, filling in details from the context:
# '{example}'
# """).strip()
# def get_summarization_seed_prompt() -> str:
# prompts = [
# dedent("""
# Please generate a single chat message instructing an LLM to summarize part of the context.
# Make sure the instruction is very explicit about the section of the context that you want to summarize.
# Include details (ids, names, titles, dates, etc.) that make it clear what you are asking about.
# """).strip(),
# dedent("""
# Please generate a single chat message instructing an LLM to summarize a section.
# Make sure the instruction is explicit about the section that should be summarized and the document it is from.
# """).strip(),
# ]
# prompt = random.choice(prompts)
# return prompt
def get_question_seed_prompt() -> str:
prompts = [
(
"Generate a question for an LLM that will test its knowledge of the information in the context above. "
"In your question be sure to include details (ids, names, titles, dates, etc.) that make it clear what you are asking about. "
"Output only a single question. Do NOT include any other text or explanation other than the question."
),
(
"Generate a message for an LLM that will test its knowledge of the information in the context above. "
"Be sure to include details (ids, names, titles, dates, etc.) in the question so that it can be answered without access to the context (i.e. closed-book setting). "
"Output only a single question. Do NOT include any other text or explanation other than the question."
),
(
"You are helping to quiz a user about the information in the context. "
"Please generate a question about the subsection of the context above. "
"Be sure to include details (ids, names, titles, dates, etc.) in the question to make it clear what you are asking about. "
"Answer only with the question, do not include any other text."
),
]
prompt = random.choice(prompts)
return prompt
def get_use_case_seed_prompt() -> str:
prompt = (
"Your primary goal is to think about practical, real-world tasks or applications that someone could achieve using the knowledge contained within the provided context. "
"Consider how a user might want to apply this information in a real-world scenario, not just recall it. Put yourself into someone else's shoes and think how you might use the information. "
"After considering potential use cases, your task will be to generate an instruction or task that reflects one of these downstream applications. "
"This instruction or task should be something a user, who has access to this context, might ask when trying to accomplish their specific goal. "
"Output only a single instruction or task. Do NOT include any other text or explanation."
)
return prompt
def get_creative_seed_prompt() -> str:
prompt = (
"You are having a creative open-ended conversation inspired by the information in the context. "
"Please generate an open question for your conversation partner to start off the discussion. "
"Answer only with the question, do not include any other text."
)
return prompt
def get_generic_seed_prompt() -> str:
return "Please generate a single chat message to begin a conversation about the information in the context. Make an open-ended request or provide a task."

108
src/ctx_to_lora/eval_utils.py
View file

@ -3,7 +3,6 @@ import logging
import os
import re
import string
import sys
import time
from argparse import Namespace
from collections import Counter, defaultdict
@ -14,7 +13,7 @@ import numpy as np
import pandas as pd
import torch
import yaml
from datasets import disable_caching, load_dataset
from datasets import disable_caching
from peft import get_peft_model
from transformers import (
PreTrainedModel,
@ -23,20 +22,16 @@ from transformers import (
Trainer,
set_seed,
)
from transformers.trainer_utils import find_executable_batch_size
from ctx_to_lora.data.collator import eval_collator, generation_collator
from ctx_to_lora.data.definitions import (
CLOSED_QA_DATASETS,
CTX_AFFIXES,
GSM8K_DATASETS,
LONGBENCH_E_TASKS,
LONGBENCH_TASKS,
MULTI_ANSWER_DATASETS,
MULTI_CHOICE_DATASETS,
)
from ctx_to_lora.data.processing import (
get_ds_kwargs,
get_tokenized_dataset,
load_answers,
)
@ -44,8 +39,6 @@ from ctx_to_lora.data.self_gen_template import SELF_QA_INTX
from ctx_to_lora.metrics import (
LENGTH_BINS,
Evaluator,
compute_gsm8k_acc,
compute_macro_f1_score,
compute_metrics,
compute_per_token_acc,
compute_perplexity,
@ -58,7 +51,6 @@ from ctx_to_lora.model_loading import (
get_model_and_tokenizer,
get_tokenizer,
)
from ctx_to_lora.modeling import hypernet
from ctx_to_lora.modeling.context_distillation import CtxDistillModel
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
from ctx_to_lora.modeling.llm_lingua import LLMLinguaModel
@ -75,81 +67,9 @@ from ctx_to_lora.utils import clear_gpu, concat_list, get_run_name, setup_loggin
logger = logging.getLogger()
# bandaid for loading old models (before restructure)
sys.modules["ctx_to_lora.modeling_utils"] = hypernet
# from https://github.com/huggingface/transformers/issues/32215#issuecomment-2386274138
class EvaluateAutoFindBatchSizeMixin:
def evaluate(self, *args, **kwargs):
print(f"Auto-finding batch size for evaluate")
# can't set self.args.eval_batch_size, it is computed from self.args.per_device_eval_batch_size
super_evaluate = super().evaluate # super cannot be called inside lambda
# assert not self.args.dataloader_persistent_workers, "dataloader_persistent_workers not supported because eval batch size does not change"
def new_evaluate(batch_size):
old_eval_batch_size = self.args.eval_batch_size
self.args.per_device_eval_batch_size = batch_size
if self.args.eval_batch_size != old_eval_batch_size:
if self.args.dataloader_persistent_workers and hasattr(
self, "_eval_dataloaders"
):
del self._eval_dataloaders # to shutdown the workers
logger.info(
f"Setting eval_batch_size to {self.args.eval_batch_size} (per device: {batch_size}) from old {old_eval_batch_size}"
)
return super_evaluate(*args, **kwargs)
# don't reset at end, but keep batch size for the next time evaluate is called
return find_executable_batch_size(
new_evaluate,
self.args.per_device_eval_batch_size,
auto_find_batch_size=self.args.auto_find_batch_size,
)()
def predict(self, *args, **kwargs):
print(f"Auto-finding batch size for predict")
# can't set self.args.eval_batch_size, it is computed from self.args.per_device_eval_batch_size
super_predict = super().predict # super cannot be called inside lambda
def new_predict(batch_size):
old_eval_batch_size = self.args.eval_batch_size
self.args.per_device_eval_batch_size = batch_size
if self.args.eval_batch_size != old_eval_batch_size:
if self.args.dataloader_persistent_workers and hasattr(
self, "_eval_dataloaders"
):
del self._eval_dataloaders # to shutdown the workers
logger.info(
f"Setting eval_batch_size to {self.args.eval_batch_size} (per device: {batch_size}) from old {old_eval_batch_size}"
)
return super_predict(*args, **kwargs)
return find_executable_batch_size(
new_predict,
self.args.per_device_eval_batch_size,
auto_find_batch_size=self.args.auto_find_batch_size,
)()
class CustomSeq2SeqTrainer(Seq2SeqTrainer, EvaluateAutoFindBatchSizeMixin):
pass
class CustomTrainer(Trainer, EvaluateAutoFindBatchSizeMixin):
pass
# ============================================================================
# Metrics and Evaluation Utilities
# ============================================================================
# from https://gist.github.com/cloneofsimo/8abd0284d4738f28f04200628f9a83f5
# https://github.com/Nordth/humanize-ai-lib/blob/main/src/humanize-string.ts
_HIDDEN_CHARS = re.compile(
r"[\u00AD\u180E\u200B-\u200F\u202A-\u202E\u2060\u2066-\u2069\uFEFF]"
)
@ -646,25 +566,6 @@ def eval_generation(
for k, v in qa_f1_metric.items():
eval_result.metrics[f"{split_name}_{k}"] = v
eval_result.metrics[f"{split_name}_num_samples_{k}"] = n
elif ds_name in GSM8K_DATASETS:
print("Computing GSM8K Accuracy")
gsm8k_acc_metric, per_sample_metric = compute_gsm8k_acc(
pred_texts, label_texts
)
for k, v in gsm8k_acc_metric.items():
eval_result.metrics[f"{split_name}_{k}"] = v
elif ds_name in MULTI_CHOICE_DATASETS:
print("Computing Multi-Choice Accuracy")
original_ds = original_datasets[ds_name]
multi_choice_acc_metric, per_sample_metric = compute_macro_f1_score(
pred_texts,
label_texts,
[x["options"] for x in original_ds["options"]],
original_ds["task"],
)
for k, v in multi_choice_acc_metric.items():
eval_result.metrics[f"{split_name}_{k}"] = v
else:
rouge_metrics, per_sample_metric = compute_rouge(pred_texts, label_texts)
for k, v in rouge_metrics.items():
@ -967,11 +868,6 @@ def evaluate(
# handling cases where there are multiple answers
if ds_name in MULTI_ANSWER_DATASETS:
answers[ds_name] = load_answers(ds_name, split)
if ds_name in MULTI_CHOICE_DATASETS:
ds_kwargs = get_ds_kwargs(ds_name, split)
original_datasets[ds_name] = load_dataset(
**ds_kwargs, trust_remote_code=True
)
print(f"Datasets: {datasets}")
print(f"Answers: {answers}")
@ -1082,7 +978,7 @@ def evaluate(
)
out.update(metrics)
else:
eval_trainer = CustomSeq2SeqTrainer(**trainer_kwargs)
eval_trainer = Seq2SeqTrainer(**trainer_kwargs)
for ds_name, ds in datasets.items():
metrics = eval_generation(
eval_trainer,

Some files were not shown because too many files have changed in this diff Show more