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

Prepare public release

Browse source
This commit is contained in:
Manan Suri 2026-06-03 01:27:45 +00:00
parent baa85db4d5
commit 7428582dea
83 changed files with 9606 additions and 14641 deletions
Download patch file Download diff file Expand all files Collapse all files

9
.gitignore vendored
View file

@ -7,6 +7,15 @@ results/
datasets/
llm-comparator/
trained_d2l/
checkpoints/
runs/
outputs/
data/video2lora/
hf_key*
wandb_key*
pat_token*
cookie.txt
www.youtube.com_cookies.txt
# Ignore data files but not scripts
/data/processed_datasets/
/data/raw_datasets/

1
LICENSE
View file

@ -1,6 +1,7 @@
MIT License
Copyright (c) 2026 Sakana AI
Copyright (c) 2026 Video2LoRA
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal

241
README.md
View file

@ -1,108 +1,161 @@
<div align="center">
<h1>Doc-to-LoRA (D2L): Learning to Instantly Internalize Contexts</h1>
:sparkles:<a href="https://pub.sakana.ai/doc-to-lora/">Interactive Web</a> |
:newspaper:<a href="https://x.com/SakanaAILabs">X</a> |
:scroll:<a href="https://arxiv.org/abs/2602.15902">Paper</a> |
:hugs:<a href="https://huggingface.co/SakanaAI">Hugging Face</a> |
:octocat:<a href="https://github.com/SakanaAI/doc-to-lora">GitHub</a>
<br>A reference implementation of Doc-to-LoRA (D2L).<br>
</div>
<div align="center">
<img height="300px" src="assets/overview_animation.gif" />
</div>
# Video2LoRA
---
Video2LoRA trains a hypernetwork that converts a video into LoRA weights for a
frozen vision-language model. The generated adapter lets the model answer later
text prompts without feeding the video tokens again.
## 🛠️ Installation
```
curl -LsSf https://astral.sh/uv/install.sh | sh
./install.sh
```
<p align="center">
<img src="assets/video2lora-animated-diagram.svg" alt="Animated Video2LoRA method diagram" width="900">
</p>
## 🤗 Pre-Trained Models
```
uv run huggingface-cli login
uv run huggingface-cli download SakanaAI/doc-to-lora --local-dir trained_d2l --include "*/"
```
## Install
## 🚀 Python API Usage
```python
# caveat: this interface only supports non-batched inputs
# for batched inference please see `src/ctx_to_lora/modeling/hypernet.py`
import torch
Install `uv`.
from ctx_to_lora.model_loading import get_tokenizer
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
# model loading
checkpoint_path = "trained_d2l/gemma_demo/checkpoint-80000/pytorch_model.bin"
state_dict = torch.load(checkpoint_path, weights_only=False)
model = ModulatedPretrainedModel.from_state_dict(
state_dict, train=False, use_sequence_packing=False
)
model.reset()
tokenizer = get_tokenizer(model.base_model.name_or_path)
# prepare data
doc = open("data/sakana_wiki.txt", "r").read()
chat = [{"role": "user", "content": "Tell me about Sakana AI."}]
chat_ids = tokenizer.apply_chat_template(
chat,
add_special_tokens=False,
return_attention_mask=False,
add_generation_prompt=True,
return_tensors="pt",
).to(model.device)
# calls after internalization will be influenced by internalized info
model.internalize(doc)
outputs = model.generate(input_ids=chat_ids, max_new_tokens=512)
print(tokenizer.decode(outputs[0]))
# remove internalized info
# model.reset()
# without internalized info, the model will halucinate
# outputs = model.generate(input_ids=chat_ids, max_new_tokens=512)
# print(tokenizer.decode(outputs[0]))
```
### 🎮 Interactive Demo
```bash
uv run demo/app.py
git clone https://github.com/MananSuri27/video2lora.git
cd video2lora
uv sync
```
<div align="center">
<h3>Video Demo</h3>
<video src="https://github.com/user-attachments/assets/16781365-5ec2-4c1c-b4f4-aeeebe3c2be5" controls autoplay muted playsinline preload="metadata" width="900"></video>
</div>
### 🧪 Experimental Scripts
To run any of the following scripts, use `uv run $PATH_TO_SCRIPT` from the root of this project.
For local video path resolution, set:
| Experiment | Data prep | Training | Evaluation | Notes |
| ------------------------------------ | ------------------------------------- | ----------------------------- | ---------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- |
| [Main experiment](scripts/main_exp/) | `scripts/main_exp/0-download_data.sh` | `scripts/main_exp/1-train.sh` | `scripts/main_exp/eval/*.sh` | Downloading data is fastest; regenerate only if you need fresh synthetic data. Evaluation scripts reproduce the main paper metrics. |
| [NIAH](scripts/niah/) | `scripts/niah/0-gen_data.sh` | `scripts/niah/1-train.sh` | `scripts/niah/2-eval.sh` | Run the scripts in order; data generation only needs to happen once |
### 🔬 Self-Generated Data Viewer
After downloading/generating the data, we can see samples of the data using this script.
```bash
uv run webui/self_gen_viewer.py
export VIDEO2LORA_DATA_ROOT=$PWD/data/video2lora
```
See more info at [webui/SELF_GEN_VIEWER.md](webui/SELF_GEN_VIEWER.md).
### 📚 Citation
```bibtex
@inproceedings{charakorn2026doctolora,
title ={Doc-to-Lo{RA}: Learning to Instantly Internalize Contexts},
author ={Rujikorn Charakorn and Edoardo Cetin and Shinnosuke Uesaka and Robert Tjarko Lange},
booktitle ={Forty-third International Conference on Machine Learning},
year ={2026},
url ={https://openreview.net/forum?id=iW1oBBO72S}
If unset, `VIDEO2LORA_DATA_ROOT` defaults to `data/video2lora`.
## Checkpoints
Download the released checkpoints from Hugging Face:
```bash
uv run huggingface-cli download MananSuri27/Video2LoRA-SmolVLM-ckpts \
--local-dir checkpoints/Video2LoRA-SmolVLM-ckpts
```
The repo contains:
```text
checkpoints/Video2LoRA-SmolVLM-ckpts/
video2lora-smolvlm2-500m-best-ce.pt
video2lora-smolvlm2-2.2b-best-ce.pt
```
## Inference
Create a JSONL manifest with one row per video. For example:
```json
{"id":"sample-0001","video_path":"/path/to/video.mp4","prompt":"Describe what is happening in this video.","task_type":"caption"}
```
```bash
uv run python -m scripts.video2lora.infer \
--checkpoint checkpoints/Video2LoRA-SmolVLM-ckpts/video2lora-smolvlm2-500m-best-ce.pt \
--manifest /path/to/manifest.jsonl \
--output outputs/tiny_generations.jsonl
```
For the 2.2B checkpoint, change `--checkpoint` to:
```bash
checkpoints/Video2LoRA-SmolVLM-ckpts/video2lora-smolvlm2-2.2b-best-ce.pt
```
The output JSONL includes the original row fields plus `prediction` and
`checkpoint`.
## Data Format
Training and inference use JSONL manifests. Each line is one example:
```json
{
"id": "sample-0001",
"video_path": "raw/finevideo/sample.mp4",
"task_type": "caption",
"prompt": "Describe what is happening in this video.",
"target_text": "A person is cooking.",
"dataset": "finevideo",
"split": "train",
"metadata": {}
}
```
Relative `video_path` values are resolved against `VIDEO2LORA_DATA_ROOT`.
Absolute paths are used as-is.
## Generate Teacher Data
The final CE training recipe uses cached teacher-generated targets. Starting
from readable manifests, generate teacher targets with:
```bash
export VIDEO2LORA_DATA_ROOT=$PWD/data/video2lora
CUDA_VISIBLE_DEVICES=0 uv run python -m scripts.video2lora.generate_finevideo_teacher_targets \
--input-manifest $VIDEO2LORA_DATA_ROOT/processed/finevideo/train.jsonl \
--output-manifest $VIDEO2LORA_DATA_ROOT/processed/finevideo/train.teacher_visual_smolvlm.jsonl \
--smolvlm-name-or-path HuggingFaceTB/SmolVLM2-2.2B-Instruct \
--per-device-batch-size 8 \
--max-frames 12 \
--video-size-longest-edge 384
```
To shard generation manually, run the same command per GPU with
`--num-shards N --shard-index R`, then merge with `--merge-only`.
## Train
The main CE training entrypoint is:
```bash
uv run accelerate launch \
--config_file accelerate_config.yaml \
--num_processes 4 \
-m scripts.video2lora.train_smolvlm_stage1 \
--smolvlm-name-or-path HuggingFaceTB/SmolVLM2-500M-Video-Instruct \
--train-manifest data/video2lora/processed/finevideo/train.teacher_visual_smolvlm.jsonl \
--val-manifest data/video2lora/processed/finevideo/val.teacher_visual_smolvlm.jsonl \
--val-core-manifest data/video2lora/processed/finevideo/val.teacher_visual_smolvlm.jsonl \
--output-dir runs/video2lora-smoke \
--per-device-batch-size 2 \
--gradient-accumulation-steps 8 \
--max-steps 1000 \
--target-modules down_proj \
--lora-r 16 \
--latent-size 512 \
--n-latent-queries 8 \
--num-blocks 9 \
--max-frames 12 \
--video-size-longest-edge 384 \
--kl-weight 0.0 \
--wandb-mode disabled
```
The self-distillation variant is also included:
```bash
uv run accelerate launch \
--config_file accelerate_config.yaml \
--num_processes 4 \
-m scripts.video2lora.train_smolvlm_stage1_selfdistill \
--smolvlm-name-or-path HuggingFaceTB/SmolVLM2-500M-Video-Instruct \
--train-manifest data/video2lora/processed/finevideo/train.jsonl \
--val-manifest data/video2lora/processed/finevideo/val.jsonl \
--val-gen-manifest data/video2lora/processed/finevideo/val_gen_100.jsonl \
--output-dir runs/video2lora-selfdistill \
--wandb-mode disabled
```
## Build FineVideo Manifests
If you have raw FineVideo metadata/videos laid out under
`$VIDEO2LORA_DATA_ROOT/raw/finevideo`, build Stage 1 manifests with:
```bash
uv run python -m scripts.video2lora.build_finevideo_stage1_manifest
```

BIN
assets/cover.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 2.3 MiB

BIN
assets/overview_animation.gif
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 4.5 MiB

840
assets/video2lora-animated-diagram.svg Normal file
View file

File diff suppressed because one or more lines are too long
Side by side

After

Width:  |  Height:  |  Size: 3.3 MiB

18
chat_templates/Qwen/Qwen3-4B-Instruct-2507.jinja
View file

@ -1,18 +0,0 @@
{%- if messages[0].role == 'system' %}
{{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
{%- endif %}
{%- for message in messages %}
{%- if message.content is string %}
{%- set content = message.content %}
{%- else %}
{%- set content = '' %}
{%- endif %}
{%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
{{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
{%- elif message.role == "assistant" %}
{{- '<|im_start|>' + message.role + '\n' }}{% generation %}{{ content + '<|im_end|>' }}{% endgeneration %}{{ '\n' }}
{%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
{{- '<|im_start|>assistant\n' }}
{%- endif %}

24
chat_templates/google/gemma-2-2b-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 %}

24
chat_templates/mistralai/Mistral-7B-Instruct-v0.2.jinja
View file

@ -1,24 +0,0 @@
{%- if messages[0]['role'] == 'system' %}
{%- set system_message = messages[0]['content'] %}
{%- set loop_messages = messages[1:] %}
{%- else %}
{%- set loop_messages = messages %}
{%- endif %}
{{- bos_token }}
{%- for message in loop_messages %}
{%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
{{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }}
{%- endif %}
{%- if message['role'] == 'user' %}
{%- if loop.first and system_message is defined %}
{{- ' [INST] ' + system_message + '\n\n' + message['content'] + ' [/INST] ' }}
{%- else %}
{{- ' [INST] ' + message['content'] + ' [/INST] ' }}
{%- endif %}
{%- elif message['role'] == 'assistant' %}
{% generation %}{{ message['content'] + eos_token }}{% endgeneration %}
{%- else %}
{{- raise_exception('Only user and assistant roles are supported, with the exception of an initial optional system message!') }}
{%- endif %}
{%- endfor %}

31
configs/main_exp/mistral/self_gen_lv1_closed_qa_1_l2l.yaml
View file

@ -1,31 +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/mistralai/Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
- self_gen/mistralai/Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_0.0/pwc_compact
- self_gen/mistralai/Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_1.0/squad_compact
- self_gen/mistralai/Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_1.0/ropes_compact
- self_gen/mistralai/Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_1.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes
- self_gen/mistralai/Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet

30
configs/main_exp/qwen/self_gen_lv1_closed_qa_1_l2l.yaml
View file

@ -1,30 +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/Qwen/Qwen3-4B-Instruct-2507_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
- self_gen/Qwen/Qwen3-4B-Instruct-2507_temp_0.0_closed_qa_prob_0.0/pwc_compact
- self_gen/Qwen/Qwen3-4B-Instruct-2507_temp_0.0_closed_qa_prob_1.0/squad_compact
- self_gen/Qwen/Qwen3-4B-Instruct-2507_temp_0.0_closed_qa_prob_1.0/ropes_compact
- self_gen/Qwen/Qwen3-4B-Instruct-2507_temp_0.0_closed_qa_prob_1.0/drop_compact
val_ds_names:
- squad
- pwc
- drop
- ropes

31
configs/main_exp/self_gen_lv1_closed_qa_1_l2l.yaml
View file

@ -1,31 +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/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
- 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

27
configs/main_exp/self_gen_lv1_closed_qa_1_no_qa_l2l.yaml
View file

@ -1,27 +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
val_ds_names:
- squad
- pwc
- drop
- ropes
- 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

14
configs/niah_exp/ctx_magic_number_32_256.yaml
View file

@ -1,14 +0,0 @@
# LoRA
lora_r: 8
lora_dropout: 0.0
target_modules:
- down_proj
# data
train_ds_names:
- ctx_magic_number_32_128
- ctx_magic_number_128_256
val_ds_names:
- ctx_magic_number_32_128
- ctx_magic_number_128_256

42
data/build_drop_compact.py
View file

@ -1,42 +0,0 @@
import gc
from datasets import Dataset, load_dataset
from tqdm import tqdm
if __name__ == "__main__":
ds_name = "ucinlp/drop"
for split in ["train", "validation"]:
ctx_qa_dict = dict()
ds = load_dataset(ds_name, split=split)
print(f"Original size: {len(ds)}")
for i, sample in tqdm(enumerate(ds)):
ctx = sample["passage"]
if ctx not in ctx_qa_dict:
ctx_qa_dict[ctx] = {"prompts": [], "responses": []}
question = sample["question"]
answer = sample["answers_spans"]["spans"][0]
ctx_qa_dict[ctx]["prompts"].append(question)
ctx_qa_dict[ctx]["responses"].append(answer)
print(f"Unique contexts: {len(ctx_qa_dict)}")
# convert ctx_qa_dict to a list of dictionaries
samples = [
{
"context": ctx,
"prompts": ctx_qa_dict[ctx]["prompts"],
"responses": ctx_qa_dict[ctx]["responses"],
}
for ctx in ctx_qa_dict
]
print(f"Sampled data: {samples[0]}")
# breakpoint()
# save to a new dataset
ds = Dataset.from_list(samples)
save_path = f"./data/raw_datasets/drop_compact/{split}/ds.parquet"
print(f"Saving dataset to {save_path}")
ds.to_parquet(save_path)
print("=" * 80)
del ds, samples, ctx_qa_dict
gc.collect()

43
data/build_pwc_compact.py
View file

@ -1,43 +0,0 @@
import gc
from datasets import Dataset, load_dataset
from tqdm import tqdm
if __name__ == "__main__":
ds_name = "sggetao/PwC"
for split in ["train", "test"]:
ctx_qa_dict = dict()
ds = load_dataset(ds_name, split=split)
print(f"Original size: {len(ds)}")
for i, sample in tqdm(enumerate(ds)):
ctx = sample["input"]
if ctx not in ctx_qa_dict:
ctx_qa_dict[ctx] = {"prompts": [], "responses": []}
# question = closed_qa_prompting(sample["prompt"])
question = sample["prompt"]
answer = sample["answer"]
ctx_qa_dict[ctx]["prompts"].append(question)
ctx_qa_dict[ctx]["responses"].append(answer)
print(f"Unique contexts: {len(ctx_qa_dict)}")
# convert ctx_qa_dict to a list of dictionaries
samples = [
{
"context": ctx,
"prompts": ctx_qa_dict[ctx]["prompts"],
"responses": ctx_qa_dict[ctx]["responses"],
}
for ctx in ctx_qa_dict
]
print(f"Sampled data: {samples[0]}")
# breakpoint()
# save to a new dataset
ds = Dataset.from_list(samples)
save_path = f"./data/raw_datasets/pwc_compact/{split}/ds.parquet"
print(f"Saving dataset to {save_path}")
ds.to_parquet(save_path)
print("=" * 80)
del ds, samples, ctx_qa_dict
gc.collect()

45
data/build_ropes_compact.py
View file

@ -1,45 +0,0 @@
import gc
from datasets import Dataset, load_dataset
from tqdm import tqdm
if __name__ == "__main__":
ds_name = "allenai/ropes"
for split in ["train", "validation"]:
ctx_qa_dict = dict()
ds = load_dataset(ds_name, split=split)
print(f"Original size: {len(ds)}")
for i, sample in tqdm(enumerate(ds)):
ctx_template = "{background}\n{situation}"
response = sample["answers"]["text"][0]
bg_txt = sample["background"]
situation_txt = sample["situation"]
ctx = ctx_template.format(background=bg_txt, situation=situation_txt)
q = sample["question"]
if ctx not in ctx_qa_dict:
ctx_qa_dict[ctx] = {"prompts": [], "responses": []}
ctx_qa_dict[ctx]["prompts"].append(q)
ctx_qa_dict[ctx]["responses"].append(response)
print(f"Unique contexts: {len(ctx_qa_dict)}")
# convert ctx_qa_dict to a list of dictionaries
samples = [
{
"context": ctx,
"prompts": ctx_qa_dict[ctx]["prompts"],
"responses": ctx_qa_dict[ctx]["responses"],
}
for ctx in ctx_qa_dict
]
print(f"Sampled data: {samples[0]}")
# breakpoint()
# save to a new dataset
ds = Dataset.from_list(samples)
save_path = f"./data/raw_datasets/ropes_compact/{split}/ds.parquet"
print(f"Saving dataset to {save_path}")
ds.to_parquet(save_path)
print("=" * 80)
del ds, samples, ctx_qa_dict
gc.collect()

42
data/build_squad_compact.py
View file

@ -1,42 +0,0 @@
import gc
from datasets import Dataset, load_dataset
from tqdm import tqdm
if __name__ == "__main__":
ds_name = "data/raw_datasets/squad"
for split in ["train", "validation"]:
ctx_qa_dict = dict()
ds = load_dataset(ds_name, split=split)
print(f"Original size: {len(ds)}")
for i, sample in tqdm(enumerate(ds)):
ctx = sample["context"]
if ctx not in ctx_qa_dict:
ctx_qa_dict[ctx] = {"prompts": [], "responses": []}
question = sample["question"]
answer = sample["answers"]["text"][0]
ctx_qa_dict[ctx]["prompts"].append(question)
ctx_qa_dict[ctx]["responses"].append(answer)
print(f"Unique contexts: {len(ctx_qa_dict)}")
# convert ctx_qa_dict to a list of dictionaries
samples = [
{
"context": ctx,
"prompts": ctx_qa_dict[ctx]["prompts"],
"responses": ctx_qa_dict[ctx]["responses"],
}
for ctx in ctx_qa_dict
]
print(f"Sampled data: {samples[0]}")
# breakpoint()
# save to a new dataset
ds = Dataset.from_list(samples)
save_path = f"./data/raw_datasets/squad_compact/{split}/ds.parquet"
print(f"Saving dataset to {save_path}")
ds.to_parquet(save_path)
print("=" * 80)
del ds, samples, ctx_qa_dict
gc.collect()

10
data/download_fineweb_edu.py
View file

@ -1,10 +0,0 @@
from huggingface_hub import snapshot_download
if __name__ == "__main__":
fw_dir = "./data/raw_datasets/fineweb_edu/"
snapshot_download(
"HuggingFaceFW/fineweb-edu",
repo_type="dataset",
local_dir=fw_dir,
allow_patterns="sample/100BT/*",
)

288
data/generate_ctx_magic_number.py
View file

@ -1,288 +0,0 @@
import argparse
import json
import math
import os
import random
# -----------------------------
# Config knobs (edit or use CLI)
# -----------------------------
TOKENS_PER_BLOCK = 40 # rough heuristic tokens per noise block
BASE_SAMPLES_PER_BIN = (
320_000 # training samples budget scaler only (val/test fixed at 1000 each)
)
RNG_SEED = 42
NOISE_BLOCK = "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again."
SPECIAL_TPL = "The special magic number is {magic_number}."
SEP = "\n" # between blocks
def save_jsonl(data: list[dict], filepath: str) -> None:
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")
essential_digits4 = lambda: f"{random.randint(0, 9_999):04d}"
def _choose_position(total_blocks: int, depth_bin: int) -> int:
"""Choose an insertion index for the special sentence within [0, total_blocks-1]
such that its relative depth falls within the depth bin [i/10, (i+1)/10).
"""
if total_blocks <= 0:
return 0
# Use floor for start and ceil for end to cover boundaries evenly
start = math.floor(total_blocks * (depth_bin / 10))
end = math.ceil(total_blocks * ((depth_bin + 1) / 10)) - 1
# clamp
start = max(0, min(start, total_blocks - 1))
end = max(start, min(end, total_blocks - 1))
return random.randint(start, end)
def _build_example(total_blocks: int, depth_bin: int) -> dict:
"""Build one example with a special line inserted among noise blocks.
total_blocks: total number of blocks in the final context (including the special one)
depth_bin: integer in [0, 9]
"""
total_blocks = max(1, total_blocks)
# Prepare blocks
magic = essential_digits4()
special_line = SPECIAL_TPL.format(magic_number=magic)
# We'll have (total_blocks - 1) noise blocks and 1 special line
noise_count = max(0, total_blocks - 1)
blocks: list[str] = [NOISE_BLOCK for _ in range(noise_count)]
insert_at = _choose_position(total_blocks, depth_bin)
# Insert special line at the desired position within the final sequence
# If noise_count == 0, we just return special
if noise_count == 0:
final_blocks = [special_line]
else:
# Compose by interleaving noise and inserting special at index
# Build a list of length `total_blocks` and fill
final_blocks = []
noise_idx = 0
for idx in range(total_blocks):
if idx == insert_at:
final_blocks.append(special_line)
else:
final_blocks.append(blocks[noise_idx])
noise_idx += 1
context = SEP.join(final_blocks)
prompt = "What is the special magic number? Reply with only the number."
response = magic
return {"context": context, "prompt": prompt, "response": response}
def generate_examples(n: int, k: int) -> list[dict]:
"""Generate n examples (all for block length k) evenly across 10 depth bins."""
if n <= 0:
return []
base = n // 10
rem = n % 10
counts = [base + (1 if i < rem else 0) for i in range(10)]
out: list[dict] = []
for depth_bin, c in enumerate(counts):
for _ in range(c):
out.append(_build_example(total_blocks=k, depth_bin=depth_bin))
random.shuffle(out)
return out
def main():
parser = argparse.ArgumentParser(
description="Generate noise-wrapped special magic number dataset (similar structure to generate_ctx_kv.py)",
)
parser.add_argument("--seed", type=int, default=RNG_SEED, help="Random seed")
parser.add_argument(
"--tokenizer-name",
type=str,
default="google/gemma-2-2b-it",
help=("Tokenizer name"),
)
parser.add_argument(
"--base-samples-per-bin",
type=int,
default=BASE_SAMPLES_PER_BIN,
help="Baseline number of TRAINING samples per token bin (scaled by bin width). Validation & test are always 1000 each.",
)
parser.add_argument(
"--out-prefix",
type=str,
default="data/raw_datasets/ctx_magic_number",
help="Output directory prefix (bin range will be appended)",
)
parser.add_argument(
"--tokens-per-block",
"--tokens-per-pair",
dest="tokens_per_block",
type=int,
default=TOKENS_PER_BLOCK,
help="Heuristic tokens per noise block 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)
# ----------------------------------------------------
# Optional: report tokenizer-based token length stats
# ----------------------------------------------------
if args.tokenizer_name:
try:
from transformers import AutoTokenizer # type: ignore
except Exception as e: # pragma: no cover
raise RuntimeError(
"Failed to import transformers. Install it or omit --tokenizer-name."
) from e
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=True)
noise_token_count = len(tokenizer(NOISE_BLOCK).input_ids)
special_example = SPECIAL_TPL.format(magic_number="0000")
special_token_count = len(tokenizer(special_example).input_ids)
print(
f"[Tokenizer: {args.tokenizer_name}] Noise block tokens: {noise_token_count} | Special line tokens: {special_token_count}"
)
tok_bins = [(32, 128), (128, 256), (256, 512), (512, 1024), (32, 1024)] + [
(1024 * i, 1024 * (i + 1)) for i in range(1, 16)
]
tok_bins += [(2**14 + 2**12 * (i), 2**14 + 2**12 * (i + 1)) for i in range(4)]
tok_bins += [(2**15 + 2**13 * (i), 2**15 + 2**13 * (i + 1)) for i in range(12)]
if args.only_first_n_bins is not None:
tok_bins = tok_bins[: args.only_first_n_bins]
if args.tokenizer_name:
max_hi = max(hi for _, hi in tok_bins)
def measure_len(k: int) -> int:
if k == 1:
ctx = SPECIAL_TPL.format(magic_number="0000")
else:
blocks = [NOISE_BLOCK] * (k - 1) + [
SPECIAL_TPL.format(magic_number="0000")
]
ctx = SEP.join(blocks)
return len(tokenizer(ctx).input_ids)
lengths: list[int] = [0]
k = 1
while True:
L = measure_len(k)
lengths.append(L)
if L >= max_hi:
break
k += 1
len_bins = []
for lo, hi in tok_bins:
k_lo = None
for kk in range(1, len(lengths)):
if lengths[kk] >= lo:
k_lo = kk
break
if k_lo is None or lengths[k_lo] >= hi:
len_bins.append((0, 0))
continue
k_hi = len(lengths)
for kk in range(k_lo, len(lengths)):
if lengths[kk] >= hi:
k_hi = kk
break
len_bins.append((k_lo, k_hi))
base_tokens = lengths[1]
delta = (lengths[2] - lengths[1]) if len(lengths) > 2 else 0
print(
f"Using tokenizer-measured block ranges. base_tokens={base_tokens} approx_delta={delta}"
)
else:
len_bins = [
(lo // args.tokens_per_block, hi // args.tokens_per_block)
for (lo, hi) in tok_bins
]
if args.dry_run:
for lb in len_bins:
if lb[1] > lb[0]:
k = max(1, lb[0])
sample = generate_examples(10, k)
print("Sample entry:")
print(json.dumps(sample[0], indent=2))
break
return
# -----------------------------------------------
# Main generation per token bin
# -----------------------------------------------
TARGET_VAL = 1000
TARGET_TEST = 1000
for len_bin, tok_bin in zip(len_bins, tok_bins):
if len_bin[1] <= len_bin[0]:
print(f"Skipping token bin {tok_bin} (no valid block counts)")
continue
k_start = max(1, len_bin[0])
k_end = max(1, len_bin[1])
k_values = list(range(k_start, k_end))
bin_size = len(k_values)
save_dir = f"{args.out_prefix}_{tok_bin[0]}_{tok_bin[1]}"
training_enabled = tok_bin[1] <= 1024 # unchanged policy
if training_enabled:
train_data: list[dict] = []
# Distribute training budget across k values.
# Scale: per_k = base_samples_per_bin / bin_size
per_k_train = max(1, args.base_samples_per_bin // max(1, bin_size))
for k in k_values:
train_data += generate_examples(per_k_train, k)
val_data: list[dict] = []
test_data: list[dict] = []
base_val = TARGET_VAL // bin_size
rem_val = TARGET_VAL % bin_size
base_test = TARGET_TEST // bin_size
rem_test = TARGET_TEST % bin_size
for idx, k in enumerate(k_values):
n_val_k = base_val + (1 if idx < rem_val else 0)
n_test_k = base_test + (1 if idx < rem_test else 0)
if n_val_k:
val_data += generate_examples(n_val_k, k)
if n_test_k:
test_data += generate_examples(n_test_k, k)
random.shuffle(val_data)
random.shuffle(test_data)
os.makedirs(save_dir, exist_ok=True)
if training_enabled:
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 training_enabled:
print(
f"Dataset generated at {save_dir} (train={len(train_data)} val={len(val_data)} test={len(test_data)})"
)
else:
print(
f"Dataset (val/test only) generated at {save_dir} (val={len(val_data)} test={len(test_data)})"
)
if __name__ == "__main__":
main()

269
data/generate_fw_edu_qa_v2.py
View file

@ -1,269 +0,0 @@
import argparse
import os
import re
from glob import glob
import pandas as pd
from datasets import Dataset, load_dataset
from vllm import LLM, SamplingParams
STOP_STRINGS = {
"google/gemma-3-12b-it": ["<eos>", "<end_of_turn>"],
}
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 suit any context.\n"
"**DO NOT** hallucinate 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 suit any context.\n\n"
"### Context ###\n"
"{context}\n\n\n"
"### Rules ###\n"
"Rules to follow when generating the questions:\n"
"1. The questions must be specific to the given context and fully answerable from information present in the given context.\n"
"2. Ask questions that are fact-seeking based on the information provided.\n"
"3. Make sure the questions are clear and unambiguous.\n"
"4. Phrases like 'based on the provided context', 'according to the context', 'in the context', etc., are **NOT ALLOWED** to appear in "
"the questions.\n"
"5. The questions should not overlap. They should be diverse, covering many aspects of the context.\n"
"6. 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"
"7. Ignore the text formatting of the context, e.g., bold, italic, underline, etc.\n"
"8. Ignore typos, spacing, and grammatical errors in the context.\n\n"
"Rules to follow when generating the answers:\n"
"1. The answers must use the (implied) information provided in the context.\n"
"2. Phrases like 'based on the provided context', 'according to the context', 'in the context', etc., are **NOT ALLOWED** to appear in "
"the answers.\n"
"3. Do not just copy words from the context. Answer the question in your own words.\n"
"4. The answers should be detailed and comprehensive. Please include additional specific details from the context.\n\n"
"Respond 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 check_should_skip(txt: str, vllm_model: str) -> bool:
"""Check if the response should be skipped based on stop strings."""
for stop in STOP_STRINGS[vllm_model]:
if stop in txt[-len(stop) :]:
return (txt.split(stop)[0], False) # Found a valid stop string
return (txt, True) # No valid stop string found, skip this response
def postprocess_qa_pairs(res_txt: str):
"""
Postprocesses the QA pairs from the response text.
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"
res_txt = remove_think(res_txt)
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 = []
n_skips = 0
if (len(questions) > 0) and (len(answers) > 0):
n_gen_pairs = min(len(questions), len(answers))
has_left_over = n_gen_pairs < len(questions) or n_gen_pairs < len(answers)
for i in range(n_gen_pairs):
response = answers[i].strip()
question = questions[i].strip()
if not response or not question:
print(f"Skipping empty question or answer at index {i}")
continue
if (not has_left_over) and (i == n_gen_pairs - 1):
response, skip = check_should_skip(response, vllm_model)
if skip:
print(f"Skipping due to missing stop string")
n_skips += 1
continue
out_q.append(question.strip())
out_a.append(response.strip())
print(f"Skipped {n_skips} responses due to missing stop strings")
return out_q, out_a
def length_filter(sample, min_len, max_len):
return min_len <= len(sample["text"]) <= max_len
def remove_think(txt):
return txt.split("</think>")[-1]
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-2-27b-it"),
help="VLLM model to use for generation",
)
parser.add_argument(
"--shard_pattern",
type=str,
required=True,
help="Pattern to match shard files (e.g., '000_0000*')",
)
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=2000,
help="Maximum length of the context to consider for generation",
)
parser.add_argument(
"--max_model_length",
type=int,
default=2**14,
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",
)
args = parser.parse_args()
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
paths = glob(
f"./data/raw_datasets/fineweb_edu/sample/100BT/{shard_pattern}.parquet"
)
split = "train[:100]" if args.debug else "train"
for path in paths:
ds = load_dataset(
"parquet",
data_files=path,
split=split,
)
ds = ds.filter(
length_filter,
fn_kwargs={"min_len": args.min_length, "max_len": args.max_length},
num_proc=8,
)
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.0,
# needed for checking if stop tokens are present
skip_special_tokens=False,
include_stop_str_in_output=True,
),
)
samples = []
for ctx, completion in zip(ctxs, completions):
questions, answers = postprocess_qa_pairs(completion.outputs[0].text)
samples.append(
{
"context": ctx,
"prompts_level_0": questions,
"responses_level_0": answers,
}
)
if args.debug:
print(f"{ctx=}")
print(f"{completion.outputs[0].text=}")
for q, a in zip(questions, answers):
print(f"{q=}")
print(f"{a=}")
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_v2/min_{args.min_length}_to_{args.max_length}/{shard_name}.parquet"
)
val_ds.to_parquet(
f"data/raw_datasets/fw_qa_v2/min_{args.min_length}_to_{args.max_length}/{shard_name}_val.parquet"
)
print(
f"Saved to data/raw_datasets/fw_qa_v2/min_{args.min_length}_to_{args.max_length}/{shard_name}.parquet"
)
print(
f"Saved to data/raw_datasets/fw_qa_v2/min_{args.min_length}_to_{args.max_length}/{shard_name}_val.parquet"
)

296
data/generate_fw_edu_qa_v2_repeat.py
View file

@ -1,296 +0,0 @@
import argparse
import gc
import os
import re
from glob import glob
from datasets import load_dataset
from vllm import LLM, SamplingParams
STOP_STRINGS = {
"google/gemma-3-12b-it": ["<eos>", "<end_of_turn>"],
}
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 suit any context.\n"
"**DO NOT** hallucinate 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 suit any context.\n\n"
"### Context ###\n"
"{context}\n\n\n"
"### Example Question-Answer Pairs ###\n"
"{qa_pairs}\n\n\n"
"### Rules ###\n"
"Rules to follow when generating the questions:\n"
"1. The questions must be specific to the given context and fully answerable from information present in *or* implied from the given context.\n"
"2. The questions must *not* be redundant with the example questions-answer pairs provided.\n"
"3. You should prioritize fact-seeking questions. Consider reversal questions, e.g., asking 'What causes X to happen?' is valid when 'Y causes X' is presented in the context.\n"
"4. If all the facts in the context are already covered by the provided examples, you must generate *more complicated* questions that require reasoning beyond simple information retrieval.\nThis includes asking about information that can be inferred, requiring synthesizing information from multiple parts of the text, or understanding relationships between concepts, events, or individuals mentioned in the context. For example, if the context says 'The Eiffel Tower was completed in 1889 after 2 years of construction', you can ask 'When did the construction of the Eiffel Tower begin?'. Here's another example: if the context says 'Alice is Bob's mother. Bob is Charlie's Dad', you can ask 'Who is Charlie's grandmother?'.\n"
"5. Phrases like 'based on the provided context', 'according to the context', 'in the context', etc., are **NOT ALLOWED** to appear in "
"the questions.\n"
"6. The questions should not overlap. They should be diverse, covering many aspects of the context.\n"
"7. 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"
"8. Ignore the text formatting of the context, e.g., bold, italic, underline, etc.\n"
"9. Ignore typos, spacing, and grammatical errors in the context.\n\n"
"Rules to follow when generating the answers:\n"
"1. The answers must use the (implied) information provided in the context.\n"
"2. Phrases like 'based on the provided context', 'according to the context', 'in the context', etc., are **NOT ALLOWED** to appear in "
"the answers.\n"
"3. Do not just copy words from the context. Answer the question in your own words.\n"
"4. The answers should be detailed and comprehensive. Please include additional specific details from the context.\n\n"
"Respond 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, example_qa_pairs, n_qa_pairs):
prompt = PROMPT_TEMPLATE.format(
context=context,
qa_pairs=example_qa_pairs,
n_qa_pairs=n_qa_pairs,
)
return prompt
def check_should_skip(txt: str, vllm_model: str) -> bool:
"""Check if the response should be skipped based on stop strings."""
for stop in STOP_STRINGS[vllm_model]:
if stop in txt[-len(stop) :]:
return (txt.split(stop)[0], False) # Found a valid stop string
return (txt, True) # No valid stop string found, skip this response
def postprocess_qa_pairs(res_txt: str):
"""
Postprocesses the QA pairs from the response text.
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"
res_txt = remove_think(res_txt)
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 = []
n_skips = 0
if (len(questions) > 0) and (len(answers) > 0):
n_gen_pairs = min(len(questions), len(answers))
has_left_over = n_gen_pairs < len(questions) or n_gen_pairs < len(answers)
for i in range(n_gen_pairs):
response = answers[i].strip()
question = questions[i].strip()
if not response or not question:
print(f"Skipping empty question or answer at index {i}")
continue
if (not has_left_over) and (i == n_gen_pairs - 1):
response, skip = check_should_skip(response, vllm_model)
if skip:
print(f"Skipping due to missing stop string")
n_skips += 1
continue
out_q.append(question.strip())
out_a.append(response.strip())
print(f"Skipped {n_skips} responses due to missing stop strings")
return out_q, out_a
def flatten_list(l):
out = []
for x in l:
out += x
return out
def remove_think(txt):
return txt.split("</think>")[-1]
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-2-27b-it"),
help="VLLM model to use for generation",
)
parser.add_argument(
"--shard_pattern",
type=str,
required=True,
help="Pattern to match shard files (e.g., '000_0000*')",
)
parser.add_argument(
"--n_qa_pairs",
type=int,
required=True,
help="Number of question-answer pairs to generate per context",
)
parser.add_argument(
"--max_model_length",
type=int,
default=2**12,
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",
)
args = parser.parse_args()
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=2**14,
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
paths = glob(f"./data/raw_datasets/fw_qa_v2/{shard_pattern}.parquet")
split = "train[:100]" if args.debug else "train"
for path in paths:
assert "_level" in path, (
"Path must contain '_level' to indicate the dataset level"
)
shard_name = path.split("/")[-1].split(".")[0].split("_debug")[0]
if "/" in shard_pattern:
shard_name = "/".join(shard_pattern.split("/")[:-1]) + "/" + shard_name
cur_level = int(shard_name.split("_level_")[-1])
next_level = cur_level + 1
ds = load_dataset(
"parquet",
data_files=path,
split=split,
)
prompt_cols = [col for col in ds.column_names if col.startswith("prompts")]
response_cols = [col for col in ds.column_names if col.startswith("responses")]
assert len(prompt_cols) > 0, "No prompt columns found in the dataset"
if len(prompt_cols) != len(response_cols):
raise ValueError(
"Number of prompt columns does not match number of response columns"
)
samples_data = []
for sample in iter(ds):
# Format existing QA pairs as examples
example_qa_pairs = ""
questions = flatten_list([sample[col] for col in prompt_cols])
answers = flatten_list([sample[col] for col in response_cols])
for i, (q, a) in enumerate(zip(questions, answers), 1):
example_qa_pairs += f"Question {i}: {q}\nAnswer {i}: {a}\n"
samples_data.append(
{"context": sample["context"], "example_qa_pairs": example_qa_pairs}
)
del ds
gc.collect()
messages = [
[
{"role": "system", "content": SYSTEM_TEMPLATE},
{
"role": "user",
"content": get_prompt(
sample["context"], sample["example_qa_pairs"], n_qa_pairs
),
},
]
for sample in samples_data
]
print(f"Generating from {len(messages)} contexts")
completions = llm.chat(
messages,
sampling_params=SamplingParams(
temperature=0.0,
# needed for checking if stop tokens are present
skip_special_tokens=False,
include_stop_str_in_output=True,
),
)
samples = []
for sample_data, completion in zip(samples_data, completions):
questions, answers = postprocess_qa_pairs(completion.outputs[0].text)
samples.append(
{
"context": sample_data["context"],
f"prompts_level_{next_level}": questions,
f"responses_level_{next_level}": answers,
}
)
if args.debug:
print(f"context={sample_data['context']}")
print(f"example_qa_pairs={sample_data['example_qa_pairs']}")
print(f"{completion.outputs[0].text=}")
for q, a in zip(questions, answers):
print(f"{q=}")
print(f"{a=}")
print()
print("=" * 80)
del samples_data
gc.collect()
print(f"Generated {len(samples)} samples")
ds = load_dataset(
"parquet",
data_files=path,
split=split,
)
ds = ds.add_column(
f"prompts_level_{next_level}",
[sample[f"prompts_level_{next_level}"] for sample in samples],
)
ds = ds.add_column(
f"responses_level_{next_level}",
[sample[f"responses_level_{next_level}"] for sample in samples],
)
shard_name_base = shard_name.split("_level_")[0]
shard_name = f"{shard_name_base}_level_{next_level}"
if args.debug:
shard_name += "_debug"
ds.to_parquet(f"data/raw_datasets/fw_qa_v2/{shard_name}.parquet")
print(f"Saved to data/raw_datasets/fw_qa_v2/{shard_name}.parquet")

387
data/gutenburg_sample.txt
View file

@ -1,387 +0,0 @@
The Project Gutenberg eBook, Addison, by William John Courthope
This eBook is for the use of anyone anywhere at no cost and with
almost no restrictions whatsoever. You may copy it, give it away or
re-use it under the terms of the Project Gutenberg License included
with this eBook or online at www.gutenberg.org
Title: Addison
Author: William John Courthope
Release Date: November 27, 2012 [eBook #41496]
Language: English
Character set encoding: ISO-8859-1
***START OF THE PROJECT GUTENBERG EBOOK ADDISON***
E-text prepared by the Online Distributed Proofreading Team
(http://www.pgdp.net) from page images generously made available by
Internet Archive (http://archive.org)
Note: Images of the original pages are available through
Internet Archive. See
http://archive.org/details/addison_00cour
Transcriber's note:
Text enclosed by underscores is in italics (_italics_).
Text enclosed by curly brackets is superscripted
(example: y{e}).
English Men of Letters
Edited by John Morley
ADDISON
by
W. J. COURTHOPE
Harper & Brothers Publishers
New York and London
1902
* * * * *
ENGLISH MEN OF LETTERS.
EDITED BY JOHN MORLEY.
JOHNSON Leslie Stephen.
GIBBON J. C. Morison.
SCOTT R. H. Hutton.
SHELLEY J. A. Symonds.
HUME T. H. Huxley.
GOLDSMITH William Black.
DEFOE William Minto.
BURNS J. C. Shairp.
SPENSER R. W. Church.
THACKERAY Anthony Trollope.
BURKE John Morley.
MILTON Mark Pattison.
HAWTHORNE Henry James, Jr.
SOUTHEY E. Dowden.
CHAUCER A. W. Ward.
BUNYAN J. A. Froude.
COWPER Goldwin Smith.
POPE Leslie Stephen.
BYRON John Nichol.
LOCKE Thomas Fowler.
WORDSWORTH F. Myers.
DRYDEN G. Saintsbury.
LANDOR Sidney Colvin.
DE QUINCEY David Masson.
LAMB Alfred Ainger.
BENTLEY R. C. Jebb.
DICKENS A. W. Ward.
GRAY E. W. Gosse.
SWIFT Leslie Stephen.
STERNE H. D. Traill.
MACAULAY J. Cotter Morison.
FIELDING Austin Dobson.
SHERIDAN Mrs. Oliphant.
ADDISON W. J. Courthope.
BACON R. W. Church.
COLERIDGE H. D. Traill.
SIR PHILIP SIDNEY J. A. Symonds.
KEATS Sidney Colvin.
CARLYLE John Nichol.
12mo, Cloth, 75 cents per volume.
_Other volumes in preparation._
PUBLISHED BY HARPER & BROTHERS, NEW YORK.
_Any of the above works will be sent by mail, postage prepaid, to any part
of the United States, Canada, or Mexico, on receipt of the price._
* * * * *
CONTENTS.
PAGE
CHAPTER I.
THE STATE OF ENGLISH SOCIETY AND LETTERS
AFTER THE RESTORATION 1
CHAPTER II.
ADDISON'S FAMILY AND EDUCATION 21
CHAPTER III.
ADDISON ON HIS TRAVELS 38
CHAPTER IV.
HIS EMPLOYMENT IN AFFAIRS OF STATE 53
CHAPTER V.
THE "TATLER" AND "SPECTATOR" 78
CHAPTER VI.
"CATO" 110
CHAPTER VII.
ADDISON'S QUARREL WITH POPE 125
CHAPTER VIII.
THE LAST YEARS OF HIS LIFE 139
CHAPTER IX.
THE GENIUS OF ADDISON 153
ADDISON.
CHAPTER I.
THE STATE OF ENGLISH SOCIETY AND LETTERS AFTER THE RESTORATION.
Of the four English men of letters whose writings most fully embody the
spirit of the eighteenth century, the one who provides the biographer with
the scantiest materials is Addison. In his _Journal to Stella_, his social
verses, and his letters to his friends, we have a vivid picture of those
relations with women and that protracted suffering which invest with such
tragic interest the history of Swift. Pope, by the publication of his own
correspondence, has enabled us, in a way that he never intended, to
understand the strange moral twist which distorted a nature by no means
devoid of noble instincts. Johnson was fortunate in the companionship of
perhaps the best biographer who ever lived. But of the real life and
character of Addison scarcely any contemporary record remains. The formal
narrative prefixed to his works by Tickell is, by that writer's own
admission, little more than a bibliography. Steele, who might have told us
more than any man about his boyhood and his manner of life in London, had
become estranged from his old friend before his death. No writer has
taken the trouble to preserve any account of the wit and wisdom that
enlivened the "little senate" at Button's. His own letters are, as a rule,
compositions as finished as his papers in the _Spectator_. Those features
in his character which excite the greatest interest have been delineated
by the hand of an enemy--an enemy who possessed an unrivalled power of
satirical portrait-painting, and was restrained by no regard for truth
from creating in the public mind such impressions about others as might
serve to heighten the favourable opinion of himself.
This absence of dramatic incident in Addison's life would lead us
naturally to conclude that he was deficient in the energy and passion
which cause a powerful nature to leave a mark upon its age. Yet such a
judgment would certainly be erroneous. Shy and reserved as he was, the
unanimous verdict of his most illustrious contemporaries is decisive as to
the respect and admiration which he excited among them. The man who could
exert so potent an influence over the mercurial Steele, who could
fascinate the haughty and cynical intellect of Swift, whose conversation,
by the admission of his satirist Pope, had in it something more charming
than that of any other man; of whom it was said that he might have been
chosen king if he wished it; such a man, though to the coarse perception
of Mandeville he might have seemed no more than "a parson in a tye-wig,"
can hardly have been deficient in force of character.
Nor would it have been possible for a writer distinguished by mere
elegance and refinement to leave a lasting impress on the literature and
society of his country. In one generation after another, men representing
opposing elements of rank, class, interest, and taste, have agreed in
acknowledging Addison's extraordinary merits. "Whoever wishes," says
Johnson--at the end of a biography strongly coloured with the
prepossessions of a semi-Jacobite Tory--"whoever wishes to attain an
English style, familiar but not coarse, and elegant but not ostentatious,
must give his days and nights to the volumes of Addison." "Such a mark of
national respect," says Macaulay, the best representative of middle-class
opinion in the present century, speaking of the statue erected to Addison
in Westminster Abbey, "was due to the unsullied statesman, to the
accomplished scholar, to the master of pure English eloquence, to the
consummate painter of life and manners. It was due, above all, to the
great satirist who alone knew how to use ridicule without abusing it; who,
without inflicting a wound, effected a great social reform, and who
reconciled wit and virtue after a long and disastrous separation, during
which wit had been led astray by profligacy, and virtue by fanaticism."
This verdict of a great critic is accepted by an age to which the grounds
of it are, perhaps, not very apparent. The author of any ideal creation--a
poem, a drama, or a novel--has an imprescriptible property in the fame of
his work. But to harmonise conflicting social elements, to bring order out
of chaos in the sphere of criticism, to form right ways of thinking about
questions of morals, taste, and breeding, are operations of which the
credit, though it is certainly to be ascribed to particular individuals,
is generally absorbed by society itself. Macaulay's eulogy is as just as
it is eloquent, but the pages of the _Spectator_ alone will hardly show
the reader why Addison should be so highly praised for having reconciled
wit with virtue. Nor, looking at him as a critic, will it appear a great
achievement to have pointed out to English society the beauties of
_Paradise Lost_, unless it be remembered that the taste of the preceding
generation still influenced Addison's contemporaries, and that in that
generation Cowley was accounted a greater poet than Milton.
To estimate Addison at his real value we must regard him as the chief
architect of Public Opinion in the eighteenth century. But here again we
are met by an initial difficulty, because it has become almost a
commonplace of contemporary criticism to represent the eighteenth century
as a period of sheer destruction. It is tacitly assumed by a school of
distinguished philosophical writers that we have arrived at a stage in the
world's history in which it is possible to take a positive and scientific
view of human affairs. As it is of course necessary that from such a
system all belief in the supernatural shall be jealously excluded, it has
not seemed impossible to write the history of Thought itself in the
eighteenth century. And in tracing the course of this supposed continuous
stream it is natural that all the great English writers of the period
should be described as in one way or another helping to pull down, or
vainly to strengthen, the theological barriers erected by centuries of
bigotry against the irresistible tide of enlightened progress.
It would be of course entirely out of place to discuss here the merits of
this new school of history. Those who consider that, whatever glimpses we
may obtain of the law and order of the universe, man is, as he always has
been and always will be, a mystery to himself, will hardly allow that the
operations of the human spirit can be traced in the dissecting-room. But
it is, in any case, obvious that to treat the great _imaginative_ writers
of any age as if they were only mechanical agents in an evolution of
thought is to do them grave injustice. Such writers are, above all things,
creative. Their first aim is to "show the very age and body of the time
his form and pressure." No work of the eighteenth century, composed in a
consciously destructive spirit, has taken its place among the acknowledged
classics of the language. Even the _Tale of a Tub_ is to be regarded as a
satire upon the aberrations of theologians from right reason, not upon the
principles of Christianity itself. The _Essay on Man_ has, no doubt,
logically a tendency towards Deism, but nobody ever read the poem for the
sake of its philosophy; and it is well known that Pope was much alarmed
when it was pointed out to him that his conclusions might be represented
as incompatible with the doctrines of revealed religion.
The truth indeed seems to be the exact converse of what is alleged by the
scientific historians. So far from the eighteenth century in England being
an age of destructive analysis, its energies were chiefly devoted to
political, social, and literary reconstruction. Whatever revolution in
faith and manners the English nation had undergone had been the work of
the two preceding centuries, and though the historic foundations of
society remained untouched, the whole form of the superstructure had been
profoundly modified.
"So tenacious are we," said Burke, towards the close of the last
century, "of our old ecclesiastical modes and fashions of institution
that very little change has been made in them since the fourteenth or
fifteenth centuries, adhering in this particular as in all else to our
old settled maxim never entirely nor at once to depart from antiquity.
We found these institutions on the whole favourable to morality and
discipline, and we thought they were susceptible of amendment without
altering the ground. We thought they were capable of receiving and
meliorating, and, above all, of preserving the accessories of science
and literature as the order of Providence should successively produce
them. And after all, with this Gothic and monkish education (for such
it is the groundwork), we may put in our claim to as ample and early
a share in all the improvements in science, in arts, and in literature
which have illuminated the modern world as any other nation in Europe.
We think one main cause of this improvement was our not despising the
patrimony of knowledge which was left us by our forefathers."
All this is, in substance, true of our political as well as our
ecclesiastical institutions. And yet, when Burke wrote, the great feudal
and mediæval structure of England had been so transformed by the Wars of
the Roses, the Reformation, the Rebellion, and the Revolution, that its
ancient outlines were barely visible. In so far, therefore, as his words
seem to imply that the social evolution he describes was produced by an
imperceptible and almost mechanical process of national instinct, the
impression they tend to create is entirely erroneous.
If we have been hitherto saved from such corruption as undermined the
republics of Italy, from the religious wars that so long enfeebled and
divided Germany, and from the Revolution that has severed modern France
from her ancient history, thanks for this are due partly, no doubt, to
favouring conditions of nature and society, but quite as much to the
genius of great individuals who prepared the mind of the nation for the
gradual assimilation of new ideas. Thus Langland and Wycliffe and their
numerous followers, long before the Reformation, had so familiarised the
minds of the people with their ideas of the Christian religion that the
Sovereign was able to assume the Headship of the Church without the shock
of a social convulsion. Fresh feelings and instincts grew up in the hearts
of whole classes of the nation without at first producing any change in
outward habits of life, and even without arousing a sense of their logical
incongruity. These mixed ideas were constantly brought before the
imagination in the works of the poets. Shakespeare abounds with passages
in which, side by side with the old feudal, monarchical, catholic, and
patriotic instincts of Englishmen, we find the sentiments of the Italian
Renaissance. Spenser conveys Puritan doctrines sometimes by the mouth of
shepherds, whose originals he had found in Theocritus and Virgil;
sometimes under allegorical forms derived from books of chivalry and the
ceremonial of the Catholic Church. Milton, the most rigidly Calvinistic of
all the English poets in his opinions, is also the most severely classical
in his style.
It was the task of Addison to carry on the reconciling traditions of our
literature. It is his praise to have accomplished his task under
conditions far more difficult than any that his predecessors had
experienced. What they had done was to give instinctive and characteristic
expression to the floating ideas of the society about them; what Addison
and his contemporaries did was to found a public opinion by a conscious
effort of reason and persuasion. Before the Civil Wars there had been at
least no visible breach in the principle of Authority in Church and State.
At the beginning of the eighteenth century constituted authority had been
recently overthrown; one king had been beheaded, another had been
expelled; the Episcopalian form of Church Government had been violently
displaced in favour of the Presbyterian, and had been with almost equal
violence restored. Whole classes of the population had been drawn into
opposing camps during the Civil War, and still stood confronting each
other with all the harsh antagonism of sentiment inherited from that
conflict. Such a bare summary alone is sufficient to indicate the nature
of the difficulties Addison had to encounter in his efforts to harmonise
public opinion; but a more detailed examination of the state of society
after the Restoration is required to place in its full light the
extraordinary merits of the success that he achieved.
There was, to begin with, a vehement opposition between town and country.
In the country the old ideas of Feudalism, modified by circumstances, but
vigorous and deep-rooted, still prevailed. True, the military system of
land-tenure had disappeared with the Restoration, but it was not so with
the relations of life, and the habits of thought and feeling which the
system had created. The features of surviving Feudalism have been
inimitably preserved for us in the character of Sir Roger de Coverley.
Living in the patriarchal fashion, in the midst of tenants and retainers,
who looked up to him as their chief, and for whose welfare and protection
he considered himself responsible, the country gentleman valued above all
things the principle of Loyalty. To the moneyed classes in the towns he
was instinctively opposed; he regarded their interests, both social and
commercial, as contrary to his own; he looked with dislike and suspicion
on the economical principles of government and conduct on which these
classes naturally rely. Even the younger sons of county families had in
Addison's day abandoned the custom, common enough in the feudal times, of
seeking their fortune in trade. Many a Will Wimble now spent his whole
life in the country, training dogs for his neighbours, fishing their
streams, making whips for their young heirs, and even garters for their
wives and daughters.[1]

8
data/sakana_wiki.txt
View file

@ -1,8 +0,0 @@
Sakana AI Co, Ltd. is a Japanese artificial intelligence company based in Tokyo.
Overview
Sakana AI's main research fields are evolution and collective intelligence of AI. The company's name is derived from the Japanese word さかな (sakana), which means fish. This represents the idea of a school of fish coming together and forming a coherent entity from simple rules, which is an analogy of collective intelligence.[2]
The company was founded by David Ha, Llion Jones and Ren Ito. Llion Jones co-authored the famous paper "Attention Is All You Need" when he was working for Google in 2017. The company raised $30M in its seed funding round from Lux Capital and Khosla Ventures.[3] The company raised approximately $200M from companies such as Mitsubishi UFJ, SMBC, Mizuho, Itochu, KDDI, Nomura and Nvidia in its series A funding round in 2024.[4]
In January 2024, Sakana AI developed a method to build new AI models by 'breeding' multiple existing models, which it sees as a means to democratise AI development, as this process does not require large computational resources.[5] Sakana AI is also developing a model called the AI Scientist, which automates the entire process of scientific research.[6] The Nikkei estimated the company's value at 19 billion yen in 2024.[7]

620
data/self_generate_qa.py
View file

@ -1,620 +0,0 @@
import argparse
import os
import random
import re
from glob import glob
import numpy as np
import yaml
from datasets import Dataset, load_dataset
from vllm import LLM, SamplingParams
from ctx_to_lora.data.definitions import (
CLOSED_QA_INTX_TEMPLATES,
RAW_DATA_DIR,
SELF_GEN_DATA_DIR,
)
from ctx_to_lora.data.processing import (
filter_none,
get_preprocessing_fn,
load_and_process_dataset,
tokenize_ctx_text,
)
from ctx_to_lora.data.self_gen_template import (
PRE_CTX,
PROMPT_TEMPLATE,
QA_PROMPT_TEMPLATE,
SELF_GEN_SYSTEM_MSG,
SELF_QA_INTX,
)
from ctx_to_lora.model_loading import get_tokenizer
from ctx_to_lora.utils import clear_gpu
STOP_STRINGS = {
"google/gemma-2-2b-it": ["<eos>", "<end_of_turn>"],
}
MODEL_CTX_LEN = {
"google/gemma-2-27b-it": 8192,
"google/gemma-2-2b-it": 8192,
"google/gemma-2-9b-it": 8192,
# qwen 4b has 256k ctx length but using lower max lengths is faster
"Qwen/Qwen3-4B-Instruct-2507": 2**13 + 2**12,
}
def truncate_middle_if_too_long(
input_ids: list[int],
max_length: int,
max_new_tokens: int = 256,
) -> list[int]:
"""
Truncate the middle of a list of tokens to fit within a maximum length.
Args:
tokens: List of token IDs
max_length: Maximum length for the truncated tokens
Returns:
List of truncated token IDs
"""
max_new_tokens_half = max_new_tokens // 2
# leave max_new_tokens for generation
half = max_length // 2 - max_new_tokens_half
if len(input_ids) > max_length:
return input_ids[:half] + input_ids[-half:]
return input_ids
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 add_closed_qa_prompt(q: str, closed_qa_prob: float = 0.1) -> str:
if random.random() <= closed_qa_prob:
q = random.choice(CLOSED_QA_INTX_TEMPLATES).format(input=q)
return q
def load_config(config_path: str) -> dict:
"""Load dataset names from YAML config file."""
with open(config_path) as f:
config = yaml.safe_load(f)
return config
def get_dataset_configs(
ds_names: list[str] | None,
config: dict | None,
split: str | None,
) -> list[tuple[str, str]]:
assert not (ds_names and config), "Cannot provide both ds_names and config"
if ds_names:
assert split, "When using ds_names, --split must be provided"
# Validate ds_names format
for ds_name in ds_names:
if not isinstance(ds_name, str):
raise ValueError(f"Invalid dataset name: {ds_name}")
return [(ds_name, split) for ds_name in ds_names]
if config:
dataset_configs = []
# Process train datasets
train_ds_names = config.get("train_ds_names", [])
# self_gen_train_ds_names = [
# (ds_name.split("/")[-1], "train")
# for ds_name in train_ds_names
# if ds_name.startswith("self_gen/")
# ]
self_gen_train_ds_names = [
(ds_name, "train")
for ds_name in train_ds_names
if ds_name.startswith("self_gen/")
]
if not self_gen_train_ds_names:
print("No self_gen datasets found in train_ds_names")
dataset_configs.extend(self_gen_train_ds_names)
# Process validation datasets
val_ds_names = config.get("val_ds_names", [])
self_gen_val_ds_names = [
(ds_name, "validation")
for ds_name in val_ds_names
if ds_name.startswith("self_gen/")
]
if not self_gen_val_ds_names:
print("No self_gen datasets found in val_ds_names")
dataset_configs.extend(self_gen_val_ds_names)
return dataset_configs
def create_messages(
ctxs: list[str],
questions: list[list[str]],
vllm_model: str,
system_template: str,
remove_qa_template: bool,
) -> list[list[dict]]:
"""Create chat messages for the model."""
# if "gemma" in vllm_model:
# gemma models do not support system messages
return [
[
{
"role": "user",
"content": (
system_template + "\n\n\n" + get_prompt(ctx, q, remove_qa_template)
).strip(),
}
]
for ctx, q_list in zip(ctxs, questions)
for q in q_list
]
# else:
# return [
# [
# {"role": "system", "content": system_template},
# {"role": "user", "content": get_prompt(ctx, q)},
# ]
# for ctx, q_list in zip(ctxs, questions)
# for q in q_list
# ]
def self_generate(
ds_name: str,
split: str,
args: argparse.Namespace,
llm: LLM,
system_template: str,
parquet_file: str | None = None,
do_truncate: bool = False,
) -> None:
"""Process a single dataset and generate QA pairs."""
shard_name = ""
# Conflict checks for ds_name-derived overrides
if ds_name is not None:
# temperature & closed_qa already handled later; add new ones
if "_temp_" in ds_name and args.temp != 0.0:
raise ValueError(
f"Multiple sources of truth for temperature: CLI arg --temp={args.temp} and dataset name contains temp specification."
)
if "_closed_qa_prob_" in ds_name and args.closed_qa_prob != 0.0:
raise ValueError(
f"Multiple sources of truth for closed_qa_prob: CLI arg --closed_qa_prob={args.closed_qa_prob} and dataset name contains closed_qa_prob specification."
)
# Base values from args
temp = args.temp
closed_qa_prob = args.closed_qa_prob
# Overrides from ds_name pattern if present
if ds_name is not None:
if "_temp_" in ds_name:
m = re.search(r"_temp_([\d.]+)", ds_name)
if m:
temp = float(m.group(1))
if "_closed_qa_prob_" in ds_name:
m = re.search(r"_closed_qa_prob_([\d.]+)", ds_name)
if m:
closed_qa_prob = float(m.group(1))
print(f"Processing dataset: {ds_name}, split: {split}")
print(f"Using temperature: {temp}")
print(f"Using closed QA prompt probability: {closed_qa_prob}")
if parquet_file:
print(f"Loading dataset from parquet file: {parquet_file}")
split = "train"
ds_name = "/".join(parquet_file.split(RAW_DATA_DIR)[-1].split("/")[:-1])
shard_name = "_" + os.path.basename(parquet_file).replace(".parquet", "")
ds = load_dataset(path="parquet", data_files=[parquet_file], split="train")
processing_fn = get_preprocessing_fn(ds_name, is_eval=False)
ds = ds.map(processing_fn, num_proc=8)
else:
ds_name = ds_name.split("/")[-1] # Extract just the dataset name
print(f"Loading dataset: {ds_name} with split: {split}")
kwargs = dict(ds_name=ds_name, split=split)
ds = load_and_process_dataset(**kwargs, num_proc=8, remove_cols=False)
print(f"Loaded dataset: {ds_name} with split: {split}")
if args.debug:
ds = ds.take(10)
ds = ds.filter(filter_none, batched=False, num_proc=8)
tk = get_tokenizer(args.vllm_model, train=True)
self_qa_intx_tokens = tk(SELF_QA_INTX, add_special_tokens=False)["input_ids"][1:]
if args.remove_qa_template:
self_qa_intx_tokens = tk("\n\n", add_special_tokens=False)["input_ids"]
n_self_qa_intx_tokens = len(self_qa_intx_tokens)
pre_ctx_tokens = tk(PRE_CTX, add_special_tokens=False)["input_ids"]
n_pre_ctx_tokens = len(pre_ctx_tokens)
sys_tokens = tk(system_template.split("\n")[0], add_special_tokens=False)[
"input_ids"
][:-1]
n_sys_tokens = len(sys_tokens)
os.environ["TOKENIZERS_PARALLELISM"] = "true"
ds = ds.map(
tokenize_ctx_text,
fn_kwargs={"tokenizer": tk},
batched=True,
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]
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}/{ds_name}/{split}/ds{shard_name}"
chunk_size = 1_000
for chunk_idx, start in enumerate(range(0, len(ctxs), chunk_size)):
print(f"Processing chunk {chunk_idx}")
chunk_ctxs = ctxs[start : start + chunk_size]
chunk_questions = questions[start : start + chunk_size]
chunk_messages = create_messages(
chunk_ctxs,
chunk_questions,
args.vllm_model,
SELF_GEN_SYSTEM_MSG,
args.remove_qa_template,
)
if do_truncate:
# we should only do this for evaluation data
tokenized_contents = tk(
[m[0]["content"] for m in chunk_messages],
add_special_tokens=False,
return_attention_mask=False,
)
tokenized_contents["input_ids"] = [
truncate_middle_if_too_long(
ids,
max_length=MODEL_CTX_LEN[args.vllm_model],
max_new_tokens=args.max_new_tokens,
)
for ids in tokenized_contents["input_ids"]
]
contents = tk.batch_decode(
tokenized_contents["input_ids"], skip_special_tokens=True
)
for c, m in zip(contents, chunk_messages):
m[0]["content"] = c
print(f"Generating from {len(chunk_messages)} contexts")
# Clear GPU memory before processing the next chunk
clear_gpu()
execute_qa_generation(
fpath + f"_{chunk_idx:04d}",
args,
llm,
temp,
tk,
self_qa_intx_tokens,
n_self_qa_intx_tokens,
sys_tokens,
n_sys_tokens,
chunk_ctxs,
ds[start : start + chunk_size]["ctx_ids"],
chunk_questions,
chunk_messages,
k,
)
def execute_qa_generation(
fpath,
args,
llm,
temp,
tk,
self_qa_intx_tokens,
n_self_qa_intx_tokens,
sys_tokens,
n_sys_tokens,
ctxs,
ctx_ids,
questions,
messages,
k,
):
completions = llm.chat(
messages,
sampling_params=SamplingParams(
max_tokens=args.max_new_tokens,
logprobs=k,
temperature=temp,
seed=42,
spaces_between_special_tokens=False,
skip_special_tokens=False,
include_stop_str_in_output=True,
),
)
self_gen_data = {
ctx: {
"ctx_ids": ctx_ids,
"input_ids": [],
"response_start_end": [],
"logprobs_vals": [],
"logprobs_indices": [],
}
for ctx, ctx_ids in zip(ctxs, ctx_ids)
}
c = 0
n_skips = 0
sys_start = None
for ctx, q_list in zip(ctxs, questions):
# self_gen_data[ctx]["ctx_ids"] = ctx_ids
for i, _ in enumerate(q_list):
# response = completions[c + i].outputs[0].text
reason = completions[c + i].outputs[0].finish_reason
if reason != "stop":
# print(f"idx: {c + i}")
print(f"finish_reason: {completions[c + i].outputs[0].finish_reason}")
print(f"Skipping due to finish_reason={reason} != 'stop'")
n_skips += 1
continue
# includes the logprob before the first response token
# but excludes the logprob from eos token
logp = completions[c + i].outputs[0].logprobs
# len = num response tokens
n_response_tokens = len(completions[c + i].outputs[0].token_ids)
logp_indices = np.empty((n_response_tokens, k), dtype=np.int32)
# float-16 is better for this range
logp_vals = np.empty((n_response_tokens, k), dtype=np.float16)
assert len(logp) == n_response_tokens, (
f"Expected {n_response_tokens} logp entries, got {len(logp)}"
)
for li, info_d in enumerate(logp):
for j, (idx, tok_info) in enumerate(info_d.items()):
logp_indices[li, j] = idx
logp_vals[li, j] = tok_info.logprob
prompt_ids = completions[c + i].prompt_token_ids # 1d list
# token_ids only includes generated tokens, not the prompt
response_token_ids = completions[c + i].outputs[0].token_ids # 1d list
all_ids = prompt_ids + response_token_ids
res_start = len(prompt_ids)
res_end = res_start + n_response_tokens
if sys_start is None:
for ii in range(len(prompt_ids) - n_sys_tokens):
if prompt_ids[ii : ii + n_sys_tokens] == sys_tokens:
# found the start of the system message
sys_start = ii
break
q_start = None
for ii in range(
len(prompt_ids) - n_self_qa_intx_tokens,
-1,
-1,
):
if prompt_ids[ii : ii + n_self_qa_intx_tokens] == self_qa_intx_tokens:
# found the start of the user input
q_start = ii + n_self_qa_intx_tokens
break
# bos + question + eos + start model turn + response + eos
input_ids = all_ids[:sys_start] + all_ids[q_start:res_end]
# relative to the input_ids
res_start = res_start - q_start + sys_start
res_end = res_start + n_response_tokens
# arrays will be saved as nested lists of numbers
self_gen_data[ctx]["input_ids"].append(input_ids)
# assume single-turn chat
self_gen_data[ctx]["response_start_end"].append((res_start, res_end))
self_gen_data[ctx]["logprobs_vals"].append(logp_vals)
self_gen_data[ctx]["logprobs_indices"].append(logp_indices)
c += i + 1
print(f"Skipped {n_skips} responses due to missing stop strings")
samples = [
{
# "context": ctx,
# "prompts": q_list,
# "responses": self_gen_data[ctx]["responses"],
"ctx_ids": self_gen_data[ctx]["ctx_ids"],
"input_ids": self_gen_data[ctx]["input_ids"],
"response_start_end": self_gen_data[ctx]["response_start_end"],
# "prompt_start_end": self_gen_data[ctx]["prompt_start_end"],
"logprobs_vals": self_gen_data[ctx]["logprobs_vals"],
"logprobs_indices": self_gen_data[ctx]["logprobs_indices"],
}
for ctx, q_list in zip(ctxs, questions)
]
if args.debug:
for sample in samples:
# print(f"context={tk.decode(sample['ctx_ids'])}")
print(f"QA={[tk.decode(ids) for ids in sample['input_ids']]}")
for input_ids, (start, end) in zip(
sample["input_ids"], sample["response_start_end"]
):
print(f"start={start}, end={end}")
print(f"response={tk.decode(input_ids[start:end])}")
print(f"logprobs_vals={[x.shape for x in sample['logprobs_vals']]}")
print(f"logprobs_indices={[x.shape for x in sample['logprobs_indices']]}")
for indices in sample["logprobs_indices"]:
print(f"logprobs_indices={indices[-1]}")
print("=" * 80)
print(f"Generated {len(samples)} samples")
# random.shuffle(samples)
# Save results
# df = pd.DataFrame(samples)
# ds_out = Dataset.from_pandas(df)
ds_out = Dataset.from_list(samples)
# fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}/{ds_name}/{split}/ds{shard_name}"
if args.debug:
fpath += "_debug"
os.makedirs(os.path.dirname(fpath), exist_ok=True)
fpath = f"{fpath}.parquet"
ds_out.to_parquet(fpath)
print(f"Saved to {fpath}")
# Cleanup
del samples, ds_out, completions, messages, ctxs, questions
clear_gpu()
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate QA pairs using VLLM")
parser.add_argument(
"--vllm_model",
type=str,
required=True,
help="VLLM model name (e.g., google/gemma-2-2b-it)",
)
parser.add_argument(
"--debug",
action="store_true",
help="Enable debug mode (process only 10 samples)",
)
# Either config file OR ds_names + split
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument(
"--config",
type=str,
help="Path to YAML config file with train_ds_names/val_ds_names",
)
group.add_argument(
"--ds_names",
type=str,
nargs="+",
help="List of dataset names/shard patterns",
)
group.add_argument(
"--glob_pattern",
type=str,
help="Glob pattern to match dataset names (e.g., 'data/raw_datasets/fw_qa_3/*')",
)
parser.add_argument(
"--split",
type=str,
help="Dataset split to use when using --ds_names (required with --ds_names)",
)
parser.add_argument(
"--temp",
type=float,
default=0.0,
help="Temperature for sampling (default: 0.0)",
)
parser.add_argument(
"--closed_qa_prob",
type=float,
default=0.0,
help="Probability of using closed QA prompt template (default: 0.0)",
)
parser.add_argument(
"--do_truncate",
action="store_true",
help="Truncate contexts to fit model context length",
)
parser.add_argument(
"--remove_qa_template",
action="store_true",
help="Remove QA template formatting from prompts",
)
parser.add_argument(
"--max_new_tokens",
type=int,
default=256,
help="Maximum number of new tokens to generate (default: 256)",
)
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
# Validate arguments
if args.ds_names and not args.split:
raise ValueError("--split is required when using --ds_names")
vllm_model = args.vllm_model
print(f"Using model: {vllm_model}")
# Setup model-specific configurations
llm_kwargs = dict(
model=vllm_model,
dtype="bfloat16",
enable_prefix_caching=True,
enable_chunked_prefill=True,
max_model_len=MODEL_CTX_LEN.get(vllm_model),
max_num_batched_tokens=16384,
max_num_seqs=32, # avoid oom when getting logprobs
)
print(f"{llm_kwargs=}")
llm = LLM(**llm_kwargs)
# Get dataset configs from config or CLI args
config = load_config(args.config) if args.config else None
if args.ds_names or args.config:
dataset_configs = get_dataset_configs(
ds_names=args.ds_names,
config=config,
split=args.split,
)
# Process each dataset
for ds_name, split in dataset_configs:
print(f"Processing dataset: {ds_name}, split: {split}")
self_generate(
ds_name, split, args, llm, SELF_GEN_SYSTEM_MSG, None, args.do_truncate
)
else:
assert args.glob_pattern, (
"glob_pattern must be provided if no ds_names or config"
)
files = glob(args.glob_pattern)
for file in files:
print(f"Processing file: {file}")
self_generate(
ds_name=None,
parquet_file=file,
split=args.split,
args=args,
llm=llm,
system_template=SELF_GEN_SYSTEM_MSG,
do_truncate=args.do_truncate,
)

685
demo/app.py
View file

@ -1,685 +0,0 @@
import os
import sys
from pathlib import Path
import gradio as gr
import torch
# Add the src directory to the path
sys.path.insert(0, str(Path(__file__).parent.parent))
from ctx_to_lora.data.processing import tokenize_ctx_text
from ctx_to_lora.model_loading import get_tokenizer
from ctx_to_lora.modeling import hypernet
sys.modules["ctx_to_lora.modeling_utils"] = hypernet
# Global state
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
modulated_model = None
chat_history = []
ctx_tokenizer = None
base_tokenizer = None
try:
DEFAULT_CONTEXT = Path("data/sakana_wiki.txt").read_text(encoding="utf-8").strip()
except FileNotFoundError:
DEFAULT_CONTEXT = ""
WARNING_MESSAGE = (
"⚠️ **Caution**: This is an educational proof-of-concept demonstration.\n"
"The model may generate inaccurate information or hallucinate facts."
)
FOOTER = """
⚠️ This model is an experimental prototype and is only available for educational and research and development purposes. It is not suitable for commercial use or in environments where failure can have significant effects (mission-critical environments).
The use of this model is at the user's own risk and its performance and results is not guaranteed in any way.
Sakana AI is not responsible for any direct or indirect loss resulting from using this model, regardless of the outcome.
"""
def load_custom_chat_template(tokenizer, model_name):
if "gemma" in model_name.lower():
template_path = "chat_templates/google/gemma-2-2b-it.jinja"
if os.path.exists(template_path):
with open(template_path) as f:
template_content = f.read()
tokenizer.chat_template = template_content
print(f"Loaded custom chat template from {template_path}")
return True
return False
def get_available_checkpoints():
trained_d2l_checkpoints = {
str(path)
for path in Path().glob("trained_d2l/**/pytorch_model.bin")
if path.is_file()
}
run_output_checkpoints = {
str(path)
for path in Path().glob("train_outputs/runs/**/pytorch_model.bin")
if path.is_file()
}
checkpoints = sorted(trained_d2l_checkpoints) + sorted(
run_output_checkpoints - trained_d2l_checkpoints
)
return checkpoints if checkpoints else ["No checkpoints found"]
def load_checkpoint(
checkpoint_path: str,
) -> tuple[str, gr.update, gr.update, gr.update, gr.update]:
global modulated_model, ctx_tokenizer, base_tokenizer, chat_history
if not checkpoint_path or checkpoint_path == "No checkpoints found":
return (
"⚠️ Please select a valid checkpoint",
gr.update(),
gr.update(),
gr.update(),
gr.update(),
)
try:
print(f"Loading checkpoint: {checkpoint_path}")
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
state_dict = torch.load(checkpoint_path, weights_only=False)
modulated_model = ModulatedPretrainedModel.from_state_dict(
state_dict,
train=False,
use_flash_attn=True,
use_sequence_packing=False,
)
modulated_model = modulated_model.to(device).to(torch.bfloat16)
modulated_model.eval()
ctx_encoder_model_name_or_path = (
modulated_model.ctx_encoder_args.ctx_encoder_model_name_or_path
or modulated_model.base_model.config.name_or_path
)
ctx_tokenizer = get_tokenizer(ctx_encoder_model_name_or_path, train=False)
base_tokenizer = get_tokenizer(
modulated_model.base_model.config.name_or_path, train=False
)
load_custom_chat_template(
base_tokenizer, modulated_model.base_model.config.name_or_path
)
chat_history = [{"role": "system", "content": ""}]
model_name = modulated_model.base_model.config.name_or_path
success_msg = (
f"✅ Successfully loaded checkpoint!\n\nBase Model: {model_name}\n\n"
"You can now add context and start chatting."
)
return (
success_msg,
gr.update(interactive=True), # msg
gr.update(interactive=True), # send_btn
gr.update(interactive=True), # system_msg
gr.update(interactive=True), # clear_btn
)
except Exception as e:
import traceback
error_msg = (
f"❌ Error loading checkpoint:\n{str(e)}\n\n{traceback.format_exc()}"
)
print(error_msg)
return (
error_msg,
gr.update(interactive=False),
gr.update(interactive=False),
gr.update(interactive=False),
gr.update(interactive=False),
)
def process_context(context: str) -> dict:
context = context.strip() if context else ""
tokenized_contexts = tokenize_ctx_text({"context": [context]}, ctx_tokenizer)
ctx_ids = tokenized_contexts["ctx_ids"]
ctx_ids = [
torch.tensor(ctx_id, dtype=torch.long, device=device) for ctx_id in ctx_ids
]
ctx_attn_mask = [torch.ones_like(ids) for ids in ctx_ids]
ctx_attn_mask = [
torch.tensor(mask, dtype=torch.long, device=device) for mask in ctx_attn_mask
]
ctx_ids = torch.nn.utils.rnn.pad_sequence(
ctx_ids,
batch_first=True,
padding_value=0,
)
ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
ctx_attn_mask,
batch_first=True,
padding_value=0,
)
return {"ctx_ids": ctx_ids, "ctx_attn_mask": ctx_attn_mask}
def add_user_message(message: str, history):
if not message.strip():
return history, ""
return history + [[message, None]], ""
def generate_response(
history: list[list[str]],
system_msg: str,
context: str,
context_scaler: float,
bias_scaler: float,
):
global modulated_model, chat_history, ctx_tokenizer, base_tokenizer
if modulated_model is None:
history[-1][1] = "Please load a checkpoint first."
yield history
return
if not history or history[-1][0] is None:
yield history
return
try:
user_message = history[-1][0]
if system_msg.strip() and chat_history[0]["role"] == "system":
chat_history[0]["content"] = system_msg.strip()
chat_history.append({"role": "user", "content": user_message})
context = context.strip() if context else ""
print(f"Processing single context with scaler: {context_scaler}")
print(f"Bias scaler: {bias_scaler}")
with torch.inference_mode(), torch.amp.autocast(str(device)):
ctx_inputs = process_context(context)
ctx_ids = ctx_inputs["ctx_ids"].to(device)
ctx_attn_mask = ctx_inputs["ctx_attn_mask"].to(device)
scalers_tensor = torch.tensor(
[context_scaler], dtype=torch.float32, device=device
)
model_inputs = base_tokenizer.apply_chat_template(
chat_history, return_tensors="pt", add_generation_prompt=True
).to(device)
print(f"Context: {context}")
print(f"Chat history: {chat_history}")
outputs = modulated_model.generate(
ctx_ids=ctx_ids,
ctx_attn_mask=ctx_attn_mask,
n_ctx_chunks=torch.tensor([len(ctx_ids)], device=ctx_ids.device),
scalers=scalers_tensor,
bias_scaler=bias_scaler,
input_ids=model_inputs,
max_new_tokens=512,
do_sample=False,
temperature=0,
)
response = base_tokenizer.decode(
outputs[0][model_inputs.shape[1] :], skip_special_tokens=True
)
chat_history.append({"role": "assistant", "content": response})
words = response.split()
partial_response = ""
for word in words:
partial_response += word + " "
history[-1][1] = partial_response.strip()
yield history
history[-1][1] = response
yield history
except Exception as e:
import traceback
error_msg = f"❌ Error: {str(e)}"
print(f"Error generating response: {str(e)}\n\n{traceback.format_exc()}")
history[-1][1] = error_msg
yield history
def reset_chat(system_msg: str):
global chat_history
chat_history = [
{"role": "system", "content": system_msg.strip() if system_msg else ""}
]
return [[None, WARNING_MESSAGE]], "Chat history reset successfully!"
custom_css = """
:root {
color-scheme: light;
}
.gradio-container {
font-family: 'Inter', sans-serif;
}
.chat-container {
border-radius: 10px;
border: 2px solid #d1d5db;
}
.context-field {
border-radius: 8px;
padding: 15px;
margin-bottom: 10px;
border: 2px solid #d1d5db;
}
.status-box {
border-radius: 8px;
padding: 15px;
margin: 10px 0;
border: 2px solid #e5e7eb;
}
.primary-button {
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
border: none;
border-radius: 6px;
color: white;
font-weight: 600;
}
.secondary-button {
background-color: #607d8b;
border: none;
border-radius: 6px;
color: white;
}
#chatbot {
height: 500px;
}
.instruction-text {
font-style: italic;
color: #666;
font-size: 0.9em;
margin-bottom: 10px;
}
.warning-box {
background-color: #fff3cd;
border: 2px solid #ffc107;
border-radius: 8px;
padding: 12px;
margin: 10px 0;
color: #856404;
font-size: 0.95em;
}
.warning-box strong {
color: #d97706;
}
.disabled-overlay {
background-color: #f5f5f5;
border: 3px dashed #999;
border-radius: 10px;
padding: 20px;
text-align: center;
color: #999;
}
.chat-disabled-notice {
border: 3px solid #f59e0b;
border-radius: 8px;
padding: 15px;
margin-bottom: 15px;
color: #92400e;
font-weight: 600;
text-align: center;
background-color: #fef3c7;
box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
}
.internalization-banner {
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
color: white;
border-radius: 10px;
padding: 20px;
margin-bottom: 15px;
text-align: center;
font-weight: 600;
box-shadow: 0 4px 8px rgba(0, 0, 0, 0.15);
border: 3px solid #5568d3;
}
.internalization-banner h3 {
margin: 0 0 10px 0;
font-size: 1.2em;
color: white;
}
.internalization-banner p {
margin: 5px 0;
font-size: 0.95em;
font-weight: 400;
color: white;
}
.context-section-header {
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
border: 3px solid #5568d3;
border-left: 6px solid #4c51bf;
padding: 15px;
margin-bottom: 15px;
border-radius: 6px;
color: white;
box-shadow: 0 2px 6px rgba(102, 126, 234, 0.3);
}
.context-section-header strong,
.context-section-header small {
color: white;
}
.chat-section-header {
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
border: 3px solid #2563eb;
border-left: 6px solid #1d4ed8;
padding: 15px;
margin-bottom: 15px;
border-radius: 6px;
color: white;
box-shadow: 0 2px 6px rgba(59, 130, 246, 0.3);
}
.chat-section-header strong,
.chat-section-header small {
color: white;
}
.panel-box {
background-color: rgba(249, 250, 251, 0.5);
border: 2px solid rgba(209, 213, 219, 0.5);
border-radius: 12px;
padding: 20px;
box-shadow: 0 4px 6px rgba(0, 0, 0, 0.05);
}
.chat-panel-box {
background-color: rgba(249, 250, 251, 0.5);
border: 2px solid rgba(209, 213, 219, 0.5);
border-radius: 12px;
padding: 20px;
box-shadow: 0 4px 6px rgba(0, 0, 0, 0.05);
}
#checkpoint-dropdown {
position: relative;
z-index: 20;
}
#checkpoint-dropdown [role="listbox"] {
z-index: 9999 !important;
}
/* Dark mode support */
.dark .panel-box,
.dark .chat-panel-box {
background-color: rgba(31, 41, 55, 0.5);
border: 2px solid rgba(75, 85, 99, 0.5);
}
.dark .context-field,
.dark .chat-container {
border-color: rgba(75, 85, 99, 0.5);
}
.dark .status-box {
border-color: rgba(55, 65, 81, 0.5);
}
"""
def create_demo():
with gr.Blocks(
title="Doc-to-LoRA Chat Interface",
theme=gr.themes.Soft(),
css=custom_css,
) as demo:
gr.Markdown(
"""
# 📜 Doc-to-LoRA Chat Interface
Load a hypernetwork checkpoint and chat with a context-modulated language model.
Add one context with a scaling parameter to influence the model's responses.
"""
)
gr.HTML(
"""
<div class="internalization-banner">
<h3>🧠 How Context Internalization Works</h3>
<p>📥 Contexts are processed by the hypernetworkto dynamically modulate the base model's parameters</p>
<p>🚫 Contexts are NOT passed as text to the base model — they influence behavior internally</p>
<p>💬 Only your chat messages (below) are sent to the language model</p>
</div>
"""
)
gr.Markdown(
"""
### 📖 Usage Instructions
1. **Load a Checkpoint**: Select a hypernetwork checkpoint from the dropdown and click "Load Checkpoint"
2. **Configure Context**:
- Enter your context information in the text field
- Adjust the scaling slider to control context influence
3. **Set Bias Scaler**: Adjust the bias scaler to control overall model behavior
4. **Start Chatting**: Once the model is loaded, type your message and press Shift+Enter or click Send
5. **Reset**: Use the "Reset Chat" button to start a new conversation
💡 **Tip**: You can use context to provide background information or specific knowledge
that should influence the model's responses.
"""
)
with gr.Row():
with gr.Column(scale=1, elem_classes="panel-box"):
gr.Markdown("### 📦 Load Checkpoint")
gr.Markdown(
"*Select a trained hypernetwork checkpoint to begin.*",
elem_classes="instruction-text",
)
checkpoint_dropdown = gr.Dropdown(
choices=get_available_checkpoints(),
label="Select Checkpoint",
value=None,
interactive=True,
elem_id="checkpoint-dropdown",
)
load_btn = gr.Button("Load Checkpoint", variant="primary", size="lg")
status_box = gr.Textbox(
label="Status",
lines=8,
interactive=False,
elem_classes="status-box",
)
gr.Markdown("---")
gr.HTML(
"""
<div class="context-section-header">
<strong>🧠 Context Internalization (Hypernetwork Input)</strong><br>
<small>This context modulates the model internally — it is NOT shown to the base model</small>
</div>
"""
)
context = gr.Textbox(
label="🧠 Context (Internalized via Hypernetwork)",
placeholder="Enter context to be internalized by the hypernetwork...",
lines=4,
value=DEFAULT_CONTEXT,
)
context_scaler = gr.Slider(
minimum=-2.0,
maximum=2.0,
step=0.01,
value=1.0,
label="Context Scaling",
)
gr.Markdown("---")
bias_scaler = gr.Slider(
minimum=-2.0,
maximum=2.0,
step=0.01,
value=1.0,
label="Bias Scaler",
info="A single scalar applied to bias parameters (independent of contexts)",
)
with gr.Column(scale=2, elem_classes="chat-panel-box"):
gr.HTML(
"""
<div class="chat-section-header">
<strong>💬 Chat Interface (Direct Input to Base Model)</strong><br>
<small>Your messages here are the ONLY text the base model sees — contexts above influence it internally</small>
</div>
"""
)
chat_status_notice = gr.HTML(
"""
<div class="chat-disabled-notice">
🔒 <strong>Chat Disabled:</strong> Please load a checkpoint first to enable chat functionality.
</div>
""",
visible=True,
)
system_msg = gr.Textbox(
label="System Message (Optional - Sent to Base Model)",
placeholder="Load a checkpoint to enable chat...",
lines=2,
interactive=False,
)
chatbot = gr.Chatbot(
label="Conversation",
show_copy_button=True,
height=500,
elem_id="chatbot",
elem_classes="chat-container",
value=[
[
None,
"🔒 Chat is currently disabled. Please load a checkpoint from the left panel to begin chatting.",
]
],
)
with gr.Row():
msg = gr.Textbox(
label="Your Message (Sent Directly to Base Model)",
placeholder="⚠️ Load a checkpoint first to start chatting...",
lines=2,
scale=4,
interactive=False,
)
send_btn = gr.Button(
"🔒 Send (Disabled)",
variant="primary",
scale=1,
interactive=False,
)
with gr.Row():
clear_btn = gr.Button(
"🔒 Reset Chat (Disabled)",
variant="secondary",
interactive=False,
)
reset_status = gr.Textbox(label="Reset Status", visible=False)
load_btn.click(
fn=load_checkpoint,
inputs=[checkpoint_dropdown],
outputs=[status_box, msg, send_btn, system_msg, clear_btn],
).then(
fn=lambda: (
gr.update(visible=False),
gr.update(
placeholder="Type your message here... (Shift+Enter for new line)"
),
gr.update(value="Send"),
gr.update(value="🔄 Reset Chat"),
gr.update(value=[[None, WARNING_MESSAGE]]),
),
outputs=[chat_status_notice, msg, send_btn, clear_btn, chatbot],
)
msg.submit(
fn=add_user_message,
inputs=[msg, chatbot],
outputs=[chatbot, msg],
).then(
fn=generate_response,
inputs=[
chatbot,
system_msg,
context,
context_scaler,
bias_scaler,
],
outputs=[chatbot],
)
send_btn.click(
fn=add_user_message,
inputs=[msg, chatbot],
outputs=[chatbot, msg],
).then(
fn=generate_response,
inputs=[
chatbot,
system_msg,
context,
context_scaler,
bias_scaler,
],
outputs=[chatbot],
)
clear_btn.click(
fn=reset_chat,
inputs=[system_msg],
outputs=[chatbot, reset_status],
)
gr.Markdown(f"---\n{FOOTER.strip()}")
return demo
if __name__ == "__main__":
demo = create_demo()
demo.launch(
server_name="0.0.0.0",
server_port=7861,
share=False,
debug=True,
)

41
examples/python_api.py
View file

@ -1,41 +0,0 @@
# caveat: this interface only supports non-batched inputs
# for batched inference please see `src/ctx_to_lora/modeling/hypernet.py`
import torch
from ctx_to_lora.model_loading import get_tokenizer
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
# model loading
checkpoint_path = "trained_d2l/gemma_demo/checkpoint-80000/pytorch_model.bin"
state_dict = torch.load(checkpoint_path, weights_only=False)
model = ModulatedPretrainedModel.from_state_dict(
state_dict, train=False, use_sequence_packing=False
)
model.reset()
tokenizer = get_tokenizer(model.base_model.name_or_path)
# prepare data
doc = open("data/sakana_wiki.txt", "r").read()
chat = [{"role": "user", "content": "Tell me about Sakana AI."}]
chat_ids = tokenizer.apply_chat_template(
chat,
add_special_tokens=False,
return_attention_mask=False,
add_generation_prompt=True,
return_tensors="pt",
).to(model.device)
# calls after internalization will be influenced by internalized info
model.internalize(doc)
outputs = model.generate(input_ids=chat_ids, max_new_tokens=512)
print(tokenizer.decode(outputs[0]))
# remove internalized info
# model.reset()
# without internalized info, the model will halucinate
# outputs = model.generate(input_ids=chat_ids, max_new_tokens=512)
# print(tokenizer.decode(outputs[0]))

24
install.sh
View file

@ -1,24 +0,0 @@
curl -LsSf https://astral.sh/uv/install.sh | sh
uv self update
uv venv --python 3.10 --seed
uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --torch-backend=cu124
uv sync
uv pip install tokenizers==0.21.0
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
# 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
# optional: needed for logging with wandb
# wandb login
# optional: dev
# uv run pre-commit install

86
pyproject.toml
View file

@ -1,50 +1,49 @@
[project]
name = "ctx-to-lora"
version = "0.0.1"
authors = [{ name = "Rujikorn Charakorn" }]
description = ""
name = "video2lora"
version = "0.1.0"
authors = [{ name = "Video2LoRA authors" }]
description = "Video-conditioned LoRA generation for vision-language models."
readme = "README.md"
requires-python = ">= 3.10"
requires-python = ">=3.10"
dependencies = [
"transformers==4.51.3",
"deepspeed==0.17.1",
"accelerate==1.6.0",
"av>=17.0.0",
"bitsandbytes>=0.46.1",
"datasets==3.6.0",
"setuptools",
"peft",
"jupyter",
"matplotlib",
"hf_transfer",
"torchmetrics",
"inflect",
"pre-commit",
"tensorboardX",
"wandb",
"fasttext-wheel",
"decord",
"einops",
"gdown>=5.1.0",
"hf_transfer",
"huggingface-hub[hf-transfer]>=0.32.0",
"jaxtyping",
"liger-kernel",
"tensorboard",
"flask",
"gradio>=4.40.0",
"pandas",
"plotly",
"rouge-score",
"vllm==0.8.5.post1",
"huggingface-hub[hf-transfer]>=0.32.0",
"matplotlib",
"num2words",
"numpy",
"opt-einsum>=3.4.0",
"kagglehub[hf-datasets]>=0.3.12",
"kaggle>=1.7.4.5",
"bitsandbytes>=0.46.1",
"google-cloud-storage>=3.2.0",
"wonderwords>=2.2.0",
"llmlingua>=0.2.2",
"pandas",
"peft",
"plotly",
"pyyaml",
"rouge-score",
"setuptools",
"tensorboard",
"tensorboardX",
"tokenizers",
"torch>=2.6.0",
"torchmetrics",
"torchvision>=0.21.0",
"transformers==4.51.3",
"wandb",
]
[build-system]
requires = ["setuptools>=61.0"]
build-backend = "setuptools.build_meta"
[tool.setuptools.packages.find]
where = ["src"]
[tool.pyright]
exclude = [
"**/node_modules",
@ -53,32 +52,19 @@ exclude = [
".venv",
".github",
".vscode",
"chat_templates",
"eval_results",
"configs",
"EditingLlama",
"icae_v2",
"lm-evaluation-harness",
"llm-comparator",
"LongBench",
"scripts",
"train_outputs",
"./data/",
"/data/",
"generated_tasks",
"assets",
"checkpoints",
"data",
"outputs",
"plots",
"runs",
"tmp",
"wandb",
".wandb",
".ruff_cache",
"assets",
]
typeCheckingMode = "off"
[tool.ruff]
line-length = 88
select = ["F401"] # remove unused imports
select = ["F401"]
ignore = ["E", "F"]
[tool.isort]

180
run_eval.py
View file

@ -1,180 +0,0 @@
import logging
from ctx_to_lora.eval_utils import run_eval
logger = logging.getLogger()
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Evaluate a checkpoint")
parser.add_argument(
"--model_name_or_path",
type=str,
default=None,
help="Evaluate a base model from HuggingFace Hub, without loading checkpoint",
)
parser.add_argument(
"--checkpoint_path",
type=str,
default=None,
help="Path to the checkpoint to evaluate",
)
parser.add_argument(
"--split",
type=str,
choices=["validation", "test"],
default="validation",
help="Which split to evaluate on",
)
parser.add_argument(
"--datasets",
type=str,
nargs="+",
help=(
"Specific datasets to evaluate on."
"If not provided, uses default from args.yaml"
),
)
parser.add_argument(
"--eval_batch_size",
type=int,
default=8,
help="Eval batch size for teacher forcing",
)
parser.add_argument(
"--eval_batch_size_gen",
type=int,
default=32,
help="Eval batch size for generation",
)
parser.add_argument(
"--max_val_samples_per_ds",
type=int,
default=-1,
help=(
"Maximum number of validation samples per dataset. "
"If -1, uses values from checkpoint config."
),
)
parser.add_argument(
"--max_test_samples_per_ds",
type=int,
default=500,
help=(
"Maximum number of validation samples per dataset. "
"If -1, uses values from checkpoint config."
),
)
parser.add_argument(
"--max_ctx_chunk_len",
type=int,
default=-1,
help="Maximum length of context chunk for evaluation",
)
parser.add_argument(
"--max_new_tokens",
type=int,
default=256,
help="Maximum number of new tokens to generate during evaluation",
)
parser.add_argument(
"--remove_context",
action="store_true",
help="Remove context when evaluating the base model.",
)
parser.add_argument(
"--use_cd",
action="store_true",
help="Use context distillation model for evaluation.",
)
parser.add_argument(
"--cd_update_iterations",
type=int,
default=20,
help="Number of update iterations for context distillation during evaluation",
)
parser.add_argument(
"--cd_use_gen_q",
action="store_true",
help="Use generated queries for context distillation training.",
)
parser.add_argument(
"--q_gen_rounds",
type=int,
default=4,
help="Number of rounds of query generation for context distillation.",
)
parser.add_argument(
"--cd_batch_size",
type=int,
default=16,
help="Batch size for context distillation.",
)
parser.add_argument(
"--use_iterative_mode",
action="store_true",
help="Use iterative mode LoRA layer-by-layer generation",
)
parser.add_argument(
"--use_llmlingua",
action="store_true",
help="Use LLMLingua compression for evaluation",
)
parser.add_argument(
"--llmlingua_compression_rate",
type=float,
default=0.9,
help="Compression rate for LLMLingua",
)
parser.add_argument(
"--use_t2l",
action="store_true",
help="Use Text-to-LoRA model for evaluation",
)
parser.add_argument(
"--add_ctx_to_input",
action="store_true",
help="Add ctx to base model's input",
)
parser.add_argument(
"--truncate_if_too_long_inp",
action="store_true",
help="Truncate input sequences that are too long",
)
parser.add_argument(
"--truncate_if_too_long_ctx",
action="store_true",
help="Truncate ctx sequences that are too long",
)
parser.add_argument(
"--gen_lora_scaling",
type=float,
default=1.0,
)
parser.add_argument(
"--flip_ctx_inp",
action="store_true",
help="Flip the order of context and input",
)
parser.add_argument(
"--use_generative_adapter",
action="store_true",
help="Use generative adapter for evaluation",
)
cli_args = vars(parser.parse_args())
if cli_args["model_name_or_path"]:
assert cli_args["max_ctx_chunk_len"] <= 0, (
f"Evaluating base model shouldn't be used with `max_ctx_chunk_len`"
)
eval_batch_size_gen = cli_args.pop("eval_batch_size_gen")
eval_batch_size = cli_args.pop("eval_batch_size")
run_eval(
**cli_args,
eval_batch_size=eval_batch_size_gen,
generative=True,
)

1
scripts/__init__.py Normal file
View file

@ -0,0 +1 @@

17
scripts/main_exp/0-download_data.py
View file

@ -1,17 +0,0 @@
from huggingface_hub import snapshot_download
if __name__ == "__main__":
self_gen_data_dir = "./data/raw_datasets/self_gen/"
snapshot_download(
"SakanaAI/self_gen_qa_d2l",
repo_type="dataset",
local_dir=self_gen_data_dir,
# we can filter based on model by using the `allow_patterns` argument
# based on https://huggingface.co/datasets/SakanaAI/self_gen_qa_d2l/tree/main
# we can use
# - `Qwen` for downloading the data for `Qwen/Qwen3-4B-Instruct-2507`
# - `google` for downloading the data for `google/gemma-2-2b-it`
# - `mistralai` for downloading the data for `mistralai/Mistral-7B-Instruct-v0.2`
#
# allow_patterns="google/*", # downloading the data for `google/gemma-2-2b-it`
)

14
scripts/main_exp/1-train.sh
View file

@ -1,14 +0,0 @@
#!/bin/bash
port=29051
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 \
--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

30
scripts/main_exp/2-train-chunk.sh
View file

@ -1,30 +0,0 @@
#!/bin/bash
port=29051
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 \
--model_name_or_path=google/gemma-2-2b-it \
--target_modules=down_proj \
--lora_r=8 \
--eval_strategy=no \
--max_qas_len=512 \
--max_qas_per_sample=1 \
--per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1 \
--max_steps=20000 \
--gradient_accumulation_steps=16 \
--max_packed_inp_len=1024 \
--max_packed_ctx_len=2048 \
--use_per_ctx_average_loss=True \
--use_kl_loss=True \
--quantize_ctx_encoder=True \
--torch_empty_cache_steps=10 \
--from_pretrained_checkpoint=train_outputs/runs/$RUN_NAME/checkpoint-80000/pytorch_model.bin \
--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"}' \
--warmup_steps=2000 \
--learning_rate=2e-5

25
scripts/main_exp/README.md
View file

@ -1,25 +0,0 @@
# D2L pipeline
### Data
You can either download the generated data (recommended, ~100 GB for each model) or generate them by youself.
Please see [`0-download_data.sh`](0-download_data.sh) for how to do model-specific data download.
```bash
# download training data for all three models (328GB)
uv run bash scripts/main_exp/0-download_data.sh
```
Generating data from scratch can take very long if not parallelized across multiple gpus.
```bash
# generate training data (takes very long if not parallelized across multiple gpus)
# optional: use the command below for generating data from scratch
# uv run bash scripts/main_exp/gen_data.sh
```
### Training
Simply run the training script once the data is ready.
```bash
# train
uv run bash scripts/main_exp/1-train.sh
```
### Evaluation
All evaluation scripts for reproducing the main results in the paper are included in [eval](eval/) directory.

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

@ -1,8 +0,0 @@
# 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 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
View file

@ -1,6 +0,0 @@
# qa
WANDB_MODE=disabled 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
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/qasper_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

1
scripts/main_exp/eval/cd_minibatch.sh
View file

@ -1 +0,0 @@
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets $1 --split test --use_cd --cd_update_iterations 50 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=5 --cd_batch_size=2

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

@ -1,6 +0,0 @@
# qa
WANDB_MODE=disabled 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
WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/qasper_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
View file

@ -1,13 +0,0 @@
# main results
# batched
WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/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/checkpoint-$step/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/checkpoint-$step/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/checkpoint-$step/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/checkpoint-$step/pytorch_model.bin --datasets squad_negative_no_passage --split test

279
scripts/main_exp/eval/imagenette_eval.py
View file

@ -1,279 +0,0 @@
import json
import os
import re
from argparse import Namespace
from difflib import SequenceMatcher
import torch
from datasets import load_dataset
from tqdm import tqdm
from transformers import AutoProcessor, Gemma3ForConditionalGeneration
from ctx_to_lora.model_loading import get_tokenizer
from ctx_to_lora.modeling.ctx_encoder import PerLayerActivations
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
from ctx_to_lora.modeling.lora_layer import apply_lora_to_layers
from ctx_to_lora.modeling.lora_merger import combine_lora
CLASS_NAMES = [
"tench",
"English springer",
"cassette player",
"chain saw",
"church",
"French horn",
"garbage truck",
"gas pump",
"golf ball",
"parachute",
]
CLASS_TO_INT = {name: i for i, name in enumerate(CLASS_NAMES)}
INPUT_TXT = f"What is in this image? Choose exactly one of the following classes: {', '.join(CLASS_NAMES)}. Response with only the correct class without any other text."
RUN_DIR = "train_outputs/runs/Oct16_02-37-04_slurm0-a3nodeset-8_94074_1d62ecb8"
def _normalize_text(text: str) -> str:
text = re.sub(r"[^a-z0-9\s]", " ", text.lower())
return " ".join(text.split())
def _normalize_compact(text: str) -> str:
return re.sub(r"[^a-z0-9]", "", text.lower())
def _build_alias_map():
alias_overrides = {
"english springer spaniel": "English springer",
"springer spaniel": "English springer",
"chainsaw": "chain saw",
"dump truck": "garbage truck",
"refuse truck": "garbage truck",
"garbage lorry": "garbage truck",
"fuel pump": "gas pump",
"gas station pump": "gas pump",
"cassette deck": "cassette player",
"cassette recorder": "cassette player",
"fish": "tench",
"tench fish": "tench",
"french horn instrument": "French horn",
"golfball": "golf ball",
"skydiving": "parachute",
"parachutist": "parachute",
}
alias_map = {}
def register(alias: str, canonical: str):
alias = _normalize_text(alias)
if alias:
alias_map[alias] = canonical
alias_map[_normalize_compact(alias)] = canonical
for name in CLASS_NAMES:
register(name, name)
register(name.replace(" ", ""), name)
register(name.replace(" ", "-"), name)
for alias, canonical in alias_overrides.items():
register(alias, canonical)
return alias_map
CLASS_ALIAS_MAP = _build_alias_map()
def pred_to_class_id(pred_txt: str) -> int:
norm_pred = _normalize_text(pred_txt)
compact_pred = _normalize_compact(pred_txt)
for alias, canonical in CLASS_ALIAS_MAP.items():
if alias and (alias in norm_pred or alias in compact_pred):
return CLASS_TO_INT[canonical]
pred_tokens = set(norm_pred.split())
best_class = None
best_token_hits = -1
for name in CLASS_NAMES:
class_tokens = set(_normalize_text(name).split())
if class_tokens and class_tokens.issubset(pred_tokens):
return CLASS_TO_INT[name]
hits = sum(token in pred_tokens for token in class_tokens)
if hits > best_token_hits:
best_token_hits = hits
best_class = name
best_ratio = -1.0
for name in CLASS_NAMES:
ratio = SequenceMatcher(None, norm_pred, _normalize_text(name)).ratio()
if ratio > best_ratio:
best_ratio = ratio
best_class = name
return CLASS_TO_INT[best_class]
def load_checkpoint():
checkpoint_path = f"{RUN_DIR}/checkpoint-80000/pytorch_model.bin"
state_dict = torch.load(checkpoint_path)
model = ModulatedPretrainedModel.from_state_dict(
state_dict,
train=False,
base_model_kwargs=dict(attn_implementation="flash_attention_2"),
use_flash_attn=True,
use_sequence_packing=False, # for generation
)
tokenizer = get_tokenizer("google/gemma-2-2b-it")
model.eval()
return model, tokenizer
def load_ctx_encoder():
model_id = "google/gemma-3-4b-it"
ctx_model = Gemma3ForConditionalGeneration.from_pretrained(
model_id, device_map="auto"
).eval()
ctx_encoder_config = Namespace(ctx_encoder_last_layer=26, keep_lm_head=True)
ctx_model.language_model = PerLayerActivations(
ctx_model.language_model, ctx_encoder_config
)
processor = AutoProcessor.from_pretrained(model_id)
return ctx_model, processor
def template_image(img, ctx_processor):
messages = [
{
"role": "system",
"content": [{"type": "text", "text": ""}],
},
{
"role": "user",
"content": [
{"type": "image", "image": img},
],
},
]
inputs = ctx_processor.apply_chat_template(
messages,
add_generation_prompt=True,
tokenize=True,
return_dict=True,
return_tensors="pt",
)
return inputs
@torch.inference_mode
def get_ctx_features(ctx_inputs, ctx_encoder):
forward_outputs = ctx_encoder(**ctx_inputs, output_hidden_states=True)
ctx_features = torch.stack(forward_outputs.hidden_states, dim=1)
return ctx_features
def generate_loras(ctx_inputs, ctx_features):
generated_loras, _ = model.hypernet.generate_weights(
ctx_features, attn_mask=torch.ones_like(ctx_inputs["input_ids"])
)
generated_loras = combine_lora(
generated_loras,
n_chunks=torch.tensor((1,), device=model.device),
lora_bias=model.hypernet.get_head_bias()
if model.hypernet.config.use_bias
else None,
)
return generated_loras
def apply_loras(model, generated_loras):
n_queries = torch.ones(1, dtype=torch.int32, device=model.device)
apply_lora_to_layers(
model.base_model,
model.hypernet.layer_indices,
generated_loras,
n_queries,
)
if __name__ == "__main__":
model, base_tokenizer = load_checkpoint()
ctx_encoder, ctx_processor = load_ctx_encoder()
ds = load_dataset("frgfm/imagenette", "full_size", split="validation")
# ds = ds.shuffle().select(range(int(0.05 * len(ds))))
input_ids = base_tokenizer.apply_chat_template(
[{"role": "user", "content": INPUT_TXT}],
add_special_tokens=False,
return_attention_mask=False,
add_generation_prompt=True,
return_tensors="pt",
).to(model.device)
preds = []
pred_txts = []
corrects = []
labels = ds["label"]
for sample in tqdm(ds):
img = sample["image"]
ctx_inputs = template_image(img, ctx_processor).to(ctx_encoder.device)
ctx_features = get_ctx_features(ctx_inputs, ctx_encoder)
generated_loras = generate_loras(ctx_inputs, ctx_features)
apply_loras(model, generated_loras)
model_outputs = model.base_model.generate(
input_ids, max_new_tokens=256, do_sample=False
)
pred_txt = base_tokenizer.decode(
model_outputs[0][len(input_ids[0]) :], skip_special_tokens=True
)
pred_txts.append(pred_txt)
preds.append(pred_to_class_id(pred_txt))
is_correct = preds[-1] == labels[len(preds) - 1]
corrects.append(is_correct)
print(
f"GT: {CLASS_NAMES[labels[len(preds) - 1]]}, Pred: {pred_txt} -> {CLASS_NAMES[preds[-1]]}, Correct: {is_correct}"
)
acc = sum(corrects) / len(corrects)
print(f"Final accuracy: {acc:4f}")
jsonl_path = os.path.join(RUN_DIR, "imagenette_eval.jsonl")
meta_path = os.path.join(RUN_DIR, "imagenette_eval.meta.json")
with open(jsonl_path, "w") as f:
for i, (pred_txt, pred_id, label_id) in enumerate(
zip(pred_txts, preds, labels)
):
f.write(
json.dumps(
{
"index": i,
"label": int(label_id),
"label_name": CLASS_NAMES[label_id],
"pred_text": pred_txt,
"pred_class_id": int(pred_id),
"pred_class_name": CLASS_NAMES[pred_id],
"correct": bool(pred_id == label_id),
}
)
+ "\n"
)
meta = {
"dataset": "frgfm/imagenette",
"subset": "full_size",
"split": "validation",
"run_dir": RUN_DIR,
"prompt": INPUT_TXT,
"accuracy": float(acc),
"num_samples": len(preds),
"class_names": CLASS_NAMES,
}
with open(meta_path, "w") as f:
json.dump(meta, f, indent=2)
print(f"Wrote samples to {jsonl_path}")
print(f"Wrote metadata to {meta_path}")

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

@ -1,5 +0,0 @@
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
View file

@ -1,4 +0,0 @@
# 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
View file

@ -1,20 +0,0 @@
# 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 13); 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 --split train --closed_qa_prob 1.0
uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names pwc_compact --split train --closed_qa_prob 0.0

29
scripts/main_exp/train-cross-enc-chunk-slurm.sh
View file

@ -1,29 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#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"
# port=29051
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 \
--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=20000 --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 --ctx_encoder_model_name_or_path=google/gemma-3-4b-it \
--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"}' \
--warmup_steps=2000 \
--learning_rate=2e-5 \
"$@"

24
scripts/main_exp/train-cross-enc-slurm.sh
View file

@ -1,24 +0,0 @@
#!/bin/bash
#SBATCH --job-name=ctxlora
#SBATCH --nodes=1
#SBATCH --partition=a3
#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"
# port=29051
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 \
--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 --ctx_encoder_model_name_or_path=google/gemma-3-4b-it \
"$@"

15
scripts/main_exp/train_no_qa.sh
View file

@ -1,15 +0,0 @@
#!/bin/bash
port=29051
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_no_qa_l2l.yaml \
--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
View file

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

42
scripts/niah/1-train.sh
View file

@ -1,42 +0,0 @@
#!/bin/bash
WANDB_MODE=disabled uv run train.py \
configs/niah_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.5 \
--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 \
--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 \
--torch_empty_cache_steps=10 \
"$@"

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

@ -1 +0,0 @@
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_3072_4096 ctx_magic_number_7168_8192 ctx_magic_number_15360_16384 ctx_magic_number_28672_32768 ctx_magic_number_57344_65536 ctx_magic_number_122880_131072 --max_ctx_chunk_len=1024 --split test --eval_batch_size_gen=4

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

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

8
scripts/niah/README.md
View file

@ -1,8 +0,0 @@
# NIAH experiment
```bash
# run the scripts in this order
# data generation is only needed to be run once
uv run bash scripts/niah/0-gen_data.sh
uv run bash scripts/niah/1-train.sh
uv run bash scripts/niah/2-eval.sh
```

39
scripts/niah/train_mistral_7b.sh
View file

@ -1,39 +0,0 @@
WANDB_MODE=disabled run uv run train.py configs/niah_exp/ctx_magic_number_32_256.yaml \
--model_name_or_path=mistralai/Mistral-7B-Instruct-v0.2 \
--num_train_epochs=1 \
--per_device_train_batch_size=-1 \
--gradient_accumulation_steps=64 \
--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=2.0 \
--use_sequence_packing=True \
--max_packed_inp_len=1024 \
--max_packed_ctx_len=1024 \
--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 \
--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 \
--torch_empty_cache_steps=10

39
scripts/niah/train_qwen3_4b.sh
View file

@ -1,39 +0,0 @@
WANDB_MODE=disabled run uv run train.py configs/niah_exp/ctx_magic_number_32_256.yaml \
--model_name_or_path=Qwen/Qwen3-4B-Instruct-2507 \
--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=0.5 \
--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 \
--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 \
--torch_empty_cache_steps=10

1
scripts/video2lora/__init__.py Normal file
View file

@ -0,0 +1 @@

249
scripts/video2lora/build_finevideo_stage1_manifest.py Normal file
View file

@ -0,0 +1,249 @@
import argparse
import os
import random
from pathlib import Path
from ctx_to_lora.data.finevideo import (
candidates_to_rows,
collect_finevideo_span_pools,
materialize_span_clips,
sample_fixed_count,
sample_mixture,
write_finevideo_manifests,
)
def parse_args():
data_root = Path(os.environ.get("VIDEO2LORA_DATA_ROOT", "data/video2lora"))
finevideo_root = data_root / "raw" / "finevideo"
processed_root = data_root / "processed" / "finevideo"
parser = argparse.ArgumentParser(
description="Build exact-span Stage 1 manifests for FineVideo with ffmpeg-cut scene/adjacent clips."
)
parser.add_argument("--train-root", default=str(finevideo_root / "train"))
parser.add_argument("--val-root", default=str(finevideo_root / "val"))
parser.add_argument(
"--clips-root",
default=str(data_root / "raw" / "finevideo_stage1_spans"),
)
parser.add_argument(
"--train-out",
default=str(processed_root / "train.jsonl"),
)
parser.add_argument(
"--val-out",
default=str(processed_root / "val.jsonl"),
)
parser.add_argument(
"--val-scene-out",
default=str(processed_root / "val_scene.jsonl"),
)
parser.add_argument(
"--val-adjacent-out",
default=str(processed_root / "val_adjacent.jsonl"),
)
parser.add_argument(
"--val-full-out",
default=str(processed_root / "val_full.jsonl"),
)
parser.add_argument(
"--val-core-out",
default=str(processed_root / "val_core.jsonl"),
)
parser.add_argument(
"--val-gen-out",
default=str(processed_root / "val_gen_100.jsonl"),
)
parser.add_argument(
"--train-target",
type=int,
default=0,
help="Target number of train examples with 60/30/10 span mix. 0 means max unique possible.",
)
parser.add_argument("--scene-ratio", type=float, default=0.60)
parser.add_argument("--adjacent-ratio", type=float, default=0.30)
parser.add_argument("--full-ratio", type=float, default=0.10)
parser.add_argument("--val-scene-size", type=int, default=384)
parser.add_argument("--val-adjacent-size", type=int, default=384)
parser.add_argument("--val-full-size", type=int, default=384)
parser.add_argument("--val-gen-size", type=int, default=100)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--num-workers", type=int, default=24)
parser.add_argument("--ffmpeg-threads", type=int, default=1)
parser.add_argument("--ffmpeg-preset", default="ultrafast")
parser.add_argument("--ffmpeg-crf", type=int, default=30)
parser.add_argument("--ffmpeg-include-audio", action="store_true")
parser.add_argument("--overwrite-clips", action="store_true")
parser.add_argument("--skip-clip-materialization", action="store_true")
return parser.parse_args()
def _stratified_val_gen_rows(
val_scene_rows: list[dict],
val_adjacent_rows: list[dict],
val_full_rows: list[dict],
*,
val_gen_size: int,
seed: int,
) -> list[dict]:
if val_gen_size <= 0:
return []
pools = [
rows
for rows in (val_scene_rows, val_adjacent_rows, val_full_rows)
if rows
]
if not pools:
return []
base = val_gen_size // 3
remainder = val_gen_size - base * 3
scene_n = base + (1 if remainder > 0 else 0)
adjacent_n = base + (1 if remainder > 1 else 0)
full_n = base
rng = random.Random(seed + 91)
out: list[dict] = []
if val_scene_rows:
out.extend(rng.sample(val_scene_rows, min(scene_n, len(val_scene_rows))))
if val_adjacent_rows:
out.extend(rng.sample(val_adjacent_rows, min(adjacent_n, len(val_adjacent_rows))))
if val_full_rows:
out.extend(rng.sample(val_full_rows, min(full_n, len(val_full_rows))))
if len(out) < val_gen_size:
selected_ids = {row["id"] for row in out}
remaining = [
row
for rows in pools
for row in rows
if row["id"] not in selected_ids
]
rng.shuffle(remaining)
out.extend(remaining[: val_gen_size - len(out)])
rng.shuffle(out)
return out[:val_gen_size]
def main():
args = parse_args()
random.seed(args.seed)
train_pools = collect_finevideo_span_pools(
args.train_root,
split="train",
clips_root=args.clips_root,
)
val_pools = collect_finevideo_span_pools(
args.val_root,
split="val",
clips_root=args.clips_root,
)
print(
"[pool] train scene/adjacent/full = "
f"{len(train_pools['scene'])}/{len(train_pools['adjacent'])}/{len(train_pools['full'])}"
)
print(
"[pool] val scene/adjacent/full = "
f"{len(val_pools['scene'])}/{len(val_pools['adjacent'])}/{len(val_pools['full'])}"
)
train_candidates = sample_mixture(
train_pools,
target_total=args.train_target,
scene_ratio=args.scene_ratio,
adjacent_ratio=args.adjacent_ratio,
full_ratio=args.full_ratio,
seed=args.seed,
)
val_scene = sample_fixed_count(
val_pools["scene"],
count=args.val_scene_size,
seed=args.seed + 1,
)
val_adjacent = sample_fixed_count(
val_pools["adjacent"],
count=args.val_adjacent_size,
seed=args.seed + 2,
)
val_full = sample_fixed_count(
val_pools["full"],
count=args.val_full_size,
seed=args.seed + 3,
)
val_candidates = list(val_scene) + list(val_adjacent) + list(val_full)
random.Random(args.seed + 4).shuffle(val_candidates)
print(f"[select] train candidates: {len(train_candidates)}")
print(
"[select] val scene/adjacent/full = "
f"{len(val_scene)}/{len(val_adjacent)}/{len(val_full)} (total={len(val_candidates)})"
)
if not args.skip_clip_materialization:
all_candidates = train_candidates + val_candidates
failures = materialize_span_clips(
all_candidates,
num_workers=args.num_workers,
overwrite=args.overwrite_clips,
ffmpeg_threads=args.ffmpeg_threads,
ffmpeg_preset=args.ffmpeg_preset,
ffmpeg_crf=args.ffmpeg_crf,
include_audio=args.ffmpeg_include_audio,
)
if failures:
print(f"[clip] failures: {len(failures)}")
print("[clip] first failures:")
for sample_id, error in failures[:20]:
print(f" {sample_id}: {error}")
failure_ids = {sample_id for sample_id, _ in failures}
train_candidates = [
candidate for candidate in train_candidates if candidate.sample_id not in failure_ids
]
val_scene = [candidate for candidate in val_scene if candidate.sample_id not in failure_ids]
val_adjacent = [
candidate for candidate in val_adjacent if candidate.sample_id not in failure_ids
]
val_full = [candidate for candidate in val_full if candidate.sample_id not in failure_ids]
val_candidates = list(val_scene) + list(val_adjacent) + list(val_full)
random.Random(args.seed + 4).shuffle(val_candidates)
train_rows = candidates_to_rows(train_candidates)
val_scene_rows = candidates_to_rows(val_scene)
val_adjacent_rows = candidates_to_rows(val_adjacent)
val_full_rows = candidates_to_rows(val_full)
val_rows = candidates_to_rows(val_candidates)
val_gen_rows = _stratified_val_gen_rows(
val_scene_rows,
val_adjacent_rows,
val_full_rows,
val_gen_size=args.val_gen_size,
seed=args.seed + 10,
)
write_finevideo_manifests(
train_rows=train_rows,
val_rows=val_rows,
val_scene_rows=val_scene_rows,
val_adjacent_rows=val_adjacent_rows,
val_full_rows=val_full_rows,
val_gen_rows=val_gen_rows,
train_out=args.train_out,
val_out=args.val_out,
val_scene_out=args.val_scene_out,
val_adjacent_out=args.val_adjacent_out,
val_full_out=args.val_full_out,
val_core_out=args.val_core_out,
val_gen_out=args.val_gen_out,
)
print(f"Wrote train rows: {len(train_rows)} -> {args.train_out}")
print(f"Wrote val rows: {len(val_rows)} -> {args.val_out}")
print(f"Wrote val_scene rows: {len(val_scene_rows)} -> {args.val_scene_out}")
print(f"Wrote val_adjacent rows: {len(val_adjacent_rows)} -> {args.val_adjacent_out}")
print(f"Wrote val_full rows: {len(val_full_rows)} -> {args.val_full_out}")
print(f"Wrote val_core rows: {min(1024, len(val_rows))} -> {args.val_core_out}")
print(f"Wrote val_gen rows: {len(val_gen_rows)} -> {args.val_gen_out}")
if __name__ == "__main__":
main()

275
scripts/video2lora/generate_finevideo_teacher_targets.py Normal file
View file

@ -0,0 +1,275 @@
import argparse
import json
import os
import time
from pathlib import Path
from typing import Any
import torch
from transformers import AutoModelForImageTextToText, AutoProcessor
from transformers.utils import is_av_available, is_cv2_available, is_decord_available
from ctx_to_lora.data.video_manifest import load_video_manifest, dump_jsonl, resolve_video_path
def parse_args():
parser = argparse.ArgumentParser(description='Generate offline teacher targets for FineVideo manifests.')
parser.add_argument('--input-manifest', required=True)
parser.add_argument('--output-manifest', required=True)
parser.add_argument('--smolvlm-name-or-path', default='HuggingFaceTB/SmolVLM2-2.2B-Instruct')
parser.add_argument('--per-device-batch-size', type=int, default=8)
parser.add_argument('--max-new-tokens', type=int, default=64)
parser.add_argument('--max-frames', type=int, default=12)
parser.add_argument('--video-size-longest-edge', type=int, default=384)
parser.add_argument('--video-fps', type=float, default=None)
parser.add_argument(
'--video-load-backend',
default='auto',
choices=('auto', 'decord', 'pyav', 'opencv', 'torchvision'),
)
parser.add_argument('--max-samples', type=int, default=None)
parser.add_argument('--flush-every', type=int, default=256)
parser.add_argument('--output-dir', default='')
parser.add_argument('--shard-index', type=int, default=0)
parser.add_argument('--num-shards', type=int, default=1)
parser.add_argument('--merge-only', action='store_true')
return parser.parse_args()
def is_skippable_video_error(exc: BaseException) -> bool:
seen = set()
cur: BaseException | None = exc
while cur is not None and id(cur) not in seen:
seen.add(id(cur))
if isinstance(cur, UnicodeDecodeError):
return True
if isinstance(cur, IndexError) and 'tuple index out of range' in str(cur).lower():
return True
err_name = cur.__class__.__name__.lower()
err_mod = cur.__class__.__module__.lower()
msg = str(cur).lower()
if 'av.' in err_mod or 'ffmpeg' in err_mod or 'pyav' in err_mod:
return True
if 'invalid total_num_frames' in msg:
return True
if 'cannot open' in msg and 'video' in msg:
return True
if 'error opening input file' in msg:
return True
if 'moov atom not found' in msg:
return True
if 'invalid data found when processing input' in msg:
return True
if 'unsupported codec' in msg:
return True
if 'decode' in err_name and 'error' in err_name:
return True
cur = cur.__cause__ or cur.__context__
return False
def prepare_generation_inputs(
processor,
messages,
device,
*,
max_frames,
video_fps,
video_size_longest_edge,
video_load_backend,
):
processor.tokenizer.padding_side = 'left'
chat_template_kwargs = dict(padding=True)
processor_name = type(processor).__name__.lower()
if 'idefics3' in processor_name:
chat_template_kwargs['num_frames'] = max_frames
if video_fps is not None:
chat_template_kwargs['video_fps'] = video_fps
else:
chat_template_kwargs['max_frames'] = max_frames
if video_fps is not None:
chat_template_kwargs['target_fps'] = video_fps
if video_load_backend == 'auto':
if is_decord_available():
video_load_backend = 'decord'
elif is_av_available():
video_load_backend = 'pyav'
elif is_cv2_available():
video_load_backend = 'opencv'
else:
video_load_backend = 'torchvision'
chat_template_kwargs['video_load_backend'] = video_load_backend
if video_size_longest_edge is not None:
video_size = {'longest_edge': video_size_longest_edge}
processor.video_size = video_size
processor.image_processor.size = video_size
prompt_inputs = processor.apply_chat_template(
messages,
add_generation_prompt=True,
tokenize=True,
return_dict=True,
return_tensors='pt',
**chat_template_kwargs,
)
moved = {}
for key, value in prompt_inputs.items():
if isinstance(value, torch.Tensor):
if value.is_floating_point():
moved[key] = value.to(device=device, dtype=torch.bfloat16)
else:
moved[key] = value.to(device=device)
else:
moved[key] = value
return moved
def chunked(seq, size):
for i in range(0, len(seq), size):
yield seq[i:i + size]
def generate_rows(
model,
processor,
tokenizer,
rows: list[dict[str, Any]],
device: str,
args,
) -> list[dict[str, Any]]:
messages = []
for row in rows:
messages.append([
{
'role': 'user',
'content': [
{'type': 'video', 'path': row['video_path']},
{'type': 'text', 'text': row['prompt']},
],
}
])
inputs = prepare_generation_inputs(
processor,
messages,
device,
max_frames=args.max_frames,
video_fps=args.video_fps,
video_size_longest_edge=args.video_size_longest_edge,
video_load_backend=args.video_load_backend,
)
generated = model.generate(
**inputs,
max_new_tokens=args.max_new_tokens,
do_sample=False,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id,
)
input_len = inputs['input_ids'].shape[1]
generated_only = generated[:, input_len:]
decoded = tokenizer.batch_decode(generated_only, skip_special_tokens=True)
out_rows = []
for row, target_text in zip(rows, decoded, strict=True):
out_row = dict(row)
out_row['target_text'] = target_text.strip() or 'Unable to determine.'
metadata = dict(out_row.get('metadata') or {})
metadata['teacher_model'] = args.smolvlm_name_or_path
metadata['teacher_generated'] = True
metadata['teacher_prompt'] = out_row['prompt']
metadata['teacher_max_new_tokens'] = args.max_new_tokens
out_row['metadata'] = metadata
out_rows.append(out_row)
return out_rows
def merge_shards(output_path: Path, tmp_dir: Path) -> int:
merged = []
for shard_file in sorted(tmp_dir.glob(f'{output_path.stem}.rank*.jsonl')):
with open(shard_file) as f:
for line in f:
line = line.strip()
if line:
merged.append(json.loads(line))
dump_jsonl(output_path, merged)
return len(merged)
def main():
args = parse_args()
output_path = Path(args.output_manifest)
tmp_dir = Path(args.output_dir) if args.output_dir else output_path.parent / f'.tmp-{output_path.stem}'
tmp_dir.mkdir(parents=True, exist_ok=True)
if args.merge_only:
merged = merge_shards(output_path, tmp_dir)
print(f'[merge] wrote {merged} rows to {output_path}', flush=True)
return
rows = load_video_manifest(args.input_manifest, max_samples=args.max_samples)
world = args.num_shards
rank = args.shard_index
if rank < 0 or rank >= world:
raise ValueError(f'Invalid shard index {rank} for num_shards={world}')
shard = rows[rank::world]
shard_path = tmp_dir / f'{output_path.stem}.rank{rank:02d}.jsonl'
processor = AutoProcessor.from_pretrained(args.smolvlm_name_or_path, trust_remote_code=True)
tokenizer = processor.tokenizer
if tokenizer.pad_token_id is None:
tokenizer.pad_token_id = tokenizer.eos_token_id
model = AutoModelForImageTextToText.from_pretrained(
args.smolvlm_name_or_path,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
).to('cuda').eval()
model.config.pad_token_id = tokenizer.pad_token_id
if getattr(model, 'generation_config', None):
model.generation_config.pad_token_id = tokenizer.pad_token_id
written = 0
skipped = 0
started = time.time()
buffer: list[dict[str, Any]] = []
with open(shard_path, 'w') as fout:
for batch_rows in chunked(shard, args.per_device_batch_size):
normalized_rows = []
for row in batch_rows:
row = dict(row)
row['video_path'] = resolve_video_path(row['video_path'])
normalized_rows.append(row)
try:
out_rows = generate_rows(model, processor, tokenizer, normalized_rows, 'cuda', args)
buffer.extend(out_rows)
written += len(out_rows)
except Exception as exc: # pylint: disable=broad-except
if len(normalized_rows) > 1:
for row in normalized_rows:
try:
out_rows = generate_rows(model, processor, tokenizer, [row], 'cuda', args)
buffer.extend(out_rows)
written += len(out_rows)
except Exception as single_exc: # pylint: disable=broad-except
if not is_skippable_video_error(single_exc):
raise
skipped += 1
print(f'[rank{rank}] skipped {row["id"]}: {type(single_exc).__name__}: {single_exc}', flush=True)
continue
if not is_skippable_video_error(exc):
raise
skipped += len(normalized_rows)
for row in normalized_rows:
print(f'[rank{rank}] skipped {row["id"]}: {type(exc).__name__}: {exc}', flush=True)
if len(buffer) >= args.flush_every:
for row in buffer:
fout.write(json.dumps(row) + '\n')
fout.flush()
buffer.clear()
elapsed = max(time.time() - started, 1e-6)
rate = written / elapsed
print(f'[rank{rank}] written={written} skipped={skipped} rate={rate:.2f} rows/s', flush=True)
for row in buffer:
fout.write(json.dumps(row) + '\n')
fout.flush()
if __name__ == '__main__':
main()

207
scripts/video2lora/infer.py Normal file
View file

@ -0,0 +1,207 @@
import argparse
import json
import re
from dataclasses import fields
from pathlib import Path
from typing import Any
import torch
from accelerate import Accelerator
from accelerate.utils import set_seed
from ctx_to_lora.data.video_manifest import dump_jsonl, load_video_manifest
from scripts.video2lora.train_smolvlm_stage1 import (
TrainArgs,
build_stage1_model,
generate_fixed_examples,
)
def parse_args():
parser = argparse.ArgumentParser(
description="Run Video2LoRA inference from a Stage 1 checkpoint."
)
parser.add_argument(
"--run-dir",
default="",
help="Directory containing train_args.json and checkpoints/.",
)
parser.add_argument(
"--checkpoint",
required=True,
help="Checkpoint filename under RUN_DIR/checkpoints/ or a direct path.",
)
parser.add_argument(
"--preset",
default="auto",
choices=("auto", "500m", "2.2b"),
help="Use built-in public checkpoint settings when train_args.json is absent.",
)
parser.add_argument("--manifest", required=True, help="JSONL video manifest.")
parser.add_argument("--output", required=True, help="Output JSONL path.")
parser.add_argument("--max-samples", type=int, default=None)
parser.add_argument("--video-size-longest-edge", type=int, default=None)
parser.add_argument("--generation-max-new-tokens", type=int, default=None)
parser.add_argument("--seed", type=int, default=None)
return parser.parse_args()
def checkpoint_path(run_dir: Path | None, checkpoint: str) -> Path:
path = Path(checkpoint)
if path.exists():
return path
if run_dir is None:
return path
return run_dir / "checkpoints" / checkpoint
def checkpoint_label(path: Path) -> str:
match = re.search(r"step-(\d+)\.pt$", path.name)
if match:
return str(int(match.group(1)))
return path.stem
def infer_preset(checkpoint: str, requested: str) -> str:
if requested != "auto":
return requested
lower = checkpoint.lower()
if "500m" in lower:
return "500m"
if "2.2b" in lower or "2p2b" in lower or "22b" in lower:
return "2.2b"
raise ValueError(
"Could not infer checkpoint preset from filename. Pass --preset 500m or --preset 2.2b."
)
def default_train_args(preset: str) -> TrainArgs:
if preset == "500m":
model_name = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
video_load_backend = "auto"
elif preset == "2.2b":
model_name = "HuggingFaceTB/SmolVLM2-2.2B-Instruct"
video_load_backend = "auto"
else:
raise ValueError(f"Unknown preset: {preset}")
return TrainArgs(
smolvlm_name_or_path=model_name,
train_manifest="",
val_manifest="",
val_core_manifest=None,
val_gen_manifest=None,
output_dir="",
per_device_batch_size=1,
eval_batch_size=1,
gradient_accumulation_steps=1,
max_steps=0,
learning_rate=0.0,
weight_decay=0.0,
warmup_ratio=0.0,
max_grad_norm=1.0,
seed=42,
log_every=1,
eval_every=1,
save_every=1,
max_train_samples=None,
max_val_samples=None,
num_workers=0,
lora_r=16,
lora_dropout=0.0,
target_modules=["down_proj"],
latent_size=512,
dropout_rate=0.0,
n_latent_queries=8,
num_blocks=9,
num_self_attn_per_block=0,
video_fps=None,
max_frames=12,
video_size_longest_edge=384,
video_load_backend=video_load_backend,
internalization_prompt="Internalize this video for later captioning.",
kl_weight=0.0,
kl_temperature=1.0,
generation_max_new_tokens=128,
legacy_sanity_manifests=[],
legacy_sanity_max_samples_per_manifest=0,
wandb_project="video2lora",
wandb_mode="disabled",
wandb_group=None,
wandb_run_name=None,
wandb_notes=None,
resume_checkpoint=None,
resume_trainer_state=None,
resume_global_step=None,
resume_ignore_scheduler_state=False,
)
def load_train_args(run_dir: Path | None, overrides: argparse.Namespace) -> TrainArgs:
args_path = run_dir / "train_args.json" if run_dir is not None else None
if args_path is None or not args_path.exists():
train_args = default_train_args(infer_preset(overrides.checkpoint, overrides.preset))
if overrides.video_size_longest_edge is not None:
train_args.video_size_longest_edge = overrides.video_size_longest_edge
if overrides.generation_max_new_tokens is not None:
train_args.generation_max_new_tokens = overrides.generation_max_new_tokens
if overrides.seed is not None:
train_args.seed = overrides.seed
return train_args
with open(args_path) as f:
raw_args: dict[str, Any] = json.load(f)
raw_args.setdefault("video_size_longest_edge", None)
raw_args.setdefault("resume_ignore_scheduler_state", False)
if overrides.video_size_longest_edge is not None:
raw_args["video_size_longest_edge"] = overrides.video_size_longest_edge
if overrides.generation_max_new_tokens is not None:
raw_args["generation_max_new_tokens"] = overrides.generation_max_new_tokens
if overrides.seed is not None:
raw_args["seed"] = overrides.seed
allowed = {field.name for field in fields(TrainArgs)}
filtered_args = {key: value for key, value in raw_args.items() if key in allowed}
missing = sorted(allowed - filtered_args.keys())
if missing:
raise ValueError(f"train_args.json is missing required fields: {missing}")
return TrainArgs(**filtered_args)
def main() -> None:
args = parse_args()
run_dir = Path(args.run_dir) if args.run_dir else None
ckpt_path = checkpoint_path(run_dir, args.checkpoint)
if not ckpt_path.exists():
raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
train_args = load_train_args(run_dir, args)
set_seed(train_args.seed)
rows = load_video_manifest(args.manifest, max_samples=args.max_samples)
accelerator = Accelerator(mixed_precision="bf16")
model, raw_model, processor, tokenizer = build_stage1_model(train_args)
state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=False)
model.load_state_dict(state_dict)
model.to(accelerator.device)
raw_model.eval()
outputs = generate_fixed_examples(
accelerator,
model,
raw_model,
processor,
tokenizer,
rows,
train_args,
)
for row in outputs:
row["checkpoint"] = checkpoint_label(ckpt_path)
dump_jsonl(args.output, outputs)
print(f"[infer] wrote {len(outputs)} rows to {args.output}")
if __name__ == "__main__":
main()

1496
scripts/video2lora/train_smolvlm_online.py Normal file
View file

File diff suppressed because it is too large Load diff

1600
scripts/video2lora/train_smolvlm_stage1.py Normal file
View file

File diff suppressed because it is too large Load diff

1235
scripts/video2lora/train_smolvlm_stage1_selfdistill.py Normal file
View file

File diff suppressed because it is too large Load diff

12
setup.py
View file

@ -1,12 +0,0 @@
# read the contents of the README file
from pathlib import Path
from setuptools import setup
this_directory = Path(__file__).parent
long_description = (this_directory / "README.md").read_text()
setup(
long_description=long_description,
long_description_content_type="text/markdown",
)

617
src/ctx_to_lora/data/finevideo.py Normal file
View file

@ -0,0 +1,617 @@
import hashlib
import json
import math
import random
import re
import subprocess
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from ctx_to_lora.data.video_manifest import (
DATA_ROOT,
dump_jsonl,
normalize_manifest_row,
relativize_video_path,
)
SCENE_PROMPTS = (
"Describe what is happening in this scene.",
"Write a short summary of this scene.",
)
ADJACENT_PROMPTS = (
"Summarize this video segment.",
"Describe the sequence of events in this segment.",
)
FULL_PROMPTS = (
"Summarize this video in 1-3 sentences.",
)
_WHITESPACE_RE = re.compile(r"\s+")
@dataclass(frozen=True)
class SpanCandidate:
sample_id: str
split: str
dataset: str
span_type: str
source_video_abs: str
output_video_abs: str
clip_start_sec: float | None
clip_end_sec: float | None
prompt: str
target_text: str
metadata: dict[str, Any]
def _read_json(path: Path) -> dict[str, Any]:
with open(path) as f:
return json.load(f)
def _clean_text(text: str) -> str:
cleaned = _WHITESPACE_RE.sub(" ", text.strip())
cleaned = cleaned.strip(" \t\r\n")
return cleaned
def _ensure_sentence(text: str) -> str:
text = _clean_text(text)
if not text:
return ""
if text[-1] not in {".", "!", "?"}:
text += "."
return text
def _choose_prompt(sample_id: str, *, span_type: str) -> str:
if span_type == "scene":
choices = SCENE_PROMPTS
elif span_type == "adjacent":
choices = ADJACENT_PROMPTS
elif span_type == "full":
choices = FULL_PROMPTS
else:
raise ValueError(f"Unknown span type: {span_type}")
digest = hashlib.sha256(sample_id.encode()).digest()
return choices[digest[0] % len(choices)]
def parse_hhmmss_to_seconds(value: str | None) -> float | None:
if not value:
return None
value = str(value).strip()
parts = value.split(":")
if len(parts) != 3:
return None
try:
hours = int(parts[0])
minutes = int(parts[1])
seconds = float(parts[2])
except ValueError:
return None
return hours * 3600.0 + minutes * 60.0 + seconds
def _scene_bounds(scene: dict[str, Any]) -> tuple[float | None, float | None]:
timestamps = scene.get("timestamps") or {}
if isinstance(timestamps, dict):
start = parse_hhmmss_to_seconds(timestamps.get("start_timestamp"))
end = parse_hhmmss_to_seconds(timestamps.get("end_timestamp"))
if start is not None and end is not None and end > start:
return start, end
return None, None
def _scene_activity_descriptions(scene: dict[str, Any]) -> list[str]:
activities = scene.get("activities") or []
out: list[str] = []
for activity in activities:
if not isinstance(activity, dict):
continue
description = _ensure_sentence(str(activity.get("description") or ""))
if description:
out.append(description)
return out
def _scene_target(scene: dict[str, Any]) -> str:
activities = _scene_activity_descriptions(scene)
if activities:
return " ".join(activities[:2])
scene_title = _ensure_sentence(str(scene.get("title") or ""))
if scene_title:
return scene_title
mood = ((scene.get("mood") or {}).get("description") or "").strip()
mood_text = _ensure_sentence(f"The scene mood is {mood.lower()}." if mood else "")
return mood_text
def _adjacent_target(scenes: list[dict[str, Any]]) -> str:
activities: list[str] = []
for scene in scenes:
activities.extend(_scene_activity_descriptions(scene))
if activities:
return " ".join(activities[:4])
titles = [
_ensure_sentence(str(scene.get("title") or ""))
for scene in scenes
]
titles = [title for title in titles if title]
return " ".join(titles[:3])
def _full_target_from_description(meta: dict[str, Any]) -> str:
description = _ensure_sentence(str((meta.get("description") or "")).strip())
return description
def _span_stem(split: str, source_stem: str, span_type: str, span_key: str) -> str:
return f"{split}-finevideo-{source_stem}-{span_type}-{span_key}"
def _clip_output_path(clips_root: Path, split: str, span_type: str, stem: str) -> Path:
digest = hashlib.sha256(stem.encode()).hexdigest()[:6]
return clips_root / split / span_type / digest / f"{stem}.mp4"
def _ffmpeg_extract_clip(
source_video: Path,
output_video: Path,
start_sec: float,
end_sec: float,
overwrite: bool,
*,
ffmpeg_threads: int = 1,
ffmpeg_preset: str = "ultrafast",
ffmpeg_crf: int = 30,
include_audio: bool = False,
) -> None:
if output_video.exists() and not overwrite:
return
output_video.parent.mkdir(parents=True, exist_ok=True)
duration = end_sec - start_sec
if duration <= 0:
raise ValueError(f"Invalid clip duration for {source_video}: {start_sec} -> {end_sec}")
command = [
"ffmpeg",
"-hide_banner",
"-loglevel",
"error",
"-threads",
str(max(1, int(ffmpeg_threads))),
"-ss",
f"{start_sec:.3f}",
"-t",
f"{duration:.3f}",
"-i",
str(source_video),
"-c:v",
"libx264",
"-preset",
ffmpeg_preset,
"-crf",
str(int(ffmpeg_crf)),
"-movflags",
"+faststart",
]
if include_audio:
command.extend(["-c:a", "aac"])
else:
command.append("-an")
command.extend(
[
"-y" if overwrite else "-n",
str(output_video),
]
)
subprocess.run(command, check=True)
def _build_scene_candidate(
*,
split: str,
source_video: Path,
source_stem: str,
scene: dict[str, Any],
scene_index: int,
clips_root: Path,
) -> SpanCandidate | None:
start, end = _scene_bounds(scene)
if start is None or end is None:
return None
target = _scene_target(scene)
if not target:
return None
span_key = f"s{scene_index:02d}"
stem = _span_stem(split, source_stem, "scene", span_key)
output_path = _clip_output_path(clips_root, split, "scene", stem)
prompt = _choose_prompt(stem, span_type="scene")
metadata = {
"dataset": "finevideo",
"span_type": "scene",
"source_video_path": relativize_video_path(str(source_video)),
"scene_index": scene_index,
"scene_id": scene.get("sceneId"),
"scene_title": scene.get("title"),
}
return SpanCandidate(
sample_id=stem,
split=split,
dataset="finevideo",
span_type="scene",
source_video_abs=str(source_video),
output_video_abs=str(output_path),
clip_start_sec=start,
clip_end_sec=end,
prompt=prompt,
target_text=target,
metadata=metadata,
)
def _build_adjacent_candidate(
*,
split: str,
source_video: Path,
source_stem: str,
scenes: list[dict[str, Any]],
start_scene_index: int,
window_size: int,
clips_root: Path,
) -> SpanCandidate | None:
window = scenes[start_scene_index : start_scene_index + window_size]
if len(window) != window_size:
return None
first_start, _ = _scene_bounds(window[0])
_, last_end = _scene_bounds(window[-1])
if first_start is None or last_end is None or last_end <= first_start:
return None
target = _adjacent_target(window)
if not target:
return None
span_key = f"a{start_scene_index:02d}w{window_size}"
stem = _span_stem(split, source_stem, "adjacent", span_key)
output_path = _clip_output_path(clips_root, split, "adjacent", stem)
prompt = _choose_prompt(stem, span_type="adjacent")
metadata = {
"dataset": "finevideo",
"span_type": "adjacent",
"source_video_path": relativize_video_path(str(source_video)),
"scene_index_start": start_scene_index,
"scene_index_end": start_scene_index + window_size - 1,
"scene_window_size": window_size,
"scene_ids": [scene.get("sceneId") for scene in window],
}
return SpanCandidate(
sample_id=stem,
split=split,
dataset="finevideo",
span_type="adjacent",
source_video_abs=str(source_video),
output_video_abs=str(output_path),
clip_start_sec=first_start,
clip_end_sec=last_end,
prompt=prompt,
target_text=target,
metadata=metadata,
)
def _build_full_candidate(
*,
split: str,
source_video: Path,
source_stem: str,
content_metadata: dict[str, Any],
youtube_title: str,
) -> SpanCandidate | None:
target = _full_target_from_description(content_metadata)
if not target:
return None
stem = _span_stem(split, source_stem, "full", "full")
prompt = _choose_prompt(stem, span_type="full")
metadata = {
"dataset": "finevideo",
"span_type": "full",
"source_video_path": relativize_video_path(str(source_video)),
"youtube_title": youtube_title,
}
return SpanCandidate(
sample_id=stem,
split=split,
dataset="finevideo",
span_type="full",
source_video_abs=str(source_video),
output_video_abs=str(source_video),
clip_start_sec=None,
clip_end_sec=None,
prompt=prompt,
target_text=target,
metadata=metadata,
)
def collect_finevideo_span_pools(
split_root: str | Path,
*,
split: str,
clips_root: str | Path,
) -> dict[str, list[SpanCandidate]]:
split_root = Path(split_root)
clips_root = Path(clips_root)
scene_pool: list[SpanCandidate] = []
adjacent_pool: list[SpanCandidate] = []
full_pool: list[SpanCandidate] = []
json_paths = sorted(split_root.rglob("*.json"))
for metadata_path in json_paths:
video_path = metadata_path.with_suffix(".mp4")
if not video_path.exists():
continue
metadata = _read_json(metadata_path)
content_metadata = metadata.get("content_metadata") or {}
scenes = content_metadata.get("scenes") or []
if not isinstance(scenes, list) or not scenes:
continue
source_stem = metadata_path.stem
youtube_title = str(metadata.get("youtube_title") or "")
for scene_index, scene in enumerate(scenes):
if not isinstance(scene, dict):
continue
candidate = _build_scene_candidate(
split=split,
source_video=video_path,
source_stem=source_stem,
scene=scene,
scene_index=scene_index,
clips_root=clips_root,
)
if candidate is not None:
scene_pool.append(candidate)
for window_size in (2, 3):
for start_scene_index in range(0, len(scenes) - window_size + 1):
candidate = _build_adjacent_candidate(
split=split,
source_video=video_path,
source_stem=source_stem,
scenes=scenes,
start_scene_index=start_scene_index,
window_size=window_size,
clips_root=clips_root,
)
if candidate is not None:
adjacent_pool.append(candidate)
full_candidate = _build_full_candidate(
split=split,
source_video=video_path,
source_stem=source_stem,
content_metadata=content_metadata,
youtube_title=youtube_title,
)
if full_candidate is not None:
full_pool.append(full_candidate)
return {
"scene": scene_pool,
"adjacent": adjacent_pool,
"full": full_pool,
}
def _counts_from_target(
target_total: int,
*,
scene_ratio: float,
adjacent_ratio: float,
full_ratio: float,
) -> tuple[int, int, int]:
scene_count = int(round(target_total * scene_ratio))
adjacent_count = int(round(target_total * adjacent_ratio))
full_count = target_total - scene_count - adjacent_count
return scene_count, adjacent_count, full_count
def max_unique_mixture_total(
*,
scene_pool_size: int,
adjacent_pool_size: int,
full_pool_size: int,
scene_ratio: float,
adjacent_ratio: float,
full_ratio: float,
) -> int:
by_scene = math.floor(scene_pool_size / scene_ratio) if scene_ratio > 0 else math.inf
by_adjacent = (
math.floor(adjacent_pool_size / adjacent_ratio) if adjacent_ratio > 0 else math.inf
)
by_full = math.floor(full_pool_size / full_ratio) if full_ratio > 0 else math.inf
return int(min(by_scene, by_adjacent, by_full))
def sample_mixture(
pools: dict[str, list[SpanCandidate]],
*,
target_total: int,
scene_ratio: float = 0.60,
adjacent_ratio: float = 0.30,
full_ratio: float = 0.10,
seed: int = 42,
) -> list[SpanCandidate]:
if target_total <= 0:
target_total = max_unique_mixture_total(
scene_pool_size=len(pools["scene"]),
adjacent_pool_size=len(pools["adjacent"]),
full_pool_size=len(pools["full"]),
scene_ratio=scene_ratio,
adjacent_ratio=adjacent_ratio,
full_ratio=full_ratio,
)
scene_count, adjacent_count, full_count = _counts_from_target(
target_total,
scene_ratio=scene_ratio,
adjacent_ratio=adjacent_ratio,
full_ratio=full_ratio,
)
if scene_count > len(pools["scene"]):
raise ValueError(
f"scene pool too small for target: need {scene_count}, have {len(pools['scene'])}"
)
if adjacent_count > len(pools["adjacent"]):
raise ValueError(
f"adjacent pool too small for target: need {adjacent_count}, have {len(pools['adjacent'])}"
)
if full_count > len(pools["full"]):
raise ValueError(
f"full pool too small for target: need {full_count}, have {len(pools['full'])}"
)
rng = random.Random(seed)
out: list[SpanCandidate] = []
out.extend(rng.sample(pools["scene"], scene_count))
out.extend(rng.sample(pools["adjacent"], adjacent_count))
out.extend(rng.sample(pools["full"], full_count))
rng.shuffle(out)
return out
def sample_fixed_count(pool: list[SpanCandidate], *, count: int, seed: int) -> list[SpanCandidate]:
if count <= 0:
return []
if count >= len(pool):
return list(pool)
rng = random.Random(seed)
return rng.sample(pool, count)
def _extract_one(
candidate: SpanCandidate,
overwrite: bool,
ffmpeg_threads: int,
ffmpeg_preset: str,
ffmpeg_crf: int,
include_audio: bool,
) -> tuple[str, str | None]:
if candidate.span_type == "full":
return candidate.sample_id, None
if candidate.clip_start_sec is None or candidate.clip_end_sec is None:
return candidate.sample_id, "missing_clip_bounds"
try:
_ffmpeg_extract_clip(
Path(candidate.source_video_abs),
Path(candidate.output_video_abs),
candidate.clip_start_sec,
candidate.clip_end_sec,
overwrite=overwrite,
ffmpeg_threads=ffmpeg_threads,
ffmpeg_preset=ffmpeg_preset,
ffmpeg_crf=ffmpeg_crf,
include_audio=include_audio,
)
return candidate.sample_id, None
except Exception as e: # pylint: disable=broad-except
return candidate.sample_id, str(e)
def materialize_span_clips(
candidates: list[SpanCandidate],
*,
num_workers: int = 16,
overwrite: bool = False,
ffmpeg_threads: int = 1,
ffmpeg_preset: str = "ultrafast",
ffmpeg_crf: int = 30,
include_audio: bool = False,
) -> list[tuple[str, str]]:
clip_candidates = [candidate for candidate in candidates if candidate.span_type != "full"]
failures: list[tuple[str, str]] = []
if not clip_candidates:
return failures
with ThreadPoolExecutor(max_workers=num_workers) as executor:
futures = [
executor.submit(
_extract_one,
candidate,
overwrite,
ffmpeg_threads,
ffmpeg_preset,
ffmpeg_crf,
include_audio,
)
for candidate in clip_candidates
]
done = 0
for future in as_completed(futures):
sample_id, error = future.result()
done += 1
if error is not None:
failures.append((sample_id, error))
if done % 500 == 0:
print(
f"[clip] done={done}/{len(clip_candidates)} failures={len(failures)}",
flush=True,
)
return failures
def candidates_to_rows(candidates: list[SpanCandidate]) -> list[dict[str, Any]]:
rows: list[dict[str, Any]] = []
for candidate in candidates:
video_rel = relativize_video_path(candidate.output_video_abs)
row = {
"id": candidate.sample_id,
"video_path": video_rel,
"task_type": "caption",
"prompt": candidate.prompt,
"target_text": candidate.target_text,
"dataset": candidate.dataset,
"split": candidate.split,
"metadata": dict(candidate.metadata),
}
if candidate.clip_start_sec is not None:
row["clip_start_sec"] = candidate.clip_start_sec
if candidate.clip_end_sec is not None:
row["clip_end_sec"] = candidate.clip_end_sec
rows.append(normalize_manifest_row(row, default_split=candidate.split))
return rows
def write_finevideo_manifests(
*,
train_rows: list[dict[str, Any]],
val_rows: list[dict[str, Any]],
val_scene_rows: list[dict[str, Any]],
val_adjacent_rows: list[dict[str, Any]],
val_full_rows: list[dict[str, Any]],
val_gen_rows: list[dict[str, Any]],
train_out: str | Path,
val_out: str | Path,
val_scene_out: str | Path,
val_adjacent_out: str | Path,
val_full_out: str | Path,
val_core_out: str | Path,
val_gen_out: str | Path,
) -> None:
dump_jsonl(train_out, train_rows)
dump_jsonl(val_out, val_rows)
dump_jsonl(val_scene_out, val_scene_rows)
dump_jsonl(val_adjacent_out, val_adjacent_rows)
dump_jsonl(val_full_out, val_full_rows)
dump_jsonl(val_core_out, val_rows[:1024])
dump_jsonl(val_gen_out, val_gen_rows)
def ensure_relative_to_data_root(path: str | Path) -> str:
path = Path(path)
if not path.is_absolute():
return str(path)
return str(path.relative_to(DATA_ROOT))

129
src/ctx_to_lora/data/video_manifest.py Normal file
View file

@ -0,0 +1,129 @@
import os
import json
from pathlib import Path
from typing import Any
DATA_ROOT = Path(os.environ.get("VIDEO2LORA_DATA_ROOT", "data/video2lora"))
CAPTION_PROMPTS = (
"Describe what is happening in this video.",
"Write a short caption for this clip.",
"Summarize this video in one sentence.",
"Describe the main event in this clip.",
)
DEFAULT_INTERNALIZATION_PROMPT = "Internalize this video for later captioning."
def resolve_video_path(video_path_str: str) -> str:
video_path = Path(video_path_str)
if not video_path.is_absolute():
video_path = DATA_ROOT / video_path
return str(video_path)
def relativize_video_path(video_path_str: str) -> str:
video_path = Path(video_path_str)
if not video_path.is_absolute():
return str(video_path)
try:
return str(video_path.relative_to(DATA_ROOT))
except ValueError:
return str(video_path)
def dump_jsonl(path: str | Path, rows: list[dict[str, Any]]) -> None:
output_path = Path(path)
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, "w") as f:
for row in rows:
f.write(json.dumps(row) + "\n")
def infer_split(row: dict[str, Any], default_split: str | None = None) -> str | None:
if row.get("split") is not None:
return str(row["split"])
sample_id = str(row.get("id", ""))
for candidate in ("train", "val", "test"):
if sample_id.startswith(f"{candidate}-") or sample_id.startswith(f"{candidate}_"):
return candidate
return default_split
def infer_dataset(row: dict[str, Any]) -> str:
metadata = row.get("metadata") or {}
if row.get("dataset") is not None:
return str(row["dataset"])
if metadata.get("dataset") is not None:
return str(metadata["dataset"])
return "unknown"
def normalize_manifest_row(
row: dict[str, Any],
*,
default_split: str | None = None,
) -> dict[str, Any]:
prompt = row.get("prompt")
target_text = row.get("target_text")
task_type = row.get("task_type")
if prompt is None and "question" in row:
prompt = row["question"]
if target_text is None and "answer" in row:
target_text = row["answer"]
if task_type is None:
task_type = "qa" if "question" in row or "answer" in row else "caption"
if prompt is None:
raise ValueError(f"Missing prompt/question in row: {row}")
if target_text is None:
raise ValueError(f"Missing target_text/answer in row: {row}")
if "video_path" not in row:
raise ValueError(f"Missing video_path in row: {row}")
metadata = dict(row.get("metadata") or {})
dataset = infer_dataset(row)
if "dataset" not in metadata:
metadata["dataset"] = dataset
normalized = {
"id": str(row.get("id", Path(str(row["video_path"])).stem)),
"video_path": relativize_video_path(str(row["video_path"])),
"task_type": str(task_type),
"prompt": str(prompt),
"target_text": str(target_text),
"dataset": dataset,
"split": infer_split(row, default_split=default_split),
"metadata": metadata,
}
for key in ("clip_start_sec", "clip_end_sec"):
if row.get(key) is not None:
normalized[key] = float(row[key])
return normalized
def load_video_manifest(
path: str | Path,
*,
max_samples: int | None = None,
default_split: str | None = None,
) -> list[dict[str, Any]]:
rows: list[dict[str, Any]] = []
with open(path) as f:
for line in f:
line = line.strip()
if not line:
continue
rows.append(
normalize_manifest_row(
json.loads(line),
default_split=default_split,
)
)
if max_samples is not None and len(rows) >= max_samples:
break
return rows

1149
src/ctx_to_lora/eval_utils.py
View file

File diff suppressed because it is too large Load diff

163
src/ctx_to_lora/metrics.py
View file

@ -1,163 +0,0 @@
from collections import defaultdict
from collections.abc import Callable
import numpy as np
import torch
from rouge_score import rouge_scorer
from transformers import EvalPrediction
LENGTH_BINS = [
# finegrain bins
(0, 2**7 - 1),
(2**7, 2**8 - 1),
(2**8, 2**9 - 1),
# coarse bins
(0, 2**9 - 1),
(2**9, 2**10 - 1),
(2**10, 2**11 - 1),
(2**11, 2**12 - 1),
(2**12, 2**13 - 1),
(0, 2**13 - 1),
(2**13, 2**14 - 1),
(2**14, 2**15 - 1),
(2**15, float("inf")),
]
def get_length_bin(length: int):
"""Get the length bin for a given length."""
for i, (start, end) in enumerate(LENGTH_BINS):
if start <= length < end:
return (start, end)
def compute_rouge(pred_texts, label_texts):
out = defaultdict(list)
scorer = rouge_scorer.RougeScorer(["rougeL"], use_stemmer=True)
for pred_text, label_text in zip(pred_texts, label_texts):
scores = scorer.score(pred_text, label_text)
for k, v in scores.items():
out[f"{k}.f1"].append(v.fmeasure)
out_mean = dict()
for k in out:
out_mean[k] = np.mean(out[k])
return out_mean, out
@torch.inference_mode()
def compute_per_token_acc(shift_logits, shift_labels, valid_masks):
indices = torch.where(valid_masks)
acc = (shift_logits.argmax(-1) == shift_labels)[indices].float()
return {
"per_token_accs": acc.flatten().tolist(),
"n_per_token_accs": valid_masks.sum().item(),
}
@torch.inference_mode()
def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
lengths = valid_masks.sum(dim=1)
is_wrong = (shift_logits.argmax(-1) != shift_labels) * valid_masks
is_correct = (shift_logits.argmax(-1) == shift_labels) * valid_masks
# NOTE: not reliable for multi-turn conversations
# ie, all tokens in the following user's turn will be correct
# still monotonically correlate with perf though
wrong_pos = torch.argmax(is_wrong, dim=1) - torch.argmax(valid_masks, dim=1)
perf = wrong_pos / lengths
# if all tokens are correct, set to 1
perf = torch.where(is_correct.sum(dim=1) == lengths, 1, perf)
return {
"prefix_matchings": perf.tolist(),
"n_prefix_matchings": valid_masks.shape[0],
}
@torch.inference_mode()
def compute_perplexity(shift_logits, shift_labels, valid_masks):
return {"perplexities_ph": [1], "n_perplexities_ph": 1}
class Evaluator:
def __init__(self, metric_fns: list[Callable]):
self.metric_fns = metric_fns
self.reset()
def reset(self):
self.accum_metrics = defaultdict(lambda: list((0,)))
self.count = defaultdict(lambda: list((0,)))
def update(self, shift_logits, shift_labels, valid_masks, lengths=None):
for metric_fn in self.metric_fns:
# overall metric
metric = metric_fn(shift_logits, shift_labels, valid_masks)
for k, v in metric.items():
key = k if not k.startswith("n_") else k[2:]
if k.startswith("n_"):
# prefix "n_" indicates the count of the metric
self.count[key].append(v)
else:
self.accum_metrics[key] += v
for start, end in LENGTH_BINS:
key_w_len = f"{key}_len_{start}-{end}"
if key_w_len not in self.accum_metrics:
# add key here so that it shows up in the output
self.accum_metrics[key_w_len] = [0]
self.count[key_w_len] = [0]
# split samples into length groups, calculate metric for each group
if lengths is not None:
for start, end in LENGTH_BINS:
logits, labels, masks = [], [], []
for logit, label, m, len in zip(
shift_logits, shift_labels, valid_masks, lengths
):
if isinstance(len, torch.Tensor):
len = len.item()
if start <= len < end:
logits.append(logit)
labels.append(label)
masks.append(m)
if not logits:
continue
metric = metric_fn(
torch.stack(logits), torch.stack(labels), torch.stack(masks)
)
for k, v in metric.items():
if k.startswith("n_"):
key = f"{k[2:]}_len_{start}-{end}"
self.count[key].append(v)
else:
key = f"{k}_len_{start}-{end}"
self.accum_metrics[key] += v
def compute(self):
# Get result across entire eval set
result = {
k: np.sum(v) / np.sum(self.count[k]) if len(v) > 1 else "None"
for k, v in self.accum_metrics.items()
}
# Reset batch statistics
self.reset()
return result
@torch.no_grad()
def compute_metrics(
eval_pred: EvalPrediction,
compute_result: bool,
evaluator: Evaluator,
) -> dict | None:
inputs = eval_pred.inputs
len_key = "ctx_ids_len" if "ctx_ids_len" in inputs else "input_ids_len"
lengths = inputs[len_key]
logits, labels = eval_pred.predictions, eval_pred.label_ids
shift_logits = logits[..., :-1, :]
shift_labels = labels[..., 1:]
valid_masks = torch.where(shift_labels != -100, 1, 0)
evaluator.update(shift_logits, shift_labels, valid_masks, lengths)
if compute_result:
return evaluator.compute()

4
src/ctx_to_lora/modeling/aggregator.py
View file

@ -112,7 +112,7 @@ class Perceiver(nn.Module):
n_latents=n_latent_queries,
intermediate_size_factor=4,
hidden_size=output_size,
attn_implementation="flash_attention_2",
attn_implementation="sdpa",
)
self.decoder_config = Idefics2PerceiverConfig(
input_size=output_size,
@ -122,7 +122,7 @@ class Perceiver(nn.Module):
n_latents=n_output_queries,
intermediate_size_factor=4,
hidden_size=output_size,
attn_implementation="flash_attention_2",
attn_implementation="sdpa",
)
self.perceiver = Idefics2Perceiver(self.config, self.decoder_config)
self.iterative_mode = False

632
src/ctx_to_lora/modeling/context_distillation.py
View file

@ -1,632 +0,0 @@
import gc
import json
import logging
import os
import re
from math import ceil
from typing import Any
import torch
from jaxtyping import Integer
from peft import PeftModel
from peft.tuners.tuners_utils import BaseTunerLayer, check_target_module_exists
from torch import Tensor, nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import ModelOutput
from ctx_to_lora.data.definitions import CTX_AFFIXES
from ctx_to_lora.data.q_generation_template import (
Q_GEN_PROMPT_TEMPLATE,
Q_GEN_PROMPT_TEMPLATE_REPEAT,
Q_GEN_SYSTEM_TEMPLATE,
STOP_STRINGS,
)
from ctx_to_lora.data.self_gen_template import SELF_QA_INTX
from ctx_to_lora.utils import log_num_train_params
logger = logging.getLogger()
def get_q_gen_prompt(context, n_qa_pairs):
prompt = Q_GEN_PROMPT_TEMPLATE.format(context=context, n_qa_pairs=n_qa_pairs)
return prompt
def get_q_gen_prompt_repeat(context, qa_pairs, n_qa_pairs):
example_qa_pairs = ""
for i, (q, a) in enumerate(qa_pairs, 1):
example_qa_pairs += f"Question {i}: {q}\nAnswer {i}: {a}\n"
prompt = Q_GEN_PROMPT_TEMPLATE_REPEAT.format(
context=context,
qa_pairs=example_qa_pairs,
n_qa_pairs=n_qa_pairs,
)
return prompt
def check_should_skip(txt: str, vllm_model: str) -> bool:
"""Check if the response should be skipped based on stop strings."""
for stop in STOP_STRINGS[vllm_model]:
if stop in txt[-len(stop) :]:
return (txt.split(stop)[0], False) # Found a valid stop string
return (txt, True) # No valid stop string found, skip this response
def postprocess_qa_pairs(res_txt: str):
"""
Postprocesses the QA pairs from the response text.
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 = []
n_skips = 0
if (len(questions) > 0) and (len(answers) > 0):
n_gen_pairs = min(len(questions), len(answers))
has_left_over = n_gen_pairs < len(questions) or n_gen_pairs < len(answers)
for i in range(n_gen_pairs):
response = answers[i].strip()
question = questions[i].strip()
if not response or not question:
print(f"Skipping empty question or answer at index {i}")
continue
if (not has_left_over) and (i == n_gen_pairs - 1):
response, skip = check_should_skip(response, "google/gemma-3-12b-it")
if skip:
print(f"Skipping due to missing stop string")
n_skips += 1
continue
out_q.append(question.strip())
out_a.append(response.strip())
print(f"Skipped {n_skips} responses due to missing stop strings")
return out_q, out_a
def build_messages(ctx_text: str, level: int, example_qa_pairs: list = None):
messages = [
{"role": "system", "content": Q_GEN_SYSTEM_TEMPLATE},
{
"role": "user",
"content": get_q_gen_prompt(ctx_text, 5)
if level == 0
else get_q_gen_prompt_repeat(ctx_text, example_qa_pairs, 5),
},
]
return messages
def get_shifted_label_pos(labels):
pos = torch.where(labels != -100)
# (batch_idx, token_idx)
return (pos[0], pos[1] - 1)
def logits_at_positions(outputs: ModelOutput, pos) -> Tensor:
logits = outputs.logits
return logits[pos[0], pos[1]]
def ctx_inp_split(
ctx_inp_ids, ctx_inp_sep_seq, pad_token_id, prefix_tokens=None, padding_side="right"
):
# Split each row in ctx_inp_ids at the first occurrence of ctx_inp_sep_seq
# Return the part after the separator for each row
batch_size = ctx_inp_ids.size(0)
sep_len = ctx_inp_sep_seq.size(0)
out_inp = []
out_ctx = []
for i in range(batch_size):
row = ctx_inp_ids[i]
# Find where the separator starts
for j in range(row.size(0) - sep_len + 1):
if torch.equal(row[j : j + sep_len], ctx_inp_sep_seq):
out_ctx.append(row[:j])
if prefix_tokens is not None:
out_inp.append(
torch.cat([prefix_tokens, row[j + sep_len :]], axis=-1)
)
else:
out_inp.append(row[j + sep_len :])
break
else:
# If separator not found
raise ValueError(f"Separator sequence not found in row {i}")
out_inp = torch.nn.utils.rnn.pad_sequence(
out_inp, batch_first=True, padding_value=pad_token_id, padding_side=padding_side
)
out_ctx = torch.nn.utils.rnn.pad_sequence(
out_ctx, batch_first=True, padding_value=pad_token_id, padding_side=padding_side
)
return out_ctx, out_inp
def get_peft_layers(model, peft_config):
out = []
for module_name, module in model.named_modules():
if not check_target_module_exists(peft_config, module_name):
continue
if not isinstance(module, BaseTunerLayer):
continue
# support just Linear layer for now
# all modules should be a leave module that is Linear layer
assert isinstance(module.base_layer, nn.Linear), (
"all modules should be a leave module that is Linear layer"
)
# this should always pass
name = module_name.split(".")[-1]
assert name in peft_config.target_modules
out.append(module)
return out
class CtxDistillModel(nn.Module):
def __init__(
self,
base_model: PeftModel,
prefix_tokens: Integer[Tensor, "n"],
ctx_inp_sep_seq: Integer[Tensor, "m"],
pad_token_id: int,
update_iterations: int,
reset: bool = True,
tokenizer=None,
q_model: PreTrainedModel | None = None,
q_tokenizer=None,
reprompt_ctx: bool = False,
lora_save_dir: str | None = None,
save_after_distill: bool = True,
q_gen_rounds: int = 4,
batch_size: int = 16,
):
super().__init__()
self.register_module("base_model", base_model)
self.register_module("q_model", q_model)
self.register_buffer("prefix_tokens", prefix_tokens)
self.register_buffer("ctx_inp_sep_seq", ctx_inp_sep_seq)
self.tokenizer = tokenizer
self.q_tokenizer = q_tokenizer
self.pad_token_id = pad_token_id
self.update_iterations = update_iterations
self.reprompt_ctx = reprompt_ctx
self.reset = reset
self.device = base_model.device
self.to(self.device)
self.q_gen_rounds = q_gen_rounds
# New save options
self.lora_save_dir = lora_save_dir
self.save_after_distill = save_after_distill
self.peft_config = base_model.peft_config["default"]
self.adapter_name = "default"
self.base_model.set_adapter("default")
for layer in get_peft_layers(self.base_model, self.peft_config):
for name, p in layer.named_parameters():
if "lora_A" in name or "lora_B" in name:
p.requires_grad = True
log_num_train_params(self.base_model)
self._init_optim()
# Mini-batch size for distillation updates
self.batch_size = batch_size
@property
def generation_config(self):
return self.base_model.generation_config
def _init_optim(self):
self.optimizer = torch.optim.AdamW(
[
p
for l in get_peft_layers(self.base_model, self.peft_config)
for p in l.parameters()
if p.requires_grad
],
lr=1e-4,
)
def reset_lora(self):
print("Resetting LoRA")
for layer in get_peft_layers(self.base_model, self.peft_config):
layer.reset_lora_parameters(self.adapter_name, init_lora_weights=True)
self._init_optim()
def save_lora(self):
"""
Save current LoRA adapter in PEFT format plus a lightweight JSON summary
for easy human inspection/manipulation.
"""
if self.lora_save_dir is None:
return
os.makedirs(self.lora_save_dir, exist_ok=True)
# Standard PEFT save (produces adapter_config.json + adapter_model.bin / safetensors)
self.base_model.save_pretrained(self.lora_save_dir)
# Human-readable summary of LoRA parameter shapes
summary = {
name: list(p.shape)
for name, p in self.base_model.named_parameters()
if ("lora_A" in name or "lora_B" in name) and p.requires_grad
}
with open(os.path.join(self.lora_save_dir, "lora_summary.json"), "w") as f:
json.dump(summary, f, indent=2)
print(f"LoRA adapter saved to {self.lora_save_dir}")
@torch.enable_grad()
def _distill_context(
self,
ctx_inp_res_ids: Integer[Tensor, "bs ctx_inp_length"],
ctx_inp_res_attention_mask: Integer[Tensor, "bs ctx_inp_length"],
teacher_labels: Integer[Tensor, "bs ctx_inp_length"],
inp_res_ids: Integer[Tensor, "bs inp_length"],
inp_res_attention_mask: Integer[Tensor, "bs inp_length"],
student_labels: Integer[Tensor, "bs inp_length"],
):
# Implements KD-style loss by computing teacher (with context) and student (no context)
# log-probs locally, using mini-batches for updates.
was_training = self.training
self.train()
num_samples = ctx_inp_res_ids.size(0)
mb = self.batch_size
num_batches = ceil(num_samples / mb)
total_steps = max(self.update_iterations * num_batches, 1)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
self.optimizer, T_max=total_steps, eta_min=0.0
)
print(
f"Starting context distillation for {self.update_iterations} epochs, "
f"mini-batch size {mb} ({num_batches} batches/epoch), "
f"cosine LR schedule over {total_steps} steps"
)
for epoch in range(self.update_iterations):
# Shuffle order each epoch for SGD
perm = torch.randperm(num_samples, device=self.device)
epoch_loss = 0.0
for b in range(num_batches):
start = b * mb
end = min(start + mb, num_samples)
indices = perm[start:end]
b_ctx_ids = ctx_inp_res_ids[indices]
b_ctx_am = ctx_inp_res_attention_mask[indices]
b_teacher_labels = teacher_labels[indices]
b_inp_ids = inp_res_ids[indices]
b_inp_am = inp_res_attention_mask[indices]
b_student_labels = student_labels[indices]
# Compute teacher distribution (top-k) for this mini-batch
with torch.no_grad(), self.base_model.disable_adapter():
t_pos = get_shifted_label_pos(b_teacher_labels)
teacher_outputs = self.base_model(
b_ctx_ids, attention_mask=b_ctx_am
)
teacher_logits = logits_at_positions(teacher_outputs, t_pos)
K = 16
topk_vals, topk_idx = teacher_logits.topk(K, dim=-1)
teacher_denom = torch.logsumexp(
teacher_logits.float(), dim=-1, keepdim=True
)
teacher_p = (topk_vals - teacher_denom).exp().detach() # [N, K]
# Student forward and update for this mini-batch
self.optimizer.zero_grad()
s_pos = get_shifted_label_pos(b_student_labels)
student_outputs = self.base_model(b_inp_ids, attention_mask=b_inp_am)
student_logits = logits_at_positions(student_outputs, s_pos)
student_denom = torch.logsumexp(
student_logits.float(), dim=-1, keepdim=True
)
selected_student_logits = student_logits.gather(-1, topk_idx)
student_logq = selected_student_logits - student_denom # [N, K]
token_losses = -(teacher_p * student_logq).sum(dim=-1) # [N]
loss = token_losses.mean()
loss.backward()
self.optimizer.step()
scheduler.step()
epoch_loss += loss.detach().item()
cur_lr = self.optimizer.param_groups[0]["lr"]
print(
f"Epoch {epoch + 1}/{self.update_iterations}, "
f"batch {b + 1}/{num_batches}: loss={epoch_loss / num_batches:.4f}, lr={cur_lr:.6e}"
)
if not was_training:
self.eval()
def generate_questions(self, *args, **kwargs):
questions = self.q_model.generate(*args, **kwargs)
return questions
def teacher_generate(self, *args, **kwargs):
# rename for separate timing
return self.base_model.generate(*args, **kwargs)
def student_generate(self, *args, **kwargs):
# rename for separate timing
return self.base_model.generate(*args, **kwargs)
def get_lora_state(self, clone: bool = True):
"""
Return a dict of current LoRA parameter tensors.
clone=True returns detached cloned tensors (safe to store).
"""
return {
name: (p.detach().clone() if clone else p)
for name, p in self.base_model.named_parameters()
if ("lora_A" in name or "lora_B" in name)
}
def generate(
self,
*model_inputs_args: Any,
distill_only: bool = False,
**model_inputs_kwargs: dict[str, Any],
):
if self.reset:
self.reset_lora()
# teacher tokens
orig_ctx_inp_ids = model_inputs_kwargs.pop("input_ids")
ctx_inp_ids = orig_ctx_inp_ids.clone()
_, orig_inp_ids = ctx_inp_split(
ctx_inp_ids,
self.ctx_inp_sep_seq,
self.pad_token_id,
self.prefix_tokens,
padding_side="left",
)
ctx_inp_attention_mask = model_inputs_kwargs.pop("attention_mask")
if self.q_model is not None:
self.q_model.to(self.base_model.device)
# Extract context-only portion after separator (remove prefix tokens from first row)
ctx_ids_full, _ = ctx_inp_split(
ctx_inp_ids, self.ctx_inp_sep_seq, self.pad_token_id
) # [bs, var_len]
ctx_ids = ctx_ids_full[0, len(self.prefix_tokens) :]
ctx_txt = self.tokenizer.decode(ctx_ids, skip_special_tokens=True)
questions = []
answers = []
# Build multiple instruction variants
for lvl in range(self.q_gen_rounds):
messages = build_messages(
ctx_txt, lvl, zip(questions, answers) if lvl > 0 else None
)
q_inputs = self.q_tokenizer.apply_chat_template(
messages,
tokenize=True,
add_special_tokens=False,
padding=False,
truncation=False,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt",
)
q_inputs = {k: v.to(self.q_model.device) for k, v in q_inputs.items()}
with torch.no_grad():
question_outputs = self.generate_questions(
input_ids=q_inputs["input_ids"],
attention_mask=q_inputs["attention_mask"],
max_new_tokens=1024,
do_sample=False,
temperature=0.0,
eos_token_id=106, # <end_of_turn> for gemma-3-12b-it
)
# Slice off the prompt portion
gen_only = question_outputs[:, q_inputs["input_ids"].shape[-1] :]
res = self.q_tokenizer.batch_decode(gen_only, skip_special_tokens=False)
gen_q_list, gen_a_list = postprocess_qa_pairs(res[0])
questions += gen_q_list
answers += gen_a_list
if len(questions) == 0:
# when q_model refuses to provide questions
# only happens with sample 116 in longbench/multifieldqa_en_e
# in this case cd just doesn't work, we fall back to zero-shot answer
print(f"Warning---no questions generated, skipping distillation")
attention_mask = torch.where(
orig_inp_ids != self.pad_token_id, 1, 0
).long()
return self.student_generate(
orig_inp_ids, attention_mask=attention_mask, **model_inputs_kwargs
)
ctx_inp_messages = [
[{"role": "user", "content": f"{ctx_txt}\n\n{SELF_QA_INTX}\n\n{q}"}]
for q in questions
]
encoded_ctx_inp = self.tokenizer.apply_chat_template(
ctx_inp_messages,
tokenize=True,
add_special_tokens=False,
padding=True,
truncation=False,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt",
)
encoded_ctx_inp = {k: v.to(self.device) for k, v in encoded_ctx_inp.items()}
ctx_inp_ids = encoded_ctx_inp["input_ids"]
ctx_inp_attention_mask = encoded_ctx_inp["attention_mask"]
self.q_model.to("cpu")
gc.collect()
torch.cuda.empty_cache()
# sample responses first
ctx_inp_res_ids = self.teacher_generate(
ctx_inp_ids,
attention_mask=ctx_inp_attention_mask,
**model_inputs_kwargs,
)
ctx_inp_res_attention_mask = torch.where(
ctx_inp_res_ids != self.pad_token_id, 1, 0
).long()
ctx_inp_res_txt = self.tokenizer.batch_decode(ctx_inp_res_ids)
bs = ctx_inp_ids.shape[0]
res_len = ctx_inp_res_ids.shape[-1] - ctx_inp_ids.shape[-1]
res_ids = ctx_inp_res_ids[:, -res_len:] # correct
pads = torch.full_like(ctx_inp_ids, self.pad_token_id)
teacher_labels = torch.cat([pads, res_ids], dim=-1)
teacher_labels = torch.where(
teacher_labels != self.pad_token_id, teacher_labels, -100
)
# student tokens
_, inp_res_ids = ctx_inp_split(
ctx_inp_res_ids,
self.ctx_inp_sep_seq,
self.pad_token_id,
self.prefix_tokens,
padding_side="left",
)
inp_res_attention_mask = torch.where(
inp_res_ids != self.pad_token_id, 1, 0
).long()
student_labels = inp_res_ids.clone()
student_labels[:, :-res_len] = -100
student_labels = torch.where(
student_labels != self.pad_token_id, student_labels, -100
)
self._distill_context(
ctx_inp_res_ids,
ctx_inp_res_attention_mask,
teacher_labels,
inp_res_ids,
inp_res_attention_mask,
student_labels,
)
# Save LoRA after distillation if requested
if distill_only:
return self.get_lora_state()
model_inputs_kwargs.pop("attention_mask", None)
model_inputs_kwargs.pop("input_ids", None)
if self.reprompt_ctx:
attention_mask = torch.where(orig_ctx_inp_ids != self.pad_token_id, 1, 0)
model_outputs = self.student_generate(
orig_ctx_inp_ids, attention_mask=attention_mask, **model_inputs_kwargs
)
else:
attention_mask = torch.where(orig_inp_ids != self.pad_token_id, 1, 0).long()
model_outputs = self.student_generate(
orig_inp_ids, attention_mask=attention_mask, **model_inputs_kwargs
)
return model_outputs
if __name__ == "__main__":
from ctx_to_lora.data.processing import load_and_process_dataset
from ctx_to_lora.model_loading import get_lora_config, get_model_and_tokenizer
model_name = "google/gemma-2-2b-it"
q_model_name = "google/gemma-3-12b-it"
peft_config = get_lora_config(
model_name, r=8, target_modules=["down_proj"], lora_dropout=0.0
)
peft_config.lora_alpha = 16
model, tokenizer = get_model_and_tokenizer(
model_name,
train=False,
requires_grad=False,
peft_config=peft_config,
)
q_model, q_tokenizer = get_model_and_tokenizer(
q_model_name,
train=False,
requires_grad=False,
peft_config=peft_config,
)
ds = load_and_process_dataset("pwc", split="train", num_proc=8)
ctx = ds[0]["context"]
inp = ds[1]["prompts"][0]
# Build a simple context/input pair separated by a unique token sequence
sep_text = SELF_QA_INTX
# ctx = "# Provided Information\nMy name is Tan."
# ctx = "# Provided Info"
prompt = f"{ctx}\n\n{sep_text}\n\n{inp}"
messages = [{"role": "user", "content": prompt}]
encoded = tokenizer.apply_chat_template(
messages,
return_tensors="pt",
return_dict=True,
add_generation_prompt=True,
)
# encoded = tokenizer(prompt, return_tensors="pt")
encoded = {k: v.to(model.device) for k, v in encoded.items()}
prefix_tokens = CTX_AFFIXES[model_name]["prefix"]
prefix_tokens = torch.tensor(prefix_tokens, dtype=torch.long)
sep_ids = (
tokenizer(sep_text.strip("\n"), add_special_tokens=False, return_tensors="pt")
.input_ids[0]
.to(model.device)
)
cd_model = CtxDistillModel(
base_model=model,
prefix_tokens=prefix_tokens,
ctx_inp_sep_seq=sep_ids,
pad_token_id=tokenizer.pad_token_id,
update_iterations=100,
q_model=q_model,
q_tokenizer=q_tokenizer,
tokenizer=tokenizer,
reprompt_ctx=False,
lora_save_dir="./saved_lora_adapter", # example save path
save_after_distill=True,
)
with torch.no_grad():
for _ in range(1):
base_model_res = model.generate(
input_ids=encoded["input_ids"],
attention_mask=encoded["attention_mask"],
max_new_tokens=256,
do_sample=False,
)
print(
f"Base model response:{tokenizer.batch_decode(base_model_res, skip_special_tokens=False)}"
)
outputs = cd_model.generate(
input_ids=encoded["input_ids"],
attention_mask=encoded["attention_mask"],
max_new_tokens=256,
do_sample=False,
)
print(
f"Student response: {tokenizer.batch_decode(outputs, skip_special_tokens=False)}"
)

70
src/ctx_to_lora/modeling/generative_adapter.py
View file

@ -1,70 +0,0 @@
import requests
import torch
def call_generate(
input_txt: str,
context_txt: str,
window_size: int | None = None,
max_new_tokens: int | None = None,
host: str = "http://127.0.0.1:8989",
timeout: int = 120,
) -> torch.Tensor:
"""Send the prompt to the API server and return the generated token tensor."""
payload: dict[str, object] = {
"input_txt": input_txt,
"context_txt": context_txt,
}
if window_size is not None:
payload["window_size"] = int(window_size)
if max_new_tokens is not None:
payload["max_new_tokens"] = int(max_new_tokens)
response = requests.post(f"{host}/generate", json=payload, timeout=timeout)
response.raise_for_status()
data = response.json()
if "output" not in data:
raise ValueError(f"Unexpected response payload: {data}")
return torch.tensor([data["output"]])
def check_server_health(host: str = "http://127.0.0.1:8989", timeout: int = 60) -> None:
"""Check if the API server is healthy and responding."""
try:
response = requests.get(f"{host}/health", timeout=timeout)
response.raise_for_status()
data = response.json()
if data.get("status") != "ok":
raise RuntimeError(f"Server is not healthy: {data}")
except requests.exceptions.ConnectionError as e:
raise RuntimeError(
f"Cannot connect to server at {host}. Is the server running?"
) from e
except requests.exceptions.Timeout as e:
raise RuntimeError(f"Server health check timed out after {timeout}s") from e
except requests.exceptions.RequestException as e:
raise RuntimeError(f"Server health check failed: {e}") from e
print("Server is healthy.")
class GenerativeAdapter(torch.nn.Module):
def __init__(self, model, tokenizer):
super().__init__()
self.base_model = model # placeholder
self.tokenizer = tokenizer
check_server_health()
@property
def generation_config(self):
return self.base_model.generation_config
def generate(self, *args, **kwargs):
ctx_ids = kwargs["ctx_ids"]
input_ids = kwargs["input_ids"]
assert ctx_ids.shape[0] == 1
assert input_ids.shape[0] == 1
context_txt = self.tokenizer.decode(ctx_ids[0])
input_txt = self.tokenizer.decode(input_ids[0])
outputs = call_generate(input_txt, context_txt)
return outputs

137
src/ctx_to_lora/modeling/hypernet.py
View file

@ -87,6 +87,16 @@ def get_hypernet_config(
aggregator_args: AggregatorArguments,
ctx_encoder_args: CtxEncoderArguments,
):
base_model_config = model.config
if getattr(base_model_config, "hidden_size", None) is not None:
base_hidden_size = base_model_config.hidden_size
elif getattr(base_model_config, "text_config", None) is not None:
base_hidden_size = base_model_config.text_config.hidden_size
else:
raise AttributeError(
f"Could not determine hidden size from config type {type(base_model_config)!r}"
)
num_modules = 0
lora_config = getattr(model, "peft_config", None)
if lora_config is not None:
@ -96,7 +106,7 @@ def get_hypernet_config(
indices = torch.arange(get_num_layers(model), device=model.device)
return HypernetConfig(
**vars(hypernet_args),
base_hidden_size=model.config.hidden_size,
base_hidden_size=base_hidden_size,
lora_config=lora_config,
layer_indices=indices,
feature_sizes=get_peft_in_out_features(model, peft_config=lora_config),
@ -466,6 +476,17 @@ class ModulatedPretrainedModel(nn.Module):
self._init_model()
self._bias_hyper_init()
def _runtime_device(
self,
ctx_ids: Tensor | None = None,
ctx_features: Tensor | None = None,
input_ids: Tensor | None = None,
) -> torch.device:
for tensor in (ctx_features, ctx_ids, input_ids):
if tensor is not None:
return tensor.device
return next(self.base_model.parameters()).device
@classmethod
def from_state_dict(
cls,
@ -533,6 +554,9 @@ class ModulatedPretrainedModel(nn.Module):
self.patch_lora_forward()
ctx_model_name = self.ctx_encoder_args.ctx_encoder_model_name_or_path
if ctx_model_name == "precomputed":
self.ctx_encoder = None
return
if ctx_model_name is None:
ctx_model_name = self.base_model.config.name_or_path
# use an explicit copy of the base model
@ -588,7 +612,9 @@ class ModulatedPretrainedModel(nn.Module):
def state_dict(self, *args, **kwargs):
# we assume ctx_encoder and base model is frozen here
if len([p for p in self.ctx_encoder.parameters() if p.requires_grad]):
if self.ctx_encoder is not None and len(
[p for p in self.ctx_encoder.parameters() if p.requires_grad]
):
raise ValueError("ctx_encoder contains trainable parameters")
if len([p for p in self.base_model.parameters() if p.requires_grad]):
raise ValueError("base model contains trainable parameters")
@ -630,43 +656,51 @@ class ModulatedPretrainedModel(nn.Module):
ctx_ids: Integer[Tensor, "bs ctx_len"],
ctx_attn_mask: Integer[Tensor, "bs ctx_len"] | None = None,
ctx_position_ids: Integer[Tensor, "bs ctx_len"] | None = None,
ctx_features: Float[Tensor, "bs ctx_len hidden_size"] | None = None,
**kwargs: Any,
):
with torch.no_grad():
ctx_encoder_kwargs = dict(
input_ids=ctx_ids,
attention_mask=ctx_attn_mask,
position_ids=ctx_position_ids,
)
if isinstance(self.ctx_encoder.base_model, ModernBertModel):
position_ids = ctx_position_ids.flatten()
indices = torch.arange(
position_ids.size(0), device=position_ids.device, dtype=torch.int32
)
# [bsz + 1]
cu_seqlens = torch.cat(
(
indices[position_ids == 0],
torch.tensor(
position_ids.size(),
device=position_ids.device,
dtype=torch.int32,
),
)
if ctx_features is None:
if self.ctx_encoder is None:
raise ValueError(
"ctx_features must be provided when ctx_encoder_model_name_or_path='precomputed'"
)
with torch.no_grad():
ctx_encoder_kwargs = dict(
input_ids=ctx_ids.squeeze(0),
cu_seqlens=cu_seqlens,
max_seqlen=position_ids.max() + 1,
attention_mask=-1,
seq_len=-1,
batch_size=-1,
input_ids=ctx_ids,
attention_mask=ctx_attn_mask,
position_ids=ctx_position_ids,
)
if isinstance(self.ctx_encoder.base_model, ModernBertModel):
position_ids = ctx_position_ids.flatten()
indices = torch.arange(
position_ids.size(0),
device=position_ids.device,
dtype=torch.int32,
)
# [bsz + 1]
cu_seqlens = torch.cat(
(
indices[position_ids == 0],
torch.tensor(
position_ids.size(),
device=position_ids.device,
dtype=torch.int32,
),
)
)
ctx_encoder_kwargs = dict(
input_ids=ctx_ids.squeeze(0),
cu_seqlens=cu_seqlens,
max_seqlen=position_ids.max() + 1,
attention_mask=-1,
seq_len=-1,
batch_size=-1,
)
ctx_features = self.ctx_encoder(**ctx_encoder_kwargs, **kwargs)
ctx_features = self.ctx_encoder(**ctx_encoder_kwargs, **kwargs)
if isinstance(self.ctx_encoder.base_model, ModernBertModel):
ctx_features = ctx_features.unsqueeze(0)
if isinstance(self.ctx_encoder.base_model, ModernBertModel):
ctx_features = ctx_features.unsqueeze(0)
if self.user_defined_scaling == 1:
return self.hypernet.generate_weights(
ctx_features, ctx_attn_mask, ctx_position_ids
@ -688,6 +722,7 @@ class ModulatedPretrainedModel(nn.Module):
ctx_ids: Integer[Tensor, "n_ctx ctx_len"] | None = None,
ctx_attn_mask: Integer[Tensor, "n_ctx ctx_len"] | None = None,
ctx_position_ids: Integer[Tensor, "n_ctx ctx_len"] | None = None,
ctx_features: Float[Tensor, "n_ctx ctx_len hidden_size"] | None = None,
n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None,
n_queries: Integer[Tensor, "n_ctx"] | None = None,
return_generated_lora: bool | None = False,
@ -697,7 +732,12 @@ class ModulatedPretrainedModel(nn.Module):
"""Forward pass of the modulated model."""
generated_loras = None
generated_layernorms = None
if ctx_ids is None and not self.use_base_input_as_ctx:
runtime_device = self._runtime_device(
ctx_ids=ctx_ids,
ctx_features=ctx_features,
input_ids=model_inputs_kwargs.get("input_ids"),
)
if ctx_ids is None and ctx_features is None and not self.use_base_input_as_ctx:
logger.warning(
(
"*" * 100,
@ -724,7 +764,10 @@ class ModulatedPretrainedModel(nn.Module):
else None
)
generated_loras, generated_layernorms = self.generate_weights(
ctx_ids, ctx_attn_mask, ctx_position_ids
ctx_ids,
ctx_attn_mask,
ctx_position_ids,
ctx_features=ctx_features,
)
if generated_loras is not None:
@ -745,9 +788,19 @@ class ModulatedPretrainedModel(nn.Module):
)
if n_queries is None:
if ctx_position_ids is None:
if not self.use_sequence_packing:
ctx_batch_size = (
ctx_ids.shape[0] if ctx_ids is not None else ctx_features.shape[0]
)
n_queries = torch.ones(
ctx_ids.shape[0], dtype=torch.int32, device=self.device
ctx_batch_size, dtype=torch.int32, device=runtime_device
)
elif ctx_position_ids is None:
ctx_batch_size = (
ctx_ids.shape[0] if ctx_ids is not None else ctx_features.shape[0]
)
n_queries = torch.ones(
ctx_batch_size, dtype=torch.int32, device=runtime_device
)
else:
# quite redundant (we do cu_seqlens many places)
@ -755,7 +808,7 @@ class ModulatedPretrainedModel(nn.Module):
n_queries = torch.ones(
(ctx_position_ids == 0).sum(),
dtype=torch.int32,
device=self.device,
device=runtime_device,
)
apply_lora_to_layers(
@ -825,6 +878,10 @@ class ModulatedPretrainedModel(nn.Module):
*model_inputs_args: Any,
**model_inputs_kwargs: dict[str, Any],
):
runtime_device = self._runtime_device(
ctx_ids=ctx_ids,
input_ids=model_inputs_kwargs.get("input_ids"),
)
generated_loras = None
generated_layernorms = None
if (
@ -840,7 +897,7 @@ class ModulatedPretrainedModel(nn.Module):
elif ctx_ids is None and self.generated_loras:
generated_loras = self.generated_loras
if n_ctx_chunks is None:
n_ctx_chunks = torch.tensor((1,), device=self.device)
n_ctx_chunks = torch.tensor((1,), device=runtime_device)
print(
"*" * 100
+ "\n\nUsing internalized LoRAs for generation\n\n"
@ -890,7 +947,7 @@ class ModulatedPretrainedModel(nn.Module):
n_queries = torch.ones(
model_inputs_kwargs["input_ids"].shape[0],
dtype=torch.int32,
device=self.device,
device=runtime_device,
)
else:
# quite redundant (we do cu_seqlens many places)
@ -898,7 +955,7 @@ class ModulatedPretrainedModel(nn.Module):
n_queries = torch.ones(
(ctx_position_ids == 0).sum(),
dtype=torch.int32,
device=self.device,
device=runtime_device,
)
apply_lora_to_layers(

160
src/ctx_to_lora/modeling/idefics2.py
View file

@ -246,7 +246,8 @@ class Idefics2PerceiverAttention(nn.Module):
def forward(
self,
latents: torch.Tensor,
context: torch.Tensor,
is_cross_attn: bool,
context: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_value: tuple[torch.Tensor] | None = None,
@ -266,9 +267,13 @@ class Idefics2PerceiverAttention(nn.Module):
use_cache (`bool`, *optional*, defaults to `False`): Whether to use past_key_value for caching.
"""
bsz, q_len, _ = latents.size()
kv_seq_len = q_len + context.size()[1]
hidden_states = torch.concat([context, latents], dim=-2)
if is_cross_attn:
if context is None:
raise ValueError("Cross-attention requires context inputs.")
hidden_states = torch.concat([context, latents], dim=-2)
else:
hidden_states = latents
kv_seq_len = hidden_states.size(1)
query_states = self.q_proj(latents)
key_states = self.k_proj(hidden_states)
@ -448,7 +453,8 @@ class Idefics2PerceiverFlashAttention2(Idefics2PerceiverAttention):
IDEFICS2_PERCEIVER_ATTENTION_CLASSES = {
# "eager": Idefics2PerceiverAttention,
"eager": Idefics2PerceiverAttention,
"sdpa": Idefics2PerceiverAttention,
"flash_attention_2": Idefics2PerceiverFlashAttention2,
}
@ -566,7 +572,7 @@ IDEFICS2_INPUTS_DOCSTRING = r"""
IDEFICS2_START_DOCSTRING,
)
class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
_supports_sdpa = False
_supports_sdpa = True
config_class = Idefics2PerceiverConfig
def __init__(self, config) -> None:
@ -612,7 +618,6 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
self.layernorm = Idefics2RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
assert self._use_flash_attention_2
def forward(
self,
@ -621,79 +626,86 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
position_ids: torch.LongTensor | None = None,
) -> torch.Tensor:
# seq embed -> bsz seq embed
if position_ids is None:
bsz = context.shape[0]
else:
if self._use_flash_attention_2 and position_ids is not None:
# flattened packed sequence
bsz = torch.where(position_ids == 0, 1, 0).sum()
else:
bsz = context.shape[0]
latents = self.latents_q.unsqueeze(0).expand((bsz, *self.latents_q.size()))
attention_mask = (
_prepare_4d_attention_mask(
attention_mask, latents.dtype, tgt_len=self.n_latents
)
if not self._use_flash_attention_2
else attention_mask
)
compressed_context = latents
cu_seq_lens_q = torch.tensor(
[self.n_latents] * (bsz + 1), device=context.device, dtype=torch.int32
) * torch.arange(bsz + 1, device=context.device, dtype=torch.int32)
max_length_q = self.n_latents
# cu_seq_lens_k = None
# max_length_k = None
if attention_mask is not None:
logger.warning_once("Using attention mask for resampler")
context, _, cu_seq_lens_k, max_length_k, _ = unpad_input(
context, attention_mask
)
context = context.unsqueeze(0)
position_ids = True # goes down flash attn path that uses cu_seq_lens
elif position_ids is not None:
logger.warning_once("Using position ids for resampler")
position_ids = position_ids.flatten()
indices = torch.arange(
position_ids.size(0), device=position_ids.device, dtype=torch.int32
)
# [bsz + 1]
cu_seq_lens_k = torch.cat(
(
indices[position_ids == 0],
torch.tensor(
position_ids.size(),
device=position_ids.device,
dtype=torch.int32,
),
if self._use_flash_attention_2:
cu_seq_lens_q = torch.tensor(
[self.n_latents] * (bsz + 1), device=context.device, dtype=torch.int32
) * torch.arange(bsz + 1, device=context.device, dtype=torch.int32)
max_length_q = self.n_latents
if attention_mask is not None:
logger.warning_once("Using attention mask for resampler")
context, _, cu_seq_lens_k, max_length_k, _ = unpad_input(
context, attention_mask
)
context = context.unsqueeze(0)
position_ids = True
elif position_ids is not None:
logger.warning_once("Using position ids for resampler")
position_ids = position_ids.flatten()
indices = torch.arange(
position_ids.size(0), device=position_ids.device, dtype=torch.int32
)
cu_seq_lens_k = torch.cat(
(
indices[position_ids == 0],
torch.tensor(
position_ids.size(),
device=position_ids.device,
dtype=torch.int32,
),
)
)
max_length_k = position_ids.max() + 1
else:
raise ValueError("either position_ids or attention_mask is required")
x_attn_kwargs = dict(
position_ids=position_ids,
cu_seq_lens_q=cu_seq_lens_q,
cu_seq_lens_k=cu_seq_lens_k,
max_length_q=max_length_q,
max_length_k=max_length_k,
)
self_attn_position_ids = torch.arange(
self.n_latents, device=context.device, dtype=torch.int32
).repeat(1, bsz)
self_attn_kwargs = dict(
position_ids=self_attn_position_ids,
cu_seq_lens_q=cu_seq_lens_q,
cu_seq_lens_k=cu_seq_lens_q,
max_length_q=max_length_q,
max_length_k=max_length_q,
)
max_length_k = position_ids.max() + 1
else:
raise ValueError("either position_ids or attention_mask is required")
x_attn_kwargs = dict(
position_ids=position_ids,
cu_seq_lens_q=cu_seq_lens_q,
cu_seq_lens_k=cu_seq_lens_k,
max_length_q=max_length_q,
max_length_k=max_length_k,
)
self_attn_position_ids = torch.arange(
self.n_latents, device=context.device, dtype=torch.int32
).repeat(1, bsz)
self_attn_kwargs = dict(
# attention_mask=self_attn_mask,
position_ids=self_attn_position_ids,
cu_seq_lens_q=cu_seq_lens_q,
cu_seq_lens_k=cu_seq_lens_q,
max_length_q=max_length_q,
max_length_k=max_length_q,
)
x_attn_mask = None
if attention_mask is not None:
context_attn_mask = _prepare_4d_attention_mask(
attention_mask, latents.dtype, tgt_len=self.n_latents
)
latent_attn_mask = torch.zeros(
(
context_attn_mask.shape[0],
1,
self.n_latents,
self.n_latents,
),
dtype=context_attn_mask.dtype,
device=context_attn_mask.device,
)
x_attn_mask = torch.cat(
[context_attn_mask, latent_attn_mask], dim=-1
)
x_attn_kwargs = {"attention_mask": x_attn_mask}
self_attn_kwargs = {"attention_mask": None}
for i, layer in enumerate(self.layers):
inp_kwargs = dict(
latents=compressed_context,
@ -738,10 +750,10 @@ class Idefics2Perceiver(Idefics2PreTrainedModel):
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
):
if position_ids is None:
bsz = context.shape[0]
else:
if self.encoder._use_flash_attention_2 and position_ids is not None:
bsz = torch.where(position_ids == 0, 1, 0).sum()
else:
bsz = context.shape[0]
projected_inputs = self.modality_projection(context)
# [bsz, n_latents, dim]

143
src/ctx_to_lora/modeling/llm_lingua.py
View file

@ -1,143 +0,0 @@
import torch
from llmlingua import PromptCompressor
from torch import nn
from ctx_to_lora.data.definitions import CTX_AFFIXES
class LLMLinguaModel(nn.Module):
def __init__(self, model, tokenizer, compression_rate):
super().__init__()
self.base_model = model
self.compressor = PromptCompressor(
model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank",
use_llmlingua2=True, # Whether to use llmlingua-2
)
model_name = self.base_model.name_or_path
self.register_buffer("prefix", torch.tensor(CTX_AFFIXES[model_name]["prefix"]))
self.register_buffer("suffix", torch.tensor(CTX_AFFIXES[model_name]["suffix"]))
self.len_prefix = len(self.prefix)
self.len_suffix = len(self.suffix)
self.tokenizer = tokenizer
self.compression_rate = compression_rate
@property
def generation_config(self):
return self.base_model.generation_config
def compress(self, prompt_txt: str, rate: float):
return self.compressor.compress_prompt(
prompt_txt, rate=rate, force_tokens=["\n", "?"]
)
def compress_tokens(self, input_ids, query_text):
bs = input_ids.shape[0]
txt = self.tokenizer.batch_decode(
input_ids[:, self.len_prefix : -self.len_suffix]
)
q_start_idx = txt[0].rfind(query_text)
ctx_txt = txt[0][:q_start_idx]
q_txt = txt[0][q_start_idx:]
compressed_txt = self.compress(ctx_txt, rate=self.compression_rate)
compressed_ids = self.tokenizer(
compressed_txt["compressed_prompt"] + "\n\n" + q_txt,
return_attention_mask=False,
add_special_tokens=False,
return_tensors="pt",
)["input_ids"].to(self.base_model.device)
out = torch.cat(
[self.prefix.expand(bs, -1), compressed_ids, self.suffix.expand(bs, -1)],
dim=-1,
)
return out
def generate(self, *args, **kwargs):
# take ctx_ids
# strip prefix and suffix
# ctx_ids is left padded
ctx_ids = kwargs["ctx_ids"][:, self.len_prefix : -self.len_suffix]
# decode ctx_ids to ctx_txt
ctx_txt = self.tokenizer.batch_decode(ctx_ids)
compressed_ctx_txt = self.compress(ctx_txt, rate=self.compression_rate)
compressed_ctx_ids = self.tokenizer(
compressed_ctx_txt["compressed_prompt"] + "\n\n",
return_attention_mask=False,
add_special_tokens=False,
return_tensors="pt",
).to(self.base_model.device)
bs = ctx_ids.shape[0]
ctx_inp_ids = torch.cat(
[
self.prefix.expand(bs, -1),
compressed_ctx_ids["input_ids"],
kwargs["input_ids"][:, self.len_prefix :],
],
dim=-1,
)
attn_mask = torch.ones_like(ctx_inp_ids)
for k in [
"ctx_ids",
"ctx_attn_mask",
"n_ctx_chunks",
"input_ids",
"attention_mask",
]:
kwargs.pop(k, None)
return self.base_model.generate(ctx_inp_ids, attention_mask=attn_mask, **kwargs)
if __name__ == "__main__":
from ctx_to_lora.model_loading import get_model_and_tokenizer
model, tokenizer = get_model_and_tokenizer(
"google/gemma-2-2b-it",
train=False,
requires_grad=False,
)
# Demo: build wrapper, create a toy context + prompt, run compression + generation.
device = "cuda"
llm = LLMLinguaModel(model, tokenizer).to(device)
# Toy context and user prompt
context_text = (
"This is a short illustrative context about large language models and compression. "
"They can reduce prompt length while preserving meaning."
)
user_prompt = (
"Summarize the context in one concise sentence." # what we want model to do
)
# Tokenize raw context (core) without special tokens
core_ctx_ids = tokenizer.apply_chat_template(
[[{"role": "user", "content": context_text}]],
tokenize=True,
add_generation_prompt=True,
return_attention_mask=False,
padding=False,
truncation=False,
return_tensors="pt",
add_special_tokens=False,
).to(device)
# Input prompt tokens (what follows the contextual block)
input_ids = tokenizer.apply_chat_template(
[[{"role": "user", "content": user_prompt}]],
tokenize=True,
add_generation_prompt=True,
return_attention_mask=False,
padding=False,
truncation=False,
return_tensors="pt",
add_special_tokens=False,
).to(device)
print("Original context length (chars):", len(context_text))
# Run generation (may vary depending on model capabilities)
output_ids = llm.generate(ctx_ids=core_ctx_ids, input_ids=input_ids)
# Decode only the tail beyond supplied input for readability
generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=False)
print(f"\nFull generated text:\n{generated_text}")

160
src/ctx_to_lora/modeling/text_to_lora.py
View file

@ -1,160 +0,0 @@
from functools import partial
import torch
from peft import PeftConfig
from torch import nn
from ctx_to_lora.modeling.lora_layer import apply_lora_to_layers, lora_forward
from ctx_to_lora.modeling.text_to_lora_impl import (
embed_texts,
get_layers,
get_peft_config,
load_hypermod,
)
from ctx_to_lora.utils import get_peft_modules
class TextToLoRA(nn.Module):
def __init__(self, model_name_or_path, prefix_tokens, device):
assert model_name_or_path == "google/gemma-2-2b-it"
super().__init__()
hypermod_dir = "trained_t2l/gemma_2b_t2l"
peft_config = get_peft_config(
PeftConfig.from_json_file(f"{hypermod_dir}/adapter_config.json")
)
# ours lora forward pass uses alpha directly
peft_config.lora_alpha = peft_config.lora_alpha / peft_config.r
self.prefix_tokens = prefix_tokens
self.device = device
(
_,
self.t2l_model,
self.base_model,
self.tokenizer,
self.emb_model,
self.emb_tokenizer,
self.task_desc_format_fn,
self.pooling_fn,
) = load_hypermod(hypermod_dir, device)
layer_indices = range(len(get_layers(self.base_model)))
self.layer_indices = torch.tensor(
layer_indices, dtype=torch.long, device=device
)
# patch base model forward pass to use lora
layers = get_layers(self.base_model)
lora_forward_fn = lora_forward
for layer_idx in self.layer_indices:
for module_info in get_peft_modules(layers[layer_idx], peft_config):
module = module_info["module"]
module.forward = partial(
lora_forward_fn,
self=module,
lora_dropout_p=peft_config.lora_dropout,
scaling=peft_config.lora_alpha,
)
@property
def generation_config(self):
return self.base_model.generation_config
def generate_weights(self, ctx_txt: str):
# generate loras
ctx_emb = embed_texts(
[ctx_txt],
self.emb_model,
self.emb_tokenizer,
self.task_desc_format_fn,
self.pooling_fn,
self.device,
)
encoder_out = self.t2l_model.task_encoder(ctx_emb)
encoded_task_emb = encoder_out["encoded_task_emb"].detach()
lora_A, lora_B = dict(), dict()
lora_dict = dict()
for target_module in self.t2l_model.target_modules:
factorized_delta_w = self.t2l_model.get_delta_weights(
self.layer_indices,
target_module,
encoded_task_emb.expand(self.layer_indices.shape[0], -1),
factorized=True,
)
# lora_A[target_module]: [n_layers, r, d_in]
# lora_A[target_module]: [n_layers, d_out, r]
lora_A[target_module], lora_B[target_module] = factorized_delta_w
# convert to lora format used by lora_forward
# dict of {module:
# {A: [bs, n_layers, r, d_inim],
# B: [bs, n_layers, r, d_outim]}}
lora_dict[target_module] = dict(
A=lora_A[target_module].unsqueeze(0),
B=lora_B[target_module].transpose(-1, -2).unsqueeze(0),
)
return lora_dict
def generate(self, *args, **kwargs):
ctx_ids_full = kwargs["ctx_ids"]
ctx_txt = self.tokenizer.decode(
ctx_ids_full[0, len(self.prefix_tokens) :], skip_special_tokens=True
)
generated_loras = self.generate_weights(ctx_txt)
apply_lora_to_layers(
self.base_model,
self.layer_indices,
generated_loras,
n_qs=torch.tensor([1], device=self.device),
position_ids=None,
)
kwargs.pop("ctx_ids", None)
kwargs.pop("ctx_attn_mask", None)
kwargs.pop("n_ctx_chunks", None)
return self.base_model.generate(*args, **kwargs)
if __name__ == "__main__":
from transformers import AutoTokenizer
from ctx_to_lora.data.definitions import CTX_AFFIXES
from ctx_to_lora.data.processing import load_and_process_dataset
model_name = "google/gemma-2-2b-it"
tokenizer = AutoTokenizer.from_pretrained(model_name)
ds = load_and_process_dataset("pwc", split="train", num_proc=8)
ctx = ds[0]["context"]
inp = ds[1]["prompts"][0]
ctx_ids = tokenizer.apply_chat_template(
[{"role": "user", "content": ctx}], return_tensors="pt", return_dict=True
)
input_ids = tokenizer.apply_chat_template(
[{"role": "user", "content": ctx + "\n\n" + inp}],
return_tensors="pt",
return_dict=True,
)
ctx_ids = {k: v.to("cuda") for k, v in ctx_ids.items()}
input_ids = {k: v.to("cuda") for k, v in input_ids.items()}
prefix_tokens = CTX_AFFIXES[model_name]["prefix"]
prefix_tokens = torch.tensor(prefix_tokens, dtype=torch.long)
t2l_model = TextToLoRA(
model_name,
prefix_tokens,
device="cuda",
)
with torch.no_grad():
for _ in range(1):
outputs = t2l_model.generate(
**input_ids,
ctx_ids=ctx_ids["input_ids"],
max_new_tokens=256,
do_sample=False,
)
print(
f"Student response: {tokenizer.batch_decode(outputs, skip_special_tokens=False)}"
)

1256
src/ctx_to_lora/modeling/text_to_lora_impl.py
View file

File diff suppressed because it is too large Load diff

92
src/ctx_to_lora/tracker/tracker.py
View file

@ -1,85 +1,25 @@
"""Unified tracking interface combining timing and CUDA memory usage.
Primary API
-----------
add_tracker(bound_method, name)
Wraps a bound instance method so that each invocation records:
- wall-clock duration (seconds) in timer.TIMER_REGISTRY[name]
- CUDA peak memory increase (bytes) in cuda_memory_tracker.MEMORY_REGISTRY[name]
(only if CUDA + torch available; otherwise memory list may stay empty / absent)
print_tracker_stats(name=None)
Convenience printer that delegates to time + memory aggregate printers.
Design
------
We implement a single wrapper (instead of nesting the individual timer & memory
wrappers) to avoid multiple layers of indirection and to ensure the measured
CUDA memory footprint reflects only the original method's body (excluding the
separate timing wrapper's slight overhead). The wrapper is idempotent: repeated
calls to add_tracker on the same method are ignored.
This file depends on sibling modules:
- tracker.timer
- tracker.cuda_memory_tracker
Both registries remain the single source of truth; no additional registry is introduced.
"""
from __future__ import annotations
from collections.abc import Callable
from time import perf_counter
from typing import Any
# Support both package (relative) and direct script execution.
try: # Package / normal import path
from .cuda_memory_tracker import ( # type: ignore
MEMORY_REGISTRY,
compute_aggregate_memory_stats,
print_aggregate_memory_stats,
print_global_memory_stats,
reset_memory_trackers,
save_memory_stats_csv,
)
from .timer import ( # type: ignore
TIMER_REGISTRY,
compute_aggregate_timer_stats,
print_aggregate_timer_stats,
print_global_timer_stats,
reset_timers,
save_timer_stats_csv,
)
except Exception: # pragma: no cover - fallback when executed directly
import pathlib
import sys
_this_file = pathlib.Path(__file__).resolve()
# project root is two levels up from tracker/ (i.e., .../src)
_src_root = _this_file.parents[2]
if str(_src_root) not in sys.path:
sys.path.insert(0, str(_src_root))
try:
from ctx_to_lora.tracker.cuda_memory_tracker import ( # type: ignore
MEMORY_REGISTRY,
compute_aggregate_memory_stats,
print_aggregate_memory_stats,
print_global_memory_stats,
reset_memory_trackers,
save_memory_stats_csv,
)
from ctx_to_lora.tracker.timer import ( # type: ignore
TIMER_REGISTRY,
compute_aggregate_timer_stats,
print_aggregate_timer_stats,
print_global_timer_stats,
reset_timers,
save_timer_stats_csv,
)
except Exception as e: # If still failing, raise a clearer error.
raise ImportError(
f"Failed to import tracking dependencies; ensure project root on PYTHONPATH. Original: {e}"
)
from .cuda_memory_tracker import ( # type: ignore
MEMORY_REGISTRY,
compute_aggregate_memory_stats,
print_aggregate_memory_stats,
print_global_memory_stats,
reset_memory_trackers,
save_memory_stats_csv,
)
from .timer import ( # type: ignore
TIMER_REGISTRY,
compute_aggregate_timer_stats,
print_aggregate_timer_stats,
print_global_timer_stats,
reset_timers,
save_timer_stats_csv,
)
try: # Optional torch import (lazy fallback if unavailable)
import torch # type: ignore

458
src/ctx_to_lora/trainer.py
View file

@ -1,458 +0,0 @@
import logging
import torch
from torch import nn
from transformers import Trainer
from transformers.trainer_pt_utils import get_parameter_names
from transformers.trainer_utils import IntervalStrategy
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
logger = logging.getLogger()
def per_ctx_loss_ce(inputs, labels, loss):
# loss still has masked out elem (0 at labels=-100)
n_queries_per_ctx = inputs["n_queries"].tolist()
position_ids = inputs["position_ids"].squeeze(0)
# account only label positions
label_mask = labels.squeeze(0) != -100
label_pos_ids = label_mask * position_ids
label_pos_ids_diff = label_pos_ids.diff(
append=torch.tensor([0], device=position_ids.device)
)
# assumes the input starts with non-assistant tokens
start_label_pos = torch.where((label_pos_ids_diff > 0) * ~label_mask)[0]
end_label_pos = torch.where((label_pos_ids_diff < 0) * label_mask)[0]
label_seq_lens = end_label_pos - start_label_pos
# these stack and split can be optimized but let's keep it simple
# mean across tokens of each q
qa_losses = torch.stack(
[
loss[start : start + llen].mean()
for start, llen in zip(start_label_pos, label_seq_lens)
]
)
# mean across queries of each ctx
per_ctx_losses = [ql.mean() for ql in torch.split(qa_losses, n_queries_per_ctx)]
# per-ctx loss
loss = torch.stack(per_ctx_losses)
return loss
def per_ctx_loss_kl(inputs, labels, loss):
# loss is compact (label indices selected)
n_queries_per_ctx = inputs["n_queries"].tolist()
position_ids = inputs["position_ids"].squeeze(0)
# account only label positions
label_mask = labels.squeeze(0) != -100
label_pos_ids = label_mask * position_ids
label_pos_ids_diff = label_pos_ids.diff(
append=torch.tensor([0], device=position_ids.device)
)
# assumes the input starts with non-assistant tokens
start_label_pos = torch.where((label_pos_ids_diff > 0) * ~label_mask)[0]
end_label_pos = torch.where((label_pos_ids_diff < 0) * label_mask)[0]
label_seq_lens = end_label_pos - start_label_pos
# find equiv start indices in the already sliced loss vector
cu_label_seq_lens = torch.cumsum(label_seq_lens, dim=0)
start_indices = torch.cat(
(
torch.tensor([0], device=cu_label_seq_lens.device),
cu_label_seq_lens[:-1],
)
)
# these stack and split can be optimized but let's keep it simple
# mean across tokens of each q
qa_losses = torch.stack(
[loss[start:end].mean() for start, end in zip(start_indices, cu_label_seq_lens)]
)
# mean across queries of each ctx
per_ctx_losses = [ql.mean() for ql in torch.split(qa_losses, n_queries_per_ctx)]
# per-ctx loss
loss = torch.stack(per_ctx_losses)
return loss
class ModulatedModelTrainer(Trainer):
# modified from the base Trainer to support per-context average loss
def get_batch_samples(self, epoch_iterator, num_batches, device):
# only used with `use_per_ctx_average_loss=True`
batch_samples = []
num_items_in_batch = None
for _ in range(num_batches):
try:
batch_samples.append(next(epoch_iterator))
except StopIteration:
break
count_num_items_in_batch = (
len(batch_samples) > 0
and "labels" in batch_samples[0]
and "n_ctx_chunks" in batch_samples[0]
)
if count_num_items_in_batch:
num_items_in_batch = dict()
num_items_in_batch["ctx"] = torch.tensor(
sum([batch["n_ctx_chunks"].numel() for batch in batch_samples])
).to(device)
# should we avg over num chunks?
# num_items_in_batch["ctx"] = sum(
# [(batch["ctx_position_ids"] == 0).sum() for batch in batch_samples]
# )
num_items_in_batch["labels"] = sum(
[(batch["labels"].ne(-100)).sum() for batch in batch_samples]
).to(device)
if num_items_in_batch is not None:
if self.args.average_tokens_across_devices:
for k in num_items_in_batch:
num_items_in_batch[k] = self.accelerator.gather(
num_items_in_batch[k]
).sum()
if torch.is_tensor(num_items_in_batch):
num_items_in_batch = num_items_in_batch.to(device)
if self.args.n_gpu > 1 and num_items_in_batch.dim() == 0:
# In the DataParallel case, convert the scalar tensor into a 1-dim tensor
num_items_in_batch = num_items_in_batch.unsqueeze(0)
return batch_samples, num_items_in_batch
class DistillationTrainer(ModulatedModelTrainer):
def __init__(self, *args, **kwargs):
self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)
self.use_per_ctx_average_loss = kwargs.pop("use_per_ctx_average_loss", False)
super().__init__(*args, **kwargs)
def compute_loss(
self, model, inputs, return_outputs=False, num_items_in_batch=None
):
# NOTE: the loss output from this fn will be ***added***
# meaning that we should always scale the loss wrt `num_items_in_batch`
# (average over the number of items in the accumulated batch)
is_train = num_items_in_batch is not None
labels = inputs.pop("labels", None)
label_pos = torch.where(labels != -100)
outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True)
if "logprobs_vals" not in inputs:
return (torch.tensor(0.0), outputs) if return_outputs else torch.tensor(0.0)
target_logp = inputs.pop("logprobs_vals").squeeze(0)
indices = inputs.pop("logprobs_indices").squeeze(0)
assert label_pos[0].shape[0] == target_logp.shape[0], (
"Label positions and target log probabilities should have the same # tokens."
f"Got : {label_pos[0].shape[0]=} and {target_logp.shape[0]=}"
)
##### KL loss
outputs_logits = outputs.logits[label_pos[0], label_pos[1] - 1] # shift back 1
logq_full_denom = torch.logsumexp(outputs_logits, dim=-1, keepdim=True) # (N,1)
selected_logits = outputs_logits.gather(1, indices) # (N,K)
# log softmax at selected indices
logq_selected = selected_logits - logq_full_denom
p = target_logp.exp()
loss = -(p * logq_selected).sum(dim=-1)
# teacher_logp = torch.full_like(outputs_logits, -torch.inf)
# teacher_logp.scatter_(1, indices, target_logp)
# # reduction = "batchmean" if num_items_in_batch is None else "sum"
# p = teacher_logp.exp()
# logq = nn.functional.log_softmax(outputs_logits, dim=-1)
# loss = -torch.sum(p * logq, dim=-1)
if self.use_per_ctx_average_loss:
loss = per_ctx_loss_kl(inputs, labels, loss)
if is_train:
if self.use_per_ctx_average_loss:
loss = loss.sum() / num_items_in_batch["ctx"]
else:
loss = loss.sum() / num_items_in_batch["labels"]
else:
# eval
loss = loss.mean()
# if reduction == "batchmean":
# loss = loss.mean()
# elif reduction == "sum":
# # loss does not scale with grad acc
# # num_items_in_batch does
# # this works for both token-avg and ctx-avg
# # loss = loss.sum() / num_items_in_batch
# `num_items_in_batch` is # tokens if `args.use_ctx_average_loss=False``
# loss = loss.sum() / num_items_in_batch
#####
##### unpack gen lora dict and compute regularization loss
l1_norm = 0
n_modules = len(gen_loras)
for module, lora in gen_loras.items():
l1_norm += lora["A"].abs().sum(0).mean() + lora["B"].abs().sum(0).mean()
l1_norm /= n_modules
if is_train:
# during eval `num_items_in_batch` will be None
l1_norm /= num_items_in_batch["ctx"]
total_loss = loss + self.gen_lora_l1_reg_coef * l1_norm
#####
scaler = self.args.gradient_accumulation_steps if is_train else 1
if self.args.average_tokens_across_devices and is_train:
total_loss *= self.accelerator.num_processes
scaler *= self.accelerator.num_processes
# rough estimate of the losses (we only log the values from one step)
if (self.state.global_step == 1 and self.args.logging_first_step) or (
self.args.logging_strategy == IntervalStrategy.STEPS
and self.state.global_step % self.state.logging_steps == 0
):
# compensate `num_items_in_batch` division
self.log(
{
"kl_loss": loss.item() * scaler,
"gen_lora_l1_norm": l1_norm.item() * scaler,
}
)
return (total_loss, outputs) if return_outputs else total_loss
def causal_lm_ce_loss(
logits,
labels,
vocab_size: int,
num_items_in_batch: torch.Tensor | None = None,
ignore_index: int = -100,
shift_labels: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
# Upcast to float if we need to compute the loss to avoid potential precision issues
logits = logits.float()
if shift_labels is None:
# Shift so that tokens < n predict n
labels = nn.functional.pad(labels, (0, 1), value=ignore_index)
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
logits = logits.view(-1, vocab_size)
shift_labels = shift_labels.view(-1)
# Enable model parallelism
shift_labels = shift_labels.to(logits.device)
# loss = fixed_cross_entropy(
# logits, shift_labels, num_items_in_batch, ignore_index, **kwargs
# )
loss = nn.functional.cross_entropy(logits, shift_labels, reduction="none")
return loss
class CrossEntropyTrainer(ModulatedModelTrainer):
def __init__(self, *args, **kwargs):
self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)
self.use_per_ctx_average_loss = kwargs.pop("use_per_ctx_average_loss", False)
super().__init__(*args, **kwargs)
def compute_loss(
self, model, inputs, return_outputs=False, num_items_in_batch=None
):
"""
How the loss is computed by Trainer.
By default, all models return the loss in the first element.
Subclass and override for custom behavior.
"""
is_train = num_items_in_batch is not None
labels = inputs.pop("labels", None)
outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True)
# [1, tot_seq_len]
logits = outputs.logits
# [tot_seq_len]
loss = causal_lm_ce_loss(logits, labels, self.model.vocab_size)
if self.use_per_ctx_average_loss:
loss = per_ctx_loss_ce(inputs, labels, loss)
if is_train:
if self.use_per_ctx_average_loss:
loss = loss.sum() / num_items_in_batch["ctx"]
else:
loss = loss.sum() / num_items_in_batch["labels"]
else:
# eval
loss = loss.mean()
#####
# if is_train:
# if self.use_per_ctx_average_loss:
# loss_kwargs["num_items_in_batch"] = num_items_in_batch["ctx"]
# else:
# loss_kwargs["num_items_in_batch"] = num_items_in_batch["labels"]
# inputs = {**inputs, **loss_kwargs}
# outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True)
# # Save past state if it exists
# if self.args.past_index >= 0:
# self._past = outputs[self.args.past_index]
# if labels is not None:
# unwrapped_model = self.accelerator.unwrap_model(model)
# if _is_peft_model(unwrapped_model):
# model_name = unwrapped_model.base_model.model._get_name()
# else:
# model_name = unwrapped_model._get_name()
# # User-defined compute_loss function
# if self.compute_loss_func is not None:
# loss = self.compute_loss_func(
# outputs, labels, num_items_in_batch=num_items_in_batch["labels"]
# )
# elif model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
# loss = self.label_smoother(outputs, labels, shift_labels=True)
# else:
# loss = self.label_smoother(outputs, labels)
# else:
# if isinstance(outputs, dict) and "loss" not in outputs:
# raise ValueError(
# "The model did not return a loss from the inputs, "
# "only the following keys: "
# f"{','.join(outputs.keys())}. "
# "For reference, the inputs it received are "
# f"{','.join(inputs.keys())}."
# )
# # We don't use .loss here since the model may return tuples instead of ModelOutput.
# loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]
#####
##### unpack gen lora dict and compute regularization loss
l1_norm = 0
n_modules = len(gen_loras)
for module, lora in gen_loras.items():
l1_norm += lora["A"].abs().sum(0).mean() + lora["B"].abs().sum(0).mean()
l1_norm /= n_modules
if is_train:
# during eval `num_items_in_batch` will be None
l1_norm /= num_items_in_batch["ctx"]
total_loss = loss + self.gen_lora_l1_reg_coef * l1_norm
#####
scaler = self.args.gradient_accumulation_steps if is_train else 1
if self.args.average_tokens_across_devices and is_train:
total_loss *= self.accelerator.num_processes
scaler *= self.accelerator.num_processes
# rough estimate of the losses (we only log the values from one step)
if (self.state.global_step == 1 and self.args.logging_first_step) or (
self.args.logging_strategy == IntervalStrategy.STEPS
and self.state.global_step % self.state.logging_steps == 0
):
# compensate `num_items_in_batch` division
self.log(
{
"ce_loss": loss.item() * scaler,
"gen_lora_l1_norm": l1_norm.item() * scaler,
}
)
return (total_loss, outputs) if return_outputs else total_loss
def get_decay_parameter_names(model) -> list[str]:
"""
Get all parameter names that weight decay will be applied to.
This function filters out parameters in two ways:
1. By layer type (nn.Embedding)
2. By parameter name patterns (containing 'bias', 'layernorm', 'rmsnorm'
or 'latents_q' [perceiver's latent queries]).
"""
decay_parameters = get_parameter_names(
model,
[nn.Embedding, nn.LayerNorm],
["scaler", "bias", "layernorm", "rmsnorm", "latents_q"],
)
return decay_parameters
def train_model(
model,
training_args,
train_dataset=None,
val_dataset=None,
train_collator=None,
compute_metrics=None,
):
checkpoint = None
if training_args.resume_from_checkpoint is not None:
checkpoint = training_args.resume_from_checkpoint
logger.info(f"Resuming from the checkpoint: {checkpoint}")
trainer_kwargs = dict(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=val_dataset,
data_collator=train_collator,
compute_metrics=compute_metrics,
)
is_modulated_model = isinstance(model, ModulatedPretrainedModel)
trainer_cls = Trainer
if is_modulated_model:
logger.info("Training with modulated model.")
trainer_cls = CrossEntropyTrainer
trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.gen_lora_l1_reg_coef
trainer_kwargs["use_per_ctx_average_loss"] = (
training_args.use_per_ctx_average_loss
)
del training_args.gen_lora_l1_reg_coef
del training_args.use_per_ctx_average_loss
if training_args.use_kl_loss:
logger.info("Training with distillation loss. Using DistillationTrainer.")
trainer_cls = DistillationTrainer
del training_args.use_kl_loss
if training_args.auto_find_batch_size:
# set the batch size to some high number
# which will be lowered by the Trainer
training_args.per_device_train_batch_size = 128
trainer = trainer_cls(**trainer_kwargs)
# if getattr(trainer, "use_per_ctx_average_loss", False):
# trainer.get_batch_samples = trainer.get_batch_samples_ctx
# MONKEY PATCH: remove embedding layers from weight decay
trainer.get_decay_parameter_names = get_decay_parameter_names
# Trainer loads the best model after training
# is done when load_best_model_at_end=True
train_result = trainer.train(resume_from_checkpoint=checkpoint)
trainer.log_metrics("train", train_result.metrics)
trainer.save_model()
# TODO: add benchmark eval?
# clear_gpu()

31
src/ctx_to_lora/utils.py
View file

@ -38,10 +38,31 @@ def evaluating(*models):
model.train(training)
def get_decoder_model(model):
visited = set()
current = model
while id(current) not in visited:
visited.add(id(current))
if hasattr(current, "model") and getattr(current, "model") is not current:
current = current.model
continue
if hasattr(current, "text_model") and getattr(current, "text_model") is not current:
current = current.text_model
continue
if hasattr(current, "language_model") and getattr(current, "language_model") is not current:
current = current.language_model
continue
break
return current
def get_layers(model):
if hasattr(model, "model"):
return get_layers(model.model)
return model.layers
model = get_decoder_model(model)
if hasattr(model, "layers"):
return model.layers
if hasattr(model, "decoder") and hasattr(model.decoder, "layers"):
return model.decoder.layers
raise AttributeError(f"Could not find decoder layers for model type {type(model)!r}")
def get_num_layers(model):
@ -49,9 +70,7 @@ def get_num_layers(model):
def get_base_model(model):
if hasattr(model, "model"):
return get_base_model(model.model)
return model
return get_decoder_model(model)
def get_num_params(model):

419
train.py
View file

@ -1,419 +0,0 @@
import contextlib
import logging
import os
from copy import deepcopy
from functools import partial
import numpy as np
import torch
import wandb
from datasets import disable_caching
from peft import PeftModel
from transformers import (
AutoConfig,
set_seed,
)
from ctx_to_lora.configs import (
AggregatorArguments,
ArgumentParser,
CtxEncoderArguments,
CtxTrainingArguments,
DataArguments,
ExperimentSetup,
HypernetArguments,
LoRAArguments,
ModelArguments,
TrainingArguments,
)
from ctx_to_lora.data.collator import ( # train_packed_collator,; DefaultDataCollator,
flatten_if_not_packed,
)
from ctx_to_lora.data.processing import get_tokenized_dataset, pack
from ctx_to_lora.metrics import (
Evaluator,
compute_metrics,
compute_per_token_acc,
compute_perplexity,
compute_prefix_matching,
)
from ctx_to_lora.model_loading import (
check_is_vision_model,
get_lora_config,
get_model_and_tokenizer,
get_tokenizer,
)
from ctx_to_lora.modeling.hypernet import (
ModulatedPretrainedModel,
get_hypernet_config,
)
from ctx_to_lora.trainer import train_model
from ctx_to_lora.utils import (
compile_linear,
extract_cli_args,
get_run_name,
log_num_train_params,
save_yaml,
setup_logging,
validate_args,
)
logger = logging.getLogger()
LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0"))
def main():
############ Argument parsing
parser = ArgumentParser(
(
DataArguments,
CtxTrainingArguments,
ModelArguments,
LoRAArguments,
TrainingArguments,
HypernetArguments,
AggregatorArguments,
CtxEncoderArguments,
)
)
(
data_args,
ctx_args,
model_args,
lora_args,
training_args,
hypernet_args,
aggregator_args,
ctx_encoder_args,
) = parser.parse()
# there shouldn't be overlap between args
validate_args(
[
data_args,
ctx_args,
model_args,
lora_args,
training_args,
hypernet_args,
aggregator_args,
ctx_encoder_args,
]
)
assert ctx_args.use_sequence_packing, (
f"Please set use_sequence_packing=True in {ctx_args}. It's faster!"
)
set_seed(training_args.seed)
checkpoint_dir = training_args.resume_from_checkpoint
# should be the same across processes
# still possible to have a name crash though
# logging_dir is just "runs/DATE_TIME_HOSTNAME"
slurm_job_id = f"_{os.getenv('SLURM_JOB_ID')}" if os.getenv("SLURM_JOB_ID") else ""
run_name = (
get_run_name(seed_str=training_args.logging_dir.strip("runs/") + slurm_job_id)
if not checkpoint_dir
else checkpoint_dir.strip("/").split("/")[-2]
)
output_dir = f"train_outputs/runs/{run_name}"
setup_logging(output_dir, debug=os.getenv("DEBUG", False))
logger.debug(f"CMD: {' '.join(os.sys.argv)}")
cli_args = extract_cli_args(os.sys.argv)
save_yaml(cli_args, f"{output_dir}/cli_args.yaml")
if "config" in cli_args:
config_name = os.path.basename(cli_args["config"]).split(".yaml")[0]
os.environ["WANDB_TAGS"] = config_name
run_name = os.path.basename(output_dir)
training_args.run_name = run_name
training_args.output_dir = output_dir
training_args.logging_dir = output_dir
if (
training_args.lr_scheduler_type == "cosine_with_min_lr"
and training_args.lr_scheduler_kwargs is None
):
training_args.lr_scheduler_kwargs = {"min_lr": 1e-7}
args = {
**vars(deepcopy(data_args)),
**vars(deepcopy(ctx_args)),
**vars(deepcopy(model_args)),
**vars(deepcopy(lora_args)),
**vars(deepcopy(training_args)),
**vars(deepcopy(hypernet_args)),
**vars(deepcopy(aggregator_args)),
**vars(deepcopy(ctx_encoder_args)),
}
args["deepspeed_plugin"] = None
logger.debug(f"args: {args}")
save_yaml(args, f"{output_dir}/args.yaml")
############ Model setup
if not ctx_args.from_pretrained_checkpoint:
model_name = model_args.model_name_or_path
base_model, tokenizer = get_model_and_tokenizer(
**vars(model_args),
train=True,
requires_grad=False,
peft_config=get_lora_config(model_name, **vars(lora_args)),
)
ctx_name = ctx_encoder_args.ctx_encoder_model_name_or_path
if ctx_name is not None:
ctx_encoder_model_config = AutoConfig.from_pretrained(
ctx_name, trust_remote_code=True
)
if ("Llama" in ctx_name and "Vision" in ctx_name) or check_is_vision_model(
ctx_name
):
ctx_encoder_model_config = ctx_encoder_model_config.text_config
ctx_tokenizer = get_tokenizer(ctx_name, train=True)
else:
ctx_name = base_model.base_model.config.name_or_path
ctx_encoder_model_config = base_model.config
ctx_tokenizer = tokenizer
if ctx_args.exp_setup == ExperimentSetup.HYPERLORA:
logger.info("Using HyperLoRA")
if not ctx_args.from_pretrained_checkpoint:
hypernet_config = get_hypernet_config(
base_model,
ctx_encoder_model_config,
hypernet_args,
aggregator_args,
ctx_encoder_args,
)
if ctx_encoder_args.layer_idx is None:
ctx_encoder_args.layer_idx = (
ctx_encoder_model_config.num_hidden_layers // 4
)
logger.info(
f"Using the first {ctx_encoder_args.layer_idx} layers"
" as the context encoder"
)
ctx_name = ctx_encoder_args.ctx_encoder_model_name_or_path
if ctx_encoder_args.ctx_encoder_last_layer is None and (
ctx_name is not None and ctx_name != base_model.name_or_path
):
logger.info(
f"Setting ctx_encoder_last_layer to {base_model.name_or_path} max layers"
f":{base_model.config.num_hidden_layers}"
)
ctx_encoder_args.ctx_encoder_last_layer = (
base_model.config.num_hidden_layers
)
model = ModulatedPretrainedModel(
base_model, hypernet_config, ctx_encoder_args
)
else:
if checkpoint_dir:
ctx_args.from_pretrained_checkpoint = (
f"{checkpoint_dir}/pytorch_model.bin"
)
logger.info(
f"Loading from checkpoint: {ctx_args.from_pretrained_checkpoint}"
)
model = ModulatedPretrainedModel.from_state_dict(
torch.load(ctx_args.from_pretrained_checkpoint, weights_only=False),
train=True,
use_flash_attn=model_args.use_flash_attn,
)
tokenizer = get_tokenizer(model.base_model.config.name_or_path, train=True)
ctx_name = model.ctx_encoder_args.ctx_encoder_model_name_or_path
if ctx_name is None:
ctx_name = model.base_model.config.name_or_path
ctx_tokenizer = get_tokenizer(ctx_name, train=True)
training_args.gen_lora_l1_reg_coef = ctx_args.gen_lora_l1_reg_coef
training_args.use_kl_loss = ctx_args.use_kl_loss
training_args.use_per_ctx_average_loss = ctx_args.use_per_ctx_average_loss
if len([p for p in model.ctx_encoder.parameters() if p.requires_grad]):
raise ValueError("ctx_encoder contains trainable parameters")
if len([p for p in model.base_model.parameters() if p.requires_grad]):
raise ValueError("base model contains trainable parameters")
model.hypernet.compile(fullgraph=True, mode="max-autotune")
else:
# activate LoRA
base_model_config = AutoConfig.from_pretrained(
model_args.model_name_or_path, trust_remote_code=True
)
base_model_config.save_pretrained(output_dir)
logger.info("Using LoRA")
model.set_adapter("default")
model = torch.compile(model)
model.train()
logger.debug(model)
log_num_train_params(model)
############ Dataset setup
logger.info("Loading dataset...")
add_ctx_to_chat = not isinstance(model, ModulatedPretrainedModel)
ctx_model_max_len = model.ctx_encoder.config.max_position_embeddings
if ctx_args.max_ctx_len > 0:
ctx_model_max_len = ctx_args.max_ctx_len
if ctx_args.max_ctx_chunk_len <= 0:
# set default chunk size to max length of the ctx encoder
ctx_args.max_ctx_chunk_len = ctx_model_max_len
if ctx_args.num_chunk_probs is not None:
ctx_args.num_chunk_probs = {
int(k): float(v) for k, v in ctx_args.num_chunk_probs.items()
}
_get_tokenized_dataset = partial(
get_tokenized_dataset,
max_qas_len=ctx_args.max_qas_len,
max_qas_per_sample=ctx_args.max_qas_per_sample,
base_model_max_len=model.base_model.config.max_position_embeddings,
tokenizer=tokenizer,
ctx_model_max_len=ctx_model_max_len,
ctx_tokenizer=ctx_tokenizer,
add_ctx_to_chat=add_ctx_to_chat,
add_negative_prompt=ctx_args.add_negative_prompt,
max_ctx_chunk_len=ctx_args.max_ctx_chunk_len,
min_ctx_chunk_len=ctx_args.min_ctx_chunk_len,
num_chunk_probs=ctx_args.num_chunk_probs,
max_ctx_chunk_num=ctx_args.max_ctx_chunk_num,
use_kl_loss=ctx_args.use_kl_loss,
)
splits = ["train"]
if training_args.eval_strategy != "no":
splits.append("validation")
tokenized_ds = {split: {} for split in splits}
for split, ds_names in zip(
splits,
[data_args.train_ds_names, data_args.val_ds_names],
):
if not ds_names:
continue
ctx_mgr = (
training_args.main_process_first()
if split == "train"
else contextlib.nullcontext()
)
with ctx_mgr:
# process and tokenize on the main process
# then other replicas can just load the cached dataset
# we dont save cache for validation ds
for ds_name in ds_names:
ds = _get_tokenized_dataset(ds_name, split)
base_name = os.path.basename(ds_name)
if ds_name.startswith("self_gen/"):
ds_name = "self_gen/" + base_name
else:
ds_name = base_name
tokenized_ds[split][ds_name] = ds
train_ds = tokenized_ds["train"]
if data_args.max_train_samples_per_ds is not None:
for ds_name, ds in train_ds.items():
if data_args.max_train_samples_per_ds >= len(ds):
continue
train_ds[ds_name] = ds.take(data_args.max_train_samples_per_ds)
logging.info(f"train_ds: {train_ds}")
val_ds = dict()
if "validation" in tokenized_ds:
n_val_samples = data_args.max_val_samples_per_ds
for ds_name, ds in tokenized_ds["validation"].items():
if ds is None:
# take some samples from train_ds
ds = train_ds[ds_name].take(n_val_samples)
train_ds[ds_name] = train_ds[ds_name].skip(n_val_samples)
val_ds[ds_name] = ds
val_indices = np.random.permutation(len(ds))[:n_val_samples]
val_ds[ds_name] = val_ds[ds_name].select(val_indices)
with training_args.main_process_first():
train_ds = pack(
train_ds,
ctx_args.max_packed_inp_len,
ctx_args.max_packed_ctx_len,
max_packed_size=-1,
seed=training_args.seed,
num_proc=30,
)
logger.info("Setting per_device_train_batch_size to 1")
training_args.per_device_train_batch_size = 1
logger.info(f"train_ds: {train_ds}")
logger.info(f"val_ds: {val_ds}")
collator = flatten_if_not_packed
if isinstance(model, ModulatedPretrainedModel):
if isinstance(model.base_model, PeftModel):
base_model = model.base_model.base_model
else:
base_model = model.base_model
if ctx_name is not None:
logger.info("Compiling ctx_encoder_model")
ctx_base_model = model.ctx_encoder.base_model
compile_linear(ctx_base_model)
elif isinstance(model, PeftModel):
base_model = model.base_model
logger.info("Compiling base_model")
base_model.compile(fullgraph=True, mode="max-autotune")
if LOCAL_RANK == 0:
wandb.init(
project=os.getenv("WANDB_PROJECT"),
name=run_name,
group=run_name,
config=args,
tags=os.getenv("WANDB_TAGS").split(","),
notes=ctx_args.notes,
resume="allow",
)
else:
wandb.init(mode="disabled")
train_model(
model,
training_args,
train_ds,
val_ds,
collator,
compute_metrics=partial(
compute_metrics,
evaluator=Evaluator(
[compute_per_token_acc, compute_prefix_matching, compute_perplexity]
),
),
)
logger.info(f"Training run finished and saved to {output_dir}")
if __name__ == "__main__":
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
os.environ["TOKENIZERS_PARALLELISM"] = "true"
os.environ["WANDB_DIR"] = ".wandb/"
os.environ["WANDB_PROJECT"] = os.getenv("WANDB_PROJECT") or "ctx_to_lora"
os.environ["WANDB_WATCH"] = ""
os.environ["WANDB_CONSOLE"] = "off"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ["OMP_NUM_THREADS"] = "23"
torch._dynamo.config.capture_scalar_outputs = True
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
if os.getenv("DEBUG", False):
disable_caching()
main()

7823
uv.lock generated
View file

File diff suppressed because it is too large Load diff

122
watcher.py
View file

@ -1,122 +0,0 @@
import itertools
import os
import time
from argparse import Namespace
from glob import glob
import wandb
import yaml
from ctx_to_lora.eval_utils import run_eval
from ctx_to_lora.utils import clear_gpu
CP_PATTERN = "train_outputs/runs/*/checkpoint*/pytorch_model.bin"
def flatten(l):
return itertools.chain.from_iterable(l)
# handmade file watcher using glob
# not using watchdog because there are too many saved files
# but we want to just watch CP_PATTERN files
class Watcher:
def __init__(self, patterns):
self.patterns = patterns
self.files = self.get_files()
self.last_files = self.files
def get_files(self):
return set(flatten(glob(pattern) for pattern in self.patterns))
def watch(self):
self.files = self.get_files()
new_files = self.files - self.last_files
return sorted(list(new_files))
def update(self, file):
if file in self.last_files:
return
self.last_files.add(file)
print(f"Added {file} to evaluated files.")
def save_state(self):
with open("watcher_state.yaml", "w") as f:
yaml.dump({"last_files": self.last_files}, f)
def load_state(self):
if not os.path.exists("watcher_state.yaml"):
return
with open("watcher_state.yaml") as f:
state = yaml.safe_load(f)
self.last_files = state["last_files"]
if __name__ == "__main__":
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1"
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
os.environ["WANDB_PROJECT"] = "ctx_to_lora"
watcher = Watcher([CP_PATTERN])
watcher.load_state()
print("Watching for new files...")
while True:
time.sleep(10)
new_files = watcher.watch()
for file in new_files:
# workaround to prevent loading incomplete checkpoints
time.sleep(20)
if not os.path.exists(file):
# cp is delete before we can read it
continue
run_dir = file.split("/checkpoint")[0]
run_name = run_dir.split("/")[-1]
print(f"Evaluating {file}")
args = Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml")))
curstep = int(file.split("checkpoint-")[1].split("/")[0])
wandb_kwargs = {
"project": os.getenv("WANDB_PROJECT"),
"group": run_name,
"name": f"{run_name}-eval",
"id": f"{run_name}-eval",
"resume": "allow",
}
wandb.init(**wandb_kwargs)
# TODO: have to change this for bigger models
eval_batch_size = 8
eval_batch_size_gen = 8
metrics = {}
# try:
# # metrics = run_eval(
# # checkpoint_path=file,
# # eval_batch_size=eval_batch_size,
# # split="validation",
# # generative=False,
# # )
# except FileNotFoundError as e:
# print(f"Error evaluating {file}: {e}. The checkpoint might be deleted.")
# continue
try:
gen_metrics = run_eval(
checkpoint_path=file,
split="validation",
eval_batch_size=eval_batch_size_gen,
max_ctx_chunk_len=args.max_ctx_chunk_len,
generative=True,
)
except FileNotFoundError as e:
print(f"The checkpoint might be deleted. Error evaluating {file}: {e}.")
gen_metrics = {}
file = ""
metrics.update(gen_metrics)
for k in metrics:
wandb.log(metrics[k], step=curstep)
wandb.finish()
print(f"Logged metrics: {metrics}")
print("=" * 80)
clear_gpu()
watcher.update(file)
watcher.save_state()

31
webui/SELF_GEN_VIEWER.md
View file

@ -1,31 +0,0 @@
# Self-Gen Data Viewer
Thanks Claude.
Running the viewer
```bash
uv run self_gen_viewer.py
```
Then open your browser and go to: **http://localhost:5001**
## Usage
1. **Select a Model Folder**: Choose from the dropdown list (e.g., `google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0`)
2. **Select a Parquet File**: Once a folder is selected, available parquet files will appear
3. **Set Number of Samples**: Adjust the sample count (default: 100, max: 1000)
4. **Click "Load Data"**: View the visualized data with context and Q&A pairs
## Data Structure
The viewer expects data in the following structure:
```
data/raw_datasets/self_gen/
├── google/
│ └── gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/
│ └── fw_qa_v2/
│ └── *.parquet
└── mistralai/
└── Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_1.0/
└── *.parquet
```

170
webui/self_gen_viewer.py
View file

@ -1,170 +0,0 @@
import traceback
from pathlib import Path
from datasets import load_dataset
from flask import Flask, jsonify, render_template, request
from transformers import AutoTokenizer
app = Flask(__name__)
# Base path for self_gen data
BASE_DATA_PATH = Path(__file__).parent.parent / "data" / "raw_datasets" / "self_gen"
# Cache for tokenizers
tokenizer_cache = {}
def get_tokenizer(model_path):
"""Get or create tokenizer with caching"""
if model_path not in tokenizer_cache:
try:
tokenizer_cache[model_path] = AutoTokenizer.from_pretrained(model_path)
except Exception as e:
print(f"Error loading tokenizer for {model_path}: {e}")
return None
return tokenizer_cache[model_path]
def discover_folders():
"""Discover all model folders in self_gen directory"""
folders = []
if not BASE_DATA_PATH.exists():
return folders
for vendor_dir in BASE_DATA_PATH.iterdir():
if vendor_dir.is_dir():
for model_dir in vendor_dir.iterdir():
if model_dir.is_dir():
rel_path = model_dir.relative_to(BASE_DATA_PATH)
folders.append(str(rel_path))
return sorted(folders)
def discover_parquet_files(folder_path):
"""Discover all parquet files in a folder"""
full_path = BASE_DATA_PATH / folder_path
parquet_files = []
if full_path.exists():
for parquet_file in full_path.glob("**/*.parquet"):
rel_path = parquet_file.relative_to(full_path)
parquet_files.append(str(rel_path))
return sorted(parquet_files)
def extract_model_name_from_folder(folder_path):
"""Extract base model name from folder path"""
# e.g., "google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0" -> "google/gemma-2-2b-it"
parts = folder_path.split("/")
if len(parts) >= 2:
vendor = parts[0]
model_part = parts[1].split("_temp_")[0]
return f"{vendor}/{model_part}"
return None
@app.route("/")
def index():
"""Main page"""
folders = discover_folders()
return render_template("self_gen_viewer.html", folders=folders)
@app.route("/api/folders")
def api_folders():
"""API endpoint to get available folders"""
folders = discover_folders()
return jsonify({"folders": folders})
@app.route("/api/parquet_files")
def api_parquet_files():
"""API endpoint to get parquet files in a folder"""
folder = request.args.get("folder", "")
if not folder:
return jsonify({"error": "No folder specified"}), 400
files = discover_parquet_files(folder)
return jsonify({"files": files})
@app.route("/api/load_data")
def api_load_data():
"""API endpoint to load and display data from a parquet file"""
folder = request.args.get("folder", "")
parquet_file = request.args.get("file", "")
num_samples = int(request.args.get("num_samples", 100))
if not folder or not parquet_file:
return jsonify({"error": "Missing parameters"}), 400
try:
# Construct full path
full_path = BASE_DATA_PATH / folder / parquet_file
if not full_path.exists():
return jsonify({"error": f"File not found: {full_path}"}), 404
# Extract model name for tokenizer
model_name = extract_model_name_from_folder(folder)
if not model_name:
return jsonify({"error": "Could not extract model name from folder"}), 400
# Load tokenizer
tokenizer = get_tokenizer(model_name)
if tokenizer is None:
return jsonify({"error": f"Could not load tokenizer for {model_name}"}), 500
# Load dataset
ds = load_dataset(
"parquet", data_files=str(full_path), split=f"train[:{num_samples}]"
)
# Process samples
samples = []
for i, sample in enumerate(ds):
processed_sample = {
"index": i,
"ctx": tokenizer.decode(sample["ctx_ids"], skip_special_tokens=False)
if "ctx_ids" in sample
else "N/A",
"questions": [],
}
# Decode input_ids if present
if "input_ids" in sample:
if isinstance(sample["input_ids"][0], list):
# Multiple Q&A pairs
processed_sample["questions"] = [
tokenizer.decode(qa, skip_special_tokens=False)
for qa in sample["input_ids"]
]
else:
# Single item
processed_sample["questions"] = [
tokenizer.decode(sample["input_ids"], skip_special_tokens=False)
]
samples.append(processed_sample)
return jsonify(
{
"success": True,
"num_samples": len(samples),
"model_name": model_name,
"file_path": str(parquet_file),
"samples": samples,
}
)
except Exception as e:
traceback.print_exc()
return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500
if __name__ == "__main__":
print(f"Data path: {BASE_DATA_PATH}")
print(f"Available folders: {discover_folders()}")
app.run(debug=True, host="0.0.0.0", port=5001)

422
webui/templates/self_gen_viewer.html
View file

@ -1,422 +0,0 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Self-Gen Data Viewer</title>
<style>
* {
margin: 0;
padding: 0;
box-sizing: border-box;
}
body {
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
min-height: 100vh;
padding: 20px;
}
.container {
max-width: 1400px;
margin: 0 auto;
background: white;
border-radius: 15px;
box-shadow: 0 10px 40px rgba(0, 0, 0, 0.3);
padding: 30px;
}
h1 {
color: #667eea;
margin-bottom: 25px;
text-align: center;
font-size: 2.5em;
}
.controls {
display: flex;
flex-direction: column;
gap: 15px;
margin-bottom: 25px;
}
.control-group {
display: flex;
flex-direction: column;
}
.samples-row {
display: grid;
grid-template-columns: 1fr auto;
gap: 15px;
align-items: end;
}
label {
font-weight: 600;
margin-bottom: 5px;
color: #333;
font-size: 0.9em;
}
select,
input {
padding: 10px;
border: 2px solid #ddd;
border-radius: 8px;
font-size: 1em;
transition: border-color 0.3s;
}
select:focus,
input:focus {
outline: none;
border-color: #667eea;
}
button {
padding: 10px 20px;
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
color: white;
border: none;
border-radius: 8px;
cursor: pointer;
font-size: 1em;
font-weight: 600;
transition: transform 0.2s, box-shadow 0.2s;
margin-top: auto;
}
button:hover {
transform: translateY(-2px);
box-shadow: 0 5px 15px rgba(102, 126, 234, 0.4);
}
button:active {
transform: translateY(0);
}
button:disabled {
background: #ccc;
cursor: not-allowed;
transform: none;
}
.info-box {
background: #f8f9fa;
padding: 15px;
border-radius: 8px;
margin-bottom: 20px;
border-left: 4px solid #667eea;
}
.info-box p {
margin: 5px 0;
color: #555;
}
.loading {
text-align: center;
padding: 40px;
color: #667eea;
font-size: 1.2em;
}
.error {
background: #fee;
color: #c33;
padding: 15px;
border-radius: 8px;
margin: 20px 0;
border-left: 4px solid #c33;
}
.sample {
background: #f8f9fa;
border: 1px solid #e0e0e0;
border-radius: 10px;
padding: 20px;
margin-bottom: 20px;
transition: box-shadow 0.3s;
}
.sample:hover {
box-shadow: 0 5px 15px rgba(0, 0, 0, 0.1);
}
.sample-header {
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
color: white;
padding: 10px 15px;
border-radius: 6px;
margin-bottom: 15px;
font-weight: 600;
}
.section {
margin-bottom: 20px;
}
.section-title {
font-weight: 600;
color: #667eea;
margin-bottom: 10px;
font-size: 1.1em;
border-bottom: 2px solid #667eea;
padding-bottom: 5px;
}
.content {
background: white;
padding: 15px;
border-radius: 6px;
border: 1px solid #e0e0e0;
white-space: pre-wrap;
word-wrap: break-word;
font-family: 'Courier New', monospace;
font-size: 0.9em;
line-height: 1.6;
max-height: 400px;
overflow-y: auto;
}
.question-item {
background: #fff;
padding: 12px;
border-radius: 6px;
border: 1px solid #ddd;
margin-bottom: 10px;
}
.question-number {
background: #667eea;
color: white;
padding: 3px 8px;
border-radius: 4px;
font-size: 0.85em;
font-weight: 600;
display: inline-block;
margin-bottom: 8px;
}
#results {
margin-top: 30px;
}
.load-more {
text-align: center;
margin-top: 20px;
}
.stats {
display: grid;
grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
gap: 15px;
margin-bottom: 20px;
}
.stat-card {
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
color: white;
padding: 15px;
border-radius: 8px;
text-align: center;
}
.stat-value {
font-size: 2em;
font-weight: 700;
}
.stat-label {
font-size: 0.9em;
opacity: 0.9;
}
::-webkit-scrollbar {
width: 8px;
height: 8px;
}
::-webkit-scrollbar-track {
background: #f1f1f1;
border-radius: 4px;
}
::-webkit-scrollbar-thumb {
background: #667eea;
border-radius: 4px;
}
::-webkit-scrollbar-thumb:hover {
background: #764ba2;
}
</style>
</head>
<body>
<div class="container">
<h1>🔍 Self-Gen Data Viewer</h1>
<div class="controls">
<div class="control-group">
<label for="folder-select">Model Folder:</label>
<select id="folder-select">
<option value="">-- Select a folder --</option>
{% for folder in folders %}
<option value="{{ folder }}">{{ folder }}</option>
{% endfor %}
</select>
</div>
<div class="control-group">
<label for="file-select">Parquet File:</label>
<select id="file-select" disabled>
<option value="">-- Select a file --</option>
</select>
</div>
<div class="samples-row">
<div class="control-group">
<label for="num-samples">Number of Samples:</label>
<input type="number" id="num-samples" value="100" min="1" max="1000" step="10">
</div>
<button id="load-btn" disabled>Load Data</button>
</div>
</div>
<div id="info" style="display: none;"></div>
<div id="results"></div>
</div>
<script>
const folderSelect = document.getElementById('folder-select');
const fileSelect = document.getElementById('file-select');
const numSamplesInput = document.getElementById('num-samples');
const loadBtn = document.getElementById('load-btn');
const infoDiv = document.getElementById('info');
const resultsDiv = document.getElementById('results');
// When folder is selected, load parquet files
folderSelect.addEventListener('change', async () => {
const folder = folderSelect.value;
fileSelect.innerHTML = '<option value="">-- Select a file --</option>';
fileSelect.disabled = true;
loadBtn.disabled = true;
if (!folder) return;
try {
const response = await fetch(`/api/parquet_files?folder=${encodeURIComponent(folder)}`);
const data = await response.json();
if (data.files && data.files.length > 0) {
data.files.forEach(file => {
const option = document.createElement('option');
option.value = file;
option.textContent = file;
fileSelect.appendChild(option);
});
fileSelect.disabled = false;
} else {
alert('No parquet files found in this folder');
}
} catch (error) {
alert('Error loading files: ' + error.message);
}
});
// Enable load button when file is selected
fileSelect.addEventListener('change', () => {
loadBtn.disabled = !fileSelect.value;
});
// Load data button
loadBtn.addEventListener('click', loadData);
async function loadData() {
const folder = folderSelect.value;
const file = fileSelect.value;
const numSamples = numSamplesInput.value;
if (!folder || !file) return;
resultsDiv.innerHTML = '<div class="loading">⏳ Loading data...</div>';
infoDiv.style.display = 'none';
try {
const response = await fetch(
`/api/load_data?folder=${encodeURIComponent(folder)}&file=${encodeURIComponent(file)}&num_samples=${numSamples}`
);
const data = await response.json();
if (data.error) {
resultsDiv.innerHTML = `<div class="error"><strong>Error:</strong> ${data.error}</div>`;
return;
}
displayData(data);
} catch (error) {
resultsDiv.innerHTML = `<div class="error"><strong>Error:</strong> ${error.message}</div>`;
}
}
function displayData(data) {
// Show info box
infoDiv.style.display = 'block';
infoDiv.innerHTML = `
<div class="info-box">
<div class="stats">
<div class="stat-card">
<div class="stat-value">${data.num_samples}</div>
<div class="stat-label">Samples</div>
</div>
<div class="stat-card">
<div class="stat-value">${data.samples.length > 0 ? data.samples[0].questions.length : 0}</div>
<div class="stat-label">Questions per Sample</div>
</div>
</div>
<p><strong>Model:</strong> ${data.model_name}</p>
<p><strong>File:</strong> ${data.file_path}</p>
</div>
`;
// Show samples
let html = '';
data.samples.forEach(sample => {
html += `
<div class="sample">
<div class="sample-header">Sample #${sample.index}</div>
<div class="section">
<div class="section-title">📝 Context</div>
<div class="content">${escapeHtml(sample.ctx)}</div>
</div>
<div class="section">
<div class="section-title">❓ Questions & Answers (${sample.questions.length})</div>
${sample.questions.map((q, i) => `
<div class="question-item">
<span class="question-number">Q&A ${i + 1}</span>
<div class="content">${escapeHtml(q)}</div>
</div>
`).join('')}
</div>
</div>
`;
});
resultsDiv.innerHTML = html;
}
function escapeHtml(text) {
const div = document.createElement('div');
div.textContent = text;
return div.innerHTML;
}
</script>
</body>
</html>