From 3b798c670c8f5ef493704a56a345c87f5c4dbbe1 Mon Sep 17 00:00:00 2001
From: 51616 <rujikorn.ch@gmail.com>
Date: Wed, 1 Oct 2025 13:11:54 +0000
Subject: [PATCH] python api + remove `generate_multi_lora` + rename config
 path + test eval scripts

---
 README.md                                     |  93 +++++-----
 .../self_gen_lv1_closed_qa_1_and_lv3_l2l.yaml |  32 ++++
 .../self_gen_lv1_closed_qa_1_l2l.yaml         |  10 -
 .../ctx_magic_number_32_256.yaml              |   0
 data/sakana_wiki.txt                          |   8 +
 data/self_generate_qa.py                      |  70 ++-----
 examples/python_api.py                        |  43 +++++
 install.sh                                    |   8 +-
 scripts/main_exp/0-download_data.py           |  11 ++
 scripts/main_exp/{train.sh => 1-train.sh}     |   0
 scripts/main_exp/README.md                    |  25 +++
 scripts/main_exp/eval/base_model.sh           |   2 +-
 scripts/main_exp/eval/base_model_test.sh      |   8 +
 scripts/main_exp/eval/cd.sh                   |   4 +-
 scripts/main_exp/eval/cd_oracle.sh            |   4 +-
 scripts/main_exp/eval/cd_oracle_test.sh       |   6 +
 scripts/main_exp/eval/cd_test.sh              |   6 +
 scripts/main_exp/eval/d2l.sh                  |  10 +-
 scripts/main_exp/eval/d2l_test.sh             |  13 ++
 scripts/main_exp/eval/llmlingua_test.sh       |   5 +
 scripts/main_exp/eval/t2l_test.sh             |   4 +
 scripts/main_exp/gen_data.sh                  |   4 +-
 scripts/main_exp/gen_data_test.sh             |  12 +-
 src/ctx_to_lora/data/preprocessing_fn.py      |  24 +++
 src/ctx_to_lora/eval_utils.py                 |   3 -
 src/ctx_to_lora/modeling/hypernet.py          | 171 ++++++++----------
 webui/app.py                                  |   2 +-
 27 files changed, 345 insertions(+), 233 deletions(-)
 create mode 100644 configs/main_exp/self_gen_lv1_closed_qa_1_and_lv3_l2l.yaml
 rename configs/{toy_exp => niah_exp}/ctx_magic_number_32_256.yaml (100%)
 create mode 100644 data/sakana_wiki.txt
 create mode 100644 examples/python_api.py
 create mode 100644 scripts/main_exp/0-download_data.py
 rename scripts/main_exp/{train.sh => 1-train.sh} (100%)
 create mode 100644 scripts/main_exp/README.md
 create mode 100755 scripts/main_exp/eval/base_model_test.sh
 create mode 100755 scripts/main_exp/eval/cd_oracle_test.sh
 create mode 100755 scripts/main_exp/eval/cd_test.sh
 create mode 100755 scripts/main_exp/eval/d2l_test.sh
 create mode 100755 scripts/main_exp/eval/llmlingua_test.sh
 create mode 100755 scripts/main_exp/eval/t2l_test.sh

diff --git a/README.md b/README.md
index 56bc8a6..b57cd52 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,5 @@
 <div align="center">
-    <h1>Ctx-to-LoRA</h1>
+    <h1>Doc-to-LoRA</h1>
     <br>
     <img height="500px" src="assets/cover.png" />
 </div>
@@ -7,56 +7,57 @@
 
 ---
 
-## 🚀 API Usage [WIP]
+## 🚀 Python API Usage
 ```python
-from ctx_to_lora.modeling import ModulatedPretrainedModel
-model = ModulatedPretrainedModel.from_state_dict(...)
+# caveat: this interface only supports non-batched inputs
+# for batched inference please see `src/ctx_to_lora/modeling/hypernet.py`
+import torch
 
-ctx_info = "..."
-query = "..."
+from ctx_to_lora.model_loading import get_tokenizer
+from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
 
-ctx_ids = model.ctx_encoder.tokenize(ctx_info)
-input_ids = model.tokenize(query)
-outputs = model.generate(ctx_ids, input_ids)
-print(model.decode(outputs))
+# model loading
+checkpoint_path = ...
+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": "Summarize what Sakana AI does."}]
+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=256)
+print(tokenizer.decode(outputs[0]))
+
+
+# remove internalized info
+model.reset()
+
+outputs = model.generate(input_ids=chat_ids, max_new_tokens=256)
+print(tokenizer.decode(outputs[0]))
 ```
 
-
-***Generate data from scratch***
-# 0. download fineweb_edu to `data/raw_datasets/fineweb_edu
-uv run data/download_fineweb_edu.py
-```
-
-1. Recursively generate more data! (depends on step 0)
+### 🎮 Interactive Demo [WIP]
 ```bash
-# run from 000 to 0013
-run uv run data/generate_fw_edu_qa_v2.py --shard_pattern "000_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048;
-run uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
+uv run webui/demo.py
 ```
 
-2. Self-generated response QA data (depends on step 0 and 1)
-```bash
-# Example commands using gemma-2-2b-it
-# self-gen data for fw_qa_v2
-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/013*_level_3*' --closed_qa_prob 1.0 # or 0.0
-
-# 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 listed in qa_short_ctx_self_gen_no_fw_qa.yaml
-uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/qa_short_ctx_self_gen_no_fw_qa.yaml
-```
-
-### Evaluation
-See [eval scripts](scripts/eval/).
-```bash
-WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 4
-
-
-# base model
-WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --eval_batch_size 2
-
-# base model w/o context
-WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets negative_nq triviaqa_retrieved squad longbench_e --split test --remove_context
-
-
+### 🧪 Experimental Scripts
+| Experiment                           | Data prep                                                                                                                           | Training                                  | Evaluation                                                                                                                               | Notes                                                                                                                               |
+| ------------------------------------ | ----------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- |
+| [Main experiment](scripts/main_exp/) | `uv run bash scripts/main_exp/0-download_data.sh`<br>`uv run bash scripts/main_exp/gen_data.sh` *(optional, rebuilds from scratch)* | `uv run bash scripts/main_exp/1-train.sh` | `uv run bash scripts/main_exp/eval/<script>.sh` (pick from `base_model.sh`, `cd.sh`, `cd_oracle.sh`, `d2l.sh`, `llmlingua.sh`, `t2l.sh`) | Downloading data is fastest; regenerate only if you need fresh synthetic data. Evaluation scripts reproduce the main paper metrics. |
+| [NIAH](scripts/niah/)                | `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`                                                                                                     | Run the scripts in order; data generation only needs to happen once                                                                 |
diff --git a/configs/main_exp/self_gen_lv1_closed_qa_1_and_lv3_l2l.yaml b/configs/main_exp/self_gen_lv1_closed_qa_1_and_lv3_l2l.yaml
new file mode 100644
index 0000000..93aaad9
--- /dev/null
+++ b/configs/main_exp/self_gen_lv1_closed_qa_1_and_lv3_l2l.yaml
@@ -0,0 +1,32 @@
+# LoRA
+lora_r: 8
+lora_dropout: 0.0
+target_modules:
+  - down_proj
+
+use_kl_loss: true
+
+ctx_encoder_type: per_layer_activations
+n_latent_queries: 8
+num_blocks: 9
+num_self_attn_per_block: 0
+
+gradient_accumulation_steps: 11
+max_packed_inp_len: 6144
+max_packed_ctx_len: 6144
+
+# data
+train_ds_names:
+  - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
+  - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_3*.parquet
+  - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/pwc_compact
+  - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/squad_compact
+  - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ropes_compact
+  - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/drop_compact
+
+val_ds_names:
+  - squad
+  - pwc
+  - drop
+  - ropes
+  - 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
diff --git a/configs/main_exp/self_gen_lv1_closed_qa_1_l2l.yaml b/configs/main_exp/self_gen_lv1_closed_qa_1_l2l.yaml
index cf18365..04db6b2 100644
--- a/configs/main_exp/self_gen_lv1_closed_qa_1_l2l.yaml
+++ b/configs/main_exp/self_gen_lv1_closed_qa_1_l2l.yaml
@@ -15,9 +15,6 @@ 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
@@ -31,11 +28,4 @@ val_ds_names:
   - pwc
   - drop
   - ropes
-  # - hotpot_qa
-  # - negative_nq
-  # - booksum
-  # - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
   - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
-  # - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/ctx_qa/validation/*.parquet
-  # - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/facts/validation/*.parquet
-
diff --git a/configs/toy_exp/ctx_magic_number_32_256.yaml b/configs/niah_exp/ctx_magic_number_32_256.yaml
similarity index 100%
rename from configs/toy_exp/ctx_magic_number_32_256.yaml
rename to configs/niah_exp/ctx_magic_number_32_256.yaml
diff --git a/data/sakana_wiki.txt b/data/sakana_wiki.txt
new file mode 100644
index 0000000..694d451
--- /dev/null
+++ b/data/sakana_wiki.txt
@@ -0,0 +1,8 @@
+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] 
\ No newline at end of file
diff --git a/data/self_generate_qa.py b/data/self_generate_qa.py
index 2f42876..aef4b91 100644
--- a/data/self_generate_qa.py
+++ b/data/self_generate_qa.py
@@ -186,21 +186,10 @@ def self_generate(
             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."
             )
-        if "_summ_prob_" in ds_name and args.summ_prob != 0.0:
-            raise ValueError(
-                f"Multiple sources of truth for summ_prob: CLI arg --summ_prob={args.summ_prob} and dataset name contains summ_prob specification."
-            )
-        if "_struct_prob_" in ds_name and args.struct_prob != 0.0:
-            raise ValueError(
-                f"Multiple sources of truth for struct_prob: CLI arg --struct_prob={args.struct_prob} and dataset name contains struct_prob specification."
-            )
 
     # Base values from args
     temp = args.temp
     closed_qa_prob = args.closed_qa_prob
-    summ_prob = args.summ_prob
-    struct_prob = args.struct_prob
-    cot_prob = args.cot_prob
 
     # Overrides from ds_name pattern if present
     if ds_name is not None:
@@ -212,21 +201,10 @@ def self_generate(
             m = re.search(r"_closed_qa_prob_([\d.]+)", ds_name)
             if m:
                 closed_qa_prob = float(m.group(1))
-        if "_summ_prob_" in ds_name:
-            m = re.search(r"_summ_prob_([\d.]+)", ds_name)
-            if m:
-                summ_prob = float(m.group(1))
-        if "_struct_prob_" in ds_name:
-            m = re.search(r"_struct_prob_([\d.]+)", ds_name)
-            if m:
-                struct_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}")
-    print(f"Using chain-of-thought prompt probability: {cot_prob}")
-    print(f"Using summarization augmentation probability: {summ_prob}")
-    print(f"Using structuring augmentation probability: {struct_prob}")
 
     if parquet_file:
         print(f"Loading dataset from parquet file: {parquet_file}")
@@ -249,7 +227,7 @@ def self_generate(
     print(f"Loaded dataset: {ds_name} with split: {split}")
 
     if args.debug:
-        ds = ds.take(50)
+        ds = ds.take(10)
 
     ds = ds.filter(filter_none, batched=False, num_proc=8)
 
@@ -445,22 +423,6 @@ def execute_qa_generation(
             # bos + question + eos + start model turn + response + eos
             input_ids = all_ids[:sys_start] + all_ids[q_start:res_end]
 
-            # TODO: save prompt logprobs with qa_start
-            # TODO: also save the prompt tokens for output alignment during training
-            # prompt_logp = completions[c + i].prompt_logprobs
-            # q_len = res_start - q_start
-            # prompt_logp_indices = np.empty((q_len, k), dtype=np.int32)
-            # # float-16 is better for this range
-            # prompt_logp_vals = np.empty((q_len, k), dtype=np.float16)
-            # for li, info_d in enumerate(prompt_logp[q_start:res_start]):
-            #     for j, (idx, tok_info) in enumerate(info_d.items()):
-            #         # vllm returns logprob for gt token first,
-            #         # which means that we might have k + 1 logprobs
-            #         if j >= k:
-            #             continue
-            #         prompt_logp_indices[li, j] = idx
-            #         prompt_logp_vals[li, j] = tok_info.logprob
-
             # relative to the input_ids
             res_start = res_start - q_start + sys_start
             res_end = res_start + n_response_tokens
@@ -472,10 +434,6 @@ def execute_qa_generation(
             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)
-            # NOTE: also have to shift back 1 for training
-            # self_gen_data[ctx]["prompt_start_end"].append((sys_start, res_start))
-            # self_gen_data[ctx]["prompt_logprobs_vals"].append(prompt_logp_vals)
-            # self_gen_data[ctx]["prompt_logprobs_indices"].append(prompt_logp_indices)
 
         c += i + 1
 
@@ -516,20 +474,22 @@ def execute_qa_generation(
     # random.shuffle(samples)
 
     # Save results
-    if not args.debug:
-        # 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}"
+    # 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}"
 
-        os.makedirs(os.path.dirname(fpath), exist_ok=True)
-        fpath = f"{fpath}.parquet"
-        ds_out.to_parquet(fpath)
-        print(f"Saved to {fpath}")
+    if args.debug:
+        fpath += "_debug"
+    os.makedirs(os.path.dirname(fpath), exist_ok=True)
 
-        # Cleanup
-        del samples, ds_out, completions, messages, ctxs, questions
-        clear_gpu()
+    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:
diff --git a/examples/python_api.py b/examples/python_api.py
new file mode 100644
index 0000000..b81c3fa
--- /dev/null
+++ b/examples/python_api.py
@@ -0,0 +1,43 @@
+# 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
+
+checkpoint_path = "../../ctx-to-lora/train_outputs/runs/Sep29_14-42-46_slurm0-a3nodeset-9_88483_1e7bb34e/checkpoint-40000/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)
+
+doc = open("data/sakana_wiki.txt").read()
+chat = [{"role": "user", "content": "Summarize what Sakana AI does."}]
+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)
+
+
+outputs = model.generate(input_ids=chat_ids, max_new_tokens=256)
+print(tokenizer.decode(outputs[0]))
+
+
+# calls after internalization will be influenced by internalized info
+model.internalize(doc)
+
+outputs = model.generate(input_ids=chat_ids, max_new_tokens=256)
+print(tokenizer.decode(outputs[0]))
+
+
+# remove internalized info
+model.reset()
+
+outputs = model.generate(input_ids=chat_ids, max_new_tokens=256)
+print(tokenizer.decode(outputs[0]))
diff --git a/install.sh b/install.sh
index f8db2b1..48afaef 100755
--- a/install.sh
+++ b/install.sh
@@ -2,6 +2,7 @@
 # module load cuda/12.4
 # module load cudnn/9.1.0
 # module load hpcx/v2.21
+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
@@ -16,12 +17,11 @@ 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
+# optional: needed for gated models
 # uv run huggingface-cli login
 
-# needed for logging with wandb
+# optional: needed for logging with wandb
 # wandb login
 
-# dev
+# optional: dev
 # uv run pre-commit install
diff --git a/scripts/main_exp/0-download_data.py b/scripts/main_exp/0-download_data.py
new file mode 100644
index 0000000..5ecfd36
--- /dev/null
+++ b/scripts/main_exp/0-download_data.py
@@ -0,0 +1,11 @@
+# TO BE ADDED: see which split is better
+
+from huggingface_hub import snapshot_download
+
+if __name__ == "__main__":
+    fw_dir = "./data/raw_datasets/self_gen/"
+    snapshot_download(
+        "SakanaAI/self_gen_qa_fw_qa",
+        repo_type="dataset",
+        local_dir=fw_dir,
+    )
diff --git a/scripts/main_exp/train.sh b/scripts/main_exp/1-train.sh
similarity index 100%
rename from scripts/main_exp/train.sh
rename to scripts/main_exp/1-train.sh
diff --git a/scripts/main_exp/README.md b/scripts/main_exp/README.md
new file mode 100644
index 0000000..72418a6
--- /dev/null
+++ b/scripts/main_exp/README.md
@@ -0,0 +1,25 @@
+# D2L pipeline
+### Data
+You can either download the generated data (recommended) or generate them by youself.
+```bash
+# download training data (recommended)
+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.
+```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.
diff --git a/scripts/main_exp/eval/base_model.sh b/scripts/main_exp/eval/base_model.sh
index 481b6ee..dc4ceef 100755
--- a/scripts/main_exp/eval/base_model.sh
+++ b/scripts/main_exp/eval/base_model.sh
@@ -2,7 +2,7 @@
 WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1
 
 # w/ truncation
-WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --truncate_if_too_long_inp
+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
diff --git a/scripts/main_exp/eval/base_model_test.sh b/scripts/main_exp/eval/base_model_test.sh
new file mode 100755
index 0000000..32b4147
--- /dev/null
+++ b/scripts/main_exp/eval/base_model_test.sh
@@ -0,0 +1,8 @@
+# no truncation
+WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --max_test_samples_per_ds 10
+
+# w/ truncation
+WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen 1 --truncate_if_too_long_inp --max_test_samples_per_ds 10
+
+# 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 --max_test_samples_per_ds 10
diff --git a/scripts/main_exp/eval/cd.sh b/scripts/main_exp/eval/cd.sh
index e27fedb..a3d4ed2 100755
--- a/scripts/main_exp/eval/cd.sh
+++ b/scripts/main_exp/eval/cd.sh
@@ -1,6 +1,6 @@
 # qa
-uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=4
+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
-uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=1
+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/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=1
diff --git a/scripts/main_exp/eval/cd_oracle.sh b/scripts/main_exp/eval/cd_oracle.sh
index d8ae957..22da250 100755
--- a/scripts/main_exp/eval/cd_oracle.sh
+++ b/scripts/main_exp/eval/cd_oracle.sh
@@ -1,6 +1,6 @@
 # qa
-uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp
+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
-uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp
+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
diff --git a/scripts/main_exp/eval/cd_oracle_test.sh b/scripts/main_exp/eval/cd_oracle_test.sh
new file mode 100755
index 0000000..87cc34b
--- /dev/null
+++ b/scripts/main_exp/eval/cd_oracle_test.sh
@@ -0,0 +1,6 @@
+# qa
+uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --max_test_samples_per_ds 10
+
+
+# longbench
+uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --max_test_samples_per_ds 10
diff --git a/scripts/main_exp/eval/cd_test.sh b/scripts/main_exp/eval/cd_test.sh
new file mode 100755
index 0000000..419ad54
--- /dev/null
+++ b/scripts/main_exp/eval/cd_test.sh
@@ -0,0 +1,6 @@
+# qa
+uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=4 --max_test_samples_per_ds 10
+
+
+# longbench
+uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=1 --max_test_samples_per_ds 10
diff --git a/scripts/main_exp/eval/d2l.sh b/scripts/main_exp/eval/d2l.sh
index 14c35e0..aec7921 100755
--- a/scripts/main_exp/eval/d2l.sh
+++ b/scripts/main_exp/eval/d2l.sh
@@ -1,13 +1,13 @@
 # main results
 # batched
-WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1
+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/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
+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/pytorch_model.bin --datasets squad --split test --eval_batch_size_gen=1 --flip_ctx_inp
+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/pytorch_model.bin --datasets squad_assistant_ctx_no_passage --split test
-WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/pytorch_model.bin --datasets squad_negative_no_passage --split test
+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
diff --git a/scripts/main_exp/eval/d2l_test.sh b/scripts/main_exp/eval/d2l_test.sh
new file mode 100755
index 0000000..c44daed
--- /dev/null
+++ b/scripts/main_exp/eval/d2l_test.sh
@@ -0,0 +1,13 @@
+# main results
+# batched
+WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ../../ctx-to-lora/train_outputs/runs/Sep29_14-42-46_slurm0-a3nodeset-9_88483_1e7bb34e/checkpoint-40000/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 --max_test_samples_per_ds 10
+
+# iterative
+WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ../../ctx-to-lora/train_outputs/runs/Sep29_14-42-46_slurm0-a3nodeset-9_88483_1e7bb34e/checkpoint-40000/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 --max_test_samples_per_ds 10
+
+# query internalization
+WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ../../ctx-to-lora/train_outputs/runs/Sep29_14-42-46_slurm0-a3nodeset-9_88483_1e7bb34e/checkpoint-40000/pytorch_model.bin --datasets squad --split test --eval_batch_size_gen=1 --flip_ctx_inp --max_test_samples_per_ds 10
+
+# replaced squad context
+WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path ../../ctx-to-lora/train_outputs/runs/Sep29_14-42-46_slurm0-a3nodeset-9_88483_1e7bb34e/checkpoint-40000/pytorch_model.bin --datasets squad_assistant_ctx_no_passage --split test --max_test_samples_per_ds 10
+WANDB_MODE=disabled uv run python run_eval.py --checkpoint_path ../../ctx-to-lora/train_outputs/runs/Sep29_14-42-46_slurm0-a3nodeset-9_88483_1e7bb34e/checkpoint-40000/pytorch_model.bin --datasets squad_negative_no_passage --split test --max_test_samples_per_ds 10
diff --git a/scripts/main_exp/eval/llmlingua_test.sh b/scripts/main_exp/eval/llmlingua_test.sh
new file mode 100755
index 0000000..1e1d0ec
--- /dev/null
+++ b/scripts/main_exp/eval/llmlingua_test.sh
@@ -0,0 +1,5 @@
+for dataset in squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e; do
+  for rate in 0.9 0.8 0.6 0.4 0.2 0.1; do
+    WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets "$dataset" --split test --eval_batch_size_gen=1 --use_llmlingua --llmlingua_compression_rate "$rate" --truncate_if_too_long_ctx --max_test_samples_per_ds 10
+  done
+done
diff --git a/scripts/main_exp/eval/t2l_test.sh b/scripts/main_exp/eval/t2l_test.sh
new file mode 100755
index 0000000..eb93b88
--- /dev/null
+++ b/scripts/main_exp/eval/t2l_test.sh
@@ -0,0 +1,4 @@
+# download t2l checkpoint
+uv run huggingface-cli download SakanaAI/text-to-lora --local-dir . --include "trained_t2l/gemma_2b_t2l"
+
+WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e --split test --eval_batch_size_gen=1 --use_t2l --max_test_samples_per_ds 10
diff --git a/scripts/main_exp/gen_data.sh b/scripts/main_exp/gen_data.sh
index 8eced94..c1c0a26 100755
--- a/scripts/main_exp/gen_data.sh
+++ b/scripts/main_exp/gen_data.sh
@@ -4,8 +4,8 @@ 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;
+  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
diff --git a/scripts/main_exp/gen_data_test.sh b/scripts/main_exp/gen_data_test.sh
index f2adb4f..9f080d5 100755
--- a/scripts/main_exp/gen_data_test.sh
+++ b/scripts/main_exp/gen_data_test.sh
@@ -4,17 +4,17 @@ uv run data/download_fineweb_edu.py
 # generate qa data
 # run from 000 to 013
 for shard_id in $(seq -f "%03g" 0 1); do
-  uv run data/generate_fw_edu_qa_v2.py --shard_pattern "${shard_id}_00000" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it --max_length=2000 --max_model_length=2048;
-  uv run data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/${shard_id}*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
+  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 --debug
+  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 --debug
 
   # 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
+  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 --debug
 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'
+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' --debug
 
 # self-gen data for other ds 
-uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names squad_compact ropes_compact drop_compact --closed_qa_prob 1.0
-uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names pwc_compact --closed_qa_prob 0.0
+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 --debug
+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 --debug
diff --git a/src/ctx_to_lora/data/preprocessing_fn.py b/src/ctx_to_lora/data/preprocessing_fn.py
index 78f6e85..c0ed117 100644
--- a/src/ctx_to_lora/data/preprocessing_fn.py
+++ b/src/ctx_to_lora/data/preprocessing_fn.py
@@ -4,6 +4,7 @@ from collections.abc import Callable
 from typing import Any
 
 from ctx_to_lora.data.definitions import CLOSED_QA_INTX_TEMPLATES, EVAL_INTX_TEMPLATES
+from ctx_to_lora.utils import concat_list
 
 logger = logging.getLogger()
 
@@ -42,6 +43,29 @@ def get_preprocessing_fn(
         # already processed data, do nothing
         return f
 
+    if "fw_qa_v2" in ds_name:
+
+        def f(sample):
+            # get questions/answers from all levels in the ds
+            q_cols = [col for col in sample.keys() if col.startswith("prompts_level")]
+            r_cols = [col for col in sample.keys() if col.startswith("responses_level")]
+            questions = concat_list([sample[col] for col in q_cols])
+            responses = concat_list([sample[col] for col in r_cols])
+            min_len = min(len(questions), len(responses))
+
+            if min_len == 0:
+                return {
+                    "context": None,
+                    "prompts": None,
+                    "responses": None,
+                }
+
+            return {
+                "context": sample["context"],
+                "prompts": questions[:min_len],
+                "responses": responses[:min_len],
+            }
+
     elif ds_name.startswith("longbench"):
 
         def f(sample):
diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py
index eb95220..60b4198 100644
--- a/src/ctx_to_lora/eval_utils.py
+++ b/src/ctx_to_lora/eval_utils.py
@@ -941,9 +941,6 @@ def evaluate(
 
     collator = generation_collator if generative else eval_collator
 
-    if max_ctx_chunk_len > 0:
-        model.generate = model.generate_with_multi_loras
-
     # if isinstance(model, CtxDistillModel):
 
     #     model.generate = partial(
diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py
index 390bcfd..5f01433 100644
--- a/src/ctx_to_lora/modeling/hypernet.py
+++ b/src/ctx_to_lora/modeling/hypernet.py
@@ -31,8 +31,10 @@ from ctx_to_lora.configs import (
     CtxEncoderArguments,
     HypernetArguments,
 )
+from ctx_to_lora.data.processing import tokenize_ctx_text
 from ctx_to_lora.model_loading import (
     get_model,
+    get_tokenizer,
 )
 from ctx_to_lora.modeling.aggregator import (
     AGGREGATOR_CLS,
@@ -458,7 +460,7 @@ class ModulatedPretrainedModel(nn.Module):
         self.user_defined_scaling = user_defined_scaling
         self.inp_compressor = inp_compressor
         self.model_accepts_loss_kwargs = True
-        self.active_adapters = []
+        self.generated_loras = None
 
         self.register_module("base_model", base_model)
         self._init_model()
@@ -496,6 +498,28 @@ class ModulatedPretrainedModel(nn.Module):
         model.load_state_dict(state_dict)
         return model
 
+    def patch_lora_forward(self):
+        layers = get_layers(self.base_model)
+
+        lora_forward_fn = (
+            lora_forward_packed if self.use_sequence_packing else lora_forward
+        )
+        for layer_idx in self.hypernet.layer_indices:
+            for module_info in get_peft_modules(layers[layer_idx], self.peft_config):
+                name = module_info["name"]
+                module = module_info["module"]
+                if getattr(module, "patched_forward", False):
+                    continue
+                logger.debug(f"Applying LoRA forward to {name}")
+                module.forward_orig = module.forward
+                module.patched_forward = True
+                module.forward = partial(
+                    lora_forward_fn,
+                    self=module,
+                    lora_dropout_p=self.peft_config.lora_dropout,
+                    scaling=self.peft_config.lora_alpha,
+                )
+
     def _init_model(self):
         # disable adapter of the base model
         # this only works with LoRA(?)
@@ -506,21 +530,7 @@ class ModulatedPretrainedModel(nn.Module):
             HyperLoRA(self.hypernet_config).to(self.device).to(torch.float32)
         )
 
-        layers = get_layers(self.base_model)
-        lora_forward_fn = (
-            lora_forward_packed if self.use_sequence_packing else lora_forward
-        )
-        for layer_idx in self.hypernet.layer_indices:
-            for module_info in get_peft_modules(layers[layer_idx], self.peft_config):
-                name = module_info["name"]
-                module = module_info["module"]
-                logger.debug(f"Applying LoRA forward to {name}")
-                module.forward = partial(
-                    lora_forward_fn,
-                    self=module,
-                    lora_dropout_p=self.peft_config.lora_dropout,
-                    scaling=self.peft_config.lora_alpha,
-                )
+        self.patch_lora_forward()
 
         ctx_model_name = self.ctx_encoder_args.ctx_encoder_model_name_or_path
         if ctx_model_name is None:
@@ -762,58 +772,6 @@ class ModulatedPretrainedModel(nn.Module):
         else:
             return model_outputs
 
-    @torch.inference_mode()
-    def generate_with_multi_loras(
-        self,
-        ctx_ids: Integer[Tensor, "bs ctx_length"],
-        ctx_attn_mask: Integer[Tensor, "bs ctx_length"] | None = None,
-        ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None,
-        n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None,
-        n_queries: Integer[Tensor, "n_ctx"] | None = None,
-        scalers: Float[Tensor, "n_ctx"] | None = None,
-        bias_scaler: float | None = None,
-        *model_inputs_args: Any,
-        **model_inputs_kwargs: dict[str, Any],
-    ):
-        generated_loras, _ = self.generate_weights(
-            ctx_ids, ctx_attn_mask, ctx_position_ids
-        )
-
-        generated_loras = combine_lora(
-            generated_loras,
-            n_ctx_chunks,
-            lora_bias=self.hypernet.get_head_bias()
-            if self.hypernet.config.use_bias
-            else None,
-            scalers=scalers,
-            bias_scaler=bias_scaler,
-        )
-
-        # apply lora hook to the base model
-        position_ids = (
-            model_inputs_kwargs["position_ids"]
-            if "position_ids" in model_inputs_kwargs
-            else None
-        )
-        if n_queries is None:
-            # assumes each group has one q
-            n_queries = torch.ones(
-                len(n_ctx_chunks), dtype=torch.int32, device=self.device
-            )
-
-        apply_lora_to_layers(
-            self.base_model,
-            self.hypernet.layer_indices,
-            generated_loras,
-            n_queries,
-            position_ids,  # TODO: remove (cant be used with generation?)
-        )
-
-        model_outputs = self.base_model.generate(
-            *model_inputs_args, **model_inputs_kwargs
-        )
-        return model_outputs
-
     def combine_lora(self, *args, **kwargs):
         # for timing
         return combine_lora(*args, **kwargs)
@@ -822,6 +780,38 @@ class ModulatedPretrainedModel(nn.Module):
         # for timing
         return apply_lora_to_layers(*args, **kwargs)
 
+    # for simple api usage
+    def internalize(self, ctx_str: str):
+        ctx_tokenizer = get_tokenizer(self.ctx_encoder.base_model.name_or_path)
+        ctx_ids = tokenize_ctx_text(dict(context=[ctx_str]), ctx_tokenizer)["ctx_ids"]
+        return self._internalize_from_ids(torch.tensor(ctx_ids, device=self.device))
+
+    def _internalize_from_ids(
+        self,
+        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,
+    ):
+        self.patch_lora_forward()
+        if ctx_attn_mask is None and ctx_position_ids is None:
+            assert ctx_ids.shape[0] == 1
+            ctx_attn_mask = torch.ones_like(ctx_ids)
+        generated_loras, generated_layernorms = self.generate_weights(
+            ctx_ids, ctx_attn_mask, ctx_position_ids
+        )
+        self.generated_loras = generated_loras
+
+    def reset(self):
+        self.generated_loras = None
+        layers = get_layers(self.base_model)
+        for layer_idx in self.hypernet.layer_indices:
+            for module_info in get_peft_modules(layers[layer_idx], self.peft_config):
+                name = module_info["name"]
+                module = module_info["module"]
+                logger.debug(f"Resetting forward for {name}")
+                module.forward = module.forward_orig
+                module.patched_forward = False
+
     @torch.inference_mode()
     def generate(
         self,
@@ -837,23 +827,22 @@ class ModulatedPretrainedModel(nn.Module):
         generated_layernorms = None
         if (
             ctx_ids is None
-            and not self.active_adapters
+            and not self.generated_loras
             and not self.use_base_input_as_ctx
         ):
-            logger.warning(
-                (
-                    "*" * 100,
-                    "\n\nNo ctx_ids provided, using the base model for generation\n\n",
-                    "*" * 100,
-                )
+            print(
+                "*" * 100
+                + "\n\nNo ctx_ids provided, using the base model for generation\n\n"
+                + "*" * 100
             )
-        if ctx_ids is None and self.active_adapters:
-            logger.info(
-                (
-                    "*" * 100,
-                    "\n\nUsing active LoRAs for generation\n\n",
-                    "*" * 100,
-                )
+        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)
+            print(
+                "*" * 100
+                + "\n\nUsing internalized LoRAs for generation\n\n"
+                + "*" * 100
             )
         else:
             if self.use_base_input_as_ctx:
@@ -895,7 +884,9 @@ class ModulatedPretrainedModel(nn.Module):
             if n_queries is None:
                 if ctx_position_ids is None:
                     n_queries = torch.ones(
-                        ctx_ids.shape[0], dtype=torch.int32, device=self.device
+                        model_inputs_kwargs["input_ids"].shape[0],
+                        dtype=torch.int32,
+                        device=self.device,
                     )
                 else:
                     # quite redundant (we do cu_seqlens many places)
@@ -913,19 +904,7 @@ class ModulatedPretrainedModel(nn.Module):
                 n_queries,
                 position_ids,
             )
-        if self.inp_compressor is not None:
-            print("Using inp compressor")
-            input_ids = model_inputs_kwargs["input_ids"]
-            print(f"input_ids before: {input_ids.shape}")
-            model_inputs_kwargs["input_ids"] = self.inp_compressor.compress_tokens(
-                input_ids,
-                query_text="Answer the following question. Only output the answer and do not output any other words.",
-            ).to(self.base_model.device)
-            # batch size always 1
-            model_inputs_kwargs["attention_mask"] = torch.ones_like(
-                model_inputs_kwargs["input_ids"]
-            )
-            print(f"input_ids after: {model_inputs_kwargs['input_ids'].shape}")
+
         model_outputs = self.base_model.generate(
             *model_inputs_args, **model_inputs_kwargs
         )
diff --git a/webui/app.py b/webui/app.py
index e674771..aa9c19e 100644
--- a/webui/app.py
+++ b/webui/app.py
@@ -675,7 +675,7 @@ def chat():
                     #         do_sample=False,
                     #     )
 
-                    outputs = modulated_model.generate_with_multi_loras(
+                    outputs = modulated_model.generate(
                         ctx_ids=ctx_ids,
                         ctx_attn_mask=ctx_attn_mask,
                         n_ctx_chunks=torch.tensor(