From 610750e6fb05f85de1d95e816867a66dc6e8f861 Mon Sep 17 00:00:00 2001
From: Rujikorn Charakorn <rujikorn.ch@gmail.com>
Date: Wed, 30 Jul 2025 18:38:51 +0900
Subject: [PATCH] refactor after distillation (#6)

---
 README.md                                     |   54 +-
 configs/context_numbers_10.yaml               |    2 -
 configs/context_numbers_10_self_gen.yaml      |    2 -
 .../qa_short_ctx_self_gen_lv3_tiny.yaml       |   10 +-
 data/self_generate_qa.py                      |   13 +-
 intx_sft.py                                   |  105 +-
 run_eval.py                                   |   20 +-
 .../short_ctx/gemma_qa_short_ctx_old_conf.sh  |    8 +-
 .../gemma_qa_short_ctx_old_conf_distill.sh    |    6 +-
 ...gemma_qa_short_ctx_old_conf_l2l_distill.sh |   42 +
 src/ctx_to_lora/configs.py                    |   16 +-
 src/ctx_to_lora/data/collator.py              |  136 +-
 src/ctx_to_lora/data/packing.py               |  107 +-
 src/ctx_to_lora/data/preprocessing_fn.py      |  280 ++++
 src/ctx_to_lora/data/processing.py            | 1222 ++---------------
 src/ctx_to_lora/eval_utils.py                 |    7 +-
 src/ctx_to_lora/modeling/hypernet.py          |  186 +--
 src/ctx_to_lora/modeling/lora_layer.py        |    1 -
 src/ctx_to_lora/trainer.py                    |  122 +-
 webui/app.py                                  |    3 +-
 20 files changed, 694 insertions(+), 1648 deletions(-)
 create mode 100644 scripts/short_ctx/gemma_qa_short_ctx_old_conf_l2l_distill.sh
 create mode 100644 src/ctx_to_lora/data/preprocessing_fn.py

diff --git a/README.md b/README.md
index 2f0dc8e..a47a1a3 100644
--- a/README.md
+++ b/README.md
@@ -31,12 +31,12 @@ print(model.decode(outputs))
 ## 🏋️ Training
 ### 🔢 HyperLoRA w/ context_numbers_10
 ```bash
-WANDB_MODE=disabled run uv run intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=1 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --load_best_model_at_end=False --add_negative_prompt=False --add_repeat_prompt=False --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208
+WANDB_MODE=disabled run uv run intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=1 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=4096 --max_packed_ctx_len=4096 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208
 ```
 
 KL loss
 ```bash
- WANDB_MODE=disabled run uv run intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=1 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --use_light_weight_lora=False --load_best_model_at_end=False --add_negative_prompt=False --add_repeat_prompt=False --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.0 --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=True --eval_steps=99999999 --use_liger_kernel=False
+WANDB_MODE=disabled run uv run intx_sft.py configs/context_numbers_10_self_gen.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=3 --per_device_train_batch_size=128 --gradient_accumulation_steps=2 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=8 --num_self_attn_per_block=0 --num_pre_head_layers=1 --lora_r=8 --eval_steps=100 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 --neftune_noise_alpha=5 --per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 --use_sequence_packing=True --max_packed_inp_len=2048 --max_packed_ctx_len=2048 --dataloader_num_workers=0 --dataloader_prefetch_factor=None --eval_on_start=False --ctx_encoder_type=early_exit --n_latent_queries=208 --use_kl_loss=True
  ```
 ### FWQA-v2 Level-0 Tiny
 ```bash
@@ -76,6 +76,11 @@ uv run python data/generate_fav_num.py
 uv run python data/generate_fav_num_big.py
 ```
 
+Self-gen for the number toy dataset
+```bash
+run uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names context_numbers_2_10 --split train
+```
+
 ### SQuAD
 ```bash
 # for some reason download directly through `load_dataset` does not work
@@ -145,7 +150,10 @@ uv run data/download_fineweb_edu.py
 1. Recursively generate more data! (depends on step 0)
 ```bash
 # run from 000 to 0013
-run uv run python 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 python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it; run uv run python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_1*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it; run uv run python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_2*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it
+run uv run python data/generate_fw_edu_qa_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 python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_0*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
+run uv run python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_1*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it;
+run uv run python data/generate_fw_edu_qa_v2_repeat.py --shard_pattern "min_0_to_2000/000*level_2*" --n_qa_pairs=5 --vllm_model=google/gemma-3-12b-it
 ```
 
 2. Self-generated response QA data (depends on step 0 and 1)
@@ -157,47 +165,11 @@ uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_
 # val split
 uv run python 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_no_fw.yaml
-uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/self_gen_qa_short_ctx_no_fw_qa.yaml
+# self-gen data for other ds listed in qa_short_ctx_self_gen_no_fw_qa.yaml
+uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/qa_short_ctx_self_gen_no_fw_qa.yaml
 ```
 
 
-<!-- ### Data Generation
-```bash
-# [WIP] there are other datasets where we used openAI to gen QA pairs
-
-# 0. download fineweb_edu to `data/raw_datasets/fineweb_edu
-uv run data/download_fineweb_edu.py
-
-# ===
-# 1 and 2 can be run in parallel (depends on step 0)
-
-# 1. generate QA data (3 QAs for each context)
-# 000_0000[0-1] to 000_0000[8-9] for small split
-vllm_model=mistralai/Mistral-Small-3.1-24B-Instruct-2503 uv run python generate_fw_edu_qa_vllm.py "00*_*" 3
-
-# 2. Augmenting pre-training data w/ paraphrasing
-# (000_00000 to 000_00009 for small split)
-vllm_model=mistralai/Mistral-Small-3.1-24B-Instruct-2503 uv run python data/generate_fw_edu_augment_vllm.py 000_00000 
-
-# ===
-# 3 and 4 can be run in parallel (depends on step 0 and 1)
-
-# 3. Pre-training data (no GPU needed) 
-# augment each context by concat'ing QAs generated in step 1 at the end
-uv run python data/generate_pretrain_from_fw_qa.py
-
-# 4. self-generate SFT data
-# each context has 3 chat instances as training data
-# the responses are not the same as the ones in the pretraining though
-# Example commands using gemma-2-2b-it
-# self-gen data for fw_qa
-# 000_0000[0-1] to 000_0000[8-9] for small split
-uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_3/00*_*'
-# self-gen data for other ds listed in qa_no_fw.yaml
-uv run python data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --config configs/qa_no_fw.yaml
-``` -->
-
 <!-- 
 ### GSM8k LoRA
 ```bash
diff --git a/configs/context_numbers_10.yaml b/configs/context_numbers_10.yaml
index 8337c97..4e79f45 100644
--- a/configs/context_numbers_10.yaml
+++ b/configs/context_numbers_10.yaml
@@ -2,8 +2,6 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
-add_repeat_prompt: false
-add_negative_prompt: false
 # eval_on_start: True
 # eval_strategy: "steps"
 # eval_steps: 500
diff --git a/configs/context_numbers_10_self_gen.yaml b/configs/context_numbers_10_self_gen.yaml
index 566ad4b..c1ba4ca 100644
--- a/configs/context_numbers_10_self_gen.yaml
+++ b/configs/context_numbers_10_self_gen.yaml
@@ -2,8 +2,6 @@ output_dir: "" # just a placeholder
 bf16: true
 model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
 label_names: ["labels"]
-add_repeat_prompt: false
-add_negative_prompt: false
 # eval_on_start: True
 # eval_strategy: "steps"
 # eval_steps: 500
diff --git a/configs/tiny_exp/qa_short_ctx_self_gen_lv3_tiny.yaml b/configs/tiny_exp/qa_short_ctx_self_gen_lv3_tiny.yaml
index 0b4e48a..e56526d 100644
--- a/configs/tiny_exp/qa_short_ctx_self_gen_lv3_tiny.yaml
+++ b/configs/tiny_exp/qa_short_ctx_self_gen_lv3_tiny.yaml
@@ -39,11 +39,11 @@ target_modules:
 
 # data
 train_ds_names:
-  - self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_3_tiny
-  - self_gen/google/gemma-2-2b-it_temp_0.0/squad_compact
-  - self_gen/google/gemma-2-2b-it_temp_0.0/pwc_compact
-  - self_gen/google/gemma-2-2b-it_temp_0.0/ropes_compact
-  - self_gen/google/gemma-2-2b-it_temp_0.0/drop_compact
+  - self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/ds_013*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
diff --git a/data/self_generate_qa.py b/data/self_generate_qa.py
index bae8205..0b6d437 100644
--- a/data/self_generate_qa.py
+++ b/data/self_generate_qa.py
@@ -239,23 +239,16 @@ def self_generate(
 
         shard_name = "_" + os.path.basename(parquet_file).replace(".parquet", "")
         ds = load_dataset(path="parquet", data_files=[parquet_file], split="train")
-        processing_fn = get_preprocessing_fn(ds_name, is_eval=False, is_pretrain=False)
+        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,
-            add_negative_prompt=False,
-            add_repeat_prompt=False,
-            repeat_prob=0,
-            is_pretrain=False,
-        )
+        kwargs = dict(ds_name=ds_name, split=split)
 
-        ds = load_and_process_dataset(**kwargs, streaming=False, num_proc=8)
+        ds = load_and_process_dataset(**kwargs, num_proc=8)
     print(f"Loaded dataset: {ds_name} with split: {split}")
 
     if args.debug:
diff --git a/intx_sft.py b/intx_sft.py
index 728551d..e49597b 100755
--- a/intx_sft.py
+++ b/intx_sft.py
@@ -32,7 +32,6 @@ from ctx_to_lora.configs import (
 )
 from ctx_to_lora.data.collator import (  # train_packed_collator,; DefaultDataCollator,
     flatten_if_not_packed,
-    train_collator,
 )
 from ctx_to_lora.data.processing import get_tokenized_dataset, pack
 from ctx_to_lora.metrics import (
@@ -44,7 +43,6 @@ from ctx_to_lora.metrics import (
 )
 from ctx_to_lora.model_loading import (
     get_lora_config,
-    get_model,
     get_model_and_tokenizer,
     get_tokenizer,
 )
@@ -125,6 +123,10 @@ def main():
         ]
     )
 
+    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
     if checkpoint_dir and not os.path.isdir(checkpoint_dir):
@@ -218,8 +220,6 @@ def main():
                 base_model,
                 hypernet_config,
                 ctx_encoder_args,
-                ctx_args.use_kl_loss,
-                ctx_args.use_sequence_packing,
             )
 
         else:
@@ -230,7 +230,6 @@ def main():
                 torch.load(ctx_args.from_pretrained_checkpoint, weights_only=False),
                 train=True,
                 use_flash_attn=model_args.use_flash_attn,
-                use_sequence_packing=ctx_args.use_sequence_packing,
             )
             tokenizer = get_tokenizer(model.base_model.config.name_or_path, train=True)
             ctx_name = model.ctx_encoder_args.ctx_encoder_model_name_or_path
@@ -267,12 +266,6 @@ def main():
     logger.info("Loading dataset...")
 
     add_ctx_to_chat = not isinstance(model, ModulatedPretrainedModel)
-    tokenizer_kwargs = {"max_length": ctx_args.max_base_len}  # not used
-    ctx_tokenizer_kwargs = {"max_length": ctx_args.max_ctx_len}  # not used for now
-
-    base_model = get_model(
-        **vars(model_args), train=False, requires_grad=False, peft_config=None
-    )
 
     _get_tokenized_dataset = partial(
         get_tokenized_dataset,
@@ -280,18 +273,10 @@ def main():
         max_qas_per_sample=ctx_args.max_qas_per_sample,
         base_model_max_len=model.base_model.config.max_position_embeddings,
         tokenizer=tokenizer,
-        tokenizer_kwargs=tokenizer_kwargs,
         ctx_model_max_len=model.ctx_encoder.config.max_position_embeddings,
         ctx_tokenizer=ctx_tokenizer,
-        ctx_tokenizer_kwargs=ctx_tokenizer_kwargs,
         add_ctx_to_chat=add_ctx_to_chat,
-        add_repeat_prompt=ctx_args.add_repeat_prompt,
-        repeat_prob=ctx_args.repeat_prob,
-        add_negative_prompt=ctx_args.add_negative_prompt,
         use_kl_loss=ctx_args.use_kl_loss,
-        set_format=None if ctx_args.use_sequence_packing else "pt",
-        base_model=base_model,
-        # streaming=data_args.streaming,
     )
     tokenized_ds = {"train": dict(), "validation": dict()}
     for split, ds_names in zip(
@@ -300,7 +285,6 @@ def main():
     ):
         if not ds_names:
             continue
-        streaming = (split == "train") and data_args.streaming
         ctx_mgr = (
             training_args.main_process_first()
             if split == "train"
@@ -311,7 +295,7 @@ def main():
             # 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, streaming=streaming)
+                ds = _get_tokenized_dataset(ds_name, split)
 
                 base_name = os.path.basename(ds_name)
                 if ds_name.startswith("self_gen/"):
@@ -337,41 +321,8 @@ def main():
             val_indices = np.random.permutation(len(ds))[:n_val_samples]
             val_ds[ds_name] = val_ds[ds_name].select(val_indices)
 
-    # if data_args.streaming:
-    # max_steps = ceil(
-    #     sum(DS_LEN[ds] for ds in train_ds)
-    #     * training_args.num_train_epochs
-    #     / training_args.per_device_train_batch_size
-    #     / training_args.gradient_accumulation_steps
-    #     / training_args.world_size
-    # )
-    # training_args.max_steps = max_steps
-
-    # # interleaving streaming datasets
-    # # simplify the probs for smaller datasets
-    # # slightly upsample those datasets
-    # probs = [
-    #     0.01 if "fw_qa" not in ds_name else 1 + 0.01 - 0.01 * len(train_ds)
-    #     for ds_name in train_ds
-    # ]
-    # train_ds = interleave_datasets(
-    #     list(train_ds.values()),
-    #     probabilities=probs,
-    #     stopping_strategy="all_exhausted",
-    #     seed=training_args.seed,
-    # )
-    # else:
-
     train_ds_len = [len(ds) for ds in train_ds.values()]
     total_len = sum(train_ds_len)
-    # max_steps = ceil(
-    #     total_len
-    #     * training_args.num_train_epochs
-    #     / training_args.per_device_train_batch_size
-    #     / training_args.gradient_accumulation_steps
-    #     / training_args.world_size
-    # )
-    # training_args.max_steps = max_steps
     train_ds = interleave_datasets(
         list(train_ds.values()),
         probabilities=get_ds_prob(train_ds_len, total_len),
@@ -380,35 +331,27 @@ def main():
     )
     logger.info(f"Train dataset length: {len(train_ds)}")
 
-    if ctx_args.use_sequence_packing:
-        with training_args.main_process_first():
-            logging.info("Packing dataset")
-            old_ds_len = len(train_ds)
-            train_ds = pack(
-                train_ds,
-                ctx_args.max_packed_inp_len,
-                ctx_args.max_packed_ctx_len,
-                max_packed_size=-1,
-                num_proc=8,
-            )
-            logger.info(f"Train dataset length: {len(train_ds)}")
-            logger.info(
-                f"Avg. # of samples per packed sequence: {old_ds_len / len(train_ds)}"
-            )
-            logger.info("Setting per_device_train_batch_size to 1")
-            training_args.per_device_train_batch_size = 1
+    with training_args.main_process_first():
+        logging.info("Packing dataset")
+        old_ds_len = len(train_ds)
+        train_ds = pack(
+            train_ds,
+            ctx_args.max_packed_inp_len,
+            ctx_args.max_packed_ctx_len,
+            max_packed_size=-1,
+            num_proc=8,
+        )
+        logger.info(f"Train dataset length: {len(train_ds)}")
+        logger.info(
+            f"Avg. # of samples per packed sequence: {old_ds_len / len(train_ds)}"
+        )
+        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
-        # DefaultDataCollator(return_tensors="pt")
-        if ctx_args.use_sequence_packing
-        else partial(train_collator, tokenizer=tokenizer)
-    )
-
-    # TODO: use SFTTrainer instead? https://huggingface.co/docs/trl/en/sft_trainer
+    collator = flatten_if_not_packed
 
     # HACK [local patch]: deepspeed model loading problem (for resume training)
     # see https://github.com/microsoft/DeepSpeed/pull/6626/files
@@ -469,12 +412,10 @@ if __name__ == "__main__":
     os.environ["TOKENIZERS_PARALLELISM"] = "true"
     os.environ["WANDB_DIR"] = ".wandb/"
     os.environ["WANDB_PROJECT"] = "ctx_to_lora"
-    os.environ["WANDB_WATCH"] = ""  # "all"
+    os.environ["WANDB_WATCH"] = ""
     os.environ["WANDB_CONSOLE"] = "off"
     os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-    # os.environ["KMP_AFFINITY"] = "disabled"  # fixing iterable dataset stuck
     os.environ["OMP_NUM_THREADS"] = "23"
-    # os.environ["HF_DATASETS_IN_MEMORY_MAX_SIZE"] = "137438953472"  # 128 TB
     torch._dynamo.config.capture_scalar_outputs = True
 
     # DOES THIS CAUSE ACCURACY DEGRADATION?
diff --git a/run_eval.py b/run_eval.py
index 6581124..19cb6f8 100644
--- a/run_eval.py
+++ b/run_eval.py
@@ -62,16 +62,16 @@ if __name__ == "__main__":
 
     eval_batch_size_gen = cli_args.pop("eval_batch_size_gen")
     eval_batch_size = cli_args.pop("eval_batch_size")
-    # run_eval(
-    #     **cli_args,
-    #     # cli_args.checkpoint_path,
-    #     # cli_args.model_name_or_path,
-    #     # cli_args.eval_batch_size,
-    #     # args,
-    #     # split=cli_args.split,
-    #     eval_batch_size=eval_batch_size,
-    #     generative=False,
-    # )
+    run_eval(
+        **cli_args,
+        # cli_args.checkpoint_path,
+        # cli_args.model_name_or_path,
+        # cli_args.eval_batch_size,
+        # args,
+        # split=cli_args.split,
+        eval_batch_size=eval_batch_size,
+        generative=False,
+    )
     run_eval(
         **cli_args,
         # cli_args.checkpoint_path,
diff --git a/scripts/short_ctx/gemma_qa_short_ctx_old_conf.sh b/scripts/short_ctx/gemma_qa_short_ctx_old_conf.sh
index 7a42174..34012a3 100644
--- a/scripts/short_ctx/gemma_qa_short_ctx_old_conf.sh
+++ b/scripts/short_ctx/gemma_qa_short_ctx_old_conf.sh
@@ -6,14 +6,14 @@
 #SBATCH --output=slurm_logs/%x-%j.out
 #SBATCH --error=slurm_logs/%x-%j.out
 
-port=$((10000 + ($SLURM_JOBID % 50000)))
+port=$((29560 + ($SLURM_JOBID % 10)))
 echo "Using port: $port"
 
 # --gradient_accumulation_steps=8 --gradient_clipping=1.0
 uv run accelerate launch --main_process_port $port \
 --num_processes=4 --gpu_ids all intx_sft.py $1 \
 --model_name_or_path=google/gemma-2-2b-it \
---num_train_epochs=5 \
+--num_train_epochs=3 \
 --per_device_train_batch_size=-1 \
 --gradient_accumulation_steps=8 \
 --per_device_eval_batch_size=64 \
@@ -29,13 +29,11 @@ uv run accelerate launch --main_process_port $port \
 --learning_rate=4e-5 \
 --lora_dropout=0.0 \
 --neftune_noise_alpha=5 \
---add_negative_prompt=False \
---add_repeat_prompt=False \
 --use_sequence_packing=True \
 --max_qas_len=2048 \
 --max_qas_per_sample=1 \
 --max_packed_inp_len=8192 \
---max_packed_ctx_len=16384 \
+--max_packed_ctx_len=8192 \
 --per_rank_gen=True \
 --per_layer_processing=True \
 --gen_lora_l1_reg_coef=0.1 \
diff --git a/scripts/short_ctx/gemma_qa_short_ctx_old_conf_distill.sh b/scripts/short_ctx/gemma_qa_short_ctx_old_conf_distill.sh
index 7a75b34..0698a50 100644
--- a/scripts/short_ctx/gemma_qa_short_ctx_old_conf_distill.sh
+++ b/scripts/short_ctx/gemma_qa_short_ctx_old_conf_distill.sh
@@ -6,14 +6,14 @@
 #SBATCH --output=slurm_logs/%x-%j.out
 #SBATCH --error=slurm_logs/%x-%j.out
 
-port=$((10000 + ($SLURM_JOBID % 50000)))
+port=$((29560 + ($SLURM_JOBID % 10)))
 echo "Using port: $port"
 
 # --gradient_accumulation_steps=8 --gradient_clipping=1.0
 uv run accelerate launch --main_process_port $port \
 --num_processes=4 --gpu_ids all intx_sft.py $1 \
 --model_name_or_path=google/gemma-2-2b-it \
---num_train_epochs=5 \
+--num_train_epochs=3 \
 --per_device_train_batch_size=-1 \
 --gradient_accumulation_steps=16 \
 --per_device_eval_batch_size=64 \
@@ -29,8 +29,6 @@ uv run accelerate launch --main_process_port $port \
 --learning_rate=4e-5 \
 --lora_dropout=0.0 \
 --neftune_noise_alpha=5 \
---add_negative_prompt=False \
---add_repeat_prompt=False \
 --use_sequence_packing=True \
 --max_qas_len=2048 \
 --max_qas_per_sample=1 \
diff --git a/scripts/short_ctx/gemma_qa_short_ctx_old_conf_l2l_distill.sh b/scripts/short_ctx/gemma_qa_short_ctx_old_conf_l2l_distill.sh
new file mode 100644
index 0000000..d293c6a
--- /dev/null
+++ b/scripts/short_ctx/gemma_qa_short_ctx_old_conf_l2l_distill.sh
@@ -0,0 +1,42 @@
+#!/bin/bash
+#SBATCH --job-name=ctxlora
+#SBATCH --nodes=1
+#SBATCH --partition=a3
+#SBATCH --gpus=4
+#SBATCH --output=slurm_logs/%x-%j.out
+#SBATCH --error=slurm_logs/%x-%j.out
+
+port=$((29560 + ($SLURM_JOBID % 10)))
+echo "Using port: $port"
+
+# --gradient_accumulation_steps=8 --gradient_clipping=1.0
+uv run accelerate launch --main_process_port $port \
+--num_processes=4 --gpu_ids all intx_sft.py $1 \
+--model_name_or_path=google/gemma-2-2b-it \
+--num_train_epochs=3 \
+--per_device_train_batch_size=-1 \
+--gradient_accumulation_steps=16 \
+--per_device_eval_batch_size=64 \
+--target_modules=down_proj \
+--num_blocks=9 \
+--num_self_attn_per_block=0 \
+--ctx_encoder_type=per_layer_activations \
+--n_latent_queries=8 \
+--num_pre_head_layers=1 \
+--lora_r=8 \
+--eval_strategy=no \
+--eval_steps=1000 \
+--save_steps=1000 \
+--learning_rate=4e-5 \
+--lora_dropout=0.0 \
+--neftune_noise_alpha=5 \
+--use_sequence_packing=True \
+--max_qas_len=2048 \
+--max_qas_per_sample=1 \
+--max_packed_inp_len=4096 \
+--max_packed_ctx_len=4096 \
+--per_rank_gen=True \
+--per_layer_processing=True \
+--gen_lora_l1_reg_coef=0.1 \
+--logging_steps=50 \
+--use_kl_loss=True
diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py
index ac28238..e7fafa7 100644
--- a/src/ctx_to_lora/configs.py
+++ b/src/ctx_to_lora/configs.py
@@ -284,7 +284,7 @@ class LoRAArguments:
         metadata={"help": ("LoRA R value.")},
     )
     lora_dropout: float | None = field(
-        default=0.02,
+        default=0.0,
         metadata={"help": ("LoRA dropout.")},
     )
     target_modules: list[str] | None = field(
@@ -355,20 +355,6 @@ class CtxTrainingArguments:
         default=None,
         metadata={"help": "Wandb notes for the experiment."},
     )
-    add_repeat_prompt: bool = field(
-        default=False,
-        metadata={"help": "Whether to add repeat prompt to the dataset."},
-    )
-    repeat_prob: float = field(
-        default=0.0,
-        metadata={
-            "help": "Probability that a context will be used with repeat prompts."
-        },
-    )
-    add_negative_prompt: bool = field(
-        default=False,
-        metadata={"help": "Whether to add negative prompt to the dataset."},
-    )
     use_kl_loss: bool = field(
         default=False,
         metadata={"help": "Whether to use KL loss."},
diff --git a/src/ctx_to_lora/data/collator.py b/src/ctx_to_lora/data/collator.py
index 68ecee8..2b77c22 100644
--- a/src/ctx_to_lora/data/collator.py
+++ b/src/ctx_to_lora/data/collator.py
@@ -5,7 +5,7 @@ from transformers.data import (
     default_data_collator,
 )
 
-from ctx_to_lora.utils import check_is_iterable, concat_list
+from ctx_to_lora.utils import check_is_iterable
 
 flattener = DataCollatorWithFlattening()
 
@@ -14,6 +14,7 @@ def flatten_if_not_packed(inp_list):
     # no padding
     sample = inp_list[0]
     n = len(inp_list)
+    # training data is packed
     if "position_ids" in sample:
         if n == 1:
             n_queries = sample.pop("n_queries")
@@ -27,16 +28,10 @@ def flatten_if_not_packed(inp_list):
 
     # for eval data (not packed) during training
     need_flatten = check_is_iterable(sample["input_ids"][0])
-    if need_flatten:
-        n_queries = torch.tensor([len(x["input_ids"]) for x in inp_list])
-        tokens = dict()
-        for k in ["input_ids", "attention_mask", "labels"]:
-            tokens[k] = concat_list([x[k] for x in inp_list])
+    assert not need_flatten, f"Validation data should not be nested."
 
-        packed_inputs = flattener(tokens, return_tensors="pt")
-    else:
-        n_queries = torch.ones(len(inp_list), dtype=torch.int32)
-        packed_inputs = flattener(inp_list, return_tensors="pt")
+    n_queries = torch.ones(len(inp_list), dtype=torch.int32)
+    packed_inputs = flattener(inp_list, return_tensors="pt")
 
     packed_inputs["n_queries"] = n_queries
     if "ctx_ids" in sample:
@@ -53,167 +48,74 @@ def flatten_if_not_packed(inp_list):
     return packed_inputs
 
 
-def train_collator(inp_list, tokenizer):
-    # input is a list of tokenized sequences
-    padding_kwargs = dict(
-        padding=True,
-        padding_side="right",
-        pad_to_multiple_of=8,
-        return_tensors="pt",
-    )
-    sample = inp_list[0]
-    ctx_ids = None
-    if "ctx_ids" in sample:
-        # pad to the longest ctx_len in the batch
-        # which can have a different length from the input_ids, attn_mask, labels
-
-        ctx_ids = [example.pop("ctx_ids") for example in inp_list]
-        ctx_ids = torch.nn.utils.rnn.pad_sequence(
-            ctx_ids,
-            batch_first=True,
-            padding_value=0,
-        )
-        ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
-        ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
-            ctx_attn_mask,
-            batch_first=True,
-            padding_value=0,
-        )
-
-    need_flatten = check_is_iterable(sample["input_ids"][0])
-
-    if need_flatten:
-        # needed for one-lora-multi-q
-        n_queries = torch.tensor([len(x["input_ids"]) for x in inp_list])
-        tokens = dict()
-        for k in ["input_ids", "attention_mask", "labels"]:
-            tokens[k] = concat_list([x[k] for x in inp_list])
-        labels = tokens.pop("labels")
-        padded_seq = tokenizer.pad(tokens, **padding_kwargs)
-    else:
-        # needed for eval data during training
-        labels = [x.pop("labels") for x in inp_list]
-        padded_seq = tokenizer.pad(inp_list, **padding_kwargs)
-        n_queries = torch.ones(len(inp_list), dtype=torch.int32)
-
-    # hacky explicit padding since the labels are not padded by default
-    labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"]
-    labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels)
-    out = {**padded_seq, "labels": labels, "n_queries": n_queries}
-    if ctx_ids is not None:
-        out["ctx_ids"] = ctx_ids
-        out["ctx_attn_mask"] = ctx_attn_mask
-    if "logprobs_vals" in sample:
-        # for training with logits
-        logprobs_vals = [x.pop("logprobs_vals") for x in inp_list]
-        # logprobs_vals = tokenizer.pad(
-        #     {"input_ids": logprobs_vals}, **padding_kwargs
-        # )["input_ids"]
-        out["logprobs_vals"] = logprobs_vals
-        out["logprobs_indices"] = [x.pop("logprobs_indices") for x in inp_list]
-
-    return out
-
-
 def eval_collator(inp_list, tokenizer):
     # only used for teacher-forcing eval
     # input is a list of tokenized sequences
-    padding_kwargs = dict(
-        padding=True,
-        padding_side="right",
-        pad_to_multiple_of=8,
-        return_tensors="pt",
-    )
+    padding_kwargs = dict(padding=True, padding_side="right", return_tensors="pt")
 
-    ctx_ids = None
-    if "ctx_ids" in inp_list[0]:
+    has_ctx_ids = "ctx_ids" in inp_list[0]
+    if has_ctx_ids:
         # pad to the longest ctx_len in the batch
         # which can have a different length from the input_ids, attn_mask, labels
 
         ctx_ids = [example.pop("ctx_ids") for example in inp_list]
+        ctx_attn_mask = [torch.ones_like(x) for x in ctx_ids]
         ctx_ids = torch.nn.utils.rnn.pad_sequence(
             ctx_ids,
             batch_first=True,
             padding_value=0,
         )
-        ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
         ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
             ctx_attn_mask,
             batch_first=True,
             padding_value=0,
         )
 
-    chat_ids = None
-    if "chat_ids" in inp_list[0]:
-        chat_ids = [x.pop("chat_ids") for x in inp_list]
-        chat_ids = torch.nn.utils.rnn.pad_sequence(
-            chat_ids,
-            batch_first=True,
-            padding_value=0,
-        )
-        chat_attn_mask = [x.pop("chat_attn_mask") for x in inp_list]
-        chat_attn_mask = torch.nn.utils.rnn.pad_sequence(
-            chat_attn_mask,
-            batch_first=True,
-            padding_value=0,
-        )
-        chat_labels = [x.pop("chat_labels") for x in inp_list]
-        chat_labels = torch.nn.utils.rnn.pad_sequence(
-            chat_labels,
-            batch_first=True,
-            padding_value=-100,
-        )
-        chat_labels = torch.where(chat_attn_mask == 0, -100, chat_labels)
+    for inp in inp_list:
+        inp["attention_mask"] = torch.ones_like(inp["input_ids"])
 
     labels = [x.pop("labels") for x in inp_list]
+    # need to pass the whole inp bc we also track the lengths (with specal keys)
     padded_seq = tokenizer.pad(inp_list, **padding_kwargs)
 
     # hacky explicit padding since the labels are not padded by default
     labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"]
     labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels)
     out = {**padded_seq, "labels": labels}
-    if ctx_ids is not None:
+
+    if has_ctx_ids:
         out["ctx_ids"] = ctx_ids
         out["ctx_attn_mask"] = ctx_attn_mask
-    if chat_ids is not None:
-        out["chat_ids"] = chat_ids
-        out["chat_attn_mask"] = chat_attn_mask
-        out["chat_labels"] = chat_labels
-    # if "input_ids_len" in inp_list[0]:
-    #     out["input_ids_len"] = [x.pop("input_ids_len") for x in inp_list]
-    # if "ctx_ids_len" in inp_list[0]:
-    #     out["ctx_ids_len"] = [x.pop("ctx_ids_len") for x in inp_list]
+
     return out
 
 
 def generation_collator(inp_list, tokenizer):
     padding_kwargs = dict(padding=True, padding_side="left", return_tensors="pt")
     input_ids = [x.pop("input_ids") for x in inp_list]
-    attn_mask = [x.pop("attention_mask") for x in inp_list]
     labels = [x.pop("labels") for x in inp_list]
     for i, label in enumerate(labels):
+        # we don't include the labels in the output during generation
         # remove the response tokens
         idx = np.argmax(label != -100)
         idx = max(1, idx)
         input_ids[i] = input_ids[i][:idx]
-        attn_mask[i] = attn_mask[i][:idx]
+    attn_mask = [torch.ones_like(x) for x in input_ids]
+
     out = tokenizer.pad(
         {"input_ids": input_ids, "attention_mask": attn_mask}, **padding_kwargs
     )
 
-    # we don't include the labels in the output
-    # since we don't want to compute the loss on the labels
-    # during generation
     if "ctx_ids" in inp_list[0]:
         # pad to the longest ctx_len in the batch
         # which can have a different length from the input_ids, attn_mask, labels
         ctx_ids = [example.pop("ctx_ids") for example in inp_list]
+        ctx_attn_mask = [torch.ones_like(x) for x in ctx_ids]
         ctx_ids = torch.nn.utils.rnn.pad_sequence(
             ctx_ids,
             batch_first=True,
             padding_value=0,
         )
-        ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
         ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
             ctx_attn_mask,
             batch_first=True,
diff --git a/src/ctx_to_lora/data/packing.py b/src/ctx_to_lora/data/packing.py
index aac38cc..cbe0349 100644
--- a/src/ctx_to_lora/data/packing.py
+++ b/src/ctx_to_lora/data/packing.py
@@ -1,6 +1,9 @@
 # based on
 # https://github.com/MeetKai/functionary/blob/aa3dbdd65f7e388f2386622606bdfeec95c2b863/functionary/train/packing/packed_dataset.py
+import json
 import logging
+import os
+import pprint
 
 import numpy as np
 
@@ -77,17 +80,13 @@ def pack_data_points_by_length(
 
 def pack_data_points_FA(
     batch: dict[str, any],
-    need_flatten: bool,
 ) -> dict[str, np.ndarray]:
     if not batch:
         raise ValueError("Batch is empty")
 
     # Pre-allocate lists with known sizes
     total_ctx_len = sum(len(x) for x in batch["ctx_ids"])
-    if not need_flatten:
-        total_inp_len = sum(len(x) for x in batch["input_ids"])
-    else:
-        total_inp_len = sum(len(y) for x in batch["input_ids"] for y in x)
+    total_inp_len = sum(len(y) for x in batch["input_ids"] for y in x)
 
     ctx_ids = np.empty(total_ctx_len, dtype=np.long)
     ctx_position_ids = np.empty(total_ctx_len, dtype=np.long)
@@ -95,9 +94,9 @@ def pack_data_points_FA(
     position_ids = np.empty(total_inp_len, dtype=np.long)
     labels = np.empty(total_inp_len, dtype=np.long)
 
-    has_logits = "logprobs_vals" in batch
+    has_logprobs = "logprobs_vals" in batch
 
-    if has_logits:
+    if has_logprobs:
         sequences = zip(
             batch["input_ids"],
             batch["labels"],
@@ -116,29 +115,26 @@ def pack_data_points_FA(
     for sample in sequences:
         input_ids_b, labels_b = sample[:2]
         inp_start = offset
-        if need_flatten:
-            # compute position_ids for each sub-list in input_ids_b
-            local_start = inp_start
-            for ids_b in input_ids_b:
-                position_ids[local_start : local_start + len(ids_b)] = np.arange(
-                    len(ids_b), dtype=np.long
-                )
-                local_start += len(ids_b)
 
-            input_ids_b = concat_list(input_ids_b)
-            labels_b = concat_list(labels_b)
+        # compute position_ids for each sub-list in input_ids_b
+        local_start = inp_start
+        for ids_b in input_ids_b:
+            position_ids[local_start : local_start + len(ids_b)] = np.arange(
+                len(ids_b), dtype=np.long
+            )
+            local_start += len(ids_b)
+
+        input_ids_b = concat_list(input_ids_b)
+        labels_b = concat_list(labels_b)
 
         inp_len = len(input_ids_b)
         inp_end = offset + inp_len
 
-        if not need_flatten:
-            position_ids[inp_start:inp_end] = np.arange(inp_len, dtype=np.long)
-
         input_ids[inp_start:inp_end] = input_ids_b
         labels[inp_start:inp_end] = labels_b
         offset += inp_len
 
-        if has_logits:
+        if has_logprobs:
             logprobs_vals_b, logprobs_indices_b = sample[2:]
             logprobs_vals_b = concat_list(logprobs_vals_b)
             logprobs_indices_b = concat_list(logprobs_indices_b)
@@ -163,7 +159,7 @@ def pack_data_points_FA(
         "position_ids": position_ids,
         "labels": labels,
     }
-    if has_logits:
+    if has_logprobs:
         out["logprobs_vals"] = logprobs_vals
         out["logprobs_indices"] = logprobs_indices
     return out
@@ -174,21 +170,17 @@ def pack_batch(
     max_packed_inp_len: int,
     max_packed_ctx_len: int,
     max_packed_size: int = -1,
+    metadata_path: str = "",
 ) -> dict[str, any]:
     need_flatten = check_is_iterable(batch["input_ids"][0][0])
-    assert need_flatten
+    assert need_flatten, (
+        f"Packing requires the input_ids to be nested "
+        f"(allowing multiple QAs per sample), but got {batch['input_ids'][0]}"
+    )
 
-    if need_flatten:
-        n_queries = [len(x) for x in batch["input_ids"]]
-        inp_lens = [[len(y) for y in x] for x in batch["input_ids"]]
-        # batch = {k: concat_list([x for x in batch[k]]) for k in keys}
-        # for k in ["input_ids", "attention_mask", "labels"]:
-        #     batch[k] = concat_list([x for x in batch[k]])
-    else:
-        n_queries = [1] * len(batch["input_ids"])
-        inp_lens = [[len(x)] for x in batch["input_ids"]]
+    n_queries = [len(x) for x in batch["input_ids"]]
+    inp_lens = [[len(y) for y in x] for x in batch["input_ids"]]
     inp_count = len(inp_lens)
-    # total_original_tokens = sum(inp_lens)
     if "ctx_ids" not in batch:
         raise ValueError("Batch must contain 'ctx_ids' and 'labels' keys")
     # we do not pad so we can just take the length of the tokens
@@ -214,8 +206,8 @@ def pack_batch(
         "labels": [],
         "n_queries": [],
     }
-    has_logits = "logprobs_vals" in batch
-    if has_logits:
+    has_logprobs = "logprobs_vals" in batch
+    if has_logprobs:
         packed_batch["logprobs_vals"] = []
         packed_batch["logprobs_indices"] = []
 
@@ -231,19 +223,19 @@ def pack_batch(
             "input_ids": batch["input_ids"][start_idx:end_idx],
             "labels": batch["labels"][start_idx:end_idx],
         }
-        if has_logits:
+        if has_logprobs:
             group_items["logprobs_vals"] = batch["logprobs_vals"][start_idx:end_idx]
             group_items["logprobs_indices"] = batch["logprobs_indices"][
                 start_idx:end_idx
             ]
-        packed_item = pack_data_points_FA(group_items, need_flatten)
+        packed_item = pack_data_points_FA(group_items)
         packed_batch["ctx_ids"].append(packed_item["ctx_ids"])
         packed_batch["ctx_position_ids"].append(packed_item["ctx_position_ids"])
         packed_batch["input_ids"].append(packed_item["input_ids"])
         packed_batch["position_ids"].append(packed_item["position_ids"])
         packed_batch["labels"].append(packed_item["labels"])
         packed_batch["n_queries"].append(n_queries[start_idx:end_idx])
-        if has_logits:
+        if has_logprobs:
             packed_batch["logprobs_vals"].append(packed_item["logprobs_vals"])
             packed_batch["logprobs_indices"].append(packed_item["logprobs_indices"])
 
@@ -272,18 +264,33 @@ def pack_batch(
         np.mean(ctx_packing_efficiency_ratios) if ctx_packing_efficiency_ratios else 0
     )
 
-    logging.debug(
-        f"Packing stats - Original samples: {inp_count}\n"
-        f"# Packed samples: {len(idx_pairs)}\n"
-        f"Avg inp packing efficiency: {avg_inp_packing_efficiency:.3f}\n"
-        f"Avg ctx packing efficiency: {avg_ctx_packing_efficiency:.3f}\n\n"
-        f"Input IDs length stats:\n"
-        f"  Avg: {np.mean(packed_inp_lens_arr):.1f}, Std: {np.std(packed_inp_lens_arr):.1f}, "
-        f"Min: {np.min(packed_inp_lens_arr)}, Max: {np.max(packed_inp_lens_arr)}\n"
-        f"Context IDs length stats:\n"
-        f"  Avg: {np.mean(packed_ctx_lens_arr):.1f}, Std: {np.std(packed_ctx_lens_arr):.1f}, "
-        f"Min: {np.min(packed_ctx_lens_arr)}, Max: {np.max(packed_ctx_lens_arr)}\n\n"
-    )
+    # Create packing statistics dictionary
+    packing_stats = {
+        "original_samples": inp_count,
+        "packed_samples": len(idx_pairs),
+        "avg_inp_packing_efficiency": float(avg_inp_packing_efficiency),
+        "avg_ctx_packing_efficiency": float(avg_ctx_packing_efficiency),
+        "input_ids_length_stats": {
+            "avg": float(np.mean(packed_inp_lens_arr)),
+            "std": float(np.std(packed_inp_lens_arr)),
+            "min": int(np.min(packed_inp_lens_arr)),
+            "max": int(np.max(packed_inp_lens_arr)),
+        },
+        "context_ids_length_stats": {
+            "avg": float(np.mean(packed_ctx_lens_arr)),
+            "std": float(np.std(packed_ctx_lens_arr)),
+            "min": int(np.min(packed_ctx_lens_arr)),
+            "max": int(np.max(packed_ctx_lens_arr)),
+        },
+    }
+
+    # Save to metadata_path if provided
+    if metadata_path:
+        os.makedirs(os.path.dirname(metadata_path), exist_ok=True)
+        with open(metadata_path, "w") as f:
+            json.dump(packing_stats, f, indent=4)
+
+    logging.debug(f"Packing stats:\n{pprint.pformat(packing_stats, indent=2)}")
 
     return packed_batch
 
diff --git a/src/ctx_to_lora/data/preprocessing_fn.py b/src/ctx_to_lora/data/preprocessing_fn.py
new file mode 100644
index 0000000..1b5c012
--- /dev/null
+++ b/src/ctx_to_lora/data/preprocessing_fn.py
@@ -0,0 +1,280 @@
+import logging
+import random
+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()
+
+
+def closed_qa_prompting(prompt: str):
+    template = random.choice(CLOSED_QA_INTX_TEMPLATES)
+    return template.format(input=prompt)
+
+
+def get_preprocessing_fn(
+    ds_name: str,
+    is_eval: bool,
+) -> Callable[[dict[str, Any]], dict[str, Any]]:
+    """
+    Get preprocessing function for a specific dataset.
+
+    Args:
+        ds_name: Name of the dataset
+
+    Returns:
+        A preprocessing function that takes and returns a dictionary
+    """
+    f = lambda x: x
+    if ds_name.startswith("self_gen") or ds_name.endswith("_compact"):
+        # 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):
+            return {
+                "context": sample["context"],
+                "prompt": sample["input"],
+                "response": sample["answers"][0],
+            }
+
+    elif ds_name == "negative_nq":
+
+        def f(sample):
+            q = sample["prompt"]
+            prompt = closed_qa_prompting(q) if not is_eval else q
+            return {
+                "context": sample["context"],
+                "prompt": prompt,
+                "response": sample["answer"],
+            }
+
+    elif ds_name == "triviaqa_retrieved":
+        # only used for eval
+        def f(sample):
+            return {
+                "context": sample["context"],
+                "prompt": sample["prompt"],
+                "response": sample["answer"],
+            }
+
+    elif ds_name == "pwc" or ds_name == "pwc_tiny":
+        # original pwc
+        def f(sample):
+            return {
+                "context": sample["input"],
+                "prompt": sample["prompt"],
+                "response": sample["answer"],
+            }
+
+    elif ds_name.startswith("hotpot_qa"):
+
+        def f(sample):
+            txt = ""
+            for p in sample["context"]["sentences"]:
+                txt += " " + "".join(p)
+
+            q = sample["question"]
+            prompt = closed_qa_prompting(q) if not is_eval else q
+
+            return {
+                "context": txt.strip(),
+                "prompt": prompt,
+                "response": sample["answer"],
+            }
+
+    elif ds_name == "squad":
+        # original squad
+        def f(sample):
+            q = sample["question"]
+            prompt = closed_qa_prompting(q) if not is_eval else q
+            return {
+                "context": sample["context"],
+                "prompt": prompt,
+                "response": sample["answers"]["text"][0],
+            }
+
+    elif ds_name == "drop":
+
+        def f(sample):
+            q = sample["question"]
+            prompt = closed_qa_prompting(q) if not is_eval else q
+            return {
+                "context": sample["passage"],
+                "prompt": prompt,
+                "response": sample["answers_spans"]["spans"][0],
+            }
+
+    elif ds_name == "ropes":
+        ctx_template = "{background}\n{situation}"
+
+        def f(sample):
+            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"]
+            q = closed_qa_prompting(q) if not is_eval else q
+            return {"context": ctx, "prompt": q, "response": response}
+
+    elif ds_name in ["narrativeqa", "quoref"]:  # , "ropes"]:
+
+        def f(sample):
+            response = sample["answers"][0]
+            if isinstance(response, list):
+                response = response[0]
+            q = sample["messages"][-1]["content"]
+            prompt = closed_qa_prompting(q) if not is_eval else q
+            return {
+                "context": sample["document"],
+                "prompt": prompt,
+                "response": response,
+            }
+
+    elif ds_name == "synthetic_convqa":
+
+        def f(sample):
+            response = sample["answers"][0]
+            if isinstance(response, list):
+                response = response[0]
+            return {
+                "context": sample["document"],
+                "prompt": sample["messages"][-1]["content"],
+                "response": response,
+            }
+
+    elif ds_name == "booksum":
+        prompt_templates = [
+            "Summarization the provided text.",
+            "# Summary",
+            "### Summary",
+            "Summary of the text",
+        ]
+        # TODO: use these templates for training
+        # analysis_templates = []
+        # summary_templates = []
+
+        def f(sample):
+            return {
+                "context": sample["chapter"],
+                "prompt": "Summarize the provided text.",
+                "response": sample["summary_text"],
+            }
+
+    elif ds_name == "gov_report":
+
+        def f(sample):
+            return {
+                "context": sample["report"],
+                "prompt": "Summarize the provided text.",
+                "response": sample["summary"],
+            }
+
+    elif "wikitext" in ds_name:
+
+        def f(sample):
+            return {
+                "context": sample["page"],
+                "prompt": "PLAECHOLDER",
+                "response": "PLAECHOLDER",
+            }
+
+    elif ds_name == "openmathintx-2":
+
+        def f(sample):
+            return {
+                "context": sample["problem"],
+                "prompt": sample["problem"],
+                "response": sample["generated_solution"],
+            }
+
+    elif "gsm8k" in ds_name:
+
+        def f(sample):
+            return {
+                "context": sample["question"],
+                "prompt": sample["question"],
+                "response": sample["answer"],
+            }
+
+    elif "opencoder-edu" in ds_name:
+
+        def f(sample):
+            return {
+                "context": sample["instruction"],
+                "prompt": sample["instruction"],
+                "response": "```python\n" + sample["code"].strip() + "\n```",
+            }
+
+    if is_eval and (ds_name in EVAL_INTX_TEMPLATES):
+        prompt_template = EVAL_INTX_TEMPLATES[ds_name]
+
+        def eval_intx_decorator(f):
+            def g(sample):
+                sample = f(sample)
+                assert "prompt" in sample, (
+                    f"Expected 'prompt' in sample, got {sample.keys()}"
+                )
+                sample["prompt"] = prompt_template.format(input=sample["prompt"])
+                return sample
+
+            return g
+
+        f = eval_intx_decorator(f)
+
+    def maybe_convert_to_list(f):
+        def g(sample):
+            sample = f(sample)
+            if "prompt" in sample:
+                sample["prompts"] = [sample.pop("prompt")]
+            if "response" in sample:
+                sample["responses"] = [sample.pop("response")]
+            return sample
+
+        return g
+
+    f = maybe_convert_to_list(f)
+
+    if "self_gen" not in ds_name:
+
+        def strip_response(f):
+            def g(sample):
+                sample = f(sample)
+                if "responses" in sample and bool(sample["responses"]):
+                    sample["responses"] = [
+                        r.strip() if isinstance(r, str) else r
+                        for r in sample["responses"]
+                    ]
+                return sample
+
+            return g
+
+        f = strip_response(f)
+
+    return f
diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py
index db30679..2b1f124 100644
--- a/src/ctx_to_lora/data/processing.py
+++ b/src/ctx_to_lora/data/processing.py
@@ -1,35 +1,54 @@
-import hashlib
 import json
 import logging
 import os
 import random
 from collections.abc import Callable
 from glob import glob
+from hashlib import sha256
 from os import path
 from typing import Any
 
 import datasets
-import numpy as np
 import torch
 from datasets import Dataset, is_caching_enabled, load_dataset
-from transformers import PreTrainedModel, PreTrainedTokenizerBase
+from transformers import PreTrainedTokenizerBase
 
 from ctx_to_lora.data.definitions import (
-    CLOSED_QA_INTX_TEMPLATES,
     DS_KWARGS,
-    EVAL_INTX_TEMPLATES,
     IGNORE_INDEX,
     RAW_DATA_DIR,
     REPEAT_PROMPTS,
     SELF_GEN_DATA_DIR,
-    SELF_GEN_SYSTEM_MSG,
     TRANSFORMED_DATA_DIR,
 )
 from ctx_to_lora.data.packing import pack_batch
-from ctx_to_lora.utils import TRAINING_TASK, concat_list
+from ctx_to_lora.data.preprocessing_fn import get_preprocessing_fn
+from ctx_to_lora.utils import check_is_iterable, concat_list
 
 logger = logging.getLogger()
 
+COLS_TO_KEEP_PREPROCESSING = [
+    "context",
+    "prompts",
+    "responses",
+    "qas",
+    "variation",
+    "logprobs_vals",
+    "logprobs_indices",
+    "input_ids",
+    "ctx_ids",
+    "response_start_end",
+]
+
+COLS_TO_KEEP_TOKENIZED = [
+    "input_ids",
+    "labels",
+    "context",
+    "ctx_ids",
+    "logprobs_vals",
+    "logprobs_indices",
+]
+
 
 def load_answers(ds_name, split):
     if ds_name.startswith("longbench"):
@@ -53,284 +72,10 @@ def load_answers(ds_name, split):
     return ds
 
 
-def closed_qa_prompting(prompt: str):
-    template = random.choice(CLOSED_QA_INTX_TEMPLATES)
-    return template.format(input=prompt)
-
-
 def get_repeat_prompt():
     return random.choice(REPEAT_PROMPTS)
 
 
-def get_preprocessing_fn(
-    ds_name: str,
-    is_eval: bool,
-    is_pretrain: bool,
-) -> Callable[[dict[str, Any]], dict[str, Any]]:
-    """
-    Get preprocessing function for a specific dataset.
-
-    Args:
-        ds_name: Name of the dataset
-
-    Returns:
-        A preprocessing function that takes and returns a dictionary
-    """
-    f = lambda x: x
-    if ds_name.startswith("self_gen") or ds_name.endswith("_compact"):
-        # 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):
-            return {
-                "context": sample["context"],
-                "prompt": sample["input"],
-                "response": sample["answers"][0],
-            }
-
-    elif ds_name == "negative_nq":
-
-        def f(sample):
-            q = sample["prompt"]
-            prompt = closed_qa_prompting(q) if not is_eval else q
-            return {
-                "context": sample["context"],
-                "prompt": prompt,
-                "response": sample["answer"],
-            }
-
-    elif ds_name == "triviaqa_retrieved":
-        # only used for eval
-        def f(sample):
-            return {
-                "context": sample["context"],
-                "prompt": sample["prompt"],
-                "response": sample["answer"],
-            }
-
-    elif ds_name == "pwc" or ds_name == "pwc_tiny":
-        # original pwc
-        def f(sample):
-            return {
-                "context": sample["input"],
-                "prompt": sample["prompt"],
-                "response": sample["answer"],
-            }
-
-    elif ds_name.startswith("hotpot_qa"):
-
-        def f(sample):
-            txt = ""
-            for p in sample["context"]["sentences"]:
-                txt += " " + "".join(p)
-
-            q = sample["question"]
-            prompt = closed_qa_prompting(q) if not is_eval else q
-
-            return {
-                "context": txt.strip(),
-                "prompt": prompt,
-                "response": sample["answer"],
-            }
-
-    elif ds_name == "squad":
-        # original squad
-        def f(sample):
-            q = sample["question"]
-            prompt = closed_qa_prompting(q) if not is_eval else q
-            return {
-                "context": sample["context"],
-                "prompt": prompt,
-                "response": sample["answers"]["text"][0],
-            }
-
-    elif ds_name == "drop":
-
-        def f(sample):
-            q = sample["question"]
-            prompt = closed_qa_prompting(q) if not is_eval else q
-            return {
-                "context": sample["passage"],
-                "prompt": prompt,
-                "response": sample["answers_spans"]["spans"][0],
-            }
-
-    elif ds_name == "ropes":
-        ctx_template = "{background}\n{situation}"
-
-        def f(sample):
-            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"]
-            q = closed_qa_prompting(q) if not is_eval else q
-            return {"context": ctx, "prompt": q, "response": response}
-
-    elif ds_name in ["narrativeqa", "quoref"]:  # , "ropes"]:
-
-        def f(sample):
-            response = sample["answers"][0]
-            if isinstance(response, list):
-                response = response[0]
-            q = sample["messages"][-1]["content"]
-            prompt = closed_qa_prompting(q) if not is_eval else q
-            return {
-                "context": sample["document"],
-                "prompt": prompt,
-                "response": response,
-            }
-
-    elif ds_name == "synthetic_convqa":
-
-        def f(sample):
-            response = sample["answers"][0]
-            if isinstance(response, list):
-                response = response[0]
-            return {
-                "context": sample["document"],
-                "prompt": sample["messages"][-1]["content"],
-                "response": response,
-            }
-
-    elif ds_name == "booksum":
-        prompt_templates = [
-            "Summarization the provided text.",
-            "# Summary",
-            "### Summary",
-            "Summary of the text",
-        ]
-        # TODO: use these templates for training
-        # analysis_templates = []
-        # summary_templates = []
-
-        def f(sample):
-            return {
-                "context": sample["chapter"],
-                "prompt": "Summarize the provided text.",
-                "response": sample["summary_text"],
-            }
-
-    elif ds_name == "gov_report":
-
-        def f(sample):
-            return {
-                "context": sample["report"],
-                "prompt": "Summarize the provided text.",
-                "response": sample["summary"],
-            }
-
-    elif "wikitext" in ds_name:
-
-        def f(sample):
-            return {
-                "context": sample["page"],
-                "prompt": "PLAECHOLDER",
-                "response": "PLAECHOLDER",
-            }
-
-    elif ds_name == "openmathintx-2":
-
-        def f(sample):
-            return {
-                "context": sample["problem"],
-                "prompt": sample["problem"],
-                "response": sample["generated_solution"],
-            }
-
-    elif "gsm8k" in ds_name:
-
-        def f(sample):
-            return {
-                "context": sample["question"],
-                "prompt": sample["question"],
-                "response": sample["answer"],
-            }
-
-    elif "opencoder-edu" in ds_name:
-
-        def f(sample):
-            return {
-                "context": sample["instruction"],
-                "prompt": sample["instruction"],
-                "response": "```python\n" + sample["code"].strip() + "\n```",
-            }
-
-    if is_eval and (ds_name in EVAL_INTX_TEMPLATES) and not is_pretrain:
-        prompt_template = EVAL_INTX_TEMPLATES[ds_name]
-
-        def eval_intx_decorator(f):
-            def g(sample):
-                sample = f(sample)
-                assert "prompt" in sample, (
-                    f"Expected 'prompt' in sample, got {sample.keys()}"
-                )
-                sample["prompt"] = prompt_template.format(input=sample["prompt"])
-                return sample
-
-            return g
-
-        f = eval_intx_decorator(f)
-
-    if not is_pretrain:
-
-        def maybe_convert_to_list(f):
-            def g(sample):
-                sample = f(sample)
-                if "prompt" in sample:
-                    sample["prompts"] = [sample.pop("prompt")]
-                if "response" in sample:
-                    sample["responses"] = [sample.pop("response")]
-                return sample
-
-            return g
-
-        f = maybe_convert_to_list(f)
-
-    if "self_gen" not in ds_name:
-
-        def strip_response(f):
-            def g(sample):
-                sample = f(sample)
-                if "responses" in sample and bool(sample["responses"]):
-                    sample["responses"] = [
-                        r.strip() if isinstance(r, str) else r
-                        for r in sample["responses"]
-                    ]
-                return sample
-
-            return g
-
-        f = strip_response(f)
-
-    return f
-
-
 def get_ds_kwargs(ds_name: str, split: str) -> dict[str, Any]:
     # custom logic for slicing iterable datasets
     take, skip = None, None
@@ -387,140 +132,12 @@ def get_ds_kwargs(ds_name: str, split: str) -> dict[str, Any]:
     return kwargs
 
 
-def validate_columns(tokenized_ds):
-    cols = ["input_ids", "attention_mask", "labels"]
-    if "ctx_ids" in tokenized_ds.column_names:
-        cols += ["ctx_ids", "ctx_attn_mask"]
-    if "chat_ids" in tokenized_ds.column_names:
-        cols += ["chat_ids", "chat_attn_mask", "chat_labels"]
-    ref_cols = set(cols)
-    assert set(tokenized_ds.column_names) == ref_cols, (
-        f"Columns mismatch: {set(tokenized_ds.column_names)} != {ref_cols}"
-    )
-
-
 def len_filter(sample, max_length: int, keys: list[str]):
     m = [len(sample[k]) <= max_length for k in keys]
     return sum(m) == len(keys)
 
 
-def len_filter_data(sample, max_length: int, keys: list[str]):
-    data = sample["data"]
-    indices = dict()
-    for k in keys:
-        indices[k] = {i for i, arr in enumerate(data[k]) if len(arr) <= max_length}
-
-    all_indices = set.intersection(*indices.values())
-    for k in data:
-        data[k] = [data[k][i] for i in all_indices]
-
-    return dict(data=data)
-
-
-# def len_filter(ds, max_length: int, keys: list[str]):
-#     for k in keys:
-#         ds = ds.select([i for i, x in enumerate(ds[k]) if len(x) <= max_length])
-#     return ds
-
-
-# def filter_long_samples(samples):
-#     return [len(ctx) < 10000 for ctx in samples["context"]]
-
-
-# def filter_long_chat(samples):
-#     return [len(chat) < 2**13 for chat in samples["chat"]]
-
-
-def add_repeat_prompt_fn(samples: dict[str, list[str]], prob: float):
-    unique_contexts = set()
-    ctxs, prompts, responses = [], [], []
-    for ctx in samples["context"]:
-        if random.random() > prob:
-            continue
-        # Only process if the context is not already in the set
-        if ctx in unique_contexts:
-            continue
-        # # remove too long contexts
-        # if len(ctx) > 1000:
-        #     continue
-        unique_contexts.add(ctx)
-        ctxs.append(ctx)
-        responses.append([ctx])
-        prompts.append([get_repeat_prompt()])
-
-    logger.debug("Adding repeat prompt...")
-    logger.debug(f"# unique contexts: {len(unique_contexts)}")
-
-    return dict(
-        context=ctxs + samples["context"],
-        responses=responses + samples["responses"],
-        prompts=prompts + samples["prompts"],
-    )
-
-
-# def add_negative_prompt_fn(samples):
-#     unique_contexts = set()
-#     ctxs, prompts, responses = [], [], []
-#     keywords = [
-#         "repeat",
-#         "rephrase",
-#         "summarize",
-#         "rewrite",
-#         "title",
-#         "keyword",
-#         "continuation",
-#     ]
-
-#     for ctx, prompt, response in zip(
-#         samples["context"], samples["prompt"], samples["response"]
-#     ):
-#         if ctx in unique_contexts:
-#             continue
-#         unique_contexts.add(ctx)
-#         if len(ctx) > 3000:
-#             continue
-#         if any(keyword in prompt for keyword in keywords):
-#             # Skip samples where the prompt contains any of the specified keywords
-#             continue
-
-#         ctxs.append(ctx)
-#         prompts.append(prompt)
-#         responses.append(response)
-
-#     logger.debug("Adding negative prompt...")
-#     logger.debug(f"# unique contexts: {len(unique_contexts)}")
-
-#     # remove one last sample if the number of samples is odd
-#     if len(ctxs) % 2 != 0:
-#         ctxs.pop()
-#         prompts.pop()
-#         responses.pop()
-
-#     # to make sure that the negative prompt/response is not the same as the original
-#     indices = list(np.random.permutation(len(ctxs))) + list(
-#         np.random.permutation(len(ctxs))
-#     )
-#     neg_ctxs, neg_prompts, neg_responses = [], [], []
-#     for idx in range(0, len(indices), 2):
-#         i = indices[idx]
-#         j = indices[idx + 1]
-#         neg_ctxs.append(ctxs[i])
-#         neg_prompts.append(ctxs[j] + "\n\n" + prompts[j])
-#         neg_responses.append(responses[j])
-
-#     return dict(
-#         context=neg_ctxs + samples["context"],
-#         prompt=neg_prompts + samples["prompt"],
-#         response=neg_responses + samples["response"],
-#     )
-
-
 def filter_none(sample):
-    # out = [True] * len(samples["context"])
-    # for k in samples:
-    #     for i, sample in enumerate(samples[k]):
-    #         if not bool(sample):
-    #             out[i] = False
     for v in sample.values():
         if v is None:
             return False
@@ -530,11 +147,6 @@ def filter_none(sample):
 def load_and_process_dataset(
     ds_name: str,
     split: str,
-    add_negative_prompt: bool,
-    add_repeat_prompt: bool,
-    repeat_prob: float,
-    is_pretrain: bool,
-    streaming: bool,
     num_proc: int,
 ):
     logger.info(f"Loading dataset {ds_name} with split {split}...")
@@ -542,87 +154,29 @@ def load_and_process_dataset(
         ds_kwargs = get_ds_kwargs(ds_name, split)
         skip = ds_kwargs.pop("skip", None)
         take = ds_kwargs.pop("take", None)
-        # only load as stream if the dataset is local (parquet)
-        load_as_stream = (
-            ("fw_qa" in ds_name)
-            and ("tiny" not in ds_name)
-            and ("mini" not in ds_name)
-            and streaming
-            and ds_kwargs["path"] == "parquet"
-        )
-        ds = load_dataset(
-            **ds_kwargs,
-            trust_remote_code=True,
-            streaming=load_as_stream,
-        )
-        if "train" in split and streaming and (not load_as_stream):
-            # if the dataset is hosted on HF, we load it first then convert to iterable
-            ds = ds.to_iterable_dataset(num_shards=128)
+        ds = load_dataset(**ds_kwargs, trust_remote_code=True)
         if skip is not None:
             ds = ds.skip(skip)
         if take is not None:
             ds = ds.take(take)
     except ValueError as e:
-        logger.warning(
-            f"Failed to load dataset {ds_name} with split {split}. Error: {e}\nSkipping..."
+        raise ValueError(
+            f"Failed to load dataset {ds_name} with split {split}. Error: {e}"
         )
-        return None
     cols_to_remove = [
-        col
-        for col in ds.column_names
-        if col
-        not in [
-            "context",
-            "prompts",
-            "responses",
-            "qas",
-            "variation",
-            "logprobs_vals",
-            "logprobs_indices",
-            "input_ids",
-            "ctx_ids",
-            "response_start_end",
-        ]
+        col for col in ds.column_names if col not in COLS_TO_KEEP_PREPROCESSING
     ]
     is_eval = split != "train"
     ds = ds.map(
-        get_preprocessing_fn(ds_name, is_eval, is_pretrain),
+        get_preprocessing_fn(ds_name, is_eval),
         remove_columns=cols_to_remove,
         num_proc=num_proc,
     )
-    # ds = ds.remove_columns(cols_to_remove)
     ds = ds.filter(
         filter_none,
         batched=False,
         num_proc=num_proc,
     )
-
-    if "train" in split:
-        # TODO: drop samples or split into multiple contexts
-        # for multi-lora training
-        # ds = ds.filter(filter_long_samples, batched=True, num_proc=num_proc)
-        if add_negative_prompt:
-            # TODO: add explicit dataset for this
-            # w/ qs that are answerable 0-shot
-            # here we add the context to the prompt
-            ds = ds.map(
-                add_negative_prompt_fn,
-                batched=True,
-                batch_size=100_000,
-                # num_proc=num_proc,
-            )
-        if (
-            add_repeat_prompt
-            and "context_numbers" not in ds_name
-            and "pretrain" not in ds_name
-        ):
-            ds = ds.map(
-                add_repeat_prompt_fn,
-                fn_kwargs={"prob": repeat_prob},
-                batched=True,
-                batch_size=100_000,
-                # num_proc=num_proc,
-            )
     return ds
 
 
@@ -633,48 +187,31 @@ def get_tokenized_dataset(
     max_qas_per_sample: int,
     base_model_max_len: int,
     tokenizer: PreTrainedTokenizerBase,
-    tokenizer_kwargs: dict[str, Any],
     ctx_model_max_len: int,
     ctx_tokenizer: PreTrainedTokenizerBase,
-    ctx_tokenizer_kwargs: dict[str, Any],
     add_ctx_to_chat: bool,
-    add_repeat_prompt: bool,
-    add_negative_prompt: bool,
     use_kl_loss: bool,
-    repeat_prob: float = 0.0,  # only used when add_repeat_prompt is True
     set_format: str | None = None,
-    streaming: bool = False,
-    base_model: PreTrainedModel | None = None,
 ) -> dict[str, Any]:
-    # assert not use_kl_loss, "KL loss is deprecated"  # TO BE REMOVED
-    if add_repeat_prompt:
-        assert repeat_prob > 0, f"add_repeat_prompt is set but repeat_prob = 0"
+    if max_qas_len > 0:
+        assert max_qas_len <= base_model_max_len, (
+            f"`max_qas_len` should be <= {base_model_max_len=}, got {max_qas_len=}"
+        )
     logger.info(f"Loading dataset {ds_name} with split {split}...")
     need_ctx_ids = not add_ctx_to_chat and bool(ctx_model_max_len)
-    is_pretrain = "pretrain" in ds_name
-    is_paraphrased = "aug_pretrain" in ds_name
 
     load_and_process_kwargs = dict(
         ds_name=ds_name,
         split=split,
-        add_negative_prompt=add_negative_prompt,
-        add_repeat_prompt=add_repeat_prompt,
-        repeat_prob=repeat_prob,
-        is_pretrain=is_pretrain,
-        streaming=streaming,
     )
     tokenize_kwargs = dict(
         max_qas_len=max_qas_len,
         max_qas_per_sample=max_qas_per_sample,
         base_model_max_len=base_model_max_len,
-        tokenizer_kwargs=tokenizer_kwargs,
         ctx_model_max_len=ctx_model_max_len,
-        ctx_tokenizer_kwargs=ctx_tokenizer_kwargs,
         add_ctx_to_chat=add_ctx_to_chat,
         use_kl_loss=use_kl_loss,
         need_ctx_ids=need_ctx_ids,
-        is_pretrain=is_pretrain,
-        is_paraphrased=is_paraphrased,
         split=split,
         set_format=set_format,
     )
@@ -684,7 +221,7 @@ def get_tokenized_dataset(
     kwargs_str += tokenizer.name_or_path + ctx_tokenizer.name_or_path
     logger.debug(f"Tokenizing dataset with kwargs: {kwargs_str}")
     kwargs_str += repr(tokenizer) + repr(ctx_tokenizer)
-    ds_hash = hashlib.sha256(kwargs_str.encode()).hexdigest()
+    ds_hash = sha256(kwargs_str.encode()).hexdigest()
     logger.debug(f"Dataset hash: {ds_hash}")
     ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}"
 
@@ -694,7 +231,7 @@ def get_tokenized_dataset(
         tokenized_ds = datasets.load_from_disk(ds_path)
         return tokenized_ds
 
-    num_proc = None if streaming and ("train" in split) else 4
+    num_proc = 4
     ds = load_and_process_dataset(
         **load_and_process_kwargs,
         num_proc=num_proc,
@@ -704,7 +241,6 @@ def get_tokenized_dataset(
         ds=ds,
         tokenizer=tokenizer,
         ctx_tokenizer=ctx_tokenizer,
-        model=base_model,
         num_proc=num_proc,
         **tokenize_kwargs,
     )
@@ -722,142 +258,49 @@ def construct_and_tokenize_ctx_qa(
     max_qas_per_sample,
     base_model_max_len,
     tokenizer,
-    tokenizer_kwargs,
     ctx_model_max_len,
     ctx_tokenizer,
-    ctx_tokenizer_kwargs,
     add_ctx_to_chat,
     use_kl_loss,
-    model,
     need_ctx_ids,
-    is_pretrain,
-    is_paraphrased,
     ds,
     split,
     set_format=None,
     num_proc=None,
 ):
-    if is_pretrain:
-        if is_paraphrased:
-            ds = ds.map(
-                build_paraphrase_pretrain,
-                # batched=True,
-                # batch_size=100_000,
-                remove_columns=["variation"],
-                num_proc=num_proc,
-            )
-        else:
-            ds = ds.map(
-                build_intx_pretrain,
-                # batched=True,
-                # batch_size=100_000,
-                num_proc=num_proc,
+    if use_kl_loss:
+        if "train" in split and "logprobs_vals" not in ds.column_names:
+            raise ValueError(
+                "`use_kl_loss` is set to True but 'logprobs_vals' column "
+                "is not present in the dataset."
             )
+    else:
+        if "logprobs_vals" in ds.column_names:
+            ds = ds.remove_columns(["logprobs_vals", "logprobs_indices"])
+
+    # for sft + chat_model, we need to convert the dataset to chat format
+    if "input_ids" in ds.column_names and "response_start_end" in ds.column_names:
+        # already tokenized dataset (e.g., self-gen qa data)
+        tokenized_ds = ds.map(get_labels_from_input_ids, num_proc=num_proc)
+    else:
+        # construct messages from prompts and responses
+        # add "messages_list" field
+        ds = ds.map(
+            convert_ctx_prompt_response_to_messages,
+            fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat},
+            num_proc=num_proc,
+        )
+        # add `input_ids`, `attention_mask`, `labels`
+        os.environ["TOKENIZERS_PARALLELISM"] = "true"
+        logging.debug("Tokenizing inputs")
         tokenized_ds = ds.map(
-            tokenize_pretrain,
-            fn_kwargs={
-                "tokenizer": tokenizer,
-                "tokenizer_kwargs": tokenizer_kwargs,
-                "add_eos_token": is_paraphrased,
-            },
+            get_sft_prompt_formatting_fn(tokenizer),
             batched=True,
             batch_size=100_000,
-            # # num_proc=num_proc,
         )
 
-    else:
-        # convert everything into a new intermediate single column format
-        # (required for tracking n_queries and
-        # ejecting entire (ctx,questions,answers) triplet after filtering)
-        # then convert back to columns after filtering
-
-        # for sft + chat_model, we need to convert the dataset to chat format
-        # add "data" field
-        if "input_ids" in ds.column_names and "response_start_end" in ds.column_names:
-            # already tokenized dataset (e.g., self-gen qa data)
-            tokenized_ds = ds.map(
-                get_labels_from_input_ids,
-                num_proc=num_proc,
-            )
-        else:
-            ds = ds.map(
-                convert_ctx_prompt_response_to_messages,
-                fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat},
-                num_proc=num_proc,
-                # remove_columns=[col for col in ds.column_names if col != "context"],
-            )
-            # add input_ids, attention_mask, labels to "data"
-            os.environ["TOKENIZERS_PARALLELISM"] = "true"
-            logging.debug("Tokenizing inputs")
-            tokenized_ds = ds.map(
-                get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer),
-                batched=True,
-                batch_size=100_000,
-                # # num_proc=num_proc,
-            )
-
-    # if "train" in split:
-    #     # base model cant handle these samples naturally
-    #     # should skip
-    #     # tokenized_ds = tokenized_ds.filter(
-    #     #     len_filter,
-    #     #     fn_kwargs={
-    #     #         "max_length": base_model_max_len,
-    #     #         "keys": ["input_ids", "labels"],
-    #     #     },
-    #     #     # batched=False,
-    #     #     # batch_size=0,
-    #     #     num_proc=num_proc,
-    #     # )
-    #     tokenized_ds = tokenized_ds.map(
-    #         len_filter_data,
-    #         fn_kwargs={
-    #             "max_length": base_model_max_len,
-    #             "keys": ["input_ids"],
-    #         },
-    #         num_proc=num_proc,
-    #     )
-
-    logger.info(tokenized_ds)
-
-    if use_kl_loss and ("train" in split):
-        if "logprobs_vals" not in ds.column_names:
-            assert model is not None, (
-                f"Base model is required for computing the "
-                "logits when `use_kl_loss` is True and logprobs are not included."
-            )
-            # add "data_logp"
-            tokenized_ds = tokenized_ds.map(
-                convert_ctx_prompt_response_to_messages,
-                fn_kwargs={"add_ctx_to_chat": True, "for_logp": True},
-                num_proc=num_proc,
-            )
-            # add logits info to "data_logp" field
-            logger.info("computing logprobs")
-            tokenized_ds = tokenized_ds.map(
-                get_sft_prompt_formatting_fn(
-                    TRAINING_TASK.COMPLETION, tokenizer, use_kl_loss, model
-                ),
-                batched=True,
-                batch_size=100_000,
-            )
-        else:
-            tokenized_ds = tokenized_ds.map(add_ctx_attn_mask, num_proc=num_proc)
-
-    keep_cols = [
-        "input_ids",
-        "attention_mask",
-        "labels",
-        "ctx_ids",
-        "ctx_attn_mask",
-        "context",
-        "data",
-        "data_logp",
-    ]
-    if use_kl_loss:
-        keep_cols += ["logprobs_vals", "logprobs_indices"]
     tokenized_ds = tokenized_ds.remove_columns(
-        [col for col in tokenized_ds.column_names if col not in keep_cols],
+        [col for col in tokenized_ds.column_names if col not in COLS_TO_KEEP_TOKENIZED],
     )
 
     if need_ctx_ids:
@@ -870,6 +313,7 @@ def construct_and_tokenize_ctx_qa(
                 fn_kwargs={"tokenizer": ctx_tokenizer},
                 batched=True,
                 batch_size=100_000,
+                remove_columns=["context"],
             )
         if "train" in split:
             # TODO: do something if ctx length is longer than the ctx model length
@@ -880,109 +324,53 @@ def construct_and_tokenize_ctx_qa(
                 num_proc=num_proc,
             )
 
-    if not is_pretrain:
-        if "train" in split:
-            logging.debug("Unpacking data")
-            tokenized_ds = tokenized_ds.map(
-                unpack_data,
-                num_proc=num_proc,
-                remove_columns=[
-                    col
-                    for col in tokenized_ds.column_names
-                    if col in ["data", "data_logp", "context"]
-                ],
-            )
-            print(tokenized_ds)
-
-            logging.debug(f"Split too long QAs with max length {max_qas_len}")
-            tokenized_ds = tokenized_ds.map(
-                split_too_long_qas,
-                fn_kwargs={
-                    "max_qas_len": max_qas_len,
-                    "max_qas_per_sample": max_qas_per_sample,
-                },
-                batched=True,
-                batch_size=100_000,
-                num_proc=4,
-            )
-        else:
-            cols_to_remove = ["data", "context"]
-            if "ctx_ids" in tokenized_ds.column_names:
-                cols_to_remove += ["ctx_ids", "ctx_attn_mask"]
-            tokenized_ds = tokenized_ds.map(
-                unpack_data_eval,
-                num_proc=num_proc,
-                remove_columns=cols_to_remove,
-                batched=True,
-                batch_size=100_000,
-            )
+    logging.info(f"Split too long QAs with max length {max_qas_len}")
+    tokenized_ds = tokenized_ds.map(
+        split_too_long_qas,
+        fn_kwargs={
+            "max_qas_len": max_qas_len,
+            "max_qas_per_sample": max_qas_per_sample,
+        },
+        batched=True,
+        batch_size=100_000,
+        num_proc=4,
+    )
+    if "train" not in split:
+        # squeeze since we always have one query per sample in eval
+        tokenized_ds = tokenized_ds.map(squeeze_tokens, num_proc=num_proc)
+        tokenized_ds = tokenized_ds.map(
+            truncate_middle_if_too_long,
+            fn_kwargs={
+                "max_length": base_model_max_len,
+                "columns": ["input_ids", "labels"],
+                "max_new_tokens": 256,
+            },
+        )
+        tokenized_ds = tokenized_ds.map(
+            add_length_info,
+            fn_kwargs={"columns": ["input_ids"]},
+        )
+        if "ctx_ids" in tokenized_ds.column_names:
             tokenized_ds = tokenized_ds.map(
                 truncate_middle_if_too_long,
                 fn_kwargs={
-                    "max_length": base_model_max_len,
-                    "columns": ["input_ids", "attention_mask", "labels"],
-                    "max_new_tokens": 256,
+                    "max_length": ctx_model_max_len,
+                    "columns": ["ctx_ids"],
+                    # cxt encoder doesnt need to add new_tokens
+                    "max_new_tokens": 0,
                 },
             )
             tokenized_ds = tokenized_ds.map(
                 add_length_info,
-                fn_kwargs={"columns": ["input_ids"]},
+                fn_kwargs={"columns": ["ctx_ids"]},
             )
-            if "ctx_ids" in tokenized_ds.column_names:
-                tokenized_ds = tokenized_ds.map(
-                    truncate_middle_if_too_long,
-                    fn_kwargs={
-                        "max_length": ctx_model_max_len,
-                        "columns": ["ctx_ids", "ctx_attn_mask"],
-                        # cxt encoder doesnt need to add new_tokens
-                        "max_new_tokens": 0,
-                    },
-                )
-                tokenized_ds = tokenized_ds.map(
-                    add_length_info,
-                    fn_kwargs={"columns": ["ctx_ids"]},
-                )
-
-    if is_pretrain:
-        tokenized_ds = tokenized_ds.remove_columns(["context", "text"])
-        if not is_paraphrased:
-            tokenized_ds = tokenized_ds.remove_columns(["qas"])
 
     if set_format:
         tokenized_ds.set_format(type=set_format)
 
-    # validate_columns(tokenized_ds)
-
     return tokenized_ds
 
 
-def build_paraphrase_pretrain(samples: dict[str, list[str]]):
-    # use "context" to predict "variation" as vice versa
-    return {
-        "text": samples["context"] + samples["variation"],
-        "context": samples["variation"] + samples["context"],
-    }
-
-
-def build_intx_pretrain(sample: dict[str, str]):
-    return {"text": sample["context"] + sample["qas"]}
-
-
-def add_ctx_attn_mask(sample: dict[str, Any]) -> dict[str, Any]:
-    """
-    Add context attention mask to the sample.
-
-    Args:
-        sample: A dictionary containing 'ctx_ids'
-
-    Returns:
-        The sample with 'ctx_attn_mask' added
-    """
-    if "ctx_ids" in sample:
-        sample["ctx_attn_mask"] = [1] * len(sample["ctx_ids"])
-    return sample
-
-
 def get_labels_from_input_ids(sample: dict[str, Any]) -> dict[str, Any]:
     """
     Extract labels from input_ids and response_start.
@@ -994,7 +382,6 @@ def get_labels_from_input_ids(sample: dict[str, Any]) -> dict[str, Any]:
         A dictionary with 'labels' field added
     """
     labels = []
-    attention_mask = []
     for input_ids_i, (start_i, end_i) in zip(
         sample["input_ids"], sample["response_start_end"]
     ):
@@ -1007,18 +394,13 @@ def get_labels_from_input_ids(sample: dict[str, Any]) -> dict[str, Any]:
             + input_ids_i[start_i:]
             + [IGNORE_INDEX] * pad_len_right
         )
-        attention_mask.append([1] * len(input_ids_i))
 
     sample["labels"] = labels
-    sample["attention_mask"] = attention_mask
     return sample
 
 
 def get_sft_prompt_formatting_fn(
-    sft_mode: TRAINING_TASK,
     tokenizer: PreTrainedTokenizerBase,
-    for_logp: bool = False,
-    model: PreTrainedModel | None = None,
 ) -> Callable[[dict[str, Any]], dict[str, Any]]:
     """
     Get a function that formats examples for supervised fine-tuning.
@@ -1033,25 +415,15 @@ def get_sft_prompt_formatting_fn(
     Raises:
         NotImplementedError: If sft_mode is not COMPLETION or tokenizer has no chat template
     """
-    if sft_mode != TRAINING_TASK.COMPLETION:
-        raise NotImplementedError(
-            f"Training task {sft_mode} not supported. "
-            "Only completion is supported for now."
-        )
 
     if tokenizer.chat_template is None:
-        raise NotImplementedError(
-            "Only chat models + SFT are supported. "
-            "Training with pre-training data is not supported yet."
-        )
-
-    data_field = "data_logp" if for_logp else "data"
+        raise NotImplementedError("Only chat models are supported")
 
     @torch.inference_mode()
     def f_intx(samples):
         # flatten all the messages into a list
         # tokenize, the pack back correctly
-        messages_list = [x for x in samples[data_field]]
+        messages_list = [x for x in samples["messages_list"]]
 
         n_queries = [len(x) for x in messages_list]
         messages = concat_list(messages_list)
@@ -1067,77 +439,27 @@ def get_sft_prompt_formatting_fn(
             return_dict=True,
         )
         labels = []
-        logprobs_vals = []
-        logprobs_indices = []
         for tok_ids, masks in zip(tokens["input_ids"], tokens["assistant_masks"]):
             o = [id_ if mask else IGNORE_INDEX for id_, mask in zip(tok_ids, masks)]
             labels.append(o)
-            if for_logp:
-                # TODO: batch compute?
-                # save only the logits for the assistant tokens
-                assert sum(masks) > 0, "No assistant tokens found in the masks"
-
-                # + 1 to also get the logits of the token before the first label token
-                num_logits_to_keep = int(len(masks) - np.argmax(masks) + 1)
-
-                # [n_labels + 1]
-                logits = model(
-                    torch.tensor([tok_ids], device=model.device),
-                    logits_to_keep=num_logits_to_keep,
-                ).logits.squeeze(0)
-
-                # include the logits of the token before the first label token
-                pre_label_idx = np.argmax(masks) - 1
-                masks[pre_label_idx] = True
-                logits = logits[
-                    torch.tensor(
-                        masks[pre_label_idx:],
-                        device=logits.device,
-                        dtype=torch.bool,
-                    )
-                ]
-
-                logits = logits[..., :-1, :]  # remove the last token (eos)
-                logprobs = torch.nn.functional.log_softmax(logits.float(), dim=-1)
-                # ignore eos
-                vals, idx = torch.topk(logprobs, k=16, dim=-1)
-                logprobs_vals.append(vals.to(torch.float16))
-                logprobs_indices.append(idx.to(torch.int32))
 
         del tokens["assistant_masks"]
         tokens["labels"] = labels
-        per_ctx_tokens = []
+        per_ctx_tokens = {"input_ids": [], "labels": []}
         i = 0
         for n in n_queries:
-            if for_logp:
-                per_ctx_data = {
-                    "logprobs_vals": logprobs_vals[i : i + n],
-                    "logprobs_indices": logprobs_indices[i : i + n],
-                }
-                # per_ctx_data = {
-                #     "logprobs_vals": logprobs_vals[i : i + n],
-                #     "logprobs_indices": logprobs_indices[i : i + n],
-                # }
-            else:
-                per_ctx_data = {
-                    "input_ids": tokens["input_ids"][i : i + n],
-                    "attention_mask": tokens["attention_mask"][i : i + n],
-                    "labels": tokens["labels"][i : i + n],
-                }
-
-            per_ctx_tokens.append(per_ctx_data)
+            per_ctx_tokens["input_ids"].append(tokens["input_ids"][i : i + n])
+            per_ctx_tokens["labels"].append(tokens["labels"][i : i + n])
             i += n
 
-        return {data_field: per_ctx_tokens}
+        return per_ctx_tokens
 
-    # return f if not apply_chat_template_fn is not None else f_intx
     return f_intx
 
 
 def convert_ctx_prompt_response_to_messages(
     example: dict[str, Any],
     add_ctx_to_chat: bool,
-    for_logp: bool = False,
 ) -> dict[str, Any]:
     """
     Convert context/prompt/response format to chat messages format.
@@ -1163,25 +485,22 @@ def convert_ctx_prompt_response_to_messages(
     system_msg = ""
     if "system_message" in example:
         system_msg = example["system_message"].strip()
-        if for_logp and not system_msg:
-            system_msg = SELF_GEN_SYSTEM_MSG
 
-    data = []
+    messages_list = []
     for prompt, response in zip(example[prompt_field], example[res_field]):
         user_msg = prompt.strip()
         if add_ctx_to_chat:
             user_msg = example["context"].strip() + "\n\n" + user_msg
 
-        data.append(
+        messages_list.append(
             [
                 {"role": "system", "content": system_msg.strip()},
                 {"role": "user", "content": user_msg.strip()},
                 {"role": "assistant", "content": response},
             ]
         )
-    out_field = "data_logp" if for_logp else "data"
 
-    return {out_field: data}
+    return {"messages_list": messages_list}
 
 
 def split_too_long_qas(
@@ -1195,10 +514,8 @@ def split_too_long_qas(
     if max_qas_len < 0 and max_qas_per_sample < 0:
         return samples
     input_ids = samples["input_ids"]
-    attention_mask = samples["attention_mask"]
     labels = samples["labels"]
     ctx_ids = samples["ctx_ids"]
-    ctx_attn_mask = samples["ctx_attn_mask"]
     target_logprobs_vals = samples.get("logprobs_vals", None)
     target_logprobs_indices = samples.get("logprobs_indices", None)
 
@@ -1225,18 +542,14 @@ def split_too_long_qas(
     # Helper function to add a batch efficiently
     def add_batch(
         inp_ids_batch,
-        attn_batch,
         labels_batch,
         ctx_id,
-        ctx_attn,
         target_vals_batch=None,
         target_indices_batch=None,
     ):
         out["input_ids"].append(inp_ids_batch)
-        out["attention_mask"].append(attn_batch)
         out["labels"].append(labels_batch)
         out["ctx_ids"].append(ctx_id)
-        out["ctx_attn_mask"].append(ctx_attn)
         if has_target_logprobs:
             out["logprobs_vals"].append(target_vals_batch)
             out["logprobs_indices"].append(target_indices_batch)
@@ -1253,19 +566,16 @@ def split_too_long_qas(
 
         # Need to split this sample
         current_ctx_id = ctx_ids[i]
-        current_ctx_attn = ctx_attn_mask[i]
         new_qas_len = 0
         new_input_ids = []
-        new_attn_mask = []
         new_labels = []
         new_target_vals = [] if has_target_logprobs else None
         new_target_indices = [] if has_target_logprobs else None
 
-        sequences = zip(input_ids[i], attention_mask[i], labels[i])
+        sequences = zip(input_ids[i], labels[i])
         if has_target_logprobs:
             sequences = zip(
                 input_ids[i],
-                attention_mask[i],
                 labels[i],
                 target_logprobs_vals[i],
                 target_logprobs_indices[i],
@@ -1273,9 +583,9 @@ def split_too_long_qas(
 
         for seq_data in sequences:
             if has_target_logprobs:
-                inp_ids, attn_mask, label, target_vals, target_indices = seq_data
+                inp_ids, label, target_vals, target_indices = seq_data
             else:
-                inp_ids, attn_mask, label = seq_data
+                inp_ids, label = seq_data
                 target_vals, target_indices = None, None
 
             inp_len = len(inp_ids)
@@ -1293,7 +603,6 @@ def split_too_long_qas(
                 # Add to current batch
                 new_qas_len += inp_len
                 new_input_ids.append(inp_ids)
-                new_attn_mask.append(attn_mask)
                 new_labels.append(label)
                 if has_target_logprobs:
                     new_target_vals.append(target_vals)
@@ -1306,10 +615,8 @@ def split_too_long_qas(
                 if new_input_ids:  # Only add non-empty batches
                     add_batch(
                         new_input_ids,
-                        new_attn_mask,
                         new_labels,
                         current_ctx_id,
-                        current_ctx_attn,
                         new_target_vals,
                         new_target_indices,
                     )
@@ -1318,7 +625,6 @@ def split_too_long_qas(
                 # Start new batch with current sequence
                 new_qas_len = inp_len
                 new_input_ids = [inp_ids]
-                new_attn_mask = [attn_mask]
                 new_labels = [label]
                 if has_target_logprobs:
                     new_target_vals = [target_vals]
@@ -1328,10 +634,8 @@ def split_too_long_qas(
         if new_input_ids:
             add_batch(
                 new_input_ids,
-                new_attn_mask,
                 new_labels,
                 current_ctx_id,
-                current_ctx_attn,
                 new_target_vals,
                 new_target_indices,
             )
@@ -1347,212 +651,44 @@ def split_too_long_qas(
     return out
 
 
-def unpack_data(sample):
-    if "data" not in sample:
-        return sample
-    data_keys = ["data"]
-    if "data_logp" in sample:
-        data_keys.append("data_logp")
-    return {k: sample[dk][k] for dk in data_keys for k in sample[dk]}
-
-
 def unpack_data_eval(samples):
     # n_queries always == 1 for eval
     data = samples["data"]
-    out = dict(input_ids=[], attention_mask=[], labels=[])
+    out = dict(input_ids=[], labels=[])
     if "ctx_ids" in samples:
         out["ctx_ids"] = []
-        out["ctx_attn_mask"] = []
     for i, d in enumerate(data):
         for tokens in zip(
             d["input_ids"],
-            d["attention_mask"],
             d["labels"],
         ):
             if "ctx_ids" in samples:
                 out["ctx_ids"].append(samples["ctx_ids"][i])
-                out["ctx_attn_mask"].append(samples["ctx_attn_mask"][i])
             out["input_ids"].append(tokens[0])
-            out["attention_mask"].append(tokens[1])
             out["labels"].append(tokens[2])
     return out
 
 
-# adapted from https://github.com/huggingface/trl/issues/632#issuecomment-1972630547
-def get_assistant_start_end_indices(
-    messages: list[dict[str, str]],
-    tokenizer: PreTrainedTokenizerBase,
-    conversation_ids: list[int],
-) -> list[tuple[int, int]]:
+def squeeze_tokens(sample: dict[str, Any]) -> dict[str, Any]:
     """
-    Get the start and end indices of assistant messages in conversation text using tokenized messages.
+    Squeeze the input_ids and labels to remove any extra dimensions.
 
     Args:
-        messages: List of message dictionaries with 'role' and 'content' keys
-        tokenizer: Tokenizer used for tokenization
-        conversation_ids: Tokenized conversation
-
+        sample: A dictionary containing 'input_ids' and 'labels'
     Returns:
-        List of (start, end) index tuples for assistant messages in the tokenized conversation
+        A dictionary with squeezed 'input_ids' and 'labels'
     """
-    assistant_ranges = []
-    current_idx = 0
-
-    for idx, message in enumerate(messages):
-        # if message["role"] == "assistant":
-        # Tokenize the assistant message content
-        tokenized_message = tokenizer.encode(
-            message["content"], add_special_tokens=False
-        )
-        # Find the start index of the assistant message in the tokenized conversation
-        for i in range(current_idx, len(conversation_ids)):
-            if conversation_ids[i : i + len(tokenized_message)] == tokenized_message:
-                start_idx = i
-                end_idx = i + len(tokenized_message)
-                current_idx = end_idx
-                if message["role"] == "assistant":
-                    if idx == (len(messages) - 1):
-                        # include everything after the last assistant message
-                        assistant_ranges.append((start_idx, len(conversation_ids)))
-                    else:
-                        assistant_ranges.append((start_idx, end_idx))
-                break
-
-    return assistant_ranges
-
-
-# def get_masked_labels(
-#     conversation_ids: dict[str, list[Any]],
-#     assistant_ranges: list[tuple[int, int]],
-# ) -> Iterator[int]:
-#     """
-#     Generate masked labels for conversation, masking non-assistant tokens.
-#     NOTE: This will also includes extra tokens between assistant and user messages
-#     in multi-turn conversations.
-#     E.g., {assistant_msg} <|eot_id|><|start_header_id|>user<|end_header_id|> {user_msg}
-#     for Llama 3 models.
-
-#     Args:
-#         conversation_ids: Dictionary with tokenized conversation info
-#         assistant_ranges: List of (start, end) indices for assistant messages
-
-#     Yields:
-#         Token ID or IGNORE_INDEX for each position
-#     """
-#     for id_, (id_s, id_e) in list(
-#         zip(conversation_ids["input_ids"], conversation_ids["offset_mapping"])
-#     ):
-#         if any(id_s >= s and id_e <= e for s, e in assistant_ranges):
-#             yield id_
-#         else:
-#             yield IGNORE_INDEX
-
-
-# def tokenize_chat_messages(
-#     example: dict[str, Any],
-#     tokenizer: PreTrainedTokenizerBase,
-#     mask_assistant_inputs: bool = True,
-#     for_kl_loss: bool = False,
-#     tokenizer_kwargs: dict[str, Any] | None = None,
-# ) -> dict[str, list[int]]:
-#     """
-#     Tokenize chat messages and optionally mask non-assistant tokens.
-
-#     Args:
-#         example: Dictionary containing 'chat' and 'messages' keys
-#         tokenizer: Tokenizer to use
-#         mask_assistant_inputs: Whether to mask non-assistant tokens
-#         for_kl_loss: change the column names to "chat_ids" and "chat_attn_mask"
-#         tokenizer_kwargs: Additional arguments to pass to tokenizer
-
-#     Returns:
-#         Dictionary containing tokenized inputs and labels
-#     """
-#     # should be used only with chat models
-#     text = example["chat"]
-#     messages = example["messages"]
-#     n_response = len([m for m in messages if m["role"] == "assistant"])
-#     if n_response != 1:
-#         raise ValueError(f"Expected 1 assistant response. Got {n_response}.")
-#     conversation_ids = tokenizer(
-#         text,
-#         # return_offsets_mapping=mask_assistant_inputs,
-#         add_special_tokens=False,
-#         truncation=False,
-#         **(tokenizer_kwargs or {}),
-#     )
-#     # if tokenizer_kwargs and (
-#     #     len(conversation_ids["input_ids"]) >= tokenizer_kwargs["max_length"]
-#     # ):
-#     #     raise ValueError(
-#     #         f"Conversation length {len(conversation_ids['input_ids'])} exceeds max length {tokenizer_kwargs['max_length']}"
-#     #     )
-
-#     if mask_assistant_inputs:
-#         assistant_ranges = get_assistant_start_end_indices(
-#             messages, tokenizer, conversation_ids["input_ids"]
-#         )
-#         # labels = get_masked_labels(conversation_ids, assistant_ranges)
-#         labels = [
-#             id_ if any(s <= i < e for s, e in assistant_ranges) else IGNORE_INDEX
-#             for i, id_ in enumerate(conversation_ids["input_ids"])
-#         ]
-#         conversation_ids["labels"] = labels
-#         # del conversation_ids["offset_mapping"]
-#     else:
-#         conversation_ids["labels"] = conversation_ids["input_ids"]
-#     if for_kl_loss:
-#         conversation_ids["chat_ids"] = conversation_ids.pop("input_ids")
-#         conversation_ids["chat_attn_mask"] = conversation_ids.pop("attention_mask")
-#         conversation_ids["chat_labels"] = conversation_ids.pop("labels")
-#     return conversation_ids
-
-
-def tokenize_pretrain(
-    samples: dict[str, Any],
-    tokenizer: PreTrainedTokenizerBase,
-    tokenizer_kwargs: dict[str, Any] | None = None,
-    add_eos_token: bool = False,
-):
-    """
-    Tokenize text for pre-training tasks.
-
-    Args:
-        samples: Dictionary containing 'text' key
-        tokenizer: Tokenizer to use
-        tokenizer_kwargs: Additional arguments to pass to tokenizer
-
-    Returns:
-        Dictionary containing tokenized inputs and labels
-    """
-    add_bos_val = tokenizer.add_bos_token
-    add_eos_val = tokenizer.add_eos_token
-    tokenizer.add_bos_token = True
-    tokenizer.add_eos_token = add_eos_token
-    tokens = tokenizer(
-        samples["text"],
-        add_special_tokens=True,
-        padding=False,
-        truncation=False,
-        **(tokenizer_kwargs or {}),
-    )
-    # if tokenizer.bos_token_id is not None:
-    #     for i in range(len(tokens["input_ids"])):
-    #         # add bos
-    #         tokens["input_ids"][i] = [tokenizer.bos_token_id] + tokens["input_ids"][i]
-    #         tokens["attention_mask"][i] = [1] + tokens["attention_mask"][i]
-    tokenizer.add_bos_token = add_bos_val
-    tokenizer.add_eos_token = add_eos_val
-
-    tokens["labels"] = tokens["input_ids"]
-    return tokens
+    first_id = sample["input_ids"][0]
+    if check_is_iterable(first_id):
+        sample["input_ids"] = first_id
+    first_label = sample["labels"][0]
+    if check_is_iterable(first_label):
+        sample["labels"] = first_label
+    return sample
 
 
 def add_length_info(sample: dict[str, any], columns: list[str]) -> dict[str, any]:
     return {f"{k}_len": len(sample[k]) for k in columns}
-    # for col in columns:
-    #     samples[f"{col}_len"] = [len(t) for t in samples[col]]
-    # return samples
 
 
 def truncate_middle_if_too_long(
@@ -1575,8 +711,6 @@ def truncate_middle_if_too_long(
     # leave max_new_tokens for generation
     half = max_length // 2 - max_new_tokens_half
     for col in columns:
-        # if len(samples[col]) <= max_length:
-        #     return samples[col]
         t = sample[col]
         sample[col] = t[:half] + t[-half:] if len(t) > max_length else t
     return sample
@@ -1599,18 +733,14 @@ def tokenize_ctx_text(
             add_generation_prompt=True,
             padding=False,
             truncation=False,
-            add_special_tokens=False,
+            add_special_tokens=False,  # special tokens are already added by the chat template
             return_dict=True,
         )
-        # special tokens are already added by the chat template
-        # tokenized_text = tokenizer(txt, truncation=False, add_special_tokens=False)
     else:
         raise NotImplementedError("Only support chat models.")
-        # txt = samples["context"]
-        # tokenized_text = tokenizer(txt, truncation=False, add_special_tokens=False)
+
     ctx_ids = tokenized_text["input_ids"]
-    ctx_attn_mask = tokenized_text["attention_mask"]
-    return dict(ctx_ids=ctx_ids, ctx_attn_mask=ctx_attn_mask)
+    return dict(ctx_ids=ctx_ids)
 
 
 def pack(
@@ -1626,12 +756,10 @@ def pack(
         max_packed_ctx_len=max_packed_ctx_len,
         max_packed_size=max_packed_size,
     )
-    ds_hash = hashlib.sha256(
-        (ds._fingerprint + json.dumps(kwargs)).encode()
-    ).hexdigest()
-    ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}"
+    ds_hash = sha256((ds._fingerprint + json.dumps(kwargs)).encode()).hexdigest()
+    ds_path = f"{TRANSFORMED_DATA_DIR}/packed_{ds_hash}"
     logger.info(
-        f"Packing dataset {ds_hash} with max_packed_inp_len={max_packed_inp_len} and max_packed_ctx_len={max_packed_ctx_len}"
+        f"Packing ds {ds_hash} with {max_packed_inp_len=} and {max_packed_ctx_len=}"
     )
     if path.exists(ds_path) and is_caching_enabled():
         logger.info(f"Loading a cached packed dataset for {ds_path}")
@@ -1643,105 +771,15 @@ def pack(
                 "max_packed_inp_len": max_packed_inp_len,
                 "max_packed_ctx_len": max_packed_ctx_len,
                 "max_packed_size": max_packed_size,
+                "metadata_path": f"{ds_path}/packing_metadata.json",
             },
             batched=True,
             batch_size=250_000,
             num_proc=num_proc,
             remove_columns=ds.column_names,
         )
-
+        # TODO: saving here is not space efficient at all...
+        # the contexts are being duplicated for each datapoint when splitting QAs
         ds.save_to_disk(ds_path, num_proc=num_proc)
+
     return ds
-
-
-# def add_logits(model, ds):
-#     # compute logits for prepared tokenized ds
-#     # must handle packed and flattened dataset
-#     # also cache!
-#     # let's do this row by row
-#     # slower for non-packed ds
-#     input = torch.tensor([ds['input_ids']], device=model.device)
-#     labels = torch.tensor([ds['labels']], device=model.device)
-#     with torch.no_grad():
-#         logits = model(input).logits
-
-#     top_logits, top_logits_positions = torch.topk(logits, args.max_logits_to_store, dim=-1)
-
-#     # Convert logits to specified precision
-#     dtype_map = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}
-#     target_dtype = dtype_map[args.logit_precision]
-#     top_logits = top_logits.to(target_dtype)
-
-#     tokens_to_mask = tok.apply_chat_template(
-#         [messages[0]],  # Just the user message
-#         tokenize=True,
-#         add_special_tokens=False,
-#         truncation=False,
-#         padding=False,
-#         add_generation_prompt=True,
-#         return_dict=True,
-#     )
-#     number_of_tokens_to_mask = len(tokens_to_mask['input_ids'])
-
-#     # Create mask: False for prompt tokens, True for response tokens
-#     responses_mask = torch.zeros(len(all_tokens['input_ids']), dtype=torch.bool)
-#     responses_mask[number_of_tokens_to_mask:] = True
-
-#     responses_masks.append(responses_mask)
-#     all_logits_values.append(top_logits.cpu())
-#     all_logits_positions.append(top_logits_positions.cpu())
-
-#     logits = model(input).logits
-
-#     top_logits, top_logits_positions = torch.topk(logits, args.max_logits_to_store, dim=-1)
-
-#     # Convert logits to specified precision
-#     dtype_map = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}
-#     target_dtype = dtype_map[args.logit_precision]
-#     top_logits = top_logits.to(target_dtype)
-
-#     tokens_to_mask = tok.apply_chat_template(
-#         [messages[0]],  # Just the user message
-#         tokenize=True,
-#         add_special_tokens=False,
-#         truncation=False,
-#         padding=False,
-#         add_generation_prompt=True,
-#         return_dict=True,
-#     )
-#     number_of_tokens_to_mask = len(tokens_to_mask['input_ids'])
-
-#     # Create mask: False for prompt tokens, True for response tokens
-#     responses_mask = torch.zeros(len(all_tokens['input_ids']), dtype=torch.bool)
-#     responses_mask[number_of_tokens_to_mask:] = True
-
-#     responses_masks.append(responses_mask)
-#     all_logits_values.append(top_logits.cpu())
-#     all_logits_positions.append(top_logits_positions.cpu())
-
-
-if __name__ == "__main__":
-    from transformers import AutoTokenizer
-
-    # model_name = "meta-llama/Llama-3.2-1B-Instruct"
-    # model_name = "meta-llama/Llama-2-7b-chat-hf"
-    model_name = "google/gemma-2-2b-it"
-    messages = [
-        {"role": "user", "content": "Hello!"},
-        {"role": "assistant", "content": "Hello!"},
-        # {"role": "user", "content": "Not too bad"},
-        # {"role": "assistant", "content": "Cooooooool"},
-    ]
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    chat = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=False, add_special_tokens=False
-    )
-    print(chat)
-    model_inputs = tokenize_chat_messages(
-        {"chat": chat, "messages": messages},
-        tokenizer,
-        # for_kl_loss=True,
-    )
-
-    print(tokenizer(chat, add_special_tokens=False))
-    print(model_inputs)
diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py
index 699e620..5494f2a 100644
--- a/src/ctx_to_lora/eval_utils.py
+++ b/src/ctx_to_lora/eval_utils.py
@@ -680,6 +680,7 @@ def evaluate(
             train=False,
             base_model_kwargs=model_kwargs,
             use_flash_attn=True,
+            use_sequence_packing=False,  # for generation
         )
     else:
         model = get_model(
@@ -724,8 +725,6 @@ def evaluate(
         if ctx_tokenizer.pad_token_id is None:
             ctx_tokenizer.pad_token_id = ctx_tokenizer.eos_token_id
 
-    tokenizer_kwargs = {}
-    ctx_tokenizer_kwargs = {}
     add_ctx_to_chat = (
         not isinstance(model, ModulatedPretrainedModel) and not args.remove_context
     )
@@ -741,13 +740,9 @@ def evaluate(
         max_qas_per_sample=1,
         base_model_max_len=model.base_model.config.max_position_embeddings,
         tokenizer=tokenizer,
-        tokenizer_kwargs=tokenizer_kwargs,
         ctx_model_max_len=ctx_model_max_len,
         ctx_tokenizer=ctx_tokenizer,
-        ctx_tokenizer_kwargs=ctx_tokenizer_kwargs,
         add_ctx_to_chat=add_ctx_to_chat,
-        add_repeat_prompt=False,
-        add_negative_prompt=False,
         use_kl_loss=False,
         set_format="pt",
     )
diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py
index 9119500..1e32560 100644
--- a/src/ctx_to_lora/modeling/hypernet.py
+++ b/src/ctx_to_lora/modeling/hypernet.py
@@ -591,19 +591,18 @@ class ModulatedPretrainedModel(nn.Module):
         base_model: PeftModel,
         hypernet_config: HypernetConfig,
         ctx_encoder_args: CtxEncoderArguments,
-        # use_kl_loss is only used for training
-        use_kl_loss: bool = False,
-        use_sequence_packing: bool = False,
         use_base_input_as_ctx: bool = False,
+        # need non-packed inputs for generation
+        use_sequence_packing: bool = True,
     ):
+        assert not use_base_input_as_ctx
         super().__init__()
         self.device = base_model.device
         self.peft_config = base_model.peft_config["default"]
         self.hypernet_config = hypernet_config
         self.ctx_encoder_args = ctx_encoder_args
-        self.use_kl_loss = use_kl_loss
-        self.use_sequence_packing = use_sequence_packing
         self.use_base_input_as_ctx = use_base_input_as_ctx
+        self.use_sequence_packing = use_sequence_packing
         self.lora_id = 1
         self.active_adapters = []
 
@@ -618,6 +617,7 @@ class ModulatedPretrainedModel(nn.Module):
         train: bool = True,
         base_model_kwargs: dict = None,
         use_flash_attn: bool = True,
+        **kwargs: Any,
     ):
         lora_config = state_dict["hypernet_config"].lora_config
         print(f"lora_config: {lora_config}")
@@ -634,7 +634,7 @@ class ModulatedPretrainedModel(nn.Module):
         if getattr(hypernet_config, "num_pre_head_layers", None) is None:
             hypernet_config.num_pre_head_layers = 4
         ctx_encoder_args = state_dict["ctx_encoder_args"]
-        model = cls(base_model, hypernet_config, ctx_encoder_args)
+        model = cls(base_model, hypernet_config, ctx_encoder_args, **kwargs)
         model.load_state_dict(state_dict)
         return model
 
@@ -650,7 +650,7 @@ class ModulatedPretrainedModel(nn.Module):
 
         layers = get_layers(self.base_model)
         lora_forward_fn = (
-            lora_forward if not self.use_sequence_packing else lora_forward_packed
+            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):
@@ -909,44 +909,57 @@ 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,
-    #     *model_inputs_args: Any,
-    #     **model_inputs_kwargs: dict[str, Any],
-    # ):
-    #     # TODO: make this persistent
-    #     generated_loras, _ = self.generate_weights(
-    #         ctx_ids, ctx_attn_mask, ctx_position_ids
-    #     )
-    #     # sum loras
-    #     for lora_at_module in generated_loras.values():
-    #         if len(lora_at_module["A"]) > 1:
-    #             print("Multiple LoRAs are generated, summing them up")
-    #         lora_at_module["A"] = torch.mean(lora_at_module["A"], dim=0, keepdim=True)
-    #         lora_at_module["B"] = torch.mean(lora_at_module["B"], dim=0, keepdim=True)
+    @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_queries: Integer[Tensor, "n_ctx"] | None = None,
+        *model_inputs_args: Any,
+        **model_inputs_kwargs: dict[str, Any],
+    ):
+        # TODO: make this persistent
+        generated_loras, _ = self.generate_weights(
+            ctx_ids, ctx_attn_mask, ctx_position_ids
+        )
+        # sum loras
+        for lora_at_module in generated_loras.values():
+            if len(lora_at_module["A"]) > 1:
+                print("Multiple LoRAs are generated, summing them up")
+            lora_at_module["A"] = torch.mean(lora_at_module["A"], dim=0, keepdim=True)
+            lora_at_module["B"] = torch.mean(lora_at_module["B"], dim=0, keepdim=True)
 
-    #     position_ids = (
-    #         model_inputs_kwargs["position_ids"]
-    #         if "position_ids" in model_inputs_kwargs
-    #         else None
-    #     )
+        # 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:
+            if ctx_position_ids is None:
+                n_queries = torch.ones(
+                    ctx_ids.shape[0], dtype=torch.int32, device=self.device
+                )
+            else:
+                # quite redundant (we do cu_seqlens many places)
+                # TODO: compute cu_seqlens here and propagate that
+                n_queries = torch.ones(
+                    (ctx_position_ids == 0).sum(), dtype=torch.int32, device=self.device
+                )
 
-    #     # apply loras
-    #     with apply_generated_loras(
-    #         self.base_model,
-    #         generated_loras,
-    #         self.hypernet.layer_indices,
-    #         position_ids,
-    #         self.training,
-    #     ):
-    #         model_outputs = self.base_model.generate(
-    #             *model_inputs_args, **model_inputs_kwargs
-    #         )
-    #     return model_outputs
+        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
 
     @torch.inference_mode()
     def generate(
@@ -1033,93 +1046,6 @@ class ModulatedPretrainedModel(nn.Module):
         return model_outputs
 
 
-# class ModulatedModelWithSharedInput(nn.Module):
-#     def __init__(
-#         self,
-#         modulated_model: ModulatedPretrainedModel,
-#         base_tokenizer: PreTrainedTokenizerBase,
-#         ctx_tokenizer: PreTrainedTokenizerBase,
-#         ctx_end_predicate: str | None = None,
-#         remove_ctx_from_base_input: bool = False,
-#     ):
-#         super().__init__()
-#         self.modulated_model = modulated_model
-#         self.base_tokenizer = base_tokenizer
-#         self.ctx_tokenizer = ctx_tokenizer
-#         self.register_module("modulated_model", self.modulated_model)
-#         self.ctx_end_predicate = ctx_end_predicate
-#         self.remove_ctx_from_base_input = remove_ctx_from_base_input
-
-#     # delegate to self.modulated_model
-#     @property
-#     def config(self):
-#         return self.modulated_model.config
-
-#     @property
-#     def device(self):
-#         return self.modulated_model.device
-
-#     def tie_weights(self):
-#         self.modulated_model.base_model.tie_weights()
-
-#     def state_dict(self, *args, **kwargs):
-#         return self.modulated_model.state_dict(*args, **kwargs)
-
-#     def forward(self, *args, **kwargs):
-#         input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0]
-#         input_txts = self.base_tokenizer.batch_decode(
-#             input_ids,
-#             skip_special_tokens=True,
-#             clean_up_tokenization_spaces=True,
-#         )
-#         ctx_txts = deepcopy(input_txts)
-#         if self.ctx_end_predicate:
-#             ctx_txts = [
-#                 txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts
-#             ]
-#         ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to(
-#             self.device
-#         )
-#         ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask
-#         if self.remove_ctx_from_base_input:
-#             raise NotImplementedError("Not implemented")
-#             # input_txts = [
-#             #     txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts
-#             # ]
-#             # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True)
-#             # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask
-#             # kwargs["input_ids"] = input_ids
-#             # kwargs["attention_mask"] = attention_mask
-#         return self.modulated_model(ctx_ids, ctx_attn_mask, *args, **kwargs)
-
-#     def generate(self, *args, **kwargs):
-#         input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0]
-#         input_txts = self.base_tokenizer.batch_decode(
-#             input_ids,
-#             skip_special_tokens=True,
-#             clean_up_tokenization_spaces=True,
-#         )
-#         ctx_txts = deepcopy(input_txts)
-#         if self.ctx_end_predicate:
-#             ctx_txts = [
-#                 txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts
-#             ]
-#         ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to(
-#             self.device
-#         )
-#         ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask
-#         if self.remove_ctx_from_base_input:
-#             raise NotImplementedError("Not implemented")
-#             # input_txts = [
-#             #     txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts
-#             # ]
-#             # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True)
-#             # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask
-#             # kwargs["input_ids"] = input_ids
-#             # kwargs["attention_mask"] = attention_mask
-#         return self.modulated_model.generate(ctx_ids, ctx_attn_mask, *args, **kwargs)
-
-
 @contextmanager
 def apply_generated_layernorm(
     base_model: nn.Module,
diff --git a/src/ctx_to_lora/modeling/lora_layer.py b/src/ctx_to_lora/modeling/lora_layer.py
index 94b7957..98ed958 100644
--- a/src/ctx_to_lora/modeling/lora_layer.py
+++ b/src/ctx_to_lora/modeling/lora_layer.py
@@ -59,7 +59,6 @@ def lora_forward_packed(
     # repeated_A = A.repeat_interleave(seq_lens, dim=0, output_size=seq_lens.sum())
     # # [tot_seq_len, d_out, r]
     # repeated_B = B.repeat_interleave(seq_lens, dim=0, output_size=seq_lens.sum())
-    # TODO: is ther a better way to implement this?
     repeated_A = A.repeat_interleave(n_qs, dim=0, output_size=tot_q)
     repeated_A = repeated_A.repeat_interleave(seq_lens, dim=0, output_size=tot_len)
 
diff --git a/src/ctx_to_lora/trainer.py b/src/ctx_to_lora/trainer.py
index b2593fb..cfe7425 100644
--- a/src/ctx_to_lora/trainer.py
+++ b/src/ctx_to_lora/trainer.py
@@ -13,25 +13,6 @@ from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
 logger = logging.getLogger()
 
 
-def add_pre_label_token(labels, is_packed):
-    """
-    Add a token before the first label token.
-    1 is used as a placeholder for the first label token.
-    We compute the loss against the teacher logits anyway.
-    """
-    if is_packed:
-        non_label_mask = torch.where(labels == -100, True, False)
-        diff = torch.diff(labels, append=torch.tensor([[-100]], device=labels.device))
-        pre_label_pos = torch.where(non_label_mask * (diff != 0))
-        labels[pre_label_pos] = 1
-    else:
-        bsz = labels.shape[0]
-        first_label_pos = torch.argmax((labels != -100).float(), axis=-1)
-        labels[torch.arange(bsz), first_label_pos - 1] = 1
-    return labels
-
-
-# TODO: test
 class DistillationTrainer(Trainer):
     def __init__(self, *args, **kwargs):
         self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)
@@ -46,71 +27,61 @@ class DistillationTrainer(Trainer):
             if num_items_in_batch is not None:
                 loss_kwargs["num_items_in_batch"] = num_items_in_batch
             inputs = {**inputs, **loss_kwargs}
-
-        labels = inputs.pop("labels", None)  # remove labels if present
-        # NOTE: have to account that we save +1 # token before the first label token
-        is_packed = "position_ids" in inputs
-        # labels = add_pre_label_token(labels, is_packed)
-        label_pos = torch.where(labels != -100)
+        use_kl = "logprobs_vals" in inputs
+        # TODO: handle when val data doesnt have logprobs
+        if use_kl:
+            labels = inputs.pop("labels", None)
+            label_pos = torch.where(labels != -100)
         # TODO: add efficient logits indexing?
         # https://github.com/huggingface/trl/pull/2773
         # + use `num_logits_to_keep` to limit the number of logits
         outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True)
-        target_logp = inputs.pop("logprobs_vals")
-        indices = inputs.pop("logprobs_indices")
 
-        # packed
-        if isinstance(target_logp, torch.Tensor):
-            target_logp = target_logp.squeeze(0)
-            indices = indices.squeeze(0)
-        # non-packed max_qas_per_sample=1
-        elif isinstance(target_logp[0], torch.Tensor):
-            target_logp = torch.cat([x[0] for x in target_logp], dim=0)
-            indices = torch.cat([x[0] for x in indices], dim=0)
-        # non-packed max_qas_per_sample>1
+        if not use_kl:
+            loss = outputs["loss"]
+
         else:
-            # convert to tensor
-            raise NotImplementedError(
-                "Non-packed data with max_qas_per_sample > 1 is not supported yet."
+            target_logp = inputs.pop("logprobs_vals")
+            indices = inputs.pop("logprobs_indices")
+
+            # packed
+            if isinstance(target_logp, torch.Tensor):
+                target_logp = target_logp.squeeze(0)
+                indices = indices.squeeze(0)
+            # non-packed max_qas_per_sample=1
+            elif isinstance(target_logp[0], torch.Tensor):
+                target_logp = torch.cat([x[0] for x in target_logp], dim=0)
+                indices = torch.cat([x[0] for x in indices], dim=0)
+            # non-packed max_qas_per_sample>1
+            else:
+                # convert to tensor
+                raise NotImplementedError(
+                    "Non-packed data with max_qas_per_sample > 1 is not supported yet."
+                )
+
+            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]=}"
             )
 
-        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]=}"
-        )
+            # TODO: per-ctx average loss?
+            ##### KL loss
+            outputs_logits = outputs.logits[
+                label_pos[0], label_pos[1] - 1
+            ]  # shift back 1
 
-        # TODO: per-ctx average loss?
-        ##### KL loss
-        outputs_logits = outputs.logits[label_pos[0], label_pos[1] - 1]  # shift back 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"
 
-        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"
-
-        # loss = torch.nn.functional.kl_div(
-        #     torch.nn.functional.log_softmax(
-        #         outputs_logits / self.distillation_temperature, dim=-1
-        #     ),
-        #     torch.nn.functional.softmax(
-        #         teacher_logits / self.distillation_temperature, dim=-1
-        #     ),
-        #     reduction=reduction,
-        #     log_target=False,
-        # ) * (self.distillation_temperature**2)
-        # loss = loss / num_items_in_batch if num_items_in_batch is not None else loss
-
-        # p = torch.nn.functional.softmax(
-        #     teacher_logits / self.distillation_temperature, dim=-1
-        # )
-        p = teacher_logp.exp()
-        logq = nn.functional.log_softmax(outputs_logits, dim=-1)
-
-        loss = -torch.sum(p * logq, dim=-1) * (self.distillation_temperature**2)
-        if reduction == "batchmean":
-            loss = loss.mean()
-        elif reduction == "sum" and num_items_in_batch is not None:
-            loss = loss.sum() / num_items_in_batch
-        #####
+            p = teacher_logp.exp()
+            logq = nn.functional.log_softmax(outputs_logits, dim=-1)
+            loss = -torch.sum(p * logq, dim=-1) * (self.distillation_temperature**2)
+            if reduction == "batchmean":
+                loss = loss.mean()
+            elif reduction == "sum" and num_items_in_batch is not None:
+                loss = loss.sum() / num_items_in_batch
+            #####
 
         ##### unpack gen lora dict and compute regularization loss
         l1_norm = 0
@@ -121,7 +92,8 @@ class DistillationTrainer(Trainer):
             self.args.logging_strategy == IntervalStrategy.STEPS
             and self.state.global_step % self.state.logging_steps == 0
         ):
-            self.log({"ce_loss": loss.item(), "gen_lora_l1_norm": l1_norm.item()})
+            loss_field = "kl_loss" if use_kl else "ce_loss"
+            self.log({loss_field: loss.item(), "gen_lora_l1_norm": l1_norm.item()})
         loss += self.gen_lora_l1_reg_coef * l1_norm
         #####
 
diff --git a/webui/app.py b/webui/app.py
index 87cbdf8..16a67dc 100644
--- a/webui/app.py
+++ b/webui/app.py
@@ -502,6 +502,7 @@ def load_checkpoint():
             state_dict,
             train=False,
             use_flash_attn=True,
+            use_sequence_packing=False,  # for generation
         )
         modulated_model = modulated_model.to(device).to(torch.bfloat16)
         modulated_model.eval()
@@ -547,7 +548,7 @@ def process_multiple_contexts(contexts, ctx_tokenizer):
     ctx_ids = [
         torch.tensor(ctx_id, dtype=torch.long, device=device) for ctx_id in ctx_ids
     ]
-    ctx_attn_mask = tokenized_contexts["ctx_attn_mask"]
+    ctx_attn_mask = [torch.ones_like(ids) for ids in ctx_ids]
     ctx_attn_mask = [
         torch.tensor(ctx_attn_mask, dtype=torch.long, device=device)
         for ctx_attn_mask in ctx_attn_mask