icml rebuttal

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -7,6 +7,7 @@ results/
 datasets/
 llm-comparator/
 trained_d2l/
+trained_t2l/
 # Ignore data files but not scripts
 /data/processed_datasets/
 /data/raw_datasets/

README.md

@@ -73,7 +73,6 @@ print(tokenizer.decode(outputs[0]))
 uv run demo/app.py
 ```
 <div align="center">
-    <h3>Video Demo</h3>
     <video src="https://github.com/user-attachments/assets/f46325b1-d040-48f0-8b3e-deba0e6218ff" controls autoplay muted playsinline preload="metadata" width="900"></video>
 </div>
 
@@ -86,6 +85,39 @@ To run any of the following scripts, use `uv run $PATH_TO_SCRIPT` from the root
 | [Main experiment](scripts/main_exp/) | `scripts/main_exp/0-download_data.sh` | `scripts/main_exp/1-train.sh` | `scripts/main_exp/eval/*.sh` | Downloading data is fastest; regenerate only if you need fresh synthetic data. Evaluation scripts reproduce the main paper metrics. |
 | [NIAH](scripts/niah/)                | `scripts/niah/0-gen_data.sh`          | `scripts/niah/1-train.sh`     | `scripts/niah/2-eval.sh`     | Run the scripts in order; data generation only needs to happen once                                                                 |
 
+`scripts/main_exp/eval/clipper.sh` adds the CLIPPER long-context benchmark to the eval pipeline via the public `chtmp223/CLIPPER` Hugging Face dataset.
+
+`scripts/main_exp/eval/rag.sh` runs a lightweight BM25-style RAG baseline over each example's `context` field. It keeps dataset-side context chunking disabled and performs retrieval chunking inside the eval-time wrapper.
+
+```bash
+WANDB_MODE=disabled uv run run_eval.py \
+  --model_name_or_path google/gemma-2-2b-it \
+  --datasets squad drop ropes \
+  --split test \
+  --eval_batch_size_gen 1 \
+  --use_rag \
+  --rag_chunk_size 256 \
+  --rag_chunk_overlap 64 \
+  --rag_top_k 4 \
+  --rag_max_retrieved_tokens 1536
+```
+
+`scripts/main_exp/eval/d2l_rag.sh` runs a hybrid mode where Doc-to-LoRA internalizes the full document while the same document is also queried with BM25-style retrieval to build a smaller prompt-side evidence block.
+
+```bash
+WANDB_MODE=disabled uv run run_eval.py \
+  --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin \
+  --datasets squad drop ropes \
+  --split test \
+  --eval_batch_size_gen 1 \
+  --use_hybrid_rag \
+  --rag_chunk_size 256 \
+  --rag_chunk_overlap 64 \
+  --rag_top_k 4 \
+  --rag_max_retrieved_tokens 1536
+```
+
+Generated JSONL outputs include compact retrieval metadata under `rag_selected_chunks` together with prompt/context token counts for debugging.
 
 ### 🔬 Self-Generated Data Viewer
 After downloading/generating the data, we can see samples of the data using this script.

data/bitter_lesson.txt

@@ -0,0 +1,22 @@
+# The Bitter Lesson
+
+## Rich Sutton
+
+### March 13, 2019
+
+The biggest lesson that can be read from 70 years of AI research is that general methods that leverage computation are ultimately the most effective, and by a large margin. The ultimate reason for this is Moore's law, or rather its generalization of continued exponentially falling cost per unit of computation. Most AI research has been conducted as if the computation available to the agent were constant (in which case leveraging human knowledge would be one of the only ways to improve performance) but, over a slightly longer time than a typical research project, massively more computation inevitably becomes available. Seeking an improvement that makes a difference in the shorter term, researchers seek to leverage their human knowledge of the domain, but the only thing that matters in the long run is the leveraging of computation. These two need not run counter to each other, but in practice they tend to. Time spent on one is time not spent on the other. There are psychological commitments to investment in one approach or the other. And the human-knowledge approach tends to complicate methods in ways that make them less suited to taking advantage of general methods leveraging computation.  There were many examples of AI researchers' belated learning of this bitter lesson, and it is instructive to review some of the most prominent.
+
+In computer chess, the methods that defeated the world champion, Kasparov, in 1997, were based on massive, deep search. At the time, this was looked upon with dismay by the majority of computer-chess researchers who had pursued methods that leveraged human understanding of the special structure of chess. When a simpler, search-based approach with special hardware and software proved vastly more effective, these human-knowledge-based chess researchers were not good losers. They said that ``brute force" search may have won this time, but it was not a general strategy, and anyway it was not how people played chess. These researchers wanted methods based on human input to win and were disappointed when they did not.
+
+A similar pattern of research progress was seen in computer Go, only delayed by a further 20 years. Enormous initial efforts went into avoiding search by taking advantage of human knowledge, or of the special features of the game, but all those efforts proved irrelevant, or worse, once search was applied effectively at scale. Also important was the use of learning by self play to learn a value function (as it was in many other games and even in chess, although learning did not play a big role in the 1997 program that first beat a world champion). Learning by self play, and learning in general, is like search in that it enables massive computation to be brought to bear. Search and learning are the two most important classes of techniques for utilizing massive amounts of computation in AI research. In computer Go, as in computer chess, researchers' initial effort was directed towards utilizing human understanding (so that less search was needed) and only much later was much greater success had by embracing search and learning.
+
+In speech recognition, there was an early competition, sponsored by DARPA, in the 1970s. Entrants included a host of special methods that took advantage of human knowledge---knowledge of words, of phonemes, of the human vocal tract, etc. On the other side were newer methods that were more statistical in nature and did much more computation, based on hidden Markov models (HMMs). Again, the statistical methods won out over the human-knowledge-based methods. This led to a major change in all of natural language processing, gradually over decades, where statistics and computation came to dominate the field. The recent rise of deep learning in speech recognition is the most recent step in this consistent direction. Deep learning methods rely even less on human knowledge, and use even more computation, together with learning on huge training sets, to produce dramatically better speech recognition systems. As in the games, researchers always tried to make systems that worked the way the researchers thought their own minds worked---they tried to put that knowledge in their systems---but it proved ultimately counterproductive, and a colossal waste of researcher's time, when, through Moore's law, massive computation became available and a means was found to put it to good use.
+
+In computer vision, there has been a similar pattern. Early methods conceived of vision as searching for edges, or generalized cylinders, or in terms of SIFT features. But today all this is discarded. Modern deep-learning neural networks use only the notions of convolution and certain kinds of invariances, and perform much better.
+
+This is a big lesson. As a field, we still have not thoroughly learned it, as we are continuing to make the same kind of mistakes. To see this, and to effectively resist it, we have to understand the appeal of these mistakes. We have to learn the bitter lesson that building in how we think we think does not work in the long run. The bitter lesson is based on the historical observations that 1) AI researchers have often tried to build knowledge into their agents, 2) this always helps in the short term, and is personally satisfying to the researcher, but 3) in the long run it plateaus and even inhibits further progress, and 4) breakthrough progress eventually arrives by an opposing approach based on scaling computation by search and learning. The eventual success is tinged with bitterness, and often incompletely digested, because it is success over a favored, human-centric approach.
+
+One thing that should be learned from the bitter lesson is the great power of general purpose methods, of methods that continue to scale with increased computation even as the available computation becomes very great. The two methods that seem to scale arbitrarily in this way are search and learning.
+
+The second general point to be learned from the bitter lesson is that the actual contents of minds are tremendously, irredeemably complex; we should stop trying to find simple ways to think about the contents of minds, such as simple ways to think about space, objects, multiple agents, or symmetries. All these are part of the arbitrary, intrinsically-complex, outside world. They are not what should be built in, as their complexity is endless; instead we should build in only the meta-methods that can find and capture this arbitrary complexity. Essential to these methods is that they can find good approximations, but the search for them should be by our methods, not by us. We want AI agents that can discover like we can, not which contain what we have discovered. Building in our discoveries only makes it harder to see how the discovering process can be done.
+

examples/python_api.py

@@ -5,7 +5,7 @@ import torch
 from ctx_to_lora.model_loading import get_tokenizer
 from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
 
-checkpoint_path = "../../ctx-to-lora/train_outputs/runs/Sep29_14-42-46_slurm0-a3nodeset-9_88483_1e7bb34e/checkpoint-40000/pytorch_model.bin"
+checkpoint_path = "trained_d2l/gemma_demo/checkpoint-80000/pytorch_model.bin"
 state_dict = torch.load(checkpoint_path, weights_only=False)
 model = ModulatedPretrainedModel.from_state_dict(
     state_dict, train=False, use_sequence_packing=False
@@ -14,8 +14,13 @@ model.reset()
 
 tokenizer = get_tokenizer(model.base_model.name_or_path)
 
-doc = open("data/sakana_wiki.txt").read()
-chat = [{"role": "user", "content": "Summarize what Sakana AI does."}]
+doc = open("data/bitter_lesson.txt").read()
+chat = [
+    {
+        "role": "user",
+        "content": "What is the bitter lesson in AI research?",
+    }
+]
 chat_ids = tokenizer.apply_chat_template(
     chat,
     add_special_tokens=False,
@@ -36,8 +41,8 @@ outputs = model.generate(input_ids=chat_ids, max_new_tokens=256)
 print(tokenizer.decode(outputs[0]))
 
 
-# remove internalized info
-model.reset()
+# # remove internalized info
+# model.reset()
 
-outputs = model.generate(input_ids=chat_ids, max_new_tokens=256)
-print(tokenizer.decode(outputs[0]))
+# outputs = model.generate(input_ids=chat_ids, max_new_tokens=256)
+# print(tokenizer.decode(outputs[0]))

pyproject.toml

@@ -39,6 +39,7 @@ dependencies = [
     "google-cloud-storage>=3.2.0",
     "wonderwords>=2.2.0",
     "llmlingua>=0.2.2",
+    "python-dateutil>=2.9.0.post0",
 ]
 
 [build-system]

run_eval.py

@@ -73,6 +73,12 @@ if __name__ == "__main__":
         default=-1,
         help="Maximum length of context chunk for evaluation",
     )
+    parser.add_argument(
+        "--ctx_chunk_overlap",
+        type=int,
+        default=0,
+        help="Number of overlapping context tokens between adjacent eval chunks",
+    )
     parser.add_argument(
         "--max_new_tokens",
         type=int,
@@ -133,6 +139,47 @@ if __name__ == "__main__":
         action="store_true",
         help="Use Text-to-LoRA model for evaluation",
     )
+    parser.add_argument(
+        "--use_rag",
+        action="store_true",
+        help="Use retrieval-augmented generation baseline over the dataset context field",
+    )
+    parser.add_argument(
+        "--use_hybrid_rag",
+        action="store_true",
+        help="Use Doc-to-LoRA internalization plus RAG-augmented prompt inputs during evaluation",
+    )
+    parser.add_argument(
+        "--rag_chunk_size",
+        type=int,
+        default=256,
+        help="Token chunk size for the RAG baseline retriever",
+    )
+    parser.add_argument(
+        "--rag_chunk_overlap",
+        type=int,
+        default=64,
+        help="Token overlap between adjacent RAG baseline chunks",
+    )
+    parser.add_argument(
+        "--rag_top_k",
+        type=int,
+        default=4,
+        help="Maximum number of retrieved chunks to include in the RAG prompt",
+    )
+    parser.add_argument(
+        "--rag_max_retrieved_tokens",
+        type=int,
+        default=1536,
+        help="Maximum total retrieved-context tokens to include in the RAG prompt",
+    )
+    parser.add_argument(
+        "--rag_retrieval_mode",
+        type=str,
+        choices=["bm25"],
+        default="bm25",
+        help="Retriever backend for the RAG baseline",
+    )
     parser.add_argument(
         "--add_ctx_to_input",
         action="store_true",
@@ -163,6 +210,14 @@ if __name__ == "__main__":
         action="store_true",
         help="Use generative adapter for evaluation",
     )
+    parser.add_argument(
+        "--log_dense_lora_magnitude",
+        action="store_true",
+        help=(
+            "Compute dense LoRA-update magnitude stats during evaluation and save "
+            "them alongside generated text."
+        ),
+    )
 
     cli_args = vars(parser.parse_args())
 

scripts/main_exp/eval/base_model.sh

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

scripts/main_exp/eval/base_model_long_ctx.sh

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

scripts/main_exp/eval/cd.sh

@@ -3,4 +3,4 @@ WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it
 
 
 # longbench
-WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/qasper_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=1
+WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/qasper_e longbench/gov_report_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp --cd_use_gen_q --q_gen_rounds=1

scripts/main_exp/eval/cd_oracle.sh

@@ -3,4 +3,4 @@ WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it
 
 
 # longbench
-WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/qasper_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp
+WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets longbench/multifieldqa_en_e longbench/2wikimqa_e longbench/qasper_e longbench/gov_report_e --split test --use_cd --cd_update_iterations 300 --eval_batch_size_gen=1 --truncate_if_too_long_inp

scripts/main_exp/eval/clipper.sh

@@ -0,0 +1,8 @@
+# plain base-model baseline truncate
+WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets oolong-synth --split test --eval_batch_size_gen 1 --truncate_if_too_long_inp --max_test_samples_per_ds 500
+
+# Context-free sanity check.
+WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets oolong-synth --split test --eval_batch_size_gen 1 --remove_context --max_test_samples_per_ds 500
+
+# Iterative D2L variant.
+WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets oolong-synth --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1 --max_test_samples_per_ds 500 --use_iterative_mode

scripts/main_exp/eval/d2l-chunking-abalation.sh

@@ -0,0 +1,10 @@
+# main results
+# batched
+
+for max_ctx_chunk_len in 256 512 1024 2048 4096 8192; do
+    WANDB_MODE=disabled uv run run_eval.py --checkpoint_path trained_d2l/gemma_2b_d2l/checkpoint-20000/pytorch_model.bin --datasets longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e --split test --max_ctx_chunk_len $max_ctx_chunk_len --eval_batch_size_gen 1
+done
+
+
+# # iterative
+# WANDB_MODE=disabled uv run run_eval.py --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e --split test --max_ctx_chunk_len 8192 --eval_batch_size_gen 1 --use_iterative_mode

scripts/main_exp/eval/d2l-plot-chunking-abalation.py

@@ -0,0 +1,248 @@
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+from pathlib import Path
+
+import matplotlib
+import matplotlib as mpl
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+sys.path.append(str(Path(__file__).resolve().parents[2] / "niah"))
+
+from plot_ablation_utils import (  # noqa: E402
+    choose_latest_runs_by_chunk,
+    ensure_output_path,
+    get_performance_value,
+    load_results_json,
+    load_script_text,
+    parse_checkpoint_path,
+    parse_eval_results_root,
+)
+
+LATTE_STYLE = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+DATASET_LABELS = {
+    "longbench/qasper_e": "Qasper",
+    "longbench/2wikimqa_e": "2WikiMQA",
+    "longbench/multifieldqa_en_e": "MultiFieldQA-EN",
+    "longbench/gov_report_e": "GovReport",
+}
+REFERENCE_PERFORMANCE = {
+    "longbench/2wikimqa_e": 0.3387,
+    "longbench/multifieldqa_en_e": 0.3938,
+    "longbench/qasper_e": 0.3839,
+}
+METRIC_CANDIDATES = (
+    "qa_f1_score",
+    "rougeL.f1",
+    "qa_f1",
+    "f1",
+    "accuracy",
+)
+
+
+def chunk_size_tick_label(chunk_size: int) -> str:
+    return rf"$2^{{{int(chunk_size).bit_length() - 1}}}$"
+
+
+def approx_chunk_label(value: float) -> str:
+    rounded = round(value)
+    if abs(value - rounded) < 0.05:
+        return str(int(rounded))
+    return f"{value:.1f}"
+
+
+def configure_plot_style() -> None:
+    plt.style.use(["ggplot", LATTE_STYLE])
+    plt.rcParams["axes.facecolor"] = "F7FBFC"
+    plt.rcParams["figure.facecolor"] = "white"
+    plt.rcParams["savefig.facecolor"] = "white"
+    plt.rcParams["grid.color"] = "cccccc"
+    plt.rcParams["grid.linewidth"] = 1
+    plt.rcParams["axes.edgecolor"] = "ccd0da"
+    plt.rcParams["legend.facecolor"] = "white"
+    plt.rcParams["legend.fontsize"] = 14
+    plt.rcParams["xtick.labelsize"] = 13
+    plt.rcParams["ytick.labelsize"] = 13
+    plt.rcParams["axes.labelweight"] = "bold"
+    plt.rcParams["axes.prop_cycle"] = mpl.cycler(color=plt.cm.tab10.colors)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Plot one D2L chunking-ablation figure per dataset.",
+    )
+    parser.add_argument(
+        "--eval-script",
+        type=Path,
+        default=Path(__file__).with_name("d2l-chunking-abalation.sh"),
+        help="Path to the eval sweep shell script.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=None,
+        help="Optional directory for output PNGs. Defaults to eval-results-*/plots/.",
+    )
+    return parser
+
+
+def parse_checkpoint_from_script(script_text: str) -> Path:
+    try:
+        return parse_checkpoint_path(script_text)
+    except ValueError:
+        match = re.search(r"--checkpoint_path\s+(\S+)", script_text)
+        if match is None:
+            raise
+        return Path(match.group(1))
+
+
+def infer_metric_name(run_dir: Path, dataset_name: str) -> str:
+    result_json = load_results_json(run_dir, dataset_name)
+    for metric_name in METRIC_CANDIDATES:
+        if f"test_{dataset_name}_{metric_name}" in result_json:
+            return metric_name
+    raise ValueError(f"Could not infer a metric for dataset {dataset_name}.")
+
+
+def dataset_label(dataset_name: str) -> str:
+    return DATASET_LABELS.get(dataset_name, dataset_name.split("/")[-1])
+
+
+def parse_datasets_from_script(script_text: str) -> list[str]:
+    match = re.search(r"--datasets\s+(.+?)\s+--split", script_text, flags=re.DOTALL)
+    if match is None:
+        raise ValueError("Could not parse datasets from the eval script.")
+    return match.group(1).split()
+
+
+def mean_context_length(run_dir: Path, dataset_name: str) -> float:
+    path = run_dir / f"test_{dataset_name}_generated_text.jsonl"
+    total = 0.0
+    count = 0
+    with path.open(encoding="utf-8") as handle:
+        for line in handle:
+            record = json.loads(line)
+            ctx_ids_len = record.get("ctx_ids_len")
+            if isinstance(ctx_ids_len, (int, float)):
+                total += float(ctx_ids_len)
+                count += 1
+    if count == 0:
+        raise ValueError(f"Could not find ctx_ids_len values in {path}.")
+    return total / count
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    configure_plot_style()
+
+    script_text = load_script_text(args.eval_script)
+    checkpoint_path = parse_checkpoint_from_script(script_text)
+    eval_results_root = parse_eval_results_root(checkpoint_path)
+    expected_datasets = parse_datasets_from_script(script_text)
+    selected_runs = choose_latest_runs_by_chunk(eval_results_root, expected_datasets)
+
+    if not selected_runs:
+        raise SystemExit(f"No matching eval runs found under {eval_results_root}.")
+
+    output_dir = (
+        eval_results_root / "plots" if args.output_dir is None else args.output_dir
+    )
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for dataset_name in expected_datasets:
+        metric_name = infer_metric_name(selected_runs[0].run_dir, dataset_name)
+        reference_performance = REFERENCE_PERFORMANCE.get(dataset_name)
+        chunk_sizes: list[int] = []
+        approx_num_chunks: list[float] = []
+        normalized_performance_values: list[float] = []
+
+        for run_info in selected_runs:
+            chunk_size = run_info.chunk_size
+            if chunk_size == 65536:
+                continue
+            performance = get_performance_value(
+                run_info.run_dir,
+                dataset_name,
+                metric_name=metric_name,
+            )
+            if performance is None:
+                continue
+
+            avg_context_length = mean_context_length(run_info.run_dir, dataset_name)
+            chunk_sizes.append(chunk_size)
+            approx_num_chunks.append(avg_context_length / chunk_size)
+            plot_value = (
+                performance / reference_performance
+                if reference_performance not in (None, 0)
+                else performance
+            )
+            normalized_performance_values.append(plot_value)
+
+        if not chunk_sizes:
+            continue
+
+        fig, ax_perf = plt.subplots(
+            1,
+            1,
+            figsize=(7.5, 5.5),
+            constrained_layout=False,
+        )
+        fig.subplots_adjust(bottom=0.18, top=0.86)
+
+        ax_perf.plot(
+            chunk_sizes,
+            normalized_performance_values,
+            marker="o",
+            linewidth=2,
+            color="#1f77b4",
+            markeredgecolor="black",
+            markeredgewidth=1,
+        )
+        for chunk_size, performance, approx_chunks in zip(
+            chunk_sizes,
+            normalized_performance_values,
+            approx_num_chunks,
+        ):
+            ax_perf.annotate(
+                approx_chunk_label(approx_chunks),
+                xy=(chunk_size, performance),
+                xytext=(0, 7),
+                textcoords="offset points",
+                ha="center",
+                va="bottom",
+                fontsize=11,
+                color="#4c4f69",
+            )
+
+        ax_perf.grid(True, alpha=0.25)
+        ax_perf.set_xscale("log", base=2)
+        ax_perf.set_xticks(chunk_sizes)
+        ax_perf.set_xticklabels(
+            [chunk_size_tick_label(chunk_size) for chunk_size in chunk_sizes]
+        )
+        ax_perf.xaxis.set_minor_locator(plt.NullLocator())
+
+        pretty_name = dataset_label(dataset_name)
+        ax_perf.set_title(f"Retrieval Performance\n({pretty_name})", fontweight="bold")
+        ax_perf.set_xlabel("Chunk Size (tokens)")
+        ax_perf.set_ylabel(
+            "Normalized Performance"
+            if reference_performance not in (None, 0)
+            else "Performance"
+        )
+        ax_perf.set_ylim(bottom=0.0)
+
+        output_path = ensure_output_path(
+            output_dir / f"d2l-chunking-abalation-{dataset_name.split('/')[-1]}.png"
+        )
+        fig.savefig(output_path, dpi=200, bbox_inches="tight")
+        print(f"Saved plot to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/main_exp/eval/d2l-plot-context-length-8192.py

@@ -0,0 +1,389 @@
+from __future__ import annotations
+
+import argparse
+import math
+import re
+import sys
+from pathlib import Path
+
+import matplotlib
+import matplotlib as mpl
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+sys.path.append(str(Path(__file__).resolve().parents[2] / "niah"))
+
+from plot_ablation_utils import (  # noqa: E402
+    choose_latest_runs_by_chunk,
+    load_results_json,
+    load_script_text,
+    parse_checkpoint_path,
+    parse_eval_results_root,
+)
+
+LATTE_STYLE = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+DATASET_LABELS = {
+    "longbench/qasper_e": "Qasper",
+    "longbench/2wikimqa_e": "2WikiMQA",
+    "longbench/multifieldqa_en_e": "MultiFieldQA-EN",
+    "longbench/gov_report_e": "GovReport",
+}
+BASE_MODEL_EVAL_ROOT = Path("eval_results/google/gemma-2-2b-it")
+METRIC_CANDIDATES = (
+    "qa_f1_score",
+    "rougeL.f1",
+    "qa_f1",
+    "f1",
+    "accuracy",
+)
+LEN_KEY_RE = re.compile(r"^test_(.+)_(.+)_len_([0-9]+)-([0-9]+|inf)$")
+MIN_CONTEXT_LENGTH = 4096
+MAX_CONTEXT_LENGTH = 16384
+
+
+def configure_plot_style() -> None:
+    plt.style.use(["ggplot", LATTE_STYLE])
+    plt.rcParams["axes.facecolor"] = "F7FBFC"
+    plt.rcParams["figure.facecolor"] = "white"
+    plt.rcParams["savefig.facecolor"] = "white"
+    plt.rcParams["grid.color"] = "cccccc"
+    plt.rcParams["grid.linewidth"] = 1
+    plt.rcParams["axes.edgecolor"] = "ccd0da"
+    plt.rcParams["legend.facecolor"] = "white"
+    plt.rcParams["legend.fontsize"] = 14
+    plt.rcParams["xtick.labelsize"] = 13
+    plt.rcParams["ytick.labelsize"] = 13
+    plt.rcParams["axes.labelweight"] = "bold"
+    plt.rcParams["axes.prop_cycle"] = mpl.cycler(color=plt.cm.tab10.colors)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Plot LongBench performance vs context length for the 8192 chunk-size D2L eval run.",
+    )
+    parser.add_argument(
+        "--eval-script",
+        type=Path,
+        default=Path(__file__).with_name("d2l-chunking-abalation.sh"),
+        help="Path to the eval sweep shell script.",
+    )
+    parser.add_argument(
+        "--chunk-size",
+        type=int,
+        default=8192,
+        help="Chunk size to select from the sweep.",
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=None,
+        help="Optional output directory. Defaults to eval-results-*/plots/.",
+    )
+    return parser
+
+
+def parse_checkpoint_from_script(script_text: str) -> Path:
+    try:
+        return parse_checkpoint_path(script_text)
+    except ValueError:
+        match = re.search(r"--checkpoint_path\s+(\S+)", script_text)
+        if match is None:
+            raise
+        return Path(match.group(1))
+
+
+def parse_datasets_from_script(script_text: str) -> list[str]:
+    match = re.search(r"--datasets\s+(.+?)\s+--split", script_text, flags=re.DOTALL)
+    if match is None:
+        raise ValueError("Could not parse datasets from the eval script.")
+    return match.group(1).split()
+
+
+def infer_metric_name(result_json: dict[str, object], dataset_name: str) -> str:
+    for metric_name in METRIC_CANDIDATES:
+        if f"test_{dataset_name}_{metric_name}" in result_json:
+            return metric_name
+    raise ValueError(f"Could not infer a metric for dataset {dataset_name}.")
+
+
+def dataset_label(dataset_name: str) -> str:
+    return DATASET_LABELS.get(dataset_name, dataset_name.split("/")[-1])
+
+
+def choose_latest_complete_base_run(expected_datasets: list[str]) -> Path:
+    candidates: list[Path] = []
+    for run_dir in sorted(BASE_MODEL_EVAL_ROOT.iterdir()):
+        if not run_dir.is_dir():
+            continue
+        if all(
+            (run_dir / f"test_{dataset_name}_results.json").exists()
+            for dataset_name in expected_datasets
+        ):
+            candidates.append(run_dir)
+
+    if not candidates:
+        raise ValueError(
+            f"Could not find a complete base-model run under {BASE_MODEL_EVAL_ROOT}."
+        )
+    return candidates[-1]
+
+
+def extract_length_series(
+    result_json: dict[str, object],
+    dataset_name: str,
+    metric_name: str,
+) -> tuple[list[int], list[float]]:
+    xs: list[int] = []
+    ys: list[float] = []
+    metric_prefix = f"test_{dataset_name}_{metric_name}_len_"
+    count_prefix = f"test_{dataset_name}_num_samples_{metric_name}_len_"
+
+    for key, value in result_json.items():
+        if not key.startswith(metric_prefix):
+            continue
+        match = LEN_KEY_RE.match(key)
+        if match is None:
+            continue
+
+        low = int(match.group(3))
+        high_str = match.group(4)
+        if low == 0 or high_str == "inf":
+            continue
+        if value in (None, "None", "N/A"):
+            continue
+
+        high = int(high_str)
+        upper_bound = high + 1
+        if upper_bound < MIN_CONTEXT_LENGTH or upper_bound > MAX_CONTEXT_LENGTH:
+            continue
+        count_key = f"{count_prefix}{low}-{high_str}"
+        count = result_json.get(count_key)
+        if not isinstance(count, int) or count <= 0:
+            continue
+
+        xs.append(upper_bound)
+        ys.append(float(value))
+
+    pairs = sorted(zip(xs, ys), key=lambda pair: pair[0])
+    return [x for x, _ in pairs], [y for _, y in pairs]
+
+
+def tick_label(length_value: int) -> str:
+    exponent = int(round(math.log2(length_value)))
+    return rf"$2^{{{exponent}}}$"
+
+
+def to_series_map(xs: list[int], ys: list[float]) -> dict[int, float]:
+    return {x: y for x, y in zip(xs, ys)}
+
+
+def chunk_count_label(context_length: int, chunk_size: int) -> str:
+    num_chunks = max(1, math.ceil(context_length / chunk_size))
+    suffix = "" if num_chunks == 1 else "s"
+    return f"{num_chunks} chunk{suffix}"
+
+
+def chunk_label_offset(dataset_name: str) -> tuple[int, int]:
+    if dataset_name == "longbench/2wikimqa_e":
+        return (0, -15)
+    return (0, 11)
+
+
+def chunk_label_position(
+    dataset_name: str,
+    index: int,
+    total: int,
+) -> tuple[tuple[int, int], str]:
+    dx, dy = chunk_label_offset(dataset_name)
+    ha = "center"
+    if total > 1 and index == 0:
+        dx += 2
+    elif total > 1 and index == total - 1:
+        dx -= 2
+    return (dx, dy), ha
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    configure_plot_style()
+
+    script_text = load_script_text(args.eval_script)
+    checkpoint_path = parse_checkpoint_from_script(script_text)
+    eval_results_root = parse_eval_results_root(checkpoint_path)
+    expected_datasets = parse_datasets_from_script(script_text)
+    selected_runs = choose_latest_runs_by_chunk(eval_results_root, expected_datasets)
+    base_run_dir = choose_latest_complete_base_run(expected_datasets)
+
+    target_run = next(
+        (
+            run_info
+            for run_info in selected_runs
+            if run_info.chunk_size == args.chunk_size
+        ),
+        None,
+    )
+    if target_run is None:
+        raise SystemExit(
+            f"Could not find a run with chunk size {args.chunk_size} under {eval_results_root}."
+        )
+
+    output_dir = eval_results_root / "plots" if args.output is None else args.output
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    color_map = plt.get_cmap("tab10", 2)
+    for dataset_name in expected_datasets:
+        result_json = load_results_json(target_run.run_dir, dataset_name)
+        metric_name = infer_metric_name(result_json, dataset_name)
+        xs, ys = extract_length_series(result_json, dataset_name, metric_name)
+        base_result_json = load_results_json(base_run_dir, dataset_name)
+        base_xs, base_ys = extract_length_series(
+            base_result_json, dataset_name, metric_name
+        )
+
+        d2l_series = to_series_map(xs, ys)
+        base_series = to_series_map(base_xs, base_ys)
+        normalized_xs = [
+            x for x in sorted(set(d2l_series) & set(base_series)) if base_series[x] != 0
+        ]
+        all_xs = sorted(set(xs) | set(base_xs))
+        if not all_xs:
+            continue
+
+        fig, (ax_raw, ax_norm) = plt.subplots(
+            1,
+            2,
+            figsize=(14.5, 5.8),
+            constrained_layout=False,
+        )
+        fig.subplots_adjust(bottom=0.16, top=0.84, wspace=0.28)
+
+        if xs:
+            ax_raw.plot(
+                xs,
+                ys,
+                marker="o",
+                linewidth=2,
+                label="D2L",
+                markeredgecolor="black",
+                markeredgewidth=1,
+                color=color_map(0),
+            )
+            for i, (x, y) in enumerate(zip(xs, ys)):
+                raw_label_offset, raw_ha = chunk_label_position(
+                    dataset_name, i, len(xs)
+                )
+                ax_raw.annotate(
+                    chunk_count_label(x, args.chunk_size),
+                    (x, y),
+                    textcoords="offset points",
+                    xytext=raw_label_offset,
+                    ha=raw_ha,
+                    fontsize=8,
+                    color="#6c6f85",
+                )
+        if base_xs:
+            ax_raw.plot(
+                base_xs,
+                base_ys,
+                linestyle="--",
+                marker="o",
+                linewidth=2,
+                label="Base",
+                color=color_map(1),
+                alpha=0.9,
+                markeredgecolor="black",
+                markeredgewidth=1,
+            )
+
+        ax_raw.set_xscale("log", base=2)
+        ax_raw.set_xticks(all_xs)
+        ax_raw.set_xticklabels([tick_label(x) for x in all_xs])
+        ax_raw.xaxis.set_minor_locator(plt.NullLocator())
+        ax_raw.grid(True, alpha=0.25)
+        ax_raw.set_title("Raw Performance", fontweight="bold")
+        ax_raw.set_xlabel("Context Length (tokens)")
+        ax_raw.set_ylabel("QA F1 Score")
+        ax_raw.set_ylim(-0.05, 1.05)
+        ax_raw.legend(loc="lower left", frameon=True, fancybox=True)
+
+        if normalized_xs:
+            normalized_d2l_ys = [d2l_series[x] / base_series[x] for x in normalized_xs]
+            normalized_base_ys = [1.0 for _ in normalized_xs]
+            ax_norm.plot(
+                normalized_xs,
+                normalized_d2l_ys,
+                marker="o",
+                linewidth=2,
+                label="D2L",
+                markeredgecolor="black",
+                markeredgewidth=1,
+                color=color_map(0),
+            )
+            for i, (x, y) in enumerate(zip(normalized_xs, normalized_d2l_ys)):
+                norm_label_offset, norm_ha = chunk_label_position(
+                    dataset_name,
+                    i,
+                    len(normalized_xs),
+                )
+                ax_norm.annotate(
+                    chunk_count_label(x, args.chunk_size),
+                    (x, y),
+                    textcoords="offset points",
+                    xytext=norm_label_offset,
+                    ha=norm_ha,
+                    fontsize=8,
+                    color="#6c6f85",
+                )
+            ax_norm.plot(
+                normalized_xs,
+                normalized_base_ys,
+                linestyle="--",
+                marker="o",
+                linewidth=2,
+                label="Base",
+                color=color_map(1),
+                alpha=0.9,
+                markeredgecolor="black",
+                markeredgewidth=1,
+            )
+        else:
+            ax_norm.text(
+                0.5,
+                0.5,
+                "No overlapping bins for normalization.",
+                ha="center",
+                va="center",
+                transform=ax_norm.transAxes,
+            )
+
+        norm_ticks = normalized_xs if normalized_xs else all_xs
+        ax_norm.set_xscale("log", base=2)
+        ax_norm.set_xticks(norm_ticks)
+        ax_norm.set_xticklabels([tick_label(x) for x in norm_ticks])
+        ax_norm.xaxis.set_minor_locator(plt.NullLocator())
+        ax_norm.grid(True, alpha=0.25)
+        ax_norm.set_title("Normalized Performance", fontweight="bold")
+        ax_norm.set_xlabel("Context Length (tokens)")
+        ax_norm.set_ylabel("Normalized Performance")
+        ax_norm.set_ylim(bottom=0.0)
+        if normalized_xs:
+            ax_norm.legend(loc="lower left", frameon=True, fancybox=True)
+
+        fig.suptitle(
+            f"Performance Grouped by Length\n({dataset_label(dataset_name)}, chunk size = {args.chunk_size})",
+            fontweight="bold",
+            fontsize=16,
+        )
+
+        output_path = (
+            output_dir
+            / f"d2l-context-length-{args.chunk_size}-{dataset_name.split('/')[-1]}.png"
+        )
+        fig.savefig(output_path, dpi=200, bbox_inches="tight")
+        print(f"Saved plot to {output_path}")
+        plt.close(fig)
+
+
+if __name__ == "__main__":
+    main()

scripts/main_exp/eval/d2l-plot-rag-vs-hybrid.py

@@ -0,0 +1,740 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import re
+import shlex
+from collections import defaultdict, deque
+from dataclasses import dataclass
+from pathlib import Path
+
+import matplotlib
+import matplotlib as mpl
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+LATTE_STYLE = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+BASE_MODEL_EVAL_ROOT = Path("eval_results/google/gemma-2-2b-it")
+DATASET_LABELS = {
+    "oolong-synth": "Oolong-Synth",
+    "longbench/qasper_e": "Qasper",
+    "longbench/2wikimqa_e": "2WikiMQA",
+    "longbench/multifieldqa_en_e": "MultiFieldQA-EN",
+    "longbench/gov_report_e": "GovReport",
+}
+METRIC_CANDIDATES = (
+    "oolong_score",
+    "qa_f1_score",
+    "rougeL.f1",
+    "qa_f1",
+    "f1",
+    "accuracy",
+)
+METRIC_LABELS = {
+    "oolong_score": "Oolong Score",
+    "qa_f1_score": "QA F1 Score",
+    "rougeL.f1": "ROUGE-L F1",
+    "qa_f1": "QA F1",
+    "f1": "F1",
+    "accuracy": "Accuracy",
+}
+COMPARISON_TOLERANCE = 0.2
+CMD_RE = re.compile(r"CMD:\s+(.*)$", re.MULTILINE)
+BASELINE_EXCLUDE_FLAGS = {
+    "--use_rag",
+    "--use_hybrid_rag",
+    "--use_cd",
+    "--use_llmlingua",
+    "--use_t2l",
+}
+NO_CONTEXT_COLUMN = "No context"
+CSV_SUBCOLUMNS = (NO_CONTEXT_COLUMN, "top-1", "top-4")
+
+
+@dataclass(frozen=True)
+class RagRunInfo:
+    run_dir: Path
+    run_name: str
+    datasets: tuple[str, ...]
+    is_hybrid: bool
+    top_k: int
+    chunk_size: int
+    chunk_overlap: int
+    max_retrieved_tokens: int
+
+
+def configure_plot_style() -> None:
+    plt.style.use(["ggplot", LATTE_STYLE])
+    plt.rcParams["axes.facecolor"] = "F7FBFC"
+    plt.rcParams["figure.facecolor"] = "white"
+    plt.rcParams["savefig.facecolor"] = "white"
+    plt.rcParams["grid.color"] = "cccccc"
+    plt.rcParams["grid.linewidth"] = 1
+    plt.rcParams["axes.edgecolor"] = "ccd0da"
+    plt.rcParams["legend.facecolor"] = "white"
+    plt.rcParams["legend.fontsize"] = 12
+    plt.rcParams["xtick.labelsize"] = 13
+    plt.rcParams["ytick.labelsize"] = 13
+    plt.rcParams["axes.labelweight"] = "bold"
+    plt.rcParams["axes.prop_cycle"] = mpl.cycler(color=plt.cm.tab10.colors)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Plot one raw-RAG vs hybrid-RAG comparison figure per dataset.",
+    )
+    parser.add_argument(
+        "--rag-root",
+        type=Path,
+        default=Path("eval_results/google/gemma-2-2b-it"),
+        help="Directory containing raw RAG eval runs.",
+    )
+    parser.add_argument(
+        "--hybrid-root",
+        type=Path,
+        default=Path("trained_d2l/gemma_2b_d2l/eval-results-20000"),
+        help="Directory containing hybrid RAG eval runs.",
+    )
+    parser.add_argument(
+        "--rag-chunk-size",
+        type=int,
+        default=256,
+        help="Filter runs by rag chunk size.",
+    )
+    parser.add_argument(
+        "--rag-chunk-overlap",
+        type=int,
+        default=64,
+        help="Filter runs by rag chunk overlap.",
+    )
+    parser.add_argument(
+        "--rag-max-retrieved-tokens",
+        type=int,
+        default=1536,
+        help="Filter runs by rag max retrieved tokens.",
+    )
+    parser.add_argument(
+        "--datasets",
+        nargs="*",
+        default=None,
+        help="Optional dataset allowlist. Defaults to all datasets with matched raw/hybrid runs.",
+    )
+    parser.add_argument(
+        "--top-ks",
+        nargs="*",
+        type=int,
+        default=None,
+        help="Optional top-k allowlist. Defaults to all matched top-k values.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=None,
+        help="Optional output directory. Defaults to HYBRID_ROOT/plots.",
+    )
+    return parser
+
+
+def _extract_flag_value(tokens: list[str], flag: str) -> str | None:
+    for index, token in enumerate(tokens):
+        if token == flag and index + 1 < len(tokens):
+            return tokens[index + 1]
+        if token.startswith(f"{flag}="):
+            return token.split("=", 1)[1]
+    return None
+
+
+def _extract_datasets(tokens: list[str]) -> tuple[str, ...]:
+    if "--datasets" not in tokens:
+        return ()
+
+    index = tokens.index("--datasets") + 1
+    datasets: list[str] = []
+    while index < len(tokens) and not tokens[index].startswith("--"):
+        datasets.append(tokens[index])
+        index += 1
+    return tuple(datasets)
+
+
+def _has_max_test_samples_flag(tokens: list[str]) -> bool:
+    return _extract_flag_value(tokens, "--max_test_samples_per_ds") is not None
+
+
+def parse_rag_run_info(run_dir: Path) -> RagRunInfo | None:
+    log_path = run_dir / "debug.log"
+    if not log_path.exists():
+        return None
+
+    log_text = log_path.read_text(encoding="utf-8", errors="replace")
+    match = CMD_RE.search(log_text)
+    if match is None:
+        return None
+
+    command = match.group(1).strip()
+    try:
+        tokens = shlex.split(command)
+    except ValueError:
+        tokens = command.split()
+
+    use_rag = "--use_rag" in tokens
+    use_hybrid_rag = "--use_hybrid_rag" in tokens
+    if use_rag == use_hybrid_rag:
+        return None
+    if _has_max_test_samples_flag(tokens):
+        return None
+
+    datasets = _extract_datasets(tokens)
+    if not datasets:
+        return None
+
+    top_k_value = _extract_flag_value(tokens, "--rag_top_k")
+    chunk_size_value = _extract_flag_value(tokens, "--rag_chunk_size")
+    chunk_overlap_value = _extract_flag_value(tokens, "--rag_chunk_overlap")
+    max_tokens_value = _extract_flag_value(tokens, "--rag_max_retrieved_tokens")
+    if None in (
+        top_k_value,
+        chunk_size_value,
+        chunk_overlap_value,
+        max_tokens_value,
+    ):
+        return None
+
+    return RagRunInfo(
+        run_dir=run_dir,
+        run_name=run_dir.name,
+        datasets=datasets,
+        is_hybrid=use_hybrid_rag,
+        top_k=int(top_k_value),
+        chunk_size=int(chunk_size_value),
+        chunk_overlap=int(chunk_overlap_value),
+        max_retrieved_tokens=int(max_tokens_value),
+    )
+
+
+def parse_plain_base_run_datasets(run_dir: Path) -> tuple[str, ...] | None:
+    log_path = run_dir / "debug.log"
+    if not log_path.exists():
+        return None
+
+    log_text = log_path.read_text(encoding="utf-8", errors="replace")
+    match = CMD_RE.search(log_text)
+    if match is None:
+        return None
+
+    command = match.group(1).strip()
+    try:
+        tokens = shlex.split(command)
+    except ValueError:
+        tokens = command.split()
+
+    if "--model_name_or_path" not in tokens:
+        return None
+    if any(flag in tokens for flag in BASELINE_EXCLUDE_FLAGS):
+        return None
+    if _has_max_test_samples_flag(tokens):
+        return None
+
+    datasets = _extract_datasets(tokens)
+    if not datasets:
+        return None
+    return datasets
+
+
+def parse_plain_d2l_run_datasets(run_dir: Path) -> tuple[str, ...] | None:
+    log_path = run_dir / "debug.log"
+    if not log_path.exists():
+        return None
+
+    log_text = log_path.read_text(encoding="utf-8", errors="replace")
+    match = CMD_RE.search(log_text)
+    if match is None:
+        return None
+
+    command = match.group(1).strip()
+    try:
+        tokens = shlex.split(command)
+    except ValueError:
+        tokens = command.split()
+
+    if "--checkpoint_path" not in tokens:
+        return None
+    if any(flag in tokens for flag in BASELINE_EXCLUDE_FLAGS):
+        return None
+    if "--add_ctx_to_input" in tokens:
+        return None
+    if _has_max_test_samples_flag(tokens):
+        return None
+    max_ctx_chunk_len = _extract_flag_value(tokens, "--max_ctx_chunk_len")
+    if max_ctx_chunk_len != "8192":
+        return None
+
+    datasets = _extract_datasets(tokens)
+    if not datasets:
+        return None
+    return datasets
+
+
+def result_json_path(run_dir: Path, dataset_name: str) -> Path:
+    return run_dir / f"test_{dataset_name}_results.json"
+
+
+def load_results_json(run_dir: Path, dataset_name: str) -> dict[str, object]:
+    return json.loads(
+        result_json_path(run_dir, dataset_name).read_text(encoding="utf-8")
+    )
+
+
+def infer_metric_name(result_json: dict[str, object], dataset_name: str) -> str:
+    for metric_name in METRIC_CANDIDATES:
+        if f"test_{dataset_name}_{metric_name}" in result_json:
+            return metric_name
+    raise ValueError(f"Could not infer a metric for dataset {dataset_name}.")
+
+
+def metric_label(metric_name: str) -> str:
+    return METRIC_LABELS.get(metric_name, metric_name.replace("_", " ").title())
+
+
+def dataset_label(dataset_name: str) -> str:
+    return DATASET_LABELS.get(dataset_name, dataset_name.split("/")[-1])
+
+
+def select_latest_runs(
+    root: Path,
+    *,
+    want_hybrid: bool,
+    allowed_datasets: set[str] | None,
+    allowed_top_ks: set[int] | None,
+    rag_chunk_size: int,
+    rag_chunk_overlap: int,
+    rag_max_retrieved_tokens: int,
+) -> dict[tuple[str, int], RagRunInfo]:
+    selected: dict[tuple[str, int], RagRunInfo] = {}
+    if not root.exists():
+        return selected
+
+    for run_dir in sorted(root.iterdir()):
+        if not run_dir.is_dir():
+            continue
+
+        run_info = parse_rag_run_info(run_dir)
+        if run_info is None or run_info.is_hybrid != want_hybrid:
+            continue
+        if run_info.chunk_size != rag_chunk_size:
+            continue
+        if run_info.chunk_overlap != rag_chunk_overlap:
+            continue
+        if run_info.max_retrieved_tokens != rag_max_retrieved_tokens:
+            continue
+        if allowed_top_ks is not None and run_info.top_k not in allowed_top_ks:
+            continue
+
+        for dataset_name in run_info.datasets:
+            if allowed_datasets is not None and dataset_name not in allowed_datasets:
+                continue
+            if not result_json_path(run_info.run_dir, dataset_name).exists():
+                continue
+
+            key = (dataset_name, run_info.top_k)
+            current = selected.get(key)
+            if current is None or run_info.run_name > current.run_name:
+                selected[key] = run_info
+
+    return selected
+
+
+def get_metric_value(
+    run_dir: Path, dataset_name: str, metric_name: str
+) -> float | None:
+    result_json = load_results_json(run_dir, dataset_name)
+    value = result_json.get(f"test_{dataset_name}_{metric_name}")
+    if value in (None, "None", "N/A"):
+        return None
+    if isinstance(value, (int, float)):
+        return float(value)
+    return None
+
+
+def select_latest_plain_base_runs(
+    root: Path,
+    datasets: set[str],
+) -> dict[str, Path]:
+    selected: dict[str, Path] = {}
+    if not root.exists():
+        return selected
+
+    for run_dir in sorted(root.iterdir()):
+        if not run_dir.is_dir():
+            continue
+
+        run_datasets = parse_plain_base_run_datasets(run_dir)
+        if run_datasets is None:
+            continue
+
+        for dataset_name in run_datasets:
+            if dataset_name not in datasets:
+                continue
+            if not result_json_path(run_dir, dataset_name).exists():
+                continue
+            selected[dataset_name] = run_dir
+
+    return selected
+
+
+def select_latest_plain_d2l_runs(
+    root: Path,
+    datasets: set[str],
+) -> dict[str, Path]:
+    selected: dict[str, Path] = {}
+    if not root.exists():
+        return selected
+
+    for run_dir in sorted(root.iterdir()):
+        if not run_dir.is_dir():
+            continue
+
+        run_datasets = parse_plain_d2l_run_datasets(run_dir)
+        if run_datasets is None:
+            continue
+
+        for dataset_name in run_datasets:
+            if dataset_name not in datasets:
+                continue
+            if not result_json_path(run_dir, dataset_name).exists():
+                continue
+            selected[dataset_name] = run_dir
+
+    return selected
+
+
+CSV_TOP_KS = (1, 4)
+
+
+def csv_output_path(output_dir: Path) -> Path:
+    return output_dir / "d2l-rag-vs-hybrid.csv"
+
+
+def confusion_output_path_for_dataset(output_dir: Path, dataset_name: str) -> Path:
+    return output_dir / f"d2l-rag-vs-hybrid-confusion-{dataset_name.split('/')[-1]}.png"
+
+
+def generated_text_path(run_dir: Path, dataset_name: str) -> Path:
+    return run_dir / f"test_{dataset_name}_generated_text.jsonl"
+
+
+def sample_key(sample: dict[str, object], dataset_name: str) -> str:
+    if dataset_name == "oolong-synth":
+        return "oolong:" + str(sample.get("oolong_id"))
+
+    input_text = str(sample.get("input", ""))
+    label_text = str(sample.get("label", ""))
+    ctx_ids_len = str(sample.get("ctx_ids_len", ""))
+    return f"{dataset_name}\n{input_text}\n{label_text}\n{ctx_ids_len}"
+
+
+def sample_metric_value(sample: dict[str, object], metric_name: str) -> float:
+    value = sample.get(metric_name)
+    if not isinstance(value, (int, float)):
+        return 0.0
+    return float(value)
+
+
+def load_metric_by_sample(
+    run_dir: Path,
+    dataset_name: str,
+    metric_name: str,
+) -> dict[str, deque[float]]:
+    path = generated_text_path(run_dir, dataset_name)
+    grouped: dict[str, deque[float]] = defaultdict(deque)
+    with path.open(encoding="utf-8") as handle:
+        for line in handle:
+            sample = json.loads(line)
+            grouped[sample_key(sample, dataset_name)].append(
+                sample_metric_value(sample, metric_name)
+            )
+    return grouped
+
+
+def relative_bucket(score: float, other_score: float) -> int:
+    delta = score - other_score
+    if delta > COMPARISON_TOLERANCE:
+        return 2
+    if delta < -COMPARISON_TOLERANCE:
+        return 0
+    return 1
+
+
+def build_confusion_counts(
+    raw_run_dir: Path,
+    hybrid_run_dir: Path,
+    dataset_name: str,
+    metric_name: str,
+) -> list[list[int]]:
+    raw_grouped = load_metric_by_sample(raw_run_dir, dataset_name, metric_name)
+    hybrid_grouped = load_metric_by_sample(hybrid_run_dir, dataset_name, metric_name)
+    raw_counts = {key: len(values) for key, values in raw_grouped.items()}
+    hybrid_counts = {key: len(values) for key, values in hybrid_grouped.items()}
+    if raw_counts != hybrid_counts:
+        raise ValueError(f"Could not align raw and hybrid samples for {dataset_name}.")
+
+    matrix = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
+    for key in sorted(raw_grouped):
+        raw_values = raw_grouped[key]
+        hybrid_values = hybrid_grouped[key]
+        while raw_values:
+            raw_score = raw_values.popleft()
+            hybrid_score = hybrid_values.popleft()
+            raw_bucket = relative_bucket(raw_score, hybrid_score)
+            hybrid_bucket = relative_bucket(hybrid_score, raw_score)
+            matrix[raw_bucket][hybrid_bucket] += 1
+    return matrix
+
+
+def plot_confusion_matrices(
+    output_dir: Path,
+    dataset_name: str,
+    metric_name: str,
+    paired_points: list[tuple[int, float, float]],
+    raw_runs: dict[tuple[str, int], RagRunInfo],
+    hybrid_runs: dict[tuple[str, int], RagRunInfo],
+) -> None:
+    ncols = len(paired_points)
+    fig, axes = plt.subplots(
+        1,
+        ncols,
+        figsize=(4.8 * ncols, 4.6),
+        constrained_layout=False,
+    )
+    if ncols == 1:
+        axes = [axes]
+    fig.subplots_adjust(bottom=0.16, top=0.80, wspace=0.35)
+
+    for ax, (top_k, _, _) in zip(axes, paired_points):
+        matrix = build_confusion_counts(
+            raw_runs[(dataset_name, top_k)].run_dir,
+            hybrid_runs[(dataset_name, top_k)].run_dir,
+            dataset_name,
+            metric_name,
+        )
+        image = ax.imshow(matrix, cmap="Blues")
+        total = sum(sum(row) for row in matrix)
+        max_cell = max(max(row) for row in matrix) if total else 0
+
+        for raw_index in range(3):
+            for hybrid_index in range(3):
+                count = matrix[raw_index][hybrid_index]
+                pct = 100.0 * count / total if total else 0.0
+                color = "white" if count > max_cell / 2 else "#4c4f69"
+                ax.text(
+                    hybrid_index,
+                    raw_index,
+                    f"{count}\n({pct:.1f}%)",
+                    ha="center",
+                    va="center",
+                    color=color,
+                    fontsize=11,
+                    fontweight="bold" if count == max_cell and count > 0 else None,
+                )
+
+        ax.set_xticks([0, 1, 2])
+        ax.set_yticks([0, 1, 2])
+        ax.set_xticklabels(["More\nIncorrect", "Same", "More\nCorrect"])
+        ax.set_yticklabels(["More\nIncorrect", "Same", "More\nCorrect"])
+        ax.set_xlabel("Hybrid RAG")
+        if ax is axes[0]:
+            ax.set_ylabel("Raw RAG")
+        ax.set_title(f"top-k = {top_k}", fontweight="bold")
+        ax.grid(False)
+
+        for spine in ax.spines.values():
+            spine.set_visible(True)
+            spine.set_color("#ccd0da")
+
+        fig.colorbar(image, ax=ax, fraction=0.046, pad=0.04)
+
+    fig.suptitle(
+        f"Relative Score Agreement\n({dataset_label(dataset_name)}, delta threshold = {COMPARISON_TOLERANCE:.1f})",
+        fontweight="bold",
+        fontsize=16,
+    )
+    output_path = confusion_output_path_for_dataset(output_dir, dataset_name)
+    fig.savefig(output_path, dpi=200, bbox_inches="tight")
+    print(f"Saved plot to {output_path}")
+    plt.close(fig)
+
+
+def write_summary_csv(
+    output_dir: Path,
+    datasets: list[str],
+    summary: dict[str, dict[str, dict[str, float | str | None]]],
+) -> None:
+    output_path = csv_output_path(output_dir)
+    with output_path.open("w", encoding="utf-8", newline="") as handle:
+        writer = csv.writer(handle)
+        top_header = [""]
+        sub_header = ["Method"]
+        for dataset_name in datasets:
+            pretty_name = dataset_label(dataset_name)
+            top_header.extend([pretty_name] * len(CSV_SUBCOLUMNS))
+            sub_header.extend(CSV_SUBCOLUMNS)
+
+        writer.writerow(top_header)
+        writer.writerow(sub_header)
+
+        for method_name in ("RAG", "RAG + D2L", "D2L"):
+            row = [method_name]
+            for dataset_name in datasets:
+                method_values = summary.get(method_name, {}).get(dataset_name, {})
+                for subcolumn in CSV_SUBCOLUMNS:
+                    value = method_values.get(subcolumn)
+                    if isinstance(value, (int, float)):
+                        row.append(f"{float(value):.4f}")
+                    elif value is None:
+                        row.append("")
+                    else:
+                        row.append(str(value))
+            writer.writerow(row)
+
+    print(f"Saved CSV to {output_path}")
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    configure_plot_style()
+
+    allowed_datasets = set(args.datasets) if args.datasets else None
+    allowed_top_ks = set(args.top_ks) if args.top_ks else None
+    output_dir = (
+        args.hybrid_root / "plots" if args.output_dir is None else args.output_dir
+    )
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    raw_runs = select_latest_runs(
+        args.rag_root,
+        want_hybrid=False,
+        allowed_datasets=allowed_datasets,
+        allowed_top_ks=allowed_top_ks,
+        rag_chunk_size=args.rag_chunk_size,
+        rag_chunk_overlap=args.rag_chunk_overlap,
+        rag_max_retrieved_tokens=args.rag_max_retrieved_tokens,
+    )
+    hybrid_runs = select_latest_runs(
+        args.hybrid_root,
+        want_hybrid=True,
+        allowed_datasets=allowed_datasets,
+        allowed_top_ks=allowed_top_ks,
+        rag_chunk_size=args.rag_chunk_size,
+        rag_chunk_overlap=args.rag_chunk_overlap,
+        rag_max_retrieved_tokens=args.rag_max_retrieved_tokens,
+    )
+
+    shared_keys = set(raw_runs) & set(hybrid_runs)
+    if not shared_keys:
+        raise SystemExit("Could not find any matched raw-RAG and hybrid-RAG runs.")
+
+    datasets = sorted({dataset_name for dataset_name, _ in shared_keys})
+    base_runs = select_latest_plain_base_runs(BASE_MODEL_EVAL_ROOT, set(datasets))
+    d2l_runs = select_latest_plain_d2l_runs(args.hybrid_root, set(datasets))
+    summary: dict[str, dict[str, dict[str, float | str | None]]] = {
+        "RAG": {},
+        "RAG + D2L": {},
+        "D2L": {},
+    }
+    for dataset_name in datasets:
+        base_run_dir = base_runs.get(dataset_name)
+        if base_run_dir is None:
+            continue
+
+        shared_top_ks = sorted(
+            top_k
+            for candidate_dataset, top_k in shared_keys
+            if candidate_dataset == dataset_name
+        )
+        if not shared_top_ks:
+            continue
+
+        raw_result_json = load_results_json(
+            raw_runs[(dataset_name, shared_top_ks[0])].run_dir, dataset_name
+        )
+        hybrid_result_json = load_results_json(
+            hybrid_runs[(dataset_name, shared_top_ks[0])].run_dir,
+            dataset_name,
+        )
+        metric_name = infer_metric_name(raw_result_json, dataset_name)
+        hybrid_metric_name = infer_metric_name(hybrid_result_json, dataset_name)
+        if metric_name != hybrid_metric_name:
+            raise ValueError(
+                f"Metric mismatch for {dataset_name}: raw uses {metric_name}, hybrid uses {hybrid_metric_name}."
+            )
+        base_result_json = load_results_json(base_run_dir, dataset_name)
+        base_metric_name = infer_metric_name(base_result_json, dataset_name)
+        if metric_name != base_metric_name:
+            raise ValueError(
+                f"Metric mismatch for {dataset_name}: comparison uses {metric_name}, baseline uses {base_metric_name}."
+            )
+        base_value = get_metric_value(base_run_dir, dataset_name, metric_name)
+        if base_value in (None, 0):
+            continue
+
+        d2l_run_dir = d2l_runs.get(dataset_name)
+        d2l_value = None
+        if d2l_run_dir is not None:
+            d2l_result_json = load_results_json(d2l_run_dir, dataset_name)
+            d2l_metric_name = infer_metric_name(d2l_result_json, dataset_name)
+            if d2l_metric_name == metric_name:
+                candidate_value = get_metric_value(
+                    d2l_run_dir, dataset_name, metric_name
+                )
+                if candidate_value is not None:
+                    d2l_value = candidate_value / base_value
+
+        paired_points: list[tuple[int, float, float]] = []
+        for top_k in shared_top_ks:
+            raw_value = get_metric_value(
+                raw_runs[(dataset_name, top_k)].run_dir, dataset_name, metric_name
+            )
+            hybrid_value = get_metric_value(
+                hybrid_runs[(dataset_name, top_k)].run_dir,
+                dataset_name,
+                metric_name,
+            )
+            if raw_value is None or hybrid_value is None:
+                continue
+            paired_points.append(
+                (top_k, raw_value / base_value, hybrid_value / base_value)
+            )
+
+        if not paired_points:
+            continue
+
+        summary["RAG"][dataset_name] = {
+            NO_CONTEXT_COLUMN: "N/A",
+            **{f"top-{top_k}": raw_value for top_k, raw_value, _ in paired_points},
+        }
+        summary["RAG + D2L"][dataset_name] = {
+            NO_CONTEXT_COLUMN: "N/A",
+            **{
+                f"top-{top_k}": hybrid_value for top_k, _, hybrid_value in paired_points
+            },
+        }
+        summary["D2L"][dataset_name] = {
+            NO_CONTEXT_COLUMN: d2l_value,
+            "top-1": "N/A",
+            "top-4": "N/A",
+        }
+        plot_confusion_matrices(
+            output_dir,
+            dataset_name,
+            metric_name,
+            paired_points,
+            raw_runs,
+            hybrid_runs,
+        )
+
+    write_summary_csv(output_dir, datasets, summary)
+
+
+if __name__ == "__main__":
+    main()

scripts/main_exp/eval/d2l.sh

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

scripts/main_exp/eval/d2l_rag.sh

@@ -0,0 +1,13 @@
+# Hybrid mode: Doc-to-LoRA internalizes the full document and BM25-style RAG adds
+# retrieved evidence to the prompt-side input.
+
+WANDB_MODE=disabled uv run run_eval.py \
+  --checkpoint_path train_outputs/runs/$RUN_NAME/checkpoint-$step/pytorch_model.bin \
+  --datasets squad drop ropes longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e \
+  --split test \
+  --eval_batch_size_gen=1 \
+  --use_hybrid_rag \
+  --rag_chunk_size 256 \
+  --rag_chunk_overlap 64 \
+  --rag_top_k 4 \
+  --rag_max_retrieved_tokens 1536

scripts/main_exp/eval/llmlingua.sh

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

scripts/main_exp/eval/llmlingua_oolong.sh

@@ -0,0 +1,3 @@
+for rate in 0.9 0.8 0.6 0.4 0.2 0.1; do
+  WANDB_MODE=disabled uv run run_eval.py --model_name_or_path google/gemma-2-2b-it --datasets oolong-synth --split test --eval_batch_size_gen=1 --use_llmlingua --llmlingua_compression_rate "$rate" --truncate_if_too_long_ctx
+done

scripts/main_exp/eval/rag.sh

@@ -0,0 +1,26 @@
+# BM25-style per-example RAG over the dataset `context` field.
+# Keep `--max_ctx_chunk_len` disabled here; the RAG wrapper does its own chunking.
+
+for top_k in 1 4; do
+  WANDB_MODE=disabled uv run run_eval.py \
+    --model_name_or_path google/gemma-2-2b-it \
+    --datasets oolong-synth longbench/qasper_e longbench/2wikimqa_e longbench/multifieldqa_en_e longbench/gov_report_e \
+    --split test \
+    --eval_batch_size_gen=1 \
+    --use_rag \
+    --rag_chunk_size 256 \
+    --rag_chunk_overlap 64 \
+    --rag_top_k 4 \
+    --rag_max_retrieved_tokens 1536
+
+  WANDB_MODE=disabled run uv run run_eval.py \
+      --checkpoint_path trained_d2l/gemma_2b_d2l/checkpoint-20000/pytorch_model.bin \
+      --datasets longbench/qasper_e \
+      --split test \
+      --eval_batch_size_gen 1 \
+      --use_hybrid_rag \
+      --rag_chunk_size 256 \
+      --rag_chunk_overlap 64 \
+      --rag_top_k 4 \
+      --rag_max_retrieved_tokens 1536 --use_iterative_mode
+done

scripts/main_exp/eval/t2l.sh

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

scripts/niah/3-eval-vary-chunk-size.sh

@@ -0,0 +1,6 @@
+CHECKPOINT_PATH=train_outputs/runs/Oct03_10-58-26_slurm0-a3nodeset-8_91076_c4dac553/checkpoint-1482/pytorch_model.bin
+CTX_CHUNK_OVERLAP=${CTX_CHUNK_OVERLAP:-0}
+
+for ctx_chunk_len in 32 64 128 256 512 1024 2048 4096; do 
+     WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_32_1024 ctx_magic_number_1024_2048 ctx_magic_number_3072_4096 ctx_magic_number_7168_8192 ctx_magic_number_15360_16384 ctx_magic_number_28672_32768 ctx_magic_number_57344_65536 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
+done

scripts/niah/3-plot-vary-chunk-size.py

@@ -0,0 +1,258 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import matplotlib
+import matplotlib as mpl
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from plot_ablation_utils import (
+    choose_latest_runs_by_chunk,
+    dataset_context_length,
+    ensure_output_path,
+    get_dense_values,
+    get_performance_value,
+    load_script_text,
+    parse_checkpoint_path,
+    parse_eval_results_root,
+    parse_expected_datasets,
+)
+
+LATTE_STYLE = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+
+
+def configure_plot_style() -> None:
+    plt.style.use(["ggplot", LATTE_STYLE])
+    plt.rcParams["axes.facecolor"] = "F7FBFC"
+    plt.rcParams["figure.facecolor"] = "white"
+    plt.rcParams["savefig.facecolor"] = "white"
+    plt.rcParams["grid.color"] = "cccccc"
+    plt.rcParams["grid.linewidth"] = 1
+    plt.rcParams["axes.edgecolor"] = "ccd0da"
+    plt.rcParams["legend.facecolor"] = "white"
+    plt.rcParams["legend.fontsize"] = 14
+    plt.rcParams["xtick.labelsize"] = 13
+    plt.rcParams["ytick.labelsize"] = 13
+    plt.rcParams["axes.labelweight"] = "bold"
+    plt.rcParams["axes.prop_cycle"] = mpl.cycler(color=plt.cm.tab10.colors)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Plot the context-length ablation from 3-eval-vary-chunk-size.sh.",
+    )
+    parser.add_argument(
+        "--eval-script",
+        type=Path,
+        default=Path(__file__).with_name("3-eval-vary-chunk-size.sh"),
+        help="Path to the eval sweep shell script.",
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=None,
+        help="Optional output PNG path. Defaults to eval-results-*/plots/3-vary-chunk-size.png.",
+    )
+    parser.add_argument(
+        "--metric",
+        default="rougeL.f1",
+        help="Performance metric to read from each *_results.json file.",
+    )
+    parser.add_argument(
+        "--dense-metric",
+        default="mean",
+        help="Dense LoRA metric to read when available.",
+    )
+    parser.add_argument(
+        "--dense-std-metric",
+        default="std",
+        help="Dense LoRA std metric to use for the shaded area when available.",
+    )
+    parser.add_argument(
+        "--dense-max-metric",
+        default="max",
+        help="Dense LoRA max metric to plot as a separate line when available.",
+    )
+    return parser
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    configure_plot_style()
+
+    script_text = load_script_text(args.eval_script)
+    checkpoint_path = parse_checkpoint_path(script_text)
+    eval_results_root = parse_eval_results_root(checkpoint_path)
+    expected_datasets = parse_expected_datasets(script_text)
+    selected_runs = choose_latest_runs_by_chunk(eval_results_root, expected_datasets)
+
+    if not selected_runs:
+        raise SystemExit(f"No matching eval runs found under {eval_results_root}.")
+
+    output_path = (
+        eval_results_root / "plots" / "3-vary-chunk-size.png"
+        if args.output is None
+        else args.output
+    )
+    output_path = ensure_output_path(output_path)
+
+    context_lengths = [
+        dataset_context_length(dataset_name) for dataset_name in expected_datasets
+    ]
+    color_map = plt.get_cmap("tab10", len(selected_runs))
+
+    fig, (ax_perf, ax_dense, ax_dense_max) = plt.subplots(
+        1,
+        3,
+        figsize=(20, 5.5),
+        constrained_layout=False,
+    )
+    fig.subplots_adjust(bottom=0.30, top=0.88, wspace=0.28)
+
+    has_any_dense_mean_values = False
+    has_any_dense_max_values = False
+    for index, run_info in enumerate(selected_runs):
+        color = color_map(index)
+        performance_xs: list[int] = []
+        performance_ys: list[float] = []
+        dense_mean_xs: list[int] = []
+        dense_mean_ys: list[float] = []
+        dense_std_lows: list[float] = []
+        dense_std_highs: list[float] = []
+        dense_max_xs: list[int] = []
+        dense_max_ys: list[float] = []
+
+        for dataset_name in expected_datasets:
+            context_length = dataset_context_length(dataset_name)
+            performance = get_performance_value(
+                run_info.run_dir,
+                dataset_name,
+                metric_name=args.metric,
+            )
+            if performance is not None:
+                performance_xs.append(context_length)
+                performance_ys.append(performance)
+
+            dense_stats = get_dense_values(
+                run_info.run_dir,
+                dataset_name,
+                metric_names=(
+                    args.dense_metric,
+                    args.dense_std_metric,
+                    args.dense_max_metric,
+                ),
+            )
+            dense_mean = dense_stats[args.dense_metric]
+            dense_std = dense_stats[args.dense_std_metric]
+            dense_max = dense_stats[args.dense_max_metric]
+            if dense_mean is not None:
+                dense_mean_xs.append(context_length)
+                dense_mean_ys.append(dense_mean)
+            if dense_max is not None:
+                dense_max_xs.append(context_length)
+                dense_max_ys.append(dense_max)
+
+        label = str(run_info.chunk_size)
+        ax_perf.plot(
+            performance_xs,
+            performance_ys,
+            marker="o",
+            linewidth=2,
+            color=color,
+            label=label,
+            markeredgecolor="black",
+            markeredgewidth=1,
+        )
+        if dense_mean_xs:
+            has_any_dense_mean_values = True
+            ax_dense.plot(
+                dense_mean_xs,
+                dense_mean_ys,
+                marker="o",
+                linewidth=2,
+                color=color,
+                markeredgecolor="black",
+                markeredgewidth=1,
+            )
+        if dense_max_xs:
+            has_any_dense_max_values = True
+            ax_dense_max.plot(
+                dense_max_xs,
+                dense_max_ys,
+                marker="o",
+                linewidth=2,
+                color=color,
+                markeredgecolor="black",
+                markeredgewidth=1,
+            )
+
+    for axis in (ax_perf, ax_dense, ax_dense_max):
+        axis.set_xscale("log", base=2)
+        axis.set_xticks(context_lengths)
+        axis.set_xticklabels(
+            [rf"$2^{{{int(value).bit_length() - 1}}}$" for value in context_lengths],
+            rotation=0,
+            ha="center",
+        )
+        axis.grid(True, alpha=0.25)
+        axis.xaxis.set_minor_locator(plt.NullLocator())
+
+    ax_perf.set_title("Retrieval Performance", fontweight="bold")
+    ax_perf.set_xlabel("Haystack Length (tokens)")
+    ax_perf.set_ylabel("ROUGE-L F1 Score")
+    ax_perf.set_ylim(-0.05, 1.05)
+
+    ax_dense.set_title("LoRA Mean Magnitude", fontweight="bold")
+    ax_dense.set_xlabel("Haystack Length (tokens)")
+    ax_dense.set_ylabel("Mean Per-Element Magnitude")
+    if has_any_dense_mean_values:
+        ax_dense.set_ylim(bottom=0.0)
+    else:
+        ax_dense.text(
+            0.5,
+            0.5,
+            "Dense LoRA stats are missing for these runs.",
+            ha="center",
+            va="center",
+            transform=ax_dense.transAxes,
+        )
+
+    ax_dense_max.set_title("LoRA Max Magnitude", fontweight="bold")
+    ax_dense_max.set_xlabel("Haystack Length (tokens)")
+    ax_dense_max.set_ylabel("Maximum Per-Element Magnitude")
+    if has_any_dense_max_values:
+        ax_dense_max.set_ylim(bottom=0.0)
+    else:
+        ax_dense_max.text(
+            0.5,
+            0.5,
+            "Dense LoRA max stats are missing for these runs.",
+            ha="center",
+            va="center",
+            transform=ax_dense_max.transAxes,
+        )
+
+    handles, labels = ax_perf.get_legend_handles_labels()
+    chunk_legend = fig.legend(
+        handles,
+        labels,
+        loc="lower center",
+        bbox_to_anchor=(0.5, -0.005),
+        ncol=min(6, len(labels)),
+        frameon=True,
+        fancybox=True,
+        title="Chunk size",
+        title_fontsize=16,
+    )
+    chunk_legend.get_frame().set_edgecolor("#bdbdbd")
+    chunk_legend.get_frame().set_linewidth(1.0)
+    chunk_legend.get_frame().set_alpha(1.0)
+
+    fig.savefig(output_path, dpi=200, bbox_inches="tight")
+    print(f"Saved plot to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/niah/4-eval-vary-num-chunks-overlap.sh

@@ -0,0 +1,6 @@
+CHECKPOINT_PATH=train_outputs/runs/Oct03_10-58-26_slurm0-a3nodeset-8_91076_c4dac553/checkpoint-1482/pytorch_model.bin
+CTX_CHUNK_OVERLAP=8
+
+for ctx_chunk_len in 16 32 64 128 256 512 1024; do 
+     WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_32_1024 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
+done

scripts/niah/4-eval-vary-num-chunks.sh

@@ -0,0 +1,26 @@
+CHECKPOINT_PATH=train_outputs/runs/Oct03_10-58-26_slurm0-a3nodeset-8_91076_c4dac553/checkpoint-1482/pytorch_model.bin
+CTX_CHUNK_OVERLAP=0
+
+for ctx_chunk_len in 128 256 512 1024 2048 4096 8192 16384 32768; do 
+     WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_28672_32768 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
+done
+
+for ctx_chunk_len in 64 128 256 512 1024 2048 4096 8192 16384; do 
+     WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_15360_16384 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
+done
+
+for ctx_chunk_len in 32 64 128 256 512 1024 2048 4096 8192; do 
+     WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_7168_8192 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
+done
+
+for ctx_chunk_len in 16 32 64 128 256 512 1024 2048 4096; do 
+     WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_3072_4096 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
+done
+
+for ctx_chunk_len in 8 16 32 64 128 256 512 1024 2048; do 
+     WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_1024_2048 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
+done
+
+for ctx_chunk_len in 4 8 16 32 64 128 256 512 1024; do 
+     WANDB_MODE=disabled uv run run_eval.py --checkpoint_path ${CHECKPOINT_PATH} --datasets ctx_magic_number_32_1024 --max_ctx_chunk_len=${ctx_chunk_len} --ctx_chunk_overlap=${CTX_CHUNK_OVERLAP} --split test --eval_batch_size_gen=4 --log_dense_lora_magnitude
+done

scripts/niah/4-plot-vary-num-chunks-overlap.py

@@ -0,0 +1,247 @@
+from __future__ import annotations
+
+import argparse
+import math
+from pathlib import Path
+
+import matplotlib
+import matplotlib as mpl
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from plot_ablation_utils import (
+    choose_latest_runs_by_dataset_chunk_and_overlap,
+    dataset_context_length,
+    ensure_output_path,
+    get_performance_value,
+    load_script_text,
+    parse_checkpoint_path,
+    parse_eval_results_root,
+    parse_expected_datasets,
+)
+
+LATTE_STYLE = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+
+
+def configure_plot_style() -> None:
+    plt.style.use(["ggplot", LATTE_STYLE])
+    plt.rcParams["axes.facecolor"] = "F7FBFC"
+    plt.rcParams["figure.facecolor"] = "white"
+    plt.rcParams["savefig.facecolor"] = "white"
+    plt.rcParams["grid.color"] = "cccccc"
+    plt.rcParams["grid.linewidth"] = 1
+    plt.rcParams["axes.edgecolor"] = "ccd0da"
+    plt.rcParams["legend.facecolor"] = "white"
+    plt.rcParams["legend.fontsize"] = 14
+    plt.rcParams["xtick.labelsize"] = 13
+    plt.rcParams["ytick.labelsize"] = 13
+    plt.rcParams["axes.labelweight"] = "bold"
+    plt.rcParams["axes.prop_cycle"] = mpl.cycler(color=plt.cm.tab10.colors)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Compare overlapping vs non-overlapping chunking for the 4-* NIAH sweeps.",
+    )
+    parser.add_argument(
+        "--non-overlap-script",
+        type=Path,
+        default=Path(__file__).with_name("4-eval-vary-num-chunks.sh"),
+        help="Path to the non-overlapping eval sweep shell script.",
+    )
+    parser.add_argument(
+        "--overlap-script",
+        type=Path,
+        default=Path(__file__).with_name("4-eval-vary-num-chunks-overlap.sh"),
+        help="Path to the overlapping eval sweep shell script.",
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=None,
+        help="Optional output PNG path. Defaults to eval-results-*/plots/4-vary-num-chunks-overlap-comparison.png.",
+    )
+    parser.add_argument(
+        "--metric",
+        default="rougeL.f1",
+        help="Performance metric to read from each *_results.json file.",
+    )
+    parser.add_argument(
+        "--non-overlap-value",
+        type=int,
+        default=0,
+        help="Overlap value used for the non-overlapping sweep.",
+    )
+    parser.add_argument(
+        "--overlap-value",
+        type=int,
+        default=8,
+        help="Overlap value used for the overlapping sweep.",
+    )
+    return parser
+
+
+def tick_label(chunk_size: int) -> str:
+    exponent = math.log2(chunk_size)
+    if exponent.is_integer():
+        return rf"$2^{{{int(exponent)}}}$"
+    return rf"$2^{{{exponent:.1f}}}$"
+
+
+def dataset_title(dataset_name: str) -> str:
+    context_length = dataset_context_length(dataset_name)
+    exponent = int(math.log2(context_length))
+    return rf"Haystack Length $2^{{{exponent}}}$"
+
+
+def line_label(overlap_value: int) -> str:
+    if overlap_value == 0:
+        return "Non-overlapping"
+    return f"Overlap = {overlap_value}"
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    configure_plot_style()
+
+    non_overlap_text = load_script_text(args.non_overlap_script)
+    overlap_text = load_script_text(args.overlap_script)
+
+    checkpoint_path = parse_checkpoint_path(non_overlap_text)
+    eval_results_root = parse_eval_results_root(checkpoint_path)
+    overlap_eval_results_root = parse_eval_results_root(
+        parse_checkpoint_path(overlap_text)
+    )
+    if overlap_eval_results_root != eval_results_root:
+        raise SystemExit(
+            "The overlap and non-overlap scripts point to different eval-results roots."
+        )
+
+    non_overlap_datasets = set(parse_expected_datasets(non_overlap_text))
+    overlap_datasets = set(parse_expected_datasets(overlap_text))
+    expected_datasets = sorted(
+        non_overlap_datasets & overlap_datasets,
+        key=dataset_context_length,
+    )
+    if not expected_datasets:
+        raise SystemExit("No shared datasets were found between the two sweep scripts.")
+
+    non_overlap_runs = choose_latest_runs_by_dataset_chunk_and_overlap(
+        eval_results_root,
+        expected_datasets,
+        ctx_chunk_overlap=args.non_overlap_value,
+    )
+    overlap_runs = choose_latest_runs_by_dataset_chunk_and_overlap(
+        eval_results_root,
+        expected_datasets,
+        ctx_chunk_overlap=args.overlap_value,
+    )
+
+    output_path = (
+        eval_results_root / "plots" / "4-vary-num-chunks-overlap-comparison.png"
+        if args.output is None
+        else args.output
+    )
+    output_path = ensure_output_path(output_path)
+
+    num_datasets = len(expected_datasets)
+    fig, axes = plt.subplots(
+        1,
+        num_datasets,
+        figsize=(6.4 * num_datasets, 5.5),
+        squeeze=False,
+        constrained_layout=False,
+    )
+    fig.subplots_adjust(bottom=0.26, top=0.84, wspace=0.28)
+
+    color_map = plt.get_cmap("tab10", 2)
+    legend_handles = []
+    legend_labels = []
+
+    for axis, dataset_name in zip(axes[0], expected_datasets):
+        plotted_anything = False
+        dataset_chunk_sizes: set[int] = set()
+
+        for overlap_value, run_infos, color in (
+            (
+                args.non_overlap_value,
+                non_overlap_runs.get(dataset_name, []),
+                color_map(0),
+            ),
+            (args.overlap_value, overlap_runs.get(dataset_name, []), color_map(1)),
+        ):
+            chunk_sizes: list[int] = []
+            performance_values: list[float] = []
+            for run_info in run_infos:
+                performance = get_performance_value(
+                    run_info.run_dir,
+                    dataset_name,
+                    metric_name=args.metric,
+                )
+                if performance is None:
+                    continue
+                chunk_sizes.append(run_info.chunk_size)
+                performance_values.append(performance)
+                dataset_chunk_sizes.add(run_info.chunk_size)
+
+            if not performance_values:
+                continue
+
+            plotted_anything = True
+            (line,) = axis.plot(
+                chunk_sizes,
+                performance_values,
+                marker="o",
+                linewidth=2,
+                color=color,
+                label=line_label(overlap_value),
+                markeredgecolor="black",
+                markeredgewidth=1,
+            )
+            if line.get_label() not in legend_labels:
+                legend_handles.append(line)
+                legend_labels.append(line.get_label())
+
+        if not plotted_anything:
+            axis.text(
+                0.5,
+                0.5,
+                "No matching runs found.",
+                ha="center",
+                va="center",
+                transform=axis.transAxes,
+            )
+        else:
+            sorted_chunk_sizes = sorted(dataset_chunk_sizes)
+            axis.set_xscale("log", base=2)
+            axis.set_xticks(sorted_chunk_sizes)
+            axis.set_xticklabels([tick_label(value) for value in sorted_chunk_sizes])
+            axis.xaxis.set_minor_locator(plt.NullLocator())
+
+        axis.grid(True, alpha=0.25)
+        axis.set_title(dataset_title(dataset_name), fontweight="bold")
+        axis.set_xlabel("Chunk Size")
+        axis.set_ylabel("ROUGE-L F1 Score")
+        axis.set_ylim(-0.05, 1.05)
+
+    if legend_handles:
+        legend = fig.legend(
+            legend_handles,
+            legend_labels,
+            loc="lower center",
+            bbox_to_anchor=(0.5, 0.01),
+            ncol=len(legend_labels),
+            frameon=True,
+            fancybox=True,
+        )
+        legend.get_frame().set_edgecolor("#bdbdbd")
+        legend.get_frame().set_linewidth(1.0)
+        legend.get_frame().set_alpha(1.0)
+
+    fig.suptitle("Overlapping vs Non-overlapping Chunking", fontweight="bold")
+    fig.savefig(output_path, dpi=200, bbox_inches="tight")
+    print(f"Saved plot to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/niah/4-plot-vary-num-chunks.py

@@ -0,0 +1,277 @@
+from __future__ import annotations
+
+import argparse
+import math
+from pathlib import Path
+
+import matplotlib
+import matplotlib as mpl
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from plot_ablation_utils import (
+    choose_latest_runs_by_dataset_and_chunk,
+    dataset_context_length,
+    ensure_output_path,
+    get_dense_values,
+    get_performance_value,
+    load_script_text,
+    parse_checkpoint_path,
+    parse_eval_results_root,
+    parse_expected_datasets,
+)
+
+LATTE_STYLE = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
+
+
+def configure_plot_style() -> None:
+    plt.style.use(["ggplot", LATTE_STYLE])
+    plt.rcParams["axes.facecolor"] = "F7FBFC"
+    plt.rcParams["figure.facecolor"] = "white"
+    plt.rcParams["savefig.facecolor"] = "white"
+    plt.rcParams["grid.color"] = "cccccc"
+    plt.rcParams["grid.linewidth"] = 1
+    plt.rcParams["axes.edgecolor"] = "ccd0da"
+    plt.rcParams["legend.facecolor"] = "white"
+    plt.rcParams["legend.fontsize"] = 14
+    plt.rcParams["xtick.labelsize"] = 13
+    plt.rcParams["ytick.labelsize"] = 13
+    plt.rcParams["axes.labelweight"] = "bold"
+    plt.rcParams["axes.prop_cycle"] = mpl.cycler(color=plt.cm.tab10.colors)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Plot the num-chunks ablation from 4-eval-vary-num-chunks.sh.",
+    )
+    parser.add_argument(
+        "--eval-script",
+        type=Path,
+        default=Path(__file__).with_name("4-eval-vary-num-chunks.sh"),
+        help="Path to the eval sweep shell script.",
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=None,
+        help="Optional output PNG path. Defaults to eval-results-*/plots/4-vary-num-chunks.png.",
+    )
+    parser.add_argument(
+        "--metric",
+        default="rougeL.f1",
+        help="Performance metric to read from each *_results.json file.",
+    )
+    parser.add_argument(
+        "--dense-metric",
+        default="mean",
+        help="Dense LoRA metric to read when available.",
+    )
+    parser.add_argument(
+        "--dense-std-metric",
+        default="std",
+        help="Dense LoRA std metric to use for the shaded area when available.",
+    )
+    parser.add_argument(
+        "--dense-max-metric",
+        default="max",
+        help="Dense LoRA max metric to plot as a separate line when available.",
+    )
+    return parser
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    configure_plot_style()
+
+    script_text = load_script_text(args.eval_script)
+    checkpoint_path = parse_checkpoint_path(script_text)
+    eval_results_root = parse_eval_results_root(checkpoint_path)
+    expected_datasets = parse_expected_datasets(script_text)
+    selected_runs_by_dataset = choose_latest_runs_by_dataset_and_chunk(
+        eval_results_root,
+        expected_datasets,
+    )
+
+    if not any(selected_runs_by_dataset.values()):
+        raise SystemExit(f"No matching eval runs found under {eval_results_root}.")
+
+    output_path = (
+        eval_results_root / "plots" / "4-vary-num-chunks.png"
+        if args.output is None
+        else args.output
+    )
+    output_path = ensure_output_path(output_path)
+
+    fig, (ax_perf, ax_dense, ax_dense_max) = plt.subplots(
+        1,
+        3,
+        figsize=(20, 5.5),
+        constrained_layout=False,
+    )
+    fig.subplots_adjust(bottom=0.27, top=0.88, wspace=0.28)
+
+    color_map = plt.get_cmap("tab10", len(expected_datasets))
+    all_num_chunks: set[int] = set()
+    has_any_dense_mean_values = False
+    has_any_dense_max_values = False
+
+    for index, dataset_name in enumerate(expected_datasets):
+        color = color_map(index)
+        context_length = dataset_context_length(dataset_name)
+        label = rf"$2^{{{int(context_length).bit_length() - 1}}}$"
+
+        num_chunks: list[int] = []
+        performance_values: list[float] = []
+        dense_mean_xs: list[int] = []
+        dense_mean_ys: list[float] = []
+        dense_max_xs: list[int] = []
+        dense_max_ys: list[float] = []
+
+        for run_info in selected_runs_by_dataset.get(dataset_name, []):
+            if run_info.chunk_size == 24:
+                continue
+
+            approx_num_chunks = max(1, context_length // run_info.chunk_size)
+            performance = get_performance_value(
+                run_info.run_dir,
+                dataset_name,
+                metric_name=args.metric,
+            )
+            if performance is None:
+                continue
+
+            num_chunks.append(approx_num_chunks)
+            performance_values.append(performance)
+            all_num_chunks.add(approx_num_chunks)
+
+            dense_stats = get_dense_values(
+                run_info.run_dir,
+                dataset_name,
+                metric_names=(
+                    args.dense_metric,
+                    args.dense_std_metric,
+                    args.dense_max_metric,
+                ),
+            )
+            dense_mean = dense_stats[args.dense_metric]
+            dense_max = dense_stats[args.dense_max_metric]
+            if dense_mean is not None:
+                dense_mean_xs.append(approx_num_chunks)
+                dense_mean_ys.append(dense_mean)
+            if dense_max is not None:
+                dense_max_xs.append(approx_num_chunks)
+                dense_max_ys.append(dense_max)
+
+        if performance_values:
+            ax_perf.plot(
+                num_chunks,
+                performance_values,
+                marker="o",
+                linewidth=2,
+                color=color,
+                label=label,
+                markeredgecolor="black",
+                markeredgewidth=1,
+            )
+        if dense_mean_ys:
+            has_any_dense_mean_values = True
+            ax_dense.plot(
+                dense_mean_xs,
+                dense_mean_ys,
+                marker="o",
+                linewidth=2,
+                color=color,
+                label=label,
+                markeredgecolor="black",
+                markeredgewidth=1,
+            )
+        if dense_max_ys:
+            has_any_dense_max_values = True
+            ax_dense_max.plot(
+                dense_max_xs,
+                dense_max_ys,
+                marker="o",
+                linewidth=2,
+                color=color,
+                label=label,
+                markeredgecolor="black",
+                markeredgewidth=1,
+            )
+
+    if not all_num_chunks:
+        raise SystemExit("No performance values were found for the selected runs.")
+
+    sorted_num_chunks = sorted(all_num_chunks)
+    for axis in (ax_perf, ax_dense, ax_dense_max):
+        axis.set_xscale("log", base=2)
+        axis.set_xticks(sorted_num_chunks)
+        axis.set_xticklabels(
+            [
+                rf"$2^{{{math.log2(value):.1f}}}$"
+                if not math.log2(value).is_integer()
+                else rf"$2^{{{int(math.log2(value))}}}$"
+                for value in sorted_num_chunks
+            ],
+            rotation=0,
+            ha="center",
+        )
+        axis.grid(True, alpha=0.25)
+        axis.xaxis.set_minor_locator(plt.NullLocator())
+
+    ax_perf.set_title("Retrieval Performance", fontweight="bold")
+    ax_perf.set_xlabel("Number of Chunks")
+    ax_perf.set_ylabel("ROUGE-L F1 Score")
+    ax_perf.set_ylim(-0.05, 1.05)
+
+    ax_dense.set_title("LoRA Mean Magnitude", fontweight="bold")
+    ax_dense.set_xlabel("Number of Chunks")
+    ax_dense.set_ylabel("Mean Per-Element Magnitude")
+    if has_any_dense_mean_values:
+        ax_dense.set_ylim(bottom=0.0)
+    else:
+        ax_dense.text(
+            0.5,
+            0.5,
+            "Dense LoRA stats are missing for these runs.",
+            ha="center",
+            va="center",
+            transform=ax_dense.transAxes,
+        )
+
+    ax_dense_max.set_title("LoRA Max Magnitude", fontweight="bold")
+    ax_dense_max.set_xlabel("Number of Chunks")
+    ax_dense_max.set_ylabel("Maximum Per-Element Magnitude")
+    if has_any_dense_max_values:
+        ax_dense_max.set_ylim(bottom=0.0)
+    else:
+        ax_dense_max.text(
+            0.5,
+            0.5,
+            "Dense LoRA max stats are missing for these runs.",
+            ha="center",
+            va="center",
+            transform=ax_dense_max.transAxes,
+        )
+
+    handles, labels = ax_perf.get_legend_handles_labels()
+    haystack_legend = fig.legend(
+        handles,
+        labels,
+        loc="lower center",
+        bbox_to_anchor=(0.5, -0.005),
+        ncol=min(6, len(labels)),
+        frameon=True,
+        fancybox=True,
+        title="Haystack length",
+        title_fontsize=16,
+    )
+    haystack_legend.get_frame().set_edgecolor("#bdbdbd")
+    haystack_legend.get_frame().set_linewidth(1.0)
+    haystack_legend.get_frame().set_alpha(1.0)
+
+    fig.savefig(output_path, dpi=200, bbox_inches="tight")
+    print(f"Saved plot to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/niah/plot_ablation_utils.py

@@ -0,0 +1,319 @@
+from __future__ import annotations
+
+import json
+import re
+import shlex
+from dataclasses import dataclass
+from pathlib import Path
+
+DATASET_RE = re.compile(r"ctx_magic_number_(\d+)_(\d+)")
+CHECKPOINT_RE = re.compile(r"^CHECKPOINT_PATH=(.+)$", re.MULTILINE)
+CMD_RE = re.compile(r"CMD:\s+(.*)$", re.MULTILINE)
+
+
+@dataclass(frozen=True)
+class RunInfo:
+    run_dir: Path
+    run_name: str
+    chunk_size: int
+    ctx_chunk_overlap: int
+    datasets: tuple[str, ...]
+    has_dense_logging: bool
+
+
+def load_script_text(script_path: Path) -> str:
+    return script_path.read_text(encoding="utf-8")
+
+
+def parse_checkpoint_path(script_text: str) -> Path:
+    match = CHECKPOINT_RE.search(script_text)
+    if match is None:
+        raise ValueError("Could not find CHECKPOINT_PATH in the eval script.")
+    return Path(match.group(1).strip().strip("\"'"))
+
+
+def parse_expected_datasets(script_text: str) -> list[str]:
+    datasets = {match.group(0) for match in DATASET_RE.finditer(script_text)}
+    if not datasets:
+        raise ValueError(
+            "Could not find any ctx_magic_number datasets in the eval script."
+        )
+    return sorted(datasets, key=dataset_context_length)
+
+
+def parse_eval_results_root(checkpoint_path: Path) -> Path:
+    checkpoint_match = re.search(r"checkpoint-(\d+)", checkpoint_path.as_posix())
+    if checkpoint_match is None:
+        raise ValueError(f"Could not parse checkpoint step from {checkpoint_path}.")
+    checkpoint_step = checkpoint_match.group(1)
+    return checkpoint_path.parent.parent / f"eval-results-{checkpoint_step}"
+
+
+def parse_dataset_context_bounds(dataset_name: str) -> tuple[int, int]:
+    match = DATASET_RE.fullmatch(dataset_name)
+    if match is None:
+        raise ValueError(f"Unsupported dataset name: {dataset_name}")
+    return int(match.group(1)), int(match.group(2))
+
+
+def dataset_context_length(dataset_name: str) -> int:
+    return parse_dataset_context_bounds(dataset_name)[1]
+
+
+def _extract_flag_value(tokens: list[str], flag: str) -> str | None:
+    for index, token in enumerate(tokens):
+        if token == flag and index + 1 < len(tokens):
+            return tokens[index + 1]
+        if token.startswith(f"{flag}="):
+            return token.split("=", 1)[1]
+    return None
+
+
+def _extract_datasets(tokens: list[str]) -> tuple[str, ...]:
+    if "--datasets" not in tokens:
+        return ()
+
+    index = tokens.index("--datasets") + 1
+    datasets: list[str] = []
+    while index < len(tokens) and not tokens[index].startswith("--"):
+        datasets.append(tokens[index])
+        index += 1
+    return tuple(datasets)
+
+
+def parse_run_info(run_dir: Path) -> RunInfo | None:
+    log_path = run_dir / "debug.log"
+    if not log_path.exists():
+        return None
+
+    log_text = log_path.read_text(encoding="utf-8", errors="replace")
+    match = CMD_RE.search(log_text)
+    if match is None:
+        return None
+
+    command = match.group(1).strip()
+    try:
+        tokens = shlex.split(command)
+    except ValueError:
+        tokens = command.split()
+
+    chunk_size_value = _extract_flag_value(tokens, "--max_ctx_chunk_len")
+    if chunk_size_value is None:
+        return None
+
+    datasets = _extract_datasets(tokens)
+    if not datasets:
+        return None
+
+    return RunInfo(
+        run_dir=run_dir,
+        run_name=run_dir.name,
+        chunk_size=int(chunk_size_value),
+        ctx_chunk_overlap=int(_extract_flag_value(tokens, "--ctx_chunk_overlap") or 0),
+        datasets=datasets,
+        has_dense_logging="--log_dense_lora_magnitude" in tokens,
+    )
+
+
+def _run_has_expected_outputs(run_dir: Path, datasets: list[str]) -> bool:
+    return all(
+        (run_dir / f"test_{dataset_name}_results.json").exists()
+        for dataset_name in datasets
+    )
+
+
+def choose_latest_runs_by_chunk(
+    eval_results_root: Path,
+    expected_datasets: list[str],
+) -> list[RunInfo]:
+    selected: dict[int, RunInfo] = {}
+    expected_dataset_set = set(expected_datasets)
+
+    for run_dir in sorted(eval_results_root.iterdir()):
+        if not run_dir.is_dir():
+            continue
+
+        run_info = parse_run_info(run_dir)
+        if run_info is None:
+            continue
+        if set(run_info.datasets) != expected_dataset_set:
+            continue
+        if not _run_has_expected_outputs(run_dir, expected_datasets):
+            continue
+
+        current = selected.get(run_info.chunk_size)
+        if current is None:
+            selected[run_info.chunk_size] = run_info
+            continue
+
+        current_key = (int(current.has_dense_logging), current.run_name)
+        candidate_key = (int(run_info.has_dense_logging), run_info.run_name)
+        if candidate_key > current_key:
+            selected[run_info.chunk_size] = run_info
+
+    return [selected[chunk_size] for chunk_size in sorted(selected)]
+
+
+def choose_latest_runs_by_dataset_and_chunk(
+    eval_results_root: Path,
+    expected_datasets: list[str],
+) -> dict[str, list[RunInfo]]:
+    selected: dict[tuple[str, int], RunInfo] = {}
+    expected_dataset_set = set(expected_datasets)
+
+    for run_dir in sorted(eval_results_root.iterdir()):
+        if not run_dir.is_dir():
+            continue
+
+        run_info = parse_run_info(run_dir)
+        if run_info is None or len(run_info.datasets) != 1:
+            continue
+
+        dataset_name = run_info.datasets[0]
+        if dataset_name not in expected_dataset_set:
+            continue
+        if not _run_has_expected_outputs(run_dir, [dataset_name]):
+            continue
+
+        key = (dataset_name, run_info.chunk_size)
+        current = selected.get(key)
+        if current is None:
+            selected[key] = run_info
+            continue
+
+        current_key = (int(current.has_dense_logging), current.run_name)
+        candidate_key = (int(run_info.has_dense_logging), run_info.run_name)
+        if candidate_key > current_key:
+            selected[key] = run_info
+
+    grouped: dict[str, list[RunInfo]] = {
+        dataset_name: [] for dataset_name in expected_datasets
+    }
+    for (dataset_name, _), run_info in selected.items():
+        grouped[dataset_name].append(run_info)
+
+    for dataset_name, run_infos in grouped.items():
+        grouped[dataset_name] = sorted(
+            run_infos, key=lambda run_info: run_info.chunk_size
+        )
+
+    return grouped
+
+
+def choose_latest_runs_by_dataset_chunk_and_overlap(
+    eval_results_root: Path,
+    expected_datasets: list[str],
+    ctx_chunk_overlap: int,
+) -> dict[str, list[RunInfo]]:
+    selected: dict[tuple[str, int], RunInfo] = {}
+    expected_dataset_set = set(expected_datasets)
+
+    for run_dir in sorted(eval_results_root.iterdir()):
+        if not run_dir.is_dir():
+            continue
+
+        run_info = parse_run_info(run_dir)
+        if run_info is None or len(run_info.datasets) != 1:
+            continue
+        if run_info.ctx_chunk_overlap != ctx_chunk_overlap:
+            continue
+
+        dataset_name = run_info.datasets[0]
+        if dataset_name not in expected_dataset_set:
+            continue
+        if not _run_has_expected_outputs(run_dir, [dataset_name]):
+            continue
+
+        key = (dataset_name, run_info.chunk_size)
+        current = selected.get(key)
+        if current is None:
+            selected[key] = run_info
+            continue
+
+        current_key = (int(current.has_dense_logging), current.run_name)
+        candidate_key = (int(run_info.has_dense_logging), run_info.run_name)
+        if candidate_key > current_key:
+            selected[key] = run_info
+
+    grouped: dict[str, list[RunInfo]] = {
+        dataset_name: [] for dataset_name in expected_datasets
+    }
+    for (dataset_name, _), run_info in selected.items():
+        grouped[dataset_name].append(run_info)
+
+    for dataset_name, run_infos in grouped.items():
+        grouped[dataset_name] = sorted(
+            run_infos, key=lambda run_info: run_info.chunk_size
+        )
+
+    return grouped
+
+
+def load_results_json(run_dir: Path, dataset_name: str) -> dict[str, object]:
+    path = run_dir / f"test_{dataset_name}_results.json"
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def load_dense_stats_json(run_dir: Path, dataset_name: str) -> dict[str, object] | None:
+    path = run_dir / f"test_{dataset_name}_dense_lora_magnitude.json"
+    if not path.exists():
+        return None
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _to_float(value: object) -> float | None:
+    if value in (None, "None", "N/A"):
+        return None
+    if isinstance(value, (int, float)):
+        return float(value)
+    return None
+
+
+def get_performance_value(
+    run_dir: Path,
+    dataset_name: str,
+    metric_name: str = "rougeL.f1",
+) -> float | None:
+    result_json = load_results_json(run_dir, dataset_name)
+    return _to_float(result_json.get(f"test_{dataset_name}_{metric_name}"))
+
+
+def get_dense_value(
+    run_dir: Path,
+    dataset_name: str,
+    metric_name: str = "mean",
+) -> float | None:
+    result_json = load_results_json(run_dir, dataset_name)
+    metric_key = f"test_{dataset_name}_dense_lora_magnitude_{metric_name}"
+    value = _to_float(result_json.get(metric_key))
+    if value is not None:
+        return value
+
+    dense_stats = load_dense_stats_json(run_dir, dataset_name)
+    if dense_stats is None:
+        return None
+
+    overall = dense_stats.get("overall", {})
+    if not isinstance(overall, dict):
+        return None
+    return _to_float(overall.get(metric_name))
+
+
+def get_dense_values(
+    run_dir: Path,
+    dataset_name: str,
+    metric_names: list[str] | tuple[str, ...],
+) -> dict[str, float | None]:
+    return {
+        metric_name: get_dense_value(
+            run_dir,
+            dataset_name,
+            metric_name=metric_name,
+        )
+        for metric_name in metric_names
+    }
+
+
+def ensure_output_path(output_path: Path) -> Path:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    return output_path

scripts/sbatch_1_gpu.sh

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

src/ctx_to_lora/data/definitions.py

@@ -1,5 +1,25 @@
+from datasets import Features, Value
+
 IGNORE_INDEX = -100
 
+OOLONG_SYNTH_FEATURES = Features(
+    {
+        "id": Value("int64"),
+        "context_len": Value("int64"),
+        "dataset": Value("large_string"),
+        "context_window_text": Value("large_string"),
+        "context_window_text_with_labels": Value("large_string"),
+        "question": Value("large_string"),
+        "task_group": Value("large_string"),
+        "task": Value("large_string"),
+        "answer": Value("large_string"),
+        "answer_type": Value("large_string"),
+        "input_subset": Value("large_string"),
+        "num_labels": Value("int64"),
+        "context_window_id": Value("int64"),
+    }
+)
+
 TRANSFORMED_DATA_DIR = "data/processed_datasets"
 RAW_DATA_DIR = "data/raw_datasets/"
 SELF_GEN_DATA_DIR = f"{RAW_DATA_DIR}/self_gen/"
@@ -27,12 +47,14 @@ LONGBENCH_TASKS = [
     "longbench/qasper",
     "longbench/multifieldqa_en",
     "longbench/2wikimqa",
+    "longbench/gov_report",
 ]
 
 LONGBENCH_E_TASKS = [
     "longbench/qasper_e",
     "longbench/multifieldqa_en_e",
     "longbench/2wikimqa_e",
+    "longbench/gov_report_e",
 ]
 
 DS_KWARGS = {
@@ -125,6 +147,35 @@ DS_KWARGS = {
             split="train",
         )
     ),
+    "clipper": dict(
+        train=dict(
+            path="chtmp223/CLIPPER",
+            split="train",
+            verification_mode="no_checks",
+        ),
+        validation=dict(
+            path="chtmp223/CLIPPER",
+            split="dev",
+            verification_mode="no_checks",
+        ),
+        test=dict(
+            path="chtmp223/CLIPPER",
+            split="test",
+            verification_mode="no_checks",
+        ),
+    ),
+    "oolong-synth": dict(
+        validation=dict(
+            path="oolongbench/oolong-synth",
+            split="validation",
+            features=OOLONG_SYNTH_FEATURES,
+        ),
+        test=dict(
+            path="oolongbench/oolong-synth",
+            split="test",
+            features=OOLONG_SYNTH_FEATURES,
+        ),
+    ),
 }
 
 # add ctx_numbers
@@ -245,6 +296,7 @@ EVAL_INTX_TEMPLATES = {
     "longbench/qasper": 'Answer the question as concisely as you can, using a single phrase or sentence if possible.\nIf the question cannot be answered based on the information in the article, write "unanswerable".\nIf the question is a yes/no question, answer "yes", "no", or "unanswerable". Do not provide any explanation.\n\nQuestion: {input}',
     "longbench/multifieldqa_en": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
     "longbench/2wikimqa": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
+    "longbench/gov_report": "Write a concise summary of the report. Focus on the main issues, findings, and conclusions.",
 }
 for ds_name in LONGBENCH_E_TASKS:
     EVAL_INTX_TEMPLATES[ds_name] = EVAL_INTX_TEMPLATES[ds_name[:-2]]

src/ctx_to_lora/data/preprocessing_fn.py

@@ -1,5 +1,7 @@
+import hashlib
 import logging
 import random
+import re
 from collections.abc import Callable
 from typing import Any
 
@@ -8,6 +10,8 @@ from ctx_to_lora.utils import concat_list
 
 logger = logging.getLogger()
 
+_CLIPPER_CONTEXT_RE = re.compile(r"<context>(.*?)</context>", re.DOTALL)
+
 
 def closed_qa_prompting(prompt: str):
     template = random.choice(CLOSED_QA_INTX_TEMPLATES)
@@ -25,6 +29,20 @@ def chat_to_str(messages: list[dict[str, str]]):
     )
 
 
+def parse_clipper_user_message(user_message: str) -> tuple[str, str]:
+    match = _CLIPPER_CONTEXT_RE.search(user_message)
+    if match is None:
+        raise ValueError("Could not find <context>...</context> block in CLIPPER row.")
+
+    context = match.group(1).strip()
+    prompt_parts = [
+        user_message[: match.start()].strip(),
+        user_message[match.end() :].strip(),
+    ]
+    prompt = "\n\n".join(part for part in prompt_parts if part)
+    return context, prompt
+
+
 def get_preprocessing_fn(
     ds_name: str,
     is_eval: bool,
@@ -75,6 +93,52 @@ def get_preprocessing_fn(
                 "response": sample["answers"][0],
             }
 
+    elif ds_name == "clipper":
+
+        def f(sample):
+            messages = sample["messages"]
+            if len(messages) < 3:
+                raise ValueError(
+                    f"Expected at least 3 messages in CLIPPER row, got {len(messages)}"
+                )
+
+            context, prompt = parse_clipper_user_message(messages[1]["content"])
+            pair_id = hashlib.sha1(
+                f"{sample['true_claim']}\n{sample['false_claim']}".encode()
+            ).hexdigest()
+
+            return {
+                # "system_message": "",
+                "context": context.replace("-", " "),
+                # "prompt": prompt,
+                "prompt": "Who is Tarrano?",
+                "response": messages[2]["content"],
+                "book_name": sample["book_name"],
+                "book_title": sample["book_title"],
+                "claim_type": sample["claim_type"],
+                "clipper_pair_id": pair_id,
+                "clipper_status": sample["status"],
+                "clipper_true_claim": sample["true_claim"],
+                "clipper_false_claim": sample["false_claim"],
+            }
+
+    elif ds_name == "oolong-synth":
+
+        def f(sample):
+            return {
+                "context": sample["context_window_text"],
+                "prompt": sample["question"],
+                "response": sample["answer"],
+                "oolong_id": sample["id"],
+                "oolong_dataset": sample["dataset"],
+                "oolong_task_group": sample["task_group"],
+                "oolong_task": sample["task"],
+                "oolong_answer_type": sample["answer_type"],
+                "oolong_input_subset": sample["input_subset"],
+                "oolong_num_labels": sample["num_labels"],
+                "oolong_context_window_id": sample["context_window_id"],
+            }
+
     elif ds_name == "pwc" or ds_name == "pwc_tiny":
         # original pwc
         def f(sample):

src/ctx_to_lora/data/processing.py

@@ -12,7 +12,14 @@ from typing import Any
 import datasets
 import numpy as np
 import torch
-from datasets import Dataset, interleave_datasets, is_caching_enabled, load_dataset
+from datasets import (
+    Dataset,
+    Features,
+    Value,
+    interleave_datasets,
+    is_caching_enabled,
+    load_dataset,
+)
 from transformers import PreTrainedTokenizerBase
 
 from ctx_to_lora.data.definitions import (
@@ -29,11 +36,29 @@ from ctx_to_lora.data.self_gen_template import QA_PROMPT_TEMPLATE
 from ctx_to_lora.utils import check_is_iterable, concat_list
 
 logger = logging.getLogger()
+OOLONG_SUBSAMPLE_SIZE = 350
+OOLONG_SUBSAMPLE_SEED = 42
+OOLONG_MAX_CONTEXT_LEN = 2**16
 
 COLS_TO_KEEP_PREPROCESSING = [
     "context",
     "prompts",
     "responses",
+    "messages",
+    "true_claim",
+    "false_claim",
+    "status",
+    "claim_type",
+    "book_name",
+    "book_title",
+    "oolong_id",
+    "oolong_dataset",
+    "oolong_task_group",
+    "oolong_task",
+    "oolong_answer_type",
+    "oolong_input_subset",
+    "oolong_num_labels",
+    "oolong_context_window_id",
     "qas",
     "variation",
     "logprobs_vals",
@@ -48,11 +73,156 @@ COLS_TO_KEEP_TOKENIZED = [
     "labels",
     "context",
     "ctx_ids",
+    "book_name",
+    "book_title",
+    "claim_type",
+    "clipper_pair_id",
+    "clipper_status",
+    "clipper_true_claim",
+    "clipper_false_claim",
+    "oolong_id",
+    "oolong_dataset",
+    "oolong_task_group",
+    "oolong_task",
+    "oolong_answer_type",
+    "oolong_input_subset",
+    "oolong_num_labels",
+    "oolong_context_window_id",
     "logprobs_vals",
     "logprobs_indices",
 ]
 
 
+def get_large_string_overridden_features(
+    features: Features,
+    feature_overrides: dict[str, Any],
+) -> Features:
+    overridden_features = Features(dict(features))
+    for column, feature in feature_overrides.items():
+        if column in overridden_features:
+            overridden_features[column] = feature
+    return overridden_features
+
+
+def get_oolong_preprocessed_features(features: Features) -> Features:
+    return get_large_string_overridden_features(
+        features,
+        {
+            "context": Value("large_string"),
+            "prompts": [Value("large_string")],
+            "responses": [Value("large_string")],
+            "oolong_dataset": Value("large_string"),
+            "oolong_task_group": Value("large_string"),
+            "oolong_task": Value("large_string"),
+            "oolong_answer_type": Value("large_string"),
+            "oolong_input_subset": Value("large_string"),
+        },
+    )
+
+
+def get_oolong_messages_features(features: Features) -> Features:
+    overridden_features = get_oolong_preprocessed_features(features)
+    overridden_features["messages_list"] = [
+        [
+            {
+                "role": Value("string"),
+                "content": Value("large_string"),
+            }
+        ]
+    ]
+    return overridden_features
+
+
+def get_dataset_processing_config(ds_name: str, split: str) -> dict[str, Any]:
+    config = {
+        "num_proc": 16,
+        "input_tokenize_batch_size": 100_000,
+        "context_tokenize_batch_size": 100_000,
+        "qa_split_batch_size": 12_500,
+        "preprocess_writer_batch_size": None,
+        "messages_writer_batch_size": None,
+    }
+    if ds_name == "oolong-synth" and "train" not in split:
+        # Oolong has very large contexts, so we still need a conservative path,
+        # but the large_string fixes let us relax the earlier one-at-a-time
+        # settings to recover some throughput.
+        config.update(
+            num_proc=2,
+            input_tokenize_batch_size=512,
+            context_tokenize_batch_size=512,
+            qa_split_batch_size=512,
+            preprocess_writer_batch_size=8,
+            messages_writer_batch_size=8,
+        )
+    return config
+
+
+def stratified_subsample_dataset(
+    ds: Dataset,
+    group_column: str,
+    max_samples: int,
+    seed: int,
+) -> Dataset:
+    if len(ds) <= max_samples:
+        return ds
+
+    group_to_indices: dict[Any, list[int]] = {}
+    for idx, group_value in enumerate(ds[group_column]):
+        group_to_indices.setdefault(group_value, []).append(idx)
+
+    group_values = sorted(group_to_indices)
+    if len(group_values) > max_samples:
+        raise ValueError(
+            f"Cannot keep all {len(group_values)} groups in only {max_samples} samples."
+        )
+
+    quotas = {group_value: 0 for group_value in group_values}
+    remaining = max_samples
+    active_groups = list(group_values)
+
+    # Allocate samples in round-robin order so each context_len bucket is
+    # represented as evenly as possible, while naturally redistributing quota
+    # from exhausted small buckets into the remaining ones.
+    while remaining > 0 and active_groups:
+        next_active_groups = []
+        for group_value in active_groups:
+            if remaining == 0:
+                break
+            if quotas[group_value] < len(group_to_indices[group_value]):
+                quotas[group_value] += 1
+                remaining -= 1
+            if quotas[group_value] < len(group_to_indices[group_value]):
+                next_active_groups.append(group_value)
+        active_groups = next_active_groups
+
+    rng = np.random.default_rng(seed)
+    sampled_indices = []
+    for group_value in group_values:
+        group_indices = np.array(group_to_indices[group_value])
+        group_quota = quotas[group_value]
+        if group_quota == 0:
+            continue
+        if group_quota == len(group_indices):
+            sampled_indices.extend(group_indices.tolist())
+        else:
+            sampled_indices.extend(
+                rng.choice(group_indices, size=group_quota, replace=False).tolist()
+            )
+
+    return ds.select(sorted(sampled_indices))
+
+
+def filter_dataset_by_max_value(
+    ds: Dataset,
+    column: str,
+    max_value: int,
+) -> Dataset:
+    eligible_indices = [
+        idx for idx, value in enumerate(ds[column]) if value <= max_value
+    ]
+    return ds.select(eligible_indices)
+
+
 def get_ds_prob(train_ds_len: list[int], total_len: int):
     # if a dataset is smaller than 1%, make it 1%
     probs = [0 for _ in train_ds_len]
@@ -208,9 +378,11 @@ def load_and_process_dataset(
     split: str,
     add_negative_prompt: bool,
     num_proc: int,
+    writer_batch_size: int | None = None,
     remove_cols: bool = True,
 ):
     logger.info(f"Loading dataset {ds_name} with split {split}...")
+    is_oolong_eval = ds_name == "oolong-synth" and "train" not in split
     try:
         ds_kwargs = get_ds_kwargs(ds_name, split)
         skip = ds_kwargs.pop("skip", None)
@@ -224,6 +396,18 @@ def load_and_process_dataset(
         raise ValueError(
             f"Failed to load dataset {ds_name} with split {split}. Error: {e}"
         )
+    if is_oolong_eval:
+        ds = filter_dataset_by_max_value(
+            ds=ds,
+            column="context_len",
+            max_value=OOLONG_MAX_CONTEXT_LEN,
+        )
+        ds = stratified_subsample_dataset(
+            ds=ds,
+            group_column="context_len",
+            max_samples=OOLONG_SUBSAMPLE_SIZE,
+            seed=OOLONG_SUBSAMPLE_SEED,
+        )
     cols_to_remove = None
     if remove_cols:
         cols_to_remove = [
@@ -233,12 +417,15 @@ def load_and_process_dataset(
     ds = ds.map(
         get_preprocessing_fn(ds_name, is_eval),
         remove_columns=cols_to_remove,
-        num_proc=16,
+        num_proc=num_proc,
+        writer_batch_size=writer_batch_size,
     )
+    if is_oolong_eval:
+        ds = ds.cast(get_oolong_preprocessed_features(ds.features))
     ds = ds.filter(
         filter_none,
         batched=False,
-        num_proc=16,
+        num_proc=num_proc,
     )
     if add_negative_prompt and "context" in ds:
         ds = ds.map(
@@ -266,6 +453,7 @@ def get_tokenized_dataset(
     add_ctx_to_chat: bool,
     add_negative_prompt: bool,
     use_kl_loss: bool,
+    ctx_chunk_overlap: int = 0,
     max_new_tokens: int = 256,
     add_self_distill_template: bool = False,
     set_format: str | None = None,
@@ -273,17 +461,21 @@ def get_tokenized_dataset(
     truncate_if_too_long_ctx: bool = False,
     flip_ctx_inp: bool = False,
 ) -> dict[str, Any]:
+    if "train" in split and ctx_chunk_overlap > 0:
+        raise ValueError("`ctx_chunk_overlap` is only supported for eval splits.")
     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 = ctx_model_max_len is not None
+    processing_config = get_dataset_processing_config(ds_name, split)
 
     load_and_process_kwargs = dict(
         ds_name=ds_name,
         split=split,
         add_negative_prompt=add_negative_prompt,
+        writer_batch_size=processing_config["preprocess_writer_batch_size"],
     )
     tokenize_kwargs = dict(
         max_qas_len=max_qas_len,
@@ -292,6 +484,7 @@ def get_tokenized_dataset(
         ctx_model_max_len=ctx_model_max_len,
         add_ctx_to_chat=add_ctx_to_chat,
         max_ctx_chunk_len=max_ctx_chunk_len,
+        ctx_chunk_overlap=ctx_chunk_overlap,
         min_ctx_chunk_len=min_ctx_chunk_len,
         num_chunk_probs=num_chunk_probs,
         max_ctx_chunk_num=max_ctx_chunk_num,
@@ -320,7 +513,7 @@ def get_tokenized_dataset(
             )
         return tokenized_ds
 
-    num_proc = 4
+    num_proc = processing_config["num_proc"]
     ds = load_and_process_dataset(
         **load_and_process_kwargs,
         num_proc=num_proc,
@@ -351,6 +544,7 @@ def get_tokenized_dataset(
         ds = ds.map(unsqueeze, fn_kwargs={"column": "prompts"})
 
     tokenized_ds = construct_and_tokenize_ctx_qa(
+        ds_name=ds_name,
         ds=ds,
         tokenizer=tokenizer,
         ctx_tokenizer=ctx_tokenizer,
@@ -358,6 +552,10 @@ def get_tokenized_dataset(
         num_proc=num_proc,
         truncate_if_too_long_inp=truncate_if_too_long_inp,
         truncate_if_too_long_ctx=truncate_if_too_long_ctx,
+        input_tokenize_batch_size=processing_config["input_tokenize_batch_size"],
+        context_tokenize_batch_size=processing_config["context_tokenize_batch_size"],
+        qa_split_batch_size=processing_config["qa_split_batch_size"],
+        messages_writer_batch_size=processing_config["messages_writer_batch_size"],
         **tokenize_kwargs,
     )
 
@@ -388,39 +586,51 @@ def construct_and_tokenize_ctx_qa(
     min_ctx_chunk_len,
     num_chunk_probs,
     max_ctx_chunk_num,
+    ds_name,
     ds,
     split,
     max_new_tokens,
+    ctx_chunk_overlap=0,
     add_self_distill_template=False,
     set_format=None,
     num_proc=None,
     truncate_if_too_long_inp=False,
     truncate_if_too_long_ctx=False,
+    input_tokenize_batch_size=100_000,
+    context_tokenize_batch_size=100_000,
+    qa_split_batch_size=12_500,
+    messages_writer_batch_size=None,
 ):
     is_train = "train" in split
+    is_oolong_eval = ds_name == "oolong-synth" and not is_train
     # 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=16)
+        tokenized_ds = ds.map(get_labels_from_input_ids, num_proc=num_proc)
     else:
         # construct messages from prompts and responses
         # add "messages_list" field
+        logging.info("Converting context/prompt/response to messages format...")
         ds = ds.map(
             convert_ctx_prompt_response_to_messages,
             fn_kwargs={
                 "add_ctx_to_chat": add_ctx_to_chat,
                 "add_self_distill_template": add_self_distill_template,
             },
-            num_proc=16,
+            num_proc=num_proc,
+            features=(
+                get_oolong_messages_features(ds.features) if is_oolong_eval else None
+            ),
+            writer_batch_size=messages_writer_batch_size,
         )
         # add `input_ids`, `attention_mask`, `labels`
         os.environ["TOKENIZERS_PARALLELISM"] = "true"
-        logging.debug("Tokenizing inputs")
+        logging.info("Tokenizing inputs")
 
         tokenized_ds = ds.map(
             get_sft_prompt_formatting_fn(tokenizer),
             batched=True,
-            batch_size=100_000,
+            batch_size=input_tokenize_batch_size,
         )
 
     tokenized_ds = tokenized_ds.remove_columns(
@@ -428,7 +638,7 @@ def construct_and_tokenize_ctx_qa(
     )
     tokenized_ds = tokenized_ds.filter(
         lambda x: bool(x["input_ids"]),  # remove empty "input_ids"
-        num_proc=16,
+        num_proc=num_proc,
     )
 
     if need_ctx_ids:
@@ -441,19 +651,19 @@ def construct_and_tokenize_ctx_qa(
                 detokenize_ctx_text,
                 fn_kwargs={"tokenizer": tokenizer},
                 batched=True,
-                batch_size=100_000,
+                batch_size=context_tokenize_batch_size,
                 remove_columns=["ctx_ids"],
             )
 
         if "ctx_ids" not in tokenized_ds.column_names:
             # tokenize the ctx_text to get ctx_ids and ctx_attn_mask
             os.environ["TOKENIZERS_PARALLELISM"] = "true"
-            logging.debug("Tokenizing context")
+            logging.info("Tokenizing context")
             tokenized_ds = tokenized_ds.map(
                 tokenize_ctx_text,
                 fn_kwargs={"tokenizer": ctx_tokenizer},
                 batched=True,
-                batch_size=100_000,
+                batch_size=context_tokenize_batch_size,
                 remove_columns=["context"],
             )
 
@@ -465,6 +675,7 @@ def construct_and_tokenize_ctx_qa(
         # with some big caveats, e.g., losing order info
         split_ctx_kwargs = {
             "max_chunk_len": max_ctx_chunk_len,
+            "chunk_overlap": ctx_chunk_overlap,
             "min_chunk_len": min_ctx_chunk_len,
             "num_chunk_probs": num_chunk_probs,
             "max_num_split": max_ctx_chunk_num,
@@ -475,7 +686,7 @@ def construct_and_tokenize_ctx_qa(
         tokenized_ds = tokenized_ds.map(
             split_too_long_ctx,
             fn_kwargs=split_ctx_kwargs,
-            num_proc=16,
+            num_proc=num_proc,
         )
 
         logging.info(f"Num samples after ctx chunking: {len(tokenized_ds)}")
@@ -493,11 +704,12 @@ def construct_and_tokenize_ctx_qa(
             split_too_long_qas,
             fn_kwargs=split_qa_kwargs,
             batched=True,
-            batch_size=12_500,
-            num_proc=16,
+            batch_size=qa_split_batch_size,
+            num_proc=num_proc,
         )
     if "train" not in split:
         # squeeze since we always have one query per sample in eval
+        logger.info("Squeezing input_ids and labels for eval...")
         tokenized_ds = tokenized_ds.map(squeeze_tokens, num_proc=num_proc)
         tokenized_ds = tokenized_ds.map(
             add_length_info,
@@ -683,6 +895,7 @@ def split_too_long_ctx(
     min_chunk_len: int,
     max_num_split: int | None,
     is_train: bool,
+    chunk_overlap: int = 0,
 ) -> dict[str, Any]:
     """
     Split context into smaller chunks if it exceeds the maximum length.
@@ -698,6 +911,12 @@ def split_too_long_ctx(
 
     chunk_len = max_chunk_len
     ctx_ids = sample["ctx_ids"]
+    if chunk_overlap < 0:
+        raise ValueError("`chunk_overlap` must be >= 0.")
+    if chunk_len > 0 and chunk_overlap >= chunk_len:
+        raise ValueError("`chunk_overlap` must be smaller than `max_chunk_len`.")
+    if is_train and chunk_overlap > 0:
+        raise ValueError("`chunk_overlap` is only supported for eval splits.")
     # Early exits
     if chunk_len <= 0 and max_num_split is None:
         return {"ctx_ids": [ctx_ids], "n_ctx_chunks": 1}
@@ -732,7 +951,22 @@ def split_too_long_ctx(
         return {"ctx_ids": [ctx_ids], "n_ctx_chunks": 1}
 
     avg_len = ceil(len(ctx_ids) / n_chunks)
-    chunks = [ctx_ids[i : i + avg_len] for i in range(0, len(ctx_ids), avg_len)]
+
+    if chunk_overlap > 0:
+        if chunk_overlap >= avg_len:
+            raise ValueError(
+                "Derived eval chunk window is too small for the requested overlap: "
+                f"window_len={avg_len}, chunk_overlap={chunk_overlap}."
+            )
+
+        stride = avg_len - chunk_overlap
+        last_start = max(0, len(ctx_ids) - avg_len)
+        start_positions = list(range(0, last_start + 1, stride))
+        if start_positions[-1] != last_start:
+            start_positions.append(last_start)
+        chunks = [ctx_ids[i : i + avg_len] for i in start_positions]
+    else:
+        chunks = [ctx_ids[i : i + avg_len] for i in range(0, len(ctx_ids), avg_len)]
 
     ctx_affixes = CTX_AFFIXES[model_name_or_path]
     prefix = ctx_affixes["prefix"]

src/ctx_to_lora/eval_utils.py

@@ -1,3 +1,4 @@
+import ast
 import json
 import logging
 import os
@@ -9,6 +10,7 @@ from collections import Counter, defaultdict
 from dataclasses import fields
 from functools import partial
 
+import dateutil
 import numpy as np
 import pandas as pd
 import torch
@@ -55,6 +57,7 @@ from ctx_to_lora.modeling.context_distillation import CtxDistillModel
 from ctx_to_lora.modeling.generative_adapter import GenerativeAdapter
 from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
 from ctx_to_lora.modeling.llm_lingua import LLMLinguaModel
+from ctx_to_lora.modeling.rag import DocToLoraRAGModel, RAGModel
 from ctx_to_lora.modeling.text_to_lora import TextToLoRA
 from ctx_to_lora.tracker.tracker import (
     add_tracker,
@@ -82,6 +85,11 @@ _SQUOTES = re.compile(r"[‘’ʼ]")  # curly apostrophes → '
 _ELLIPSIS = re.compile(r"…")  # single‐char ellipsis → "..."
 _ENDASH = re.compile(r"\u2013")
 _EMDASH = re.compile(r"\u2014")
+_CLIPPER_ANSWER_TAG = re.compile(
+    r"<answer>\s*(true|false)\s*</answer>",
+    re.IGNORECASE | re.DOTALL,
+)
+_CLIPPER_BOOL = re.compile(r"\b(true|false)\b", re.IGNORECASE)
 
 
 def humanize_str(text: str) -> str:
@@ -178,6 +186,199 @@ def compute_qa_f1_score(
     ), dict(qa_f1_score=f1_scores, qa_precision=precisions, qa_recall=recalls)
 
 
+def extract_clipper_answer(text: str) -> str | None:
+    if not isinstance(text, str):
+        return None
+
+    text = humanize_str(text)
+    match = _CLIPPER_ANSWER_TAG.search(text)
+    if match is not None:
+        return match.group(1).lower()
+
+    matches = _CLIPPER_BOOL.findall(text)
+    if matches:
+        return matches[-1].lower()
+
+    return None
+
+
+def compute_clipper_metrics(
+    decoded_txts: list[dict[str, object]],
+) -> tuple[dict[str, float | int | str], dict[str, list[float]], dict[str, int]]:
+    row_accs = []
+    parse_successes = []
+    true_accs = []
+    false_accs = []
+    pair_outcomes = defaultdict(dict)
+
+    for sample in decoded_txts:
+        predicted = extract_clipper_answer(sample["generated"])
+        gold = sample.get("clipper_status")
+        if gold not in {"true", "false"}:
+            gold = extract_clipper_answer(sample.get("label", ""))
+
+        is_parsed = predicted in {"true", "false"}
+        is_correct = bool(is_parsed and gold in {"true", "false"} and predicted == gold)
+
+        row_accs.append(float(is_correct))
+        parse_successes.append(float(is_parsed))
+
+        if gold == "true":
+            true_accs.append(float(is_correct))
+        elif gold == "false":
+            false_accs.append(float(is_correct))
+
+        pair_id = sample.get("clipper_pair_id")
+        if pair_id and gold in {"true", "false"}:
+            pair_outcomes[pair_id][gold] = is_correct
+
+    pair_correct_by_id = {}
+    for pair_id, outcomes in pair_outcomes.items():
+        if {"true", "false"}.issubset(outcomes):
+            pair_correct_by_id[pair_id] = float(outcomes["true"] and outcomes["false"])
+
+    metrics = {
+        "clipper_accuracy": float(np.mean(row_accs)) if row_accs else "None",
+        "clipper_parse_rate": (
+            float(np.mean(parse_successes)) if parse_successes else "None"
+        ),
+        "clipper_true_accuracy": float(np.mean(true_accs)) if true_accs else "None",
+        "clipper_false_accuracy": float(np.mean(false_accs)) if false_accs else "None",
+        "clipper_pair_accuracy": (
+            float(np.mean(list(pair_correct_by_id.values())))
+            if pair_correct_by_id
+            else "None"
+        ),
+    }
+    per_sample_metric = {
+        "clipper_accuracy": row_accs,
+        "clipper_parse_rate": parse_successes,
+        "clipper_pair_accuracy": [
+            pair_correct_by_id.get(sample.get("clipper_pair_id"), 0.0)
+            for sample in decoded_txts
+        ],
+    }
+    counts = {
+        "clipper_accuracy": len(row_accs),
+        "clipper_parse_rate": len(parse_successes),
+        "clipper_true_accuracy": len(true_accs),
+        "clipper_false_accuracy": len(false_accs),
+        "clipper_pair_accuracy": len(pair_correct_by_id),
+    }
+    return metrics, per_sample_metric, counts
+
+
+def parse_oolong_gold_answer(answer: object) -> object | None:
+    if not isinstance(answer, str):
+        return answer
+
+    answer = humanize_str(answer).strip()
+    if not answer:
+        return None
+
+    if "datetime.date" in answer:
+        match = re.search(r"datetime\.date\((\d+),\s*(\d+),\s*(\d+)\)", answer)
+        if match is None:
+            return answer
+        year, month, day = (int(group) for group in match.groups())
+        return pd.Timestamp(year=year, month=month, day=day).date()
+
+    try:
+        parsed = ast.literal_eval(answer)
+    except (ValueError, SyntaxError):
+        return answer
+
+    if isinstance(parsed, list) and parsed:
+        return parsed[0]
+    return parsed
+
+
+def extract_oolong_answer(answer: object) -> tuple[object | None, str]:
+    if not isinstance(answer, str):
+        return None, "low"
+
+    answer = humanize_str(answer).strip()
+    if not answer:
+        return None, "low"
+
+    parse_confidence = "low"
+    if ":" not in answer:
+        candidate = answer if len(answer) < 20 else answer.split()[-1]
+        return candidate.strip("[]* "), parse_confidence
+
+    candidate = answer.split(":")[-1].strip()
+    candidate = candidate.replace("*", "")
+    candidate = candidate.replace("[", "")
+    candidate = candidate.replace("]", "")
+    parse_confidence = "med"
+
+    if any(marker in answer for marker in ("User:", "Answer:", "Date:", "Label")):
+        parse_confidence = "high"
+
+    if len(candidate) < 20:
+        parse_confidence = "vhigh"
+    elif "more common" in candidate:
+        candidate = "more common"
+    elif "less common" in candidate:
+        candidate = "less common"
+    elif "same frequency" in candidate:
+        candidate = "same frequency"
+
+    return candidate.strip(), parse_confidence
+
+
+def compute_oolong_metrics(
+    decoded_txts: list[dict[str, object]],
+) -> tuple[dict[str, float | int | str], dict[str, list[float]], dict[str, int]]:
+    scores = []
+    parse_successes = []
+
+    for sample in decoded_txts:
+        gold = parse_oolong_gold_answer(sample.get("label"))
+        predicted, parse_confidence = extract_oolong_answer(sample.get("generated"))
+        answer_type = sample.get("oolong_answer_type")
+        score = 0.0
+
+        if str(predicted) == str(gold):
+            score = 1.0
+        elif (
+            isinstance(predicted, str)
+            and predicted in {"more common", "less common", "same frequency"}
+            and predicted in str(gold)
+        ):
+            score = 1.0
+        elif answer_type == "ANSWER_TYPE.NUMERIC":
+            try:
+                score = 0.75 ** abs(int(gold) - int(predicted))
+            except (TypeError, ValueError):
+                parse_confidence = "low"
+        elif answer_type == "ANSWER_TYPE.DATE":
+            try:
+                parsed_date = dateutil.parser.parse(str(predicted)).date()
+                score = float(parsed_date == gold)
+            except (TypeError, ValueError, OverflowError):
+                parse_confidence = "low"
+
+        scores.append(float(score))
+        parse_successes.append(float(parse_confidence != "low"))
+
+    metrics = {
+        "oolong_score": float(np.mean(scores)) if scores else "None",
+        "oolong_parse_rate": (
+            float(np.mean(parse_successes)) if parse_successes else "None"
+        ),
+    }
+    per_sample_metric = {
+        "oolong_score": scores,
+        "oolong_parse_rate": parse_successes,
+    }
+    counts = {
+        "oolong_score": len(scores),
+        "oolong_parse_rate": len(parse_successes),
+    }
+    return metrics, per_sample_metric, counts
+
+
 def add_longbench_tasks(ds_names: list[str]) -> None:
     """Add longbench tasks to dataset names list."""
     if "longbench" in ds_names:
@@ -211,6 +412,48 @@ def save_generated_text(
             f.write(json.dumps(sample) + "\n")
 
 
+def save_dense_lora_magnitude_stats(
+    stats: dict[str, object],
+    output_dir: str,
+    split: str,
+) -> None:
+    """Save aggregate dense LoRA magnitude stats to JSON."""
+    os.makedirs(output_dir, exist_ok=True)
+    if "/" in split:
+        split_dir = os.path.join(output_dir, os.path.dirname(split))
+        os.makedirs(split_dir, exist_ok=True)
+
+    with open(f"{output_dir}/{split}_dense_lora_magnitude.json", "w") as f:
+        json.dump(stats, f, indent=2, sort_keys=True)
+
+
+def flatten_dense_lora_magnitude_metrics(
+    stats: dict[str, object],
+) -> dict[str, float | int]:
+    """Flatten dense LoRA stats for trainer logging."""
+    out = {}
+
+    overall = stats.get("overall", {})
+    for metric_name in ("mean", "min", "max", "std", "count"):
+        value = overall.get(metric_name)
+        if value is not None:
+            out[f"dense_lora_magnitude_{metric_name}"] = value
+
+    for metric_name in ("num_examples", "num_batches", "num_layer_matrices"):
+        value = stats.get(metric_name)
+        if value is not None:
+            out[f"dense_lora_magnitude_{metric_name}"] = value
+
+    by_module = stats.get("by_module", {})
+    for module_name, module_stats in by_module.items():
+        for metric_name in ("mean", "min", "max", "std", "count"):
+            value = module_stats.get(metric_name)
+            if value is not None:
+                out[f"dense_lora_magnitude_{module_name}_{metric_name}"] = value
+
+    return out
+
+
 # ============================================================================
 # CSV Export Utilities
 # ============================================================================
@@ -478,6 +721,23 @@ def decode_test_result(
     test_dataset, test_result, tokenizer, ctx_tokenizer
 ) -> list[dict]:
     """Decode test results into human-readable format."""
+    metadata_keys = (
+        "book_name",
+        "book_title",
+        "claim_type",
+        "clipper_pair_id",
+        "clipper_status",
+        "clipper_true_claim",
+        "clipper_false_claim",
+        "oolong_id",
+        "oolong_dataset",
+        "oolong_task_group",
+        "oolong_task",
+        "oolong_answer_type",
+        "oolong_input_subset",
+        "oolong_num_labels",
+        "oolong_context_window_id",
+    )
     out = []
     for sample, pred_toks in zip(test_dataset, test_result.predictions):
         d = dict()
@@ -511,6 +771,15 @@ def decode_test_result(
             d["context"] = ctx_tokenizer.decode(
                 concat_list(sample["ctx_ids"]), skip_special_tokens=False
             )
+        for key in metadata_keys:
+            if key in sample:
+                value = sample[key]
+                if hasattr(value, "item"):
+                    try:
+                        value = value.item()
+                    except (TypeError, ValueError):
+                        pass
+                d[key] = value
         for k in sample:
             if k.endswith("_len"):
                 d[k] = sample[k].item()
@@ -546,12 +815,29 @@ def eval_generation(
             split_name += "_no_context"
 
         clear_gpu()
+        if hasattr(eval_trainer.model, "reset_retrieval_history"):
+            eval_trainer.model.reset_retrieval_history()
+        if getattr(eval_trainer.model, "log_dense_lora_magnitude", False) and hasattr(
+            eval_trainer.model, "reset_dense_lora_magnitude_stats"
+        ):
+            eval_trainer.model.reset_dense_lora_magnitude_stats()
         eval_result = eval_trainer.predict(
             ds,
             metric_key_prefix=split_name,
             **gen_kwargs,
         )
         decoded_txts = decode_test_result(ds, eval_result, tokenizer, ctx_tokenizer)
+        if hasattr(eval_trainer.model, "get_retrieval_history"):
+            retrieval_history = eval_trainer.model.get_retrieval_history()
+            if len(retrieval_history) == len(decoded_txts):
+                for decoded, retrieval in zip(decoded_txts, retrieval_history):
+                    decoded.update(retrieval)
+            else:
+                logger.warning(
+                    "RAG retrieval history length (%s) does not match decoded samples (%s).",
+                    len(retrieval_history),
+                    len(decoded_txts),
+                )
 
         pred_texts = [txt["generated"] for txt in decoded_txts]
         label_texts = [txt["label"] for txt in decoded_txts]
@@ -560,8 +846,17 @@ def eval_generation(
         else:
             answers_list = [[txt] for txt in label_texts]
         n = len(pred_texts)
+        metric_mode = (
+            "clipper"
+            if ds_name == "clipper"
+            else "oolong"
+            if ds_name == "oolong-synth"
+            else "closed_qa"
+            if ds_name in CLOSED_QA_DATASETS
+            else "rouge"
+        )
 
-        if ds_name in CLOSED_QA_DATASETS:
+        if metric_mode == "closed_qa":
             print("Computing QA F1 Score")
             qa_f1_metric, per_sample_metric = compute_qa_f1_score(
                 pred_texts, answers_list
@@ -569,6 +864,22 @@ def eval_generation(
             for k, v in qa_f1_metric.items():
                 eval_result.metrics[f"{split_name}_{k}"] = v
                 eval_result.metrics[f"{split_name}_num_samples_{k}"] = n
+        elif metric_mode == "clipper":
+            print("Computing CLIPPER Accuracy")
+            clipper_metrics, per_sample_metric, metric_counts = compute_clipper_metrics(
+                decoded_txts
+            )
+            for k, v in clipper_metrics.items():
+                eval_result.metrics[f"{split_name}_{k}"] = v
+                eval_result.metrics[f"{split_name}_num_samples_{k}"] = metric_counts[k]
+        elif metric_mode == "oolong":
+            print("Computing Oolong Score")
+            oolong_metrics, per_sample_metric, metric_counts = compute_oolong_metrics(
+                decoded_txts
+            )
+            for k, v in oolong_metrics.items():
+                eval_result.metrics[f"{split_name}_{k}"] = v
+                eval_result.metrics[f"{split_name}_num_samples_{k}"] = metric_counts[k]
         else:
             rouge_metrics, per_sample_metric = compute_rouge(pred_texts, label_texts)
             for k, v in rouge_metrics.items():
@@ -577,8 +888,22 @@ def eval_generation(
 
         # Ensure all keys for length metrics are present, even if a bin is empty
         for low, high in LENGTH_BINS:
-            if ds_name in CLOSED_QA_DATASETS:
+            if metric_mode == "closed_qa":
                 eval_result.metrics[f"{split_name}_qa_f1_len_{low}-{high}"] = "None"
+            elif metric_mode == "clipper":
+                eval_result.metrics[
+                    f"{split_name}_clipper_accuracy_len_{low}-{high}"
+                ] = "None"
+                eval_result.metrics[
+                    f"{split_name}_clipper_pair_accuracy_len_{low}-{high}"
+                ] = "None"
+            elif metric_mode == "oolong":
+                eval_result.metrics[f"{split_name}_oolong_score_len_{low}-{high}"] = (
+                    "None"
+                )
+                eval_result.metrics[
+                    f"{split_name}_oolong_parse_rate_len_{low}-{high}"
+                ] = "None"
             else:
                 eval_result.metrics[f"{split_name}_rougeL.f1_len_{low}-{high}"] = "None"
 
@@ -596,10 +921,9 @@ def eval_generation(
 
         for group_key, data in grouped_texts.items():
             if data["count"] > 0:
-                if ds_name in CLOSED_QA_DATASETS:
+                if metric_mode == "closed_qa":
                     # Get corresponding answers for this group
                     group_answers = []
-                    group_idx = 0
                     for i, txt in enumerate(decoded_txts):
                         len_key = (
                             "ctx_ids_len" if "ctx_ids_len" in txt else "input_ids_len"
@@ -624,6 +948,50 @@ def eval_generation(
                         eval_result.metrics[
                             f"{split_name}_num_samples_{k}_len_{group_key}"
                         ] = data["count"]
+                elif metric_mode == "clipper":
+                    group_txts = []
+                    for txt in decoded_txts:
+                        len_key = (
+                            "ctx_ids_len" if "ctx_ids_len" in txt else "input_ids_len"
+                        )
+                        input_len = txt[len_key]
+                        for low, high in LENGTH_BINS:
+                            if (
+                                low <= input_len <= high
+                                and f"{low}-{high}" == group_key
+                            ):
+                                group_txts.append(txt)
+
+                    group_clipper_metrics, _, group_counts = compute_clipper_metrics(
+                        group_txts
+                    )
+                    for k, v in group_clipper_metrics.items():
+                        eval_result.metrics[f"{split_name}_{k}_len_{group_key}"] = v
+                        eval_result.metrics[
+                            f"{split_name}_num_samples_{k}_len_{group_key}"
+                        ] = group_counts[k]
+                elif metric_mode == "oolong":
+                    group_txts = []
+                    for txt in decoded_txts:
+                        len_key = (
+                            "ctx_ids_len" if "ctx_ids_len" in txt else "input_ids_len"
+                        )
+                        input_len = txt[len_key]
+                        for low, high in LENGTH_BINS:
+                            if (
+                                low <= input_len <= high
+                                and f"{low}-{high}" == group_key
+                            ):
+                                group_txts.append(txt)
+
+                    group_oolong_metrics, _, group_counts = compute_oolong_metrics(
+                        group_txts
+                    )
+                    for k, v in group_oolong_metrics.items():
+                        eval_result.metrics[f"{split_name}_{k}_len_{group_key}"] = v
+                        eval_result.metrics[
+                            f"{split_name}_num_samples_{k}_len_{group_key}"
+                        ] = group_counts[k]
                 else:
                     group_rouge_metrics, _ = compute_rouge(
                         data["generated"], data["label"]
@@ -635,6 +1003,19 @@ def eval_generation(
                             f"{split_name}_num_samples_{k}_len_{group_key}"
                         ] = data["count"]
 
+        if getattr(eval_trainer.model, "log_dense_lora_magnitude", False) and hasattr(
+            eval_trainer.model, "get_dense_lora_magnitude_stats"
+        ):
+            dense_lora_stats = eval_trainer.model.get_dense_lora_magnitude_stats()
+            dense_lora_metrics = flatten_dense_lora_magnitude_metrics(dense_lora_stats)
+            for metric_name, value in dense_lora_metrics.items():
+                eval_result.metrics[f"{split_name}_{metric_name}"] = value
+            save_dense_lora_magnitude_stats(
+                dense_lora_stats,
+                split=split_name,
+                output_dir=eval_trainer.args.output_dir,
+            )
+
         save_generated_text(
             decoded_txts,
             per_sample_metric,
@@ -679,7 +1060,23 @@ def eval_teacher_forcing(
         if remove_context:
             split_name += "_no_context"
 
+        if getattr(eval_trainer.model, "log_dense_lora_magnitude", False) and hasattr(
+            eval_trainer.model, "reset_dense_lora_magnitude_stats"
+        ):
+            eval_trainer.model.reset_dense_lora_magnitude_stats()
         metrics = eval_trainer.evaluate(ds, metric_key_prefix=split_name)
+        if getattr(eval_trainer.model, "log_dense_lora_magnitude", False) and hasattr(
+            eval_trainer.model, "get_dense_lora_magnitude_stats"
+        ):
+            dense_lora_stats = eval_trainer.model.get_dense_lora_magnitude_stats()
+            dense_lora_metrics = flatten_dense_lora_magnitude_metrics(dense_lora_stats)
+            for metric_name, value in dense_lora_metrics.items():
+                metrics[f"{split_name}_{metric_name}"] = value
+            save_dense_lora_magnitude_stats(
+                dense_lora_stats,
+                split=split_name,
+                output_dir=eval_trainer.args.output_dir,
+            )
         out[split_name] = metrics
         eval_trainer.log_metrics(split_name, metrics)
         eval_trainer.save_metrics(split_name, metrics)
@@ -707,6 +1104,7 @@ def evaluate(
     args: Namespace,
     split: str,
     max_ctx_chunk_len: int,
+    ctx_chunk_overlap: int,
     max_new_tokens: int,
     generative: bool,
 ) -> dict[str, dict]:
@@ -750,6 +1148,20 @@ def evaluate(
         add_tracker(model.generate_weights, "generate_weights")
         add_tracker(model.combine_lora, "combine_lora")
         add_tracker(model.apply_lora_to_layers, "apply_lora_to_layers")
+        model.enable_dense_lora_magnitude_logging(
+            getattr(args, "log_dense_lora_magnitude", False)
+        )
+        if getattr(args, "use_hybrid_rag", False):
+            model = DocToLoraRAGModel(
+                model=model,
+                tokenizer=tokenizer,
+                chunk_size=args.rag_chunk_size,
+                chunk_overlap=args.rag_chunk_overlap,
+                top_k=args.rag_top_k,
+                max_retrieved_tokens=args.rag_max_retrieved_tokens,
+                retrieval_mode=args.rag_retrieval_mode,
+            )
+            add_tracker(model._retrieve, "retrieve")
     else:
         model = base_model = get_model(
             model_name_or_path,
@@ -823,6 +1235,19 @@ def evaluate(
             add_tracker(model.base_model.generate, "base_model.generate")
             add_tracker(model.generate_weights, "generate_weights")
 
+        elif use_rag := getattr(args, "use_rag", False):
+            model = RAGModel(
+                model=base_model,
+                tokenizer=tokenizer,
+                chunk_size=args.rag_chunk_size,
+                chunk_overlap=args.rag_chunk_overlap,
+                top_k=args.rag_top_k,
+                max_retrieved_tokens=args.rag_max_retrieved_tokens,
+                retrieval_mode=args.rag_retrieval_mode,
+            )
+            ctx_model_max_len = base_model.config.max_position_embeddings
+            add_tracker(model._retrieve, "retrieve")
+
         elif args.use_generative_adapter:
             model = GenerativeAdapter(model=base_model, tokenizer=tokenizer)
             ctx_model_max_len = base_model.config.max_position_embeddings
@@ -849,6 +1274,7 @@ def evaluate(
         max_qas_len=-1,
         max_qas_per_sample=1,
         max_ctx_chunk_len=max_ctx_chunk_len,
+        ctx_chunk_overlap=ctx_chunk_overlap,
         min_ctx_chunk_len=-1,
         num_chunk_probs=None,
         max_ctx_chunk_num=None,
@@ -1020,6 +1446,7 @@ def run_eval(
     max_val_samples_per_ds: int = -1,
     max_test_samples_per_ds: int = -1,
     max_ctx_chunk_len: int = -1,
+    ctx_chunk_overlap: int = 0,
     remove_context: bool = False,
     max_new_tokens: int = 256,
     generative: bool = False,
@@ -1032,24 +1459,84 @@ def run_eval(
     use_llmlingua: bool = False,
     llmlingua_compression_rate: float = 0.9,
     use_t2l: bool = False,
+    use_rag: bool = False,
+    use_hybrid_rag: bool = False,
+    rag_chunk_size: int = 256,
+    rag_chunk_overlap: int = 64,
+    rag_top_k: int = 4,
+    rag_max_retrieved_tokens: int = 1536,
+    rag_retrieval_mode: str = "bm25",
     add_ctx_to_input: bool = False,
     truncate_if_too_long_inp: bool = False,
     truncate_if_too_long_ctx: bool = False,
     flip_ctx_inp: bool = False,
     gen_lora_scaling: float = 1,
     use_generative_adapter: bool = False,
+    log_dense_lora_magnitude: bool = False,
 ) -> None:
     """Run evaluation with the specified parameters."""
     assert bool(model_name_or_path) ^ bool(checkpoint_path), (
         "Either --model_name_or_path or --checkpoint_path must be provided"
     )
-    if (use_cd or use_llmlingua or use_t2l) and eval_batch_size != 1:
+    if (
+        use_cd or use_llmlingua or use_t2l or use_rag or use_hybrid_rag
+    ) and eval_batch_size != 1:
         raise ValueError("When using a baseline method, eval_batch_size must be 1.")
+    if ctx_chunk_overlap < 0:
+        raise ValueError("`ctx_chunk_overlap` must be >= 0.")
+    if max_ctx_chunk_len > 0 and ctx_chunk_overlap >= max_ctx_chunk_len:
+        raise ValueError(
+            "`ctx_chunk_overlap` must be smaller than `max_ctx_chunk_len`."
+        )
 
     if use_llmlingua and add_ctx_to_input:
         raise ValueError(
             "LLMLingua always adds compressed context to input by default."
         )
+    baseline_methods = [
+        use_cd,
+        use_llmlingua,
+        use_t2l,
+        use_rag,
+        use_hybrid_rag,
+        use_generative_adapter,
+    ]
+    if sum(bool(method) for method in baseline_methods) > 1:
+        raise ValueError("Please enable at most one baseline method at a time.")
+    if use_rag and remove_context:
+        raise ValueError(
+            "RAG baseline requires context; `remove_context` is incompatible."
+        )
+    if use_hybrid_rag and remove_context:
+        raise ValueError(
+            "Hybrid D2L+RAG requires context; `remove_context` is incompatible."
+        )
+    if use_rag and add_ctx_to_input:
+        raise ValueError(
+            "RAG baseline rebuilds its own prompt; `add_ctx_to_input` is incompatible."
+        )
+    if use_hybrid_rag and add_ctx_to_input:
+        raise ValueError(
+            "Hybrid D2L+RAG rebuilds its own prompt; `add_ctx_to_input` is incompatible."
+        )
+    if use_rag and max_ctx_chunk_len > 0:
+        raise ValueError(
+            "RAG baseline performs its own retrieval chunking; `max_ctx_chunk_len` must stay disabled."
+        )
+    if use_hybrid_rag and model_name_or_path:
+        raise ValueError(
+            "Hybrid D2L+RAG is only available when evaluating a D2L checkpoint."
+        )
+    if rag_chunk_size <= 0:
+        raise ValueError("`rag_chunk_size` must be > 0.")
+    if rag_chunk_overlap < 0 or rag_chunk_overlap >= rag_chunk_size:
+        raise ValueError(
+            "`rag_chunk_overlap` must be >= 0 and smaller than `rag_chunk_size`."
+        )
+    if rag_top_k <= 0:
+        raise ValueError("`rag_top_k` must be > 0.")
+    if rag_max_retrieved_tokens <= 0:
+        raise ValueError("`rag_max_retrieved_tokens` must be > 0.")
     if use_generative_adapter and (
         model_name_or_path != "mistralai/Mistral-7B-Instruct-v0.2"
     ):
@@ -1096,6 +1583,20 @@ def run_eval(
         if use_llmlingua:
             args.use_llmlingua = use_llmlingua
             args.llmlingua_compression_rate = llmlingua_compression_rate
+        if use_rag:
+            args.use_rag = use_rag
+            args.rag_chunk_size = rag_chunk_size
+            args.rag_chunk_overlap = rag_chunk_overlap
+            args.rag_top_k = rag_top_k
+            args.rag_max_retrieved_tokens = rag_max_retrieved_tokens
+            args.rag_retrieval_mode = rag_retrieval_mode
+        if use_hybrid_rag:
+            args.use_hybrid_rag = use_hybrid_rag
+            args.rag_chunk_size = rag_chunk_size
+            args.rag_chunk_overlap = rag_chunk_overlap
+            args.rag_top_k = rag_top_k
+            args.rag_max_retrieved_tokens = rag_max_retrieved_tokens
+            args.rag_retrieval_mode = rag_retrieval_mode
     else:
         args = Namespace(
             model_name_or_path=model_name_or_path,
@@ -1117,6 +1618,20 @@ def run_eval(
             args.llmlingua_compression_rate = llmlingua_compression_rate
         if use_t2l:
             args.use_t2l = use_t2l
+        if use_rag:
+            args.use_rag = use_rag
+            args.rag_chunk_size = rag_chunk_size
+            args.rag_chunk_overlap = rag_chunk_overlap
+            args.rag_top_k = rag_top_k
+            args.rag_max_retrieved_tokens = rag_max_retrieved_tokens
+            args.rag_retrieval_mode = rag_retrieval_mode
+        if use_hybrid_rag:
+            args.use_hybrid_rag = use_hybrid_rag
+            args.rag_chunk_size = rag_chunk_size
+            args.rag_chunk_overlap = rag_chunk_overlap
+            args.rag_top_k = rag_top_k
+            args.rag_max_retrieved_tokens = rag_max_retrieved_tokens
+            args.rag_retrieval_mode = rag_retrieval_mode
         args.use_generative_adapter = use_generative_adapter
     if max_val_samples_per_ds > 0:
         args.max_val_samples_per_ds = max_val_samples_per_ds
@@ -1127,6 +1642,7 @@ def run_eval(
     args.truncate_if_too_long_inp = truncate_if_too_long_inp
     args.truncate_if_too_long_ctx = truncate_if_too_long_ctx
     args.flip_ctx_inp = flip_ctx_inp
+    args.log_dense_lora_magnitude = log_dense_lora_magnitude
     setup_logging(args.logging_dir)
     logger.debug(f"CMD: {' '.join(os.sys.argv)}")
 
@@ -1144,6 +1660,7 @@ def run_eval(
         args,
         split,
         max_ctx_chunk_len,
+        ctx_chunk_overlap,
         max_new_tokens,
         generative=generative,
     )

src/ctx_to_lora/lora_stats.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+
+import torch
+from torch import Tensor
+
+
+@dataclass
+class RunningMagnitudeStats:
+    count: int = 0
+    total: float = 0.0
+    total_sq: float = 0.0
+    min_val: float | None = None
+    max_val: float | None = None
+
+    def update(self, values: Tensor) -> None:
+        if values.numel() == 0:
+            return
+
+        values = values.to(dtype=torch.float64)
+        batch_count = int(values.numel())
+        batch_sum = float(values.sum().item())
+        batch_sum_sq = float(values.square().sum().item())
+        batch_min = float(values.min().item())
+        batch_max = float(values.max().item())
+
+        self.count += batch_count
+        self.total += batch_sum
+        self.total_sq += batch_sum_sq
+        self.min_val = (
+            batch_min if self.min_val is None else min(self.min_val, batch_min)
+        )
+        self.max_val = (
+            batch_max if self.max_val is None else max(self.max_val, batch_max)
+        )
+
+    def to_dict(self) -> dict[str, float | int | None]:
+        if self.count == 0:
+            return {
+                "count": 0,
+                "mean": None,
+                "min": None,
+                "max": None,
+                "std": None,
+            }
+
+        mean = self.total / self.count
+        variance = max(self.total_sq / self.count - mean**2, 0.0)
+        return {
+            "count": self.count,
+            "mean": mean,
+            "min": self.min_val,
+            "max": self.max_val,
+            "std": math.sqrt(variance),
+        }
+
+
+def _update_dense_magnitude_stats_for_module(
+    stats: RunningMagnitudeStats | list[RunningMagnitudeStats],
+    A: Tensor,
+    B: Tensor,
+    scaling: float,
+    out_chunk_size: int,
+) -> None:
+    if A.ndim != 4 or B.ndim != 4:
+        raise ValueError(
+            f"Expected combined LoRA tensors with 4 dims, got {A.shape=} and {B.shape=}"
+        )
+    if A.shape[:3] != B.shape[:3]:
+        raise ValueError(f"Mismatched LoRA shapes: {A.shape=} and {B.shape=}")
+    if out_chunk_size <= 0:
+        raise ValueError("`out_chunk_size` must be > 0.")
+    if isinstance(stats, RunningMagnitudeStats):
+        stats = [stats]
+    if not torch.cuda.is_available():
+        raise RuntimeError("Dense LoRA magnitude logging requires CUDA.")
+
+    # Keep the entire computation on GPU to avoid CPU offloading.
+    A_gpu = A.detach().to(device="cuda", dtype=torch.float32)
+    B_gpu = B.detach().to(device="cuda", dtype=torch.float32)
+
+    batch_size, num_layers, _, d_in = A_gpu.shape
+    d_out = B_gpu.shape[-1]
+
+    A_flat = A_gpu.reshape(batch_size * num_layers, A_gpu.shape[-2], d_in)
+    B_flat = B_gpu.reshape(batch_size * num_layers, B_gpu.shape[-2], d_out)
+
+    for start in range(0, d_out, out_chunk_size):
+        stop = min(start + out_chunk_size, d_out)
+        B_block = B_flat[:, :, start:stop].transpose(1, 2)
+        dense_block = torch.bmm(B_block, A_flat)
+        dense_block.mul_(float(scaling))
+        abs_block = dense_block.abs()
+        for stat in stats:
+            stat.update(abs_block)
+
+
+@dataclass
+class DenseLoraMagnitudeAggregator:
+    scaling: float
+    out_chunk_size: int = 8
+    overall: RunningMagnitudeStats = field(default_factory=RunningMagnitudeStats)
+    by_module: dict[str, RunningMagnitudeStats] = field(default_factory=dict)
+    num_examples: int = 0
+    num_batches: int = 0
+    num_layer_matrices: int = 0
+
+    def reset(self) -> None:
+        self.overall = RunningMagnitudeStats()
+        self.by_module = {}
+        self.num_examples = 0
+        self.num_batches = 0
+        self.num_layer_matrices = 0
+
+    def update(self, combined_loras: dict[str, dict[str, Tensor]]) -> None:
+        if not combined_loras:
+            return
+
+        first_module = next(iter(combined_loras.values()))
+        first_A = first_module["A"]
+        batch_size, num_layers = first_A.shape[:2]
+        self.num_examples += int(batch_size)
+        self.num_batches += 1
+        self.num_layer_matrices += int(batch_size * num_layers)
+
+        for module_name, module_lora in combined_loras.items():
+            module_stats = self.by_module.setdefault(
+                module_name, RunningMagnitudeStats()
+            )
+            _update_dense_magnitude_stats_for_module(
+                [module_stats, self.overall],
+                module_lora["A"],
+                module_lora["B"],
+                scaling=self.scaling,
+                out_chunk_size=self.out_chunk_size,
+            )
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "scaling": self.scaling,
+            "out_chunk_size": self.out_chunk_size,
+            "num_examples": self.num_examples,
+            "num_batches": self.num_batches,
+            "num_layer_matrices": self.num_layer_matrices,
+            "overall": self.overall.to_dict(),
+            "by_module": {
+                module_name: stats.to_dict()
+                for module_name, stats in sorted(self.by_module.items())
+            },
+        }

src/ctx_to_lora/modeling/hypernet.py

@@ -32,6 +32,7 @@ from ctx_to_lora.configs import (
     HypernetArguments,
 )
 from ctx_to_lora.data.processing import tokenize_ctx_text
+from ctx_to_lora.lora_stats import DenseLoraMagnitudeAggregator
 from ctx_to_lora.model_loading import (
     get_model,
     get_tokenizer,
@@ -461,6 +462,10 @@ class ModulatedPretrainedModel(nn.Module):
         self.inp_compressor = inp_compressor
         self.model_accepts_loss_kwargs = True
         self.generated_loras = None
+        self.log_dense_lora_magnitude = False
+        self.dense_lora_magnitude = DenseLoraMagnitudeAggregator(
+            scaling=self.peft_config.lora_alpha
+        )
 
         self.register_module("base_model", base_model)
         self._init_model()
@@ -683,6 +688,27 @@ class ModulatedPretrainedModel(nn.Module):
     def enable_iterative_mode(self, x: bool):
         self.hypernet.enable_iterative_mode(x)
 
+    def enable_dense_lora_magnitude_logging(
+        self, x: bool, out_chunk_size: int = 8
+    ) -> None:
+        self.log_dense_lora_magnitude = x
+        self.dense_lora_magnitude.scaling = self.peft_config.lora_alpha
+        self.dense_lora_magnitude.out_chunk_size = out_chunk_size
+        self.dense_lora_magnitude.reset()
+
+    def reset_dense_lora_magnitude_stats(self) -> None:
+        self.dense_lora_magnitude.reset()
+
+    def get_dense_lora_magnitude_stats(self) -> dict[str, object]:
+        return self.dense_lora_magnitude.to_dict()
+
+    def _maybe_track_dense_lora_magnitude(
+        self, generated_loras: dict[str, dict[str, Tensor]] | None
+    ) -> None:
+        if not self.log_dense_lora_magnitude or generated_loras is None:
+            return
+        self.dense_lora_magnitude.update(generated_loras)
+
     def forward(
         self,
         ctx_ids: Integer[Tensor, "n_ctx ctx_len"] | None = None,
@@ -735,6 +761,7 @@ class ModulatedPretrainedModel(nn.Module):
                 if self.hypernet.config.use_bias
                 else None,
             )
+            self._maybe_track_dense_lora_magnitude(generated_loras)
 
             # input_ids in model_inputs_kwargs contains only
             # prompt + response (for hypernet training)
@@ -877,6 +904,7 @@ class ModulatedPretrainedModel(nn.Module):
                 scalers=scalers,
                 bias_scaler=bias_scaler,
             )
+            self._maybe_track_dense_lora_magnitude(generated_loras)
 
             # apply lora hook to the base model
             # TODO: we dont this position_ids for generation?

src/ctx_to_lora/modeling/rag.py

@@ -0,0 +1,516 @@
+import math
+import re
+from collections import Counter
+from dataclasses import asdict, dataclass
+from typing import Any
+
+import torch
+from torch import nn
+
+from ctx_to_lora.data.definitions import CTX_AFFIXES
+
+_WORD_RE = re.compile(r"[A-Za-z0-9_]+")
+
+
+@dataclass
+class RetrievedChunk:
+    text: str
+    token_ids: list[int]
+    start: int
+    end: int
+    score: float
+
+
+class _RAGMixin:
+    def _init_rag(
+        self,
+        tokenizer,
+        chunk_size: int = 256,
+        chunk_overlap: int = 64,
+        top_k: int = 4,
+        max_retrieved_tokens: int = 1536,
+        retrieval_mode: str = "bm25",
+        prompt_instruction: str | None = None,
+    ) -> None:
+        if chunk_size <= 0:
+            raise ValueError("`chunk_size` must be > 0.")
+        if chunk_overlap < 0:
+            raise ValueError("`chunk_overlap` must be >= 0.")
+        if chunk_overlap >= chunk_size:
+            raise ValueError("`chunk_overlap` must be smaller than `chunk_size`.")
+        if top_k <= 0:
+            raise ValueError("`top_k` must be > 0.")
+        if max_retrieved_tokens <= 0:
+            raise ValueError("`max_retrieved_tokens` must be > 0.")
+        if retrieval_mode != "bm25":
+            raise ValueError(
+                f"Unsupported `retrieval_mode`: {retrieval_mode}. Only 'bm25' is supported."
+            )
+
+        self.tokenizer = tokenizer
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+        self.top_k = top_k
+        self.max_retrieved_tokens = max_retrieved_tokens
+        self.retrieval_mode = retrieval_mode
+        self.prompt_instruction = (
+            prompt_instruction
+            or "Answer the question using only the retrieved context below. "
+            "If the answer is not supported by the retrieved context, say you don't know."
+        )
+
+        model_name = getattr(self.base_model, "name_or_path", None)
+        if model_name is None and hasattr(self.base_model, "config"):
+            model_name = getattr(
+                self.base_model.config,
+                "name_or_path",
+                getattr(self.base_model.config, "_name_or_path", None),
+            )
+        if model_name not in CTX_AFFIXES:
+            raise ValueError(
+                f"No context affixes configured for model '{model_name}'. "
+                "RAG wrappers only support models listed in `CTX_AFFIXES`."
+            )
+
+        affixes = CTX_AFFIXES[model_name]
+        self.prefix = affixes["prefix"]
+        self.suffix = affixes["suffix"]
+        self.len_prefix = len(self.prefix)
+        self.len_suffix = len(self.suffix)
+        self.last_retrieval = None
+        self._retrieval_history = []
+
+    @property
+    def generation_config(self):
+        return self.generation_model.generation_config
+
+    @property
+    def device(self):
+        try:
+            return next(self.base_model.parameters()).device
+        except StopIteration:
+            return torch.device("cpu")
+
+    def reset_retrieval_history(self) -> None:
+        self.last_retrieval = None
+        self._retrieval_history = []
+
+    def get_retrieval_history(self) -> list[dict[str, Any]]:
+        return list(self._retrieval_history)
+
+    def _strip_padding_and_affixes(
+        self, token_ids: torch.Tensor, attn_mask: torch.Tensor | None = None
+    ) -> list[int]:
+        if attn_mask is not None:
+            valid = token_ids[attn_mask.bool()].tolist()
+        else:
+            valid = token_ids.tolist()
+
+        if len(valid) >= self.len_prefix and valid[: self.len_prefix] == self.prefix:
+            valid = valid[self.len_prefix :]
+        if len(valid) >= self.len_suffix and valid[-self.len_suffix :] == self.suffix:
+            valid = valid[: -self.len_suffix]
+        return valid
+
+    def _decode_question(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None
+    ) -> tuple[str, list[int]]:
+        question_ids = self._strip_padding_and_affixes(input_ids, attention_mask)
+        question = self.tokenizer.decode(question_ids, skip_special_tokens=True).strip()
+        return question, question_ids
+
+    def _decode_context(
+        self,
+        ctx_ids: torch.Tensor,
+        ctx_attn_mask: torch.Tensor | None = None,
+    ) -> tuple[str, list[int]]:
+        if ctx_ids.ndim == 1:
+            ctx_ids = ctx_ids.unsqueeze(0)
+        if ctx_attn_mask is None:
+            ctx_attn_mask = torch.ones_like(ctx_ids)
+        elif ctx_attn_mask.ndim == 1:
+            ctx_attn_mask = ctx_attn_mask.unsqueeze(0)
+
+        full_ids = []
+        for chunk_ids, chunk_mask in zip(ctx_ids, ctx_attn_mask):
+            full_ids.extend(self._strip_padding_and_affixes(chunk_ids, chunk_mask))
+        context = self.tokenizer.decode(full_ids, skip_special_tokens=True).strip()
+        return context, full_ids
+
+    def _tokenize_text(self, text: str) -> list[int]:
+        return self.tokenizer(
+            text,
+            add_special_tokens=False,
+            return_attention_mask=False,
+        )["input_ids"]
+
+    def _chunk_context(self, context_text: str) -> list[RetrievedChunk]:
+        context_ids = self._tokenize_text(context_text)
+        if not context_ids:
+            return []
+        if len(context_ids) <= self.chunk_size:
+            return [
+                RetrievedChunk(
+                    text=context_text.strip(),
+                    token_ids=context_ids,
+                    start=0,
+                    end=len(context_ids),
+                    score=0.0,
+                )
+            ]
+
+        stride = self.chunk_size - self.chunk_overlap
+        chunks = []
+        for start in range(0, len(context_ids), stride):
+            end = min(start + self.chunk_size, len(context_ids))
+            chunk_ids = context_ids[start:end]
+            if not chunk_ids:
+                break
+            chunk_text = self.tokenizer.decode(
+                chunk_ids, skip_special_tokens=True
+            ).strip()
+            chunks.append(
+                RetrievedChunk(
+                    text=chunk_text,
+                    token_ids=chunk_ids,
+                    start=start,
+                    end=end,
+                    score=0.0,
+                )
+            )
+            if end == len(context_ids):
+                break
+        return chunks
+
+    def _lexical_tokenize(self, text: str) -> list[str]:
+        return _WORD_RE.findall(text.lower())
+
+    def _score_chunks_bm25(
+        self, query: str, chunks: list[RetrievedChunk]
+    ) -> list[RetrievedChunk]:
+        query_terms = self._lexical_tokenize(query)
+        if not chunks:
+            return []
+        if not query_terms:
+            return [
+                RetrievedChunk(**{**asdict(chunk), "score": 0.0})
+                for chunk in chunks[: self.top_k]
+            ]
+
+        tokenized_chunks = [self._lexical_tokenize(chunk.text) for chunk in chunks]
+        num_chunks = len(tokenized_chunks)
+        avgdl = sum(len(chunk_terms) for chunk_terms in tokenized_chunks) / max(
+            num_chunks, 1
+        )
+        doc_freq = Counter()
+        for chunk_terms in tokenized_chunks:
+            for term in set(chunk_terms):
+                doc_freq[term] += 1
+
+        k1 = 1.5
+        b = 0.75
+        ranked = []
+        for idx, (chunk, chunk_terms) in enumerate(zip(chunks, tokenized_chunks)):
+            term_freq = Counter(chunk_terms)
+            doc_len = len(chunk_terms)
+            score = 0.0
+            for term in query_terms:
+                freq = term_freq.get(term, 0)
+                if freq == 0:
+                    continue
+                df = doc_freq[term]
+                idf = math.log(1.0 + (num_chunks - df + 0.5) / (df + 0.5))
+                denom = freq + k1 * (1.0 - b + b * doc_len / max(avgdl, 1e-8))
+                score += idf * (freq * (k1 + 1.0)) / denom
+            score += 1e-4 / (1 + idx)
+            ranked.append(RetrievedChunk(**{**asdict(chunk), "score": float(score)}))
+
+        ranked.sort(key=lambda chunk: (-chunk.score, chunk.start, chunk.end))
+        return ranked
+
+    def _merge_overlapping_chunks(
+        self, context_ids: list[int], chunks: list[RetrievedChunk]
+    ) -> list[RetrievedChunk]:
+        if not chunks:
+            return []
+        chunks = sorted(chunks, key=lambda chunk: (chunk.start, chunk.end))
+        merged = [chunks[0]]
+        for chunk in chunks[1:]:
+            prev = merged[-1]
+            if chunk.start <= prev.end:
+                start = prev.start
+                end = max(prev.end, chunk.end)
+                merged_ids = context_ids[start:end]
+                merged[-1] = RetrievedChunk(
+                    text=self.tokenizer.decode(
+                        merged_ids, skip_special_tokens=True
+                    ).strip(),
+                    token_ids=merged_ids,
+                    start=start,
+                    end=end,
+                    score=max(prev.score, chunk.score),
+                )
+                continue
+            merged.append(chunk)
+        return merged
+
+    def _select_chunks_with_budget(
+        self,
+        ranked_chunks: list[RetrievedChunk],
+        context_ids: list[int],
+    ) -> list[RetrievedChunk]:
+        if not ranked_chunks:
+            return []
+
+        selected = []
+        budget_used = 0
+        for chunk in ranked_chunks:
+            if len(selected) >= self.top_k:
+                break
+            chunk_len = len(chunk.token_ids)
+            if selected and budget_used + chunk_len > self.max_retrieved_tokens:
+                continue
+            if chunk_len > self.max_retrieved_tokens:
+                truncated_ids = chunk.token_ids[: self.max_retrieved_tokens]
+                selected.append(
+                    RetrievedChunk(
+                        text=self.tokenizer.decode(
+                            truncated_ids, skip_special_tokens=True
+                        ).strip(),
+                        token_ids=truncated_ids,
+                        start=chunk.start,
+                        end=chunk.start + len(truncated_ids),
+                        score=chunk.score,
+                    )
+                )
+                budget_used = len(truncated_ids)
+                break
+
+            selected.append(chunk)
+            budget_used += chunk_len
+
+        if not selected:
+            top_chunk = ranked_chunks[0]
+            truncated_ids = top_chunk.token_ids[: self.max_retrieved_tokens]
+            selected = [
+                RetrievedChunk(
+                    text=self.tokenizer.decode(
+                        truncated_ids, skip_special_tokens=True
+                    ).strip(),
+                    token_ids=truncated_ids,
+                    start=top_chunk.start,
+                    end=top_chunk.start + len(truncated_ids),
+                    score=top_chunk.score,
+                )
+            ]
+
+        return self._merge_overlapping_chunks(context_ids, selected)
+
+    def _build_prompt_text(
+        self,
+        question: str,
+        selected_chunks: list[RetrievedChunk],
+    ) -> str:
+        sections = [self.prompt_instruction.strip()]
+        if selected_chunks:
+            sections.append("Retrieved context:")
+            for idx, chunk in enumerate(selected_chunks, start=1):
+                sections.append(f"[{idx}] {chunk.text.strip()}")
+        sections.append(f"Question:\n{question.strip()}")
+        return "\n\n".join(section for section in sections if section.strip())
+
+    def _build_rag_inputs(
+        self,
+        question: str,
+        selected_chunks: list[RetrievedChunk],
+    ) -> dict[str, torch.Tensor]:
+        prompt_text = self._build_prompt_text(question, selected_chunks)
+        messages = [
+            {"role": "system", "content": ""},
+            {"role": "user", "content": prompt_text},
+        ]
+        tokenized = self.tokenizer.apply_chat_template(
+            [messages],
+            tokenize=True,
+            add_generation_prompt=True,
+            return_attention_mask=True,
+            return_dict=True,
+            return_tensors="pt",
+            add_special_tokens=False,
+        )
+        return {k: v.to(self.device) for k, v in tokenized.items()}
+
+    def _retrieve(
+        self,
+        question: str,
+        chunks: list[RetrievedChunk],
+        context_ids: list[int],
+    ) -> tuple[list[RetrievedChunk], bool]:
+        if len(context_ids) <= self.max_retrieved_tokens:
+            full_context = RetrievedChunk(
+                text=self.tokenizer.decode(
+                    context_ids, skip_special_tokens=True
+                ).strip(),
+                token_ids=context_ids,
+                start=0,
+                end=len(context_ids),
+                score=0.0,
+            )
+            return [full_context], True
+
+        ranked = self._score_chunks_bm25(question, chunks)
+        return self._select_chunks_with_budget(ranked, context_ids), False
+
+    def _record_retrieval(
+        self,
+        question: str,
+        context_token_count: int,
+        prompt_token_count: int,
+        used_full_context: bool,
+        selected_chunks: list[RetrievedChunk],
+    ) -> None:
+        retrieval = {
+            "rag_question": question,
+            "rag_context_token_count": context_token_count,
+            "rag_prompt_token_count": prompt_token_count,
+            "rag_used_full_context": used_full_context,
+            "rag_chunk_size": self.chunk_size,
+            "rag_chunk_overlap": self.chunk_overlap,
+            "rag_top_k": self.top_k,
+            "rag_max_retrieved_tokens": self.max_retrieved_tokens,
+            "rag_retrieval_mode": self.retrieval_mode,
+            "rag_selected_chunks": [asdict(chunk) for chunk in selected_chunks],
+        }
+        self.last_retrieval = retrieval
+        self._retrieval_history.append(retrieval)
+
+    def _prepare_rag_generation_inputs(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        ctx_ids: torch.Tensor,
+        ctx_attn_mask: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        question, _ = self._decode_question(input_ids[0], attention_mask[0])
+        context_text, context_ids = self._decode_context(ctx_ids, ctx_attn_mask)
+        chunks = self._chunk_context(context_text)
+        selected_chunks, used_full_context = self._retrieve(
+            question, chunks, context_ids
+        )
+        rag_inputs = self._build_rag_inputs(question, selected_chunks)
+        prompt_token_count = int(rag_inputs["attention_mask"].sum().item())
+        self._record_retrieval(
+            question=question,
+            context_token_count=len(context_ids),
+            prompt_token_count=prompt_token_count,
+            used_full_context=used_full_context,
+            selected_chunks=selected_chunks,
+        )
+        return rag_inputs
+
+
+class RAGModel(nn.Module, _RAGMixin):
+    def __init__(
+        self,
+        model: nn.Module,
+        tokenizer,
+        chunk_size: int = 256,
+        chunk_overlap: int = 64,
+        top_k: int = 4,
+        max_retrieved_tokens: int = 1536,
+        retrieval_mode: str = "bm25",
+        prompt_instruction: str | None = None,
+    ):
+        super().__init__()
+        self.base_model = model
+        self.generation_model = model
+        self._init_rag(
+            tokenizer=tokenizer,
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+            top_k=top_k,
+            max_retrieved_tokens=max_retrieved_tokens,
+            retrieval_mode=retrieval_mode,
+            prompt_instruction=prompt_instruction,
+        )
+
+    def generate(self, *args, **kwargs):
+        input_ids = kwargs["input_ids"]
+        if input_ids.shape[0] != 1:
+            raise ValueError("RAGModel only supports batch_size=1 during evaluation.")
+
+        attention_mask = kwargs.get("attention_mask")
+        ctx_ids = kwargs.get("ctx_ids")
+        ctx_attn_mask = kwargs.get("ctx_attn_mask")
+        if ctx_ids is None:
+            raise ValueError("RAGModel requires `ctx_ids`.")
+
+        rag_inputs = self._prepare_rag_generation_inputs(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            ctx_ids=ctx_ids,
+            ctx_attn_mask=ctx_attn_mask,
+        )
+
+        for key in [
+            "ctx_ids",
+            "ctx_attn_mask",
+            "n_ctx_chunks",
+            "input_ids",
+            "attention_mask",
+        ]:
+            kwargs.pop(key, None)
+
+        return self.base_model.generate(*args, **rag_inputs, **kwargs)
+
+
+class DocToLoraRAGModel(nn.Module, _RAGMixin):
+    def __init__(
+        self,
+        model: nn.Module,
+        tokenizer,
+        chunk_size: int = 256,
+        chunk_overlap: int = 64,
+        top_k: int = 4,
+        max_retrieved_tokens: int = 1536,
+        retrieval_mode: str = "bm25",
+        prompt_instruction: str | None = None,
+    ):
+        super().__init__()
+        self.model = model
+        self.base_model = model.base_model
+        self.generation_model = model
+        self._init_rag(
+            tokenizer=tokenizer,
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+            top_k=top_k,
+            max_retrieved_tokens=max_retrieved_tokens,
+            retrieval_mode=retrieval_mode,
+            prompt_instruction=prompt_instruction,
+        )
+
+    def generate(self, *args, **kwargs):
+        input_ids = kwargs["input_ids"]
+        if input_ids.shape[0] != 1:
+            raise ValueError(
+                "DocToLoraRAGModel only supports batch_size=1 during evaluation."
+            )
+
+        attention_mask = kwargs.get("attention_mask")
+        ctx_ids = kwargs.get("ctx_ids")
+        ctx_attn_mask = kwargs.get("ctx_attn_mask")
+        if ctx_ids is None:
+            raise ValueError("DocToLoraRAGModel requires `ctx_ids`.")
+
+        rag_inputs = self._prepare_rag_generation_inputs(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            ctx_ids=ctx_ids,
+            ctx_attn_mask=ctx_attn_mask,
+        )
+
+        kwargs = dict(kwargs)
+        kwargs["input_ids"] = rag_inputs["input_ids"]
+        kwargs["attention_mask"] = rag_inputs["attention_mask"]
+        return self.model.generate(*args, **kwargs)

src/ctx_to_lora/utils.py

@@ -255,9 +255,13 @@ def compile_linear(model):
 
 def clear_gpu():
     gc.collect()
+    if not torch.cuda.is_available():
+        return
     torch.cuda.empty_cache()
-    torch.cuda.reset_max_memory_allocated()
-    torch.cuda.reset_max_memory_cached()
+    try:
+        torch.cuda.reset_peak_memory_stats()
+    except RuntimeError:
+        return
 
 
 def concat_list(l):

tests/test_chunk_overlap.py

@@ -0,0 +1,91 @@
+import inspect
+import unittest
+
+from ctx_to_lora.data.definitions import CTX_AFFIXES
+from ctx_to_lora.data.processing import (
+    construct_and_tokenize_ctx_qa,
+    split_too_long_ctx,
+)
+
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
+
+
+class ChunkOverlapTest(unittest.TestCase):
+    def test_construct_and_tokenize_ctx_qa_accepts_overlap_kwarg(self):
+        signature = inspect.signature(construct_and_tokenize_ctx_qa)
+        self.assertIn("ctx_chunk_overlap", signature.parameters)
+
+    def _strip_affixes(self, chunks):
+        prefix = CTX_AFFIXES[MODEL_NAME]["prefix"]
+        suffix = CTX_AFFIXES[MODEL_NAME]["suffix"]
+        out = []
+        for idx, chunk in enumerate(chunks):
+            if idx > 0:
+                self.assertEqual(chunk[: len(prefix)], prefix)
+                chunk = chunk[len(prefix) :]
+            if idx < len(chunks) - 1:
+                self.assertEqual(chunk[-len(suffix) :], suffix)
+                chunk = chunk[: -len(suffix)]
+            out.append(chunk)
+        return out
+
+    def test_eval_overlap_adds_shared_boundary_tokens(self):
+        ctx_ids = list(range(10))
+        out = split_too_long_ctx(
+            sample={"ctx_ids": ctx_ids},
+            model_name_or_path=MODEL_NAME,
+            num_chunk_probs=None,
+            max_chunk_len=4,
+            min_chunk_len=-1,
+            max_num_split=None,
+            is_train=False,
+            chunk_overlap=2,
+        )
+
+        stripped_chunks = self._strip_affixes(out["ctx_ids"])
+
+        self.assertEqual(out["n_ctx_chunks"], 4)
+        self.assertEqual(
+            stripped_chunks,
+            [
+                [0, 1, 2, 3],
+                [2, 3, 4, 5],
+                [4, 5, 6, 7],
+                [6, 7, 8, 9],
+            ],
+        )
+
+    def test_zero_overlap_preserves_existing_balanced_split(self):
+        ctx_ids = list(range(10))
+        out = split_too_long_ctx(
+            sample={"ctx_ids": ctx_ids},
+            model_name_or_path=MODEL_NAME,
+            num_chunk_probs=None,
+            max_chunk_len=4,
+            min_chunk_len=-1,
+            max_num_split=None,
+            is_train=False,
+            chunk_overlap=0,
+        )
+
+        stripped_chunks = self._strip_affixes(out["ctx_ids"])
+
+        self.assertEqual(out["n_ctx_chunks"], 3)
+        self.assertEqual(stripped_chunks, [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9]])
+
+    def test_overlap_is_rejected_for_train(self):
+        with self.assertRaisesRegex(ValueError, "eval splits"):
+            split_too_long_ctx(
+                sample={"ctx_ids": list(range(8))},
+                model_name_or_path=MODEL_NAME,
+                num_chunk_probs=None,
+                max_chunk_len=4,
+                min_chunk_len=-1,
+                max_num_split=None,
+                is_train=True,
+                chunk_overlap=1,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_clipper_eval.py

@@ -0,0 +1,60 @@
+from ctx_to_lora.data.preprocessing_fn import parse_clipper_user_message
+from ctx_to_lora.eval_utils import compute_clipper_metrics, extract_clipper_answer
+
+
+def test_parse_clipper_user_message_splits_context_and_prompt():
+    user_message = (
+        "You are provided with a context and a statement.\n\n"
+        "<context>Book body</context>\n\n"
+        "<statement>A claim.</statement>\n\n"
+        "<question>Is it true or false?</question>"
+    )
+
+    context, prompt = parse_clipper_user_message(user_message)
+
+    assert context == "Book body"
+    assert "<statement>A claim.</statement>" in prompt
+    assert "<context>" not in prompt
+
+
+def test_extract_clipper_answer_prefers_answer_tag():
+    text = "<explanation>Reasoning mentions false.</explanation><answer>TRUE</answer>"
+    assert extract_clipper_answer(text) == "true"
+
+
+def test_compute_clipper_metrics_reports_pair_accuracy():
+    decoded_txts = [
+        {
+            "generated": "<answer>TRUE</answer>",
+            "label": "<answer>TRUE</answer>",
+            "clipper_status": "true",
+            "clipper_pair_id": "pair-a",
+        },
+        {
+            "generated": "<answer>false</answer>",
+            "label": "<answer>FALSE</answer>",
+            "clipper_status": "false",
+            "clipper_pair_id": "pair-a",
+        },
+        {
+            "generated": "<answer>true</answer>",
+            "label": "<answer>TRUE</answer>",
+            "clipper_status": "true",
+            "clipper_pair_id": "pair-b",
+        },
+        {
+            "generated": "<answer>true</answer>",
+            "label": "<answer>FALSE</answer>",
+            "clipper_status": "false",
+            "clipper_pair_id": "pair-b",
+        },
+    ]
+
+    metrics, per_sample, counts = compute_clipper_metrics(decoded_txts)
+
+    assert metrics["clipper_accuracy"] == 0.75
+    assert metrics["clipper_true_accuracy"] == 1.0
+    assert metrics["clipper_false_accuracy"] == 0.5
+    assert metrics["clipper_pair_accuracy"] == 0.5
+    assert per_sample["clipper_pair_accuracy"] == [1.0, 1.0, 0.0, 0.0]
+    assert counts["clipper_pair_accuracy"] == 2

tests/test_dense_lora_magnitude.py

@@ -0,0 +1,249 @@
+import json
+import runpy
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from ctx_to_lora.eval_utils import eval_generation
+from ctx_to_lora.lora_stats import DenseLoraMagnitudeAggregator
+
+
+def _manual_dense_stats(combined_loras, scaling):
+    overall = []
+    by_module = {}
+    for module_name, module_lora in combined_loras.items():
+        module_values = []
+        A = module_lora["A"]
+        B = module_lora["B"]
+        for batch_idx in range(A.shape[0]):
+            for layer_idx in range(A.shape[1]):
+                dense = (
+                    torch.einsum(
+                        "ro,ri->oi",
+                        B[batch_idx, layer_idx],
+                        A[batch_idx, layer_idx],
+                    )
+                    * scaling
+                )
+                module_values.append(dense.abs().reshape(-1))
+        by_module[module_name] = torch.cat(module_values)
+        overall.append(by_module[module_name])
+    return torch.cat(overall), by_module
+
+
+class DenseLoraMagnitudeTests(unittest.TestCase):
+    @unittest.skipUnless(
+        torch.cuda.is_available(),
+        "CUDA is required for dense LoRA magnitude computation.",
+    )
+    def test_dense_lora_magnitude_aggregator_matches_manual_batched_dense_update(self):
+        scaling = 1.5
+        combined_loras = {
+            "down_proj": {
+                "A": torch.tensor(
+                    [
+                        [
+                            [[1.0, -2.0], [0.5, 3.0]],
+                            [[-1.0, 2.0], [4.0, -0.5]],
+                        ],
+                        [
+                            [[2.0, 1.0], [1.5, -1.0]],
+                            [[0.25, -3.0], [2.0, 0.75]],
+                        ],
+                    ]
+                ),
+                "B": torch.tensor(
+                    [
+                        [
+                            [[2.0, -1.0, 0.5], [1.0, 0.0, -2.0]],
+                            [[0.5, 1.5, -1.0], [2.0, -0.5, 1.0]],
+                        ],
+                        [
+                            [[-1.0, 2.0, 0.25], [0.5, -1.5, 3.0]],
+                            [[1.0, 0.0, 2.0], [-2.0, 1.0, -0.5]],
+                        ],
+                    ]
+                ),
+            },
+            "up_proj": {
+                "A": torch.tensor(
+                    [
+                        [
+                            [[0.5, 1.0], [2.0, -1.0]],
+                            [[1.0, -0.5], [0.25, 3.0]],
+                        ],
+                        [
+                            [[-2.0, 0.5], [1.0, 1.5]],
+                            [[0.75, -1.25], [2.5, 0.5]],
+                        ],
+                    ]
+                ),
+                "B": torch.tensor(
+                    [
+                        [
+                            [[1.0, -2.0], [0.5, 1.5]],
+                            [[-1.0, 2.0], [3.0, 0.25]],
+                        ],
+                        [
+                            [[2.0, 1.0], [-0.5, 0.5]],
+                            [[1.5, -1.0], [0.0, 2.0]],
+                        ],
+                    ]
+                ),
+            },
+        }
+
+        aggregator = DenseLoraMagnitudeAggregator(scaling=scaling, out_chunk_size=2)
+        aggregator.update(combined_loras)
+
+        overall_manual, by_module_manual = _manual_dense_stats(combined_loras, scaling)
+        stats = aggregator.to_dict()
+
+        self.assertEqual(stats["num_examples"], 2)
+        self.assertEqual(stats["num_batches"], 1)
+        self.assertEqual(stats["num_layer_matrices"], 4)
+
+        self.assertAlmostEqual(
+            stats["overall"]["mean"], overall_manual.mean().item(), places=6
+        )
+        self.assertAlmostEqual(
+            stats["overall"]["min"], overall_manual.min().item(), places=6
+        )
+        self.assertAlmostEqual(
+            stats["overall"]["max"], overall_manual.max().item(), places=6
+        )
+        self.assertAlmostEqual(
+            stats["overall"]["std"],
+            overall_manual.std(unbiased=False).item(),
+            places=6,
+        )
+
+        for module_name, module_values in by_module_manual.items():
+            module_stats = stats["by_module"][module_name]
+            self.assertAlmostEqual(
+                module_stats["mean"], module_values.mean().item(), places=6
+            )
+            self.assertAlmostEqual(
+                module_stats["min"], module_values.min().item(), places=6
+            )
+            self.assertAlmostEqual(
+                module_stats["max"], module_values.max().item(), places=6
+            )
+            self.assertAlmostEqual(
+                module_stats["std"], module_values.std(unbiased=False).item(), places=6
+            )
+
+    def test_eval_generation_saves_dense_lora_magnitude_json_and_metrics(self):
+        dense_stats = {
+            "overall": {"count": 8, "mean": 1.25, "min": 0.0, "max": 4.0, "std": 1.0},
+            "by_module": {
+                "down_proj": {
+                    "count": 8,
+                    "mean": 1.25,
+                    "min": 0.0,
+                    "max": 4.0,
+                    "std": 1.0,
+                }
+            },
+            "num_examples": 1,
+            "num_batches": 1,
+            "num_layer_matrices": 2,
+            "out_chunk_size": 8,
+            "scaling": 16.0,
+        }
+
+        class FakeModel:
+            def __init__(self, stats):
+                self.log_dense_lora_magnitude = True
+                self._stats = stats
+                self.reset_calls = 0
+
+            def reset_dense_lora_magnitude_stats(self):
+                self.reset_calls += 1
+
+            def get_dense_lora_magnitude_stats(self):
+                return self._stats
+
+        class FakeTrainer:
+            def __init__(self, output_dir, stats):
+                self.args = SimpleNamespace(output_dir=output_dir)
+                self.model = FakeModel(stats)
+                self.logged = []
+                self.saved = []
+
+            def predict(self, *args, **kwargs):
+                return SimpleNamespace(predictions=[[0, 1]], metrics={})
+
+            def log_metrics(self, split, metrics):
+                self.logged.append((split, metrics))
+
+            def save_metrics(self, split, metrics):
+                self.saved.append((split, metrics))
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            trainer = FakeTrainer(tmpdir, dense_stats)
+            decoded = [{"generated": "answer", "label": "answer", "ctx_ids_len": 4}]
+            dataset = [{"input_ids": [1], "labels": [-100, 1]}]
+
+            with patch(
+                "ctx_to_lora.eval_utils.decode_test_result", return_value=decoded
+            ):
+                result = eval_generation(
+                    trainer,
+                    tokenizer=None,
+                    ctx_tokenizer=None,
+                    datasets={"dummy": dataset},
+                    original_datasets={},
+                    answers={},
+                    split="test",
+                    remove_context=False,
+                    gen_kwargs={"max_new_tokens": 4, "do_sample": False},
+                )
+
+            json_path = Path(tmpdir) / "test_dummy_dense_lora_magnitude.json"
+            self.assertTrue(json_path.exists())
+            self.assertEqual(json.loads(json_path.read_text()), dense_stats)
+            self.assertEqual(trainer.model.reset_calls, 1)
+            self.assertIn("test_dummy_dense_lora_magnitude_mean", result["test_dummy"])
+            self.assertEqual(
+                result["test_dummy"]["test_dummy_dense_lora_magnitude_mean"], 1.25
+            )
+            self.assertIn(
+                "test_dummy_dense_lora_magnitude_down_proj_max", result["test_dummy"]
+            )
+
+    def test_run_eval_cli_forwards_log_dense_lora_magnitude_flag(self):
+        script_path = Path(__file__).resolve().parents[1] / "run_eval.py"
+        captured = {}
+
+        def fake_run_eval(**kwargs):
+            captured.update(kwargs)
+
+        argv = [
+            str(script_path),
+            "--model_name_or_path",
+            "google/gemma-2-2b-it",
+            "--datasets",
+            "dummy_dataset",
+            "--log_dense_lora_magnitude",
+        ]
+
+        with patch("ctx_to_lora.eval_utils.run_eval", side_effect=fake_run_eval):
+            old_argv = sys.argv
+            try:
+                sys.argv = argv
+                runpy.run_path(str(script_path), run_name="__main__")
+            finally:
+                sys.argv = old_argv
+
+        self.assertTrue(captured["log_dense_lora_magnitude"])
+        self.assertTrue(captured["generative"])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_oolong_eval.py

@@ -0,0 +1,110 @@
+from collections import Counter
+
+from datasets import Dataset
+
+from ctx_to_lora.data.preprocessing_fn import get_preprocessing_fn
+from ctx_to_lora.data.processing import (
+    OOLONG_MAX_CONTEXT_LEN,
+    filter_dataset_by_max_value,
+    stratified_subsample_dataset,
+)
+from ctx_to_lora.eval_utils import compute_oolong_metrics, extract_oolong_answer
+
+
+def test_oolong_preprocessing_maps_core_fields_and_metadata():
+    preprocess = get_preprocessing_fn("oolong-synth", is_eval=True)
+    sample = {
+        "id": 7,
+        "dataset": "spam",
+        "context_window_text": "context body",
+        "question": "Give your final answer in the form 'Label: answer'.",
+        "task_group": "classification",
+        "task": "most_freq",
+        "answer": "['ham']",
+        "answer_type": "ANSWER_TYPE.LABEL",
+        "input_subset": "all",
+        "num_labels": 2,
+        "context_window_id": 101,
+    }
+
+    processed = preprocess(sample)
+
+    assert processed["context"] == "context body"
+    assert processed["prompts"] == [sample["question"]]
+    assert processed["responses"] == ["['ham']"]
+    assert processed["oolong_task"] == "most_freq"
+    assert processed["oolong_answer_type"] == "ANSWER_TYPE.LABEL"
+    assert processed["oolong_context_window_id"] == 101
+
+
+def test_extract_oolong_answer_prefers_suffix_after_colon():
+    answer, confidence = extract_oolong_answer("Reasoning...\nLabel: [ham]")
+
+    assert answer == "ham"
+    assert confidence == "vhigh"
+
+
+def test_compute_oolong_metrics_handles_exact_numeric_and_date_scoring():
+    decoded_txts = [
+        {
+            "generated": "Label: ham",
+            "label": "['ham']",
+            "oolong_answer_type": "ANSWER_TYPE.LABEL",
+        },
+        {
+            "generated": "Answer: 8",
+            "label": "['10']",
+            "oolong_answer_type": "ANSWER_TYPE.NUMERIC",
+        },
+        {
+            "generated": "Date: 2025-02-07",
+            "label": "[datetime.date(2025, 2, 7)]",
+            "oolong_answer_type": "ANSWER_TYPE.DATE",
+        },
+    ]
+
+    metrics, per_sample, counts = compute_oolong_metrics(decoded_txts)
+
+    assert metrics["oolong_score"] == (1.0 + (0.75**2) + 1.0) / 3
+    assert metrics["oolong_parse_rate"] == 1.0
+    assert per_sample["oolong_score"] == [1.0, 0.75**2, 1.0]
+    assert counts["oolong_score"] == 3
+
+
+def test_stratified_subsample_dataset_balances_context_len_groups():
+    ds = Dataset.from_dict(
+        {
+            "context_len": [1] * 800 + [2] * 50 + [3] * 800,
+            "row_id": list(range(1650)),
+        }
+    )
+
+    sampled = stratified_subsample_dataset(
+        ds=ds,
+        group_column="context_len",
+        max_samples=1000,
+        seed=42,
+    )
+    counts = Counter(sampled["context_len"])
+
+    assert len(sampled) == 1000
+    assert counts[2] == 50
+    assert counts[1] == 475
+    assert counts[3] == 475
+
+
+def test_filter_dataset_by_max_value_drops_over_64k_context_len_rows():
+    ds = Dataset.from_dict(
+        {
+            "context_len": [1024, OOLONG_MAX_CONTEXT_LEN, OOLONG_MAX_CONTEXT_LEN + 1],
+            "row_id": [0, 1, 2],
+        }
+    )
+
+    filtered = filter_dataset_by_max_value(
+        ds=ds,
+        column="context_len",
+        max_value=OOLONG_MAX_CONTEXT_LEN,
+    )
+
+    assert filtered["context_len"] == [1024, OOLONG_MAX_CONTEXT_LEN]

tests/test_rag.py

@@ -0,0 +1,410 @@
+import json
+import runpy
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from ctx_to_lora.data.definitions import CTX_AFFIXES
+from ctx_to_lora.eval_utils import eval_generation, run_eval
+from ctx_to_lora.modeling.rag import DocToLoraRAGModel, RAGModel, RetrievedChunk
+
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
+PREFIX = CTX_AFFIXES[MODEL_NAME]["prefix"]
+SUFFIX = CTX_AFFIXES[MODEL_NAME]["suffix"]
+
+
+class FakeTokenizer:
+    def __init__(self, name_or_path=MODEL_NAME):
+        self.name_or_path = name_or_path
+        self.pad_token_id = 0
+        self.eos_token_id = 1
+        self.chat_template = "fake"
+        self._token_to_id = {}
+        self._id_to_token = {}
+        self._next_id = 1000
+
+    def _encode_words(self, text: str) -> list[int]:
+        tokens = []
+        for word in text.split():
+            if word not in self._token_to_id:
+                token_id = self._next_id
+                self._next_id += 1
+                self._token_to_id[word] = token_id
+                self._id_to_token[token_id] = word
+            tokens.append(self._token_to_id[word])
+        return tokens
+
+    def __call__(
+        self,
+        text,
+        add_special_tokens=False,
+        return_attention_mask=False,
+        return_tensors=None,
+        **kwargs,
+    ):
+        del add_special_tokens, kwargs
+        input_ids = self._encode_words(text)
+        out = {"input_ids": input_ids}
+        if return_attention_mask:
+            out["attention_mask"] = [1] * len(input_ids)
+        if return_tensors == "pt":
+            out = {k: torch.tensor([v], dtype=torch.long) for k, v in out.items()}
+        return out
+
+    def decode(self, ids, skip_special_tokens=True):
+        del skip_special_tokens
+        if isinstance(ids, torch.Tensor):
+            ids = ids.tolist()
+        words = []
+        for token_id in ids:
+            if token_id in self._id_to_token:
+                words.append(self._id_to_token[token_id])
+        return " ".join(words)
+
+    def apply_chat_template(
+        self,
+        conversations,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_attention_mask=True,
+        return_dict=True,
+        return_tensors=None,
+        add_special_tokens=False,
+        **kwargs,
+    ):
+        del tokenize, add_special_tokens, kwargs
+        messages = (
+            conversations[0]
+            if conversations and isinstance(conversations[0], list)
+            else conversations
+        )
+        user_content = [m["content"] for m in messages if m["role"] == "user"][-1]
+        token_ids = PREFIX + self._encode_words(user_content)
+        if add_generation_prompt:
+            token_ids += SUFFIX
+        out = {"input_ids": [token_ids]}
+        if return_attention_mask:
+            out["attention_mask"] = [[1] * len(token_ids)]
+        if return_tensors == "pt":
+            out = {k: torch.tensor(v, dtype=torch.long) for k, v in out.items()}
+        if return_dict:
+            return out
+        return out["input_ids"]
+
+
+class FakeBaseModel(torch.nn.Module):
+    def __init__(self, name_or_path=MODEL_NAME):
+        super().__init__()
+        self.dummy = torch.nn.Parameter(torch.zeros(1))
+        self.name_or_path = name_or_path
+        self.config = SimpleNamespace(
+            name_or_path=name_or_path,
+            _name_or_path=name_or_path,
+            max_position_embeddings=4096,
+        )
+        self.generation_config = SimpleNamespace(pad_token_id=0)
+        self.last_generate_kwargs = None
+
+    def generate(self, *args, **kwargs):
+        del args
+        self.last_generate_kwargs = kwargs
+        return kwargs["input_ids"]
+
+
+class FakeDocToLoraModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.base_model = FakeBaseModel()
+        self.generation_config = self.base_model.generation_config
+        self.last_generate_kwargs = None
+
+    def generate(self, *args, **kwargs):
+        del args
+        self.last_generate_kwargs = dict(kwargs)
+        return kwargs["input_ids"]
+
+
+def build_chat_tensor(tokenizer: FakeTokenizer, text: str) -> dict[str, torch.Tensor]:
+    token_ids = PREFIX + tokenizer._encode_words(text) + SUFFIX
+    return {
+        "input_ids": torch.tensor([token_ids], dtype=torch.long),
+        "attention_mask": torch.ones((1, len(token_ids)), dtype=torch.long),
+    }
+
+
+class RAGModelTests(unittest.TestCase):
+    def setUp(self):
+        self.tokenizer = FakeTokenizer()
+        self.base_model = FakeBaseModel()
+        self.rag = RAGModel(
+            model=self.base_model,
+            tokenizer=self.tokenizer,
+            chunk_size=4,
+            chunk_overlap=2,
+            top_k=2,
+            max_retrieved_tokens=8,
+        )
+
+    def test_chunk_context_respects_overlap(self):
+        chunks = self.rag._chunk_context("a b c d e f g h i j")
+        self.assertEqual(
+            [(chunk.start, chunk.end) for chunk in chunks],
+            [(0, 4), (2, 6), (4, 8), (6, 10)],
+        )
+        self.assertEqual(chunks[1].text, "c d e f")
+        self.assertEqual(chunks[-1].text, "g h i j")
+
+    def test_bm25_prefers_relevant_chunk(self):
+        chunks = [
+            RetrievedChunk(
+                "apple orange", self.tokenizer._encode_words("apple orange"), 0, 2, 0.0
+            ),
+            RetrievedChunk(
+                "banana pear", self.tokenizer._encode_words("banana pear"), 2, 4, 0.0
+            ),
+        ]
+        ranked = self.rag._score_chunks_bm25("apple", chunks)
+        self.assertEqual(ranked[0].text, "apple orange")
+        self.assertGreater(ranked[0].score, ranked[1].score)
+
+    def test_generate_rebuilds_prompt_from_selected_chunks(self):
+        rag = RAGModel(
+            model=self.base_model,
+            tokenizer=self.tokenizer,
+            chunk_size=3,
+            chunk_overlap=0,
+            top_k=1,
+            max_retrieved_tokens=3,
+        )
+        question_inputs = build_chat_tensor(self.tokenizer, "answer")
+        context_inputs = build_chat_tensor(
+            self.tokenizer,
+            "cat dog fish bird apple answer kiwi mango noise noise",
+        )
+
+        rag.generate(
+            input_ids=question_inputs["input_ids"],
+            attention_mask=question_inputs["attention_mask"],
+            ctx_ids=context_inputs["input_ids"],
+            ctx_attn_mask=context_inputs["attention_mask"],
+            max_new_tokens=16,
+        )
+
+        generated_input_ids = self.base_model.last_generate_kwargs["input_ids"][0]
+        generated_attention_mask = self.base_model.last_generate_kwargs[
+            "attention_mask"
+        ][0]
+        rebuilt_prompt, _ = rag._decode_question(
+            generated_input_ids,
+            generated_attention_mask,
+        )
+
+        self.assertIn("bird apple answer", rebuilt_prompt)
+        self.assertIn("Question:", rebuilt_prompt)
+        self.assertNotIn("cat dog fish", rebuilt_prompt)
+        self.assertEqual(len(rag.get_retrieval_history()), 1)
+        self.assertEqual(
+            rag.get_retrieval_history()[0]["rag_selected_chunks"][0]["text"],
+            "bird apple answer",
+        )
+
+    def test_hybrid_generate_preserves_ctx_and_rewrites_prompt_input(self):
+        doc_model = FakeDocToLoraModel()
+        hybrid = DocToLoraRAGModel(
+            model=doc_model,
+            tokenizer=self.tokenizer,
+            chunk_size=3,
+            chunk_overlap=0,
+            top_k=1,
+            max_retrieved_tokens=3,
+        )
+        question_inputs = build_chat_tensor(self.tokenizer, "answer")
+        context_inputs = build_chat_tensor(
+            self.tokenizer,
+            "cat dog fish bird apple answer kiwi mango noise noise",
+        )
+
+        hybrid.generate(
+            input_ids=question_inputs["input_ids"],
+            attention_mask=question_inputs["attention_mask"],
+            ctx_ids=context_inputs["input_ids"],
+            ctx_attn_mask=context_inputs["attention_mask"],
+            n_ctx_chunks=torch.tensor([1], dtype=torch.int32),
+            max_new_tokens=8,
+        )
+
+        forwarded = doc_model.last_generate_kwargs
+        self.assertTrue(torch.equal(forwarded["ctx_ids"], context_inputs["input_ids"]))
+        self.assertTrue(
+            torch.equal(forwarded["ctx_attn_mask"], context_inputs["attention_mask"])
+        )
+        self.assertTrue(
+            torch.equal(forwarded["n_ctx_chunks"], torch.tensor([1], dtype=torch.int32))
+        )
+        rebuilt_prompt, _ = hybrid._decode_question(
+            forwarded["input_ids"][0],
+            forwarded["attention_mask"][0],
+        )
+        self.assertIn("bird apple answer", rebuilt_prompt)
+        self.assertEqual(
+            hybrid.get_retrieval_history()[0]["rag_selected_chunks"][0]["text"],
+            "bird apple answer",
+        )
+
+    def test_run_eval_rejects_eval_chunking_with_rag(self):
+        with self.assertRaisesRegex(ValueError, "retrieval chunking"):
+            run_eval(
+                model_name_or_path="google/gemma-2-2b-it",
+                datasets=["dummy_dataset"],
+                use_rag=True,
+                max_ctx_chunk_len=128,
+                eval_batch_size=1,
+            )
+
+    def test_run_eval_rejects_hybrid_without_checkpoint(self):
+        with self.assertRaisesRegex(ValueError, "only available"):
+            run_eval(
+                model_name_or_path="google/gemma-2-2b-it",
+                datasets=["dummy_dataset"],
+                use_hybrid_rag=True,
+                eval_batch_size=1,
+            )
+
+    def test_eval_generation_saves_retrieval_metadata(self):
+        class FakeModel:
+            def __init__(self):
+                self._history = [
+                    {
+                        "rag_question": "who?",
+                        "rag_selected_chunks": [
+                            {"text": "retrieved text", "score": 1.0}
+                        ],
+                    }
+                ]
+
+            def reset_retrieval_history(self):
+                self._history = list(self._history)
+
+            def get_retrieval_history(self):
+                return self._history
+
+        class FakeTrainer:
+            def __init__(self, output_dir):
+                self.args = SimpleNamespace(output_dir=output_dir)
+                self.model = FakeModel()
+
+            def predict(self, *args, **kwargs):
+                del args, kwargs
+                return SimpleNamespace(predictions=[[0, 1]], metrics={})
+
+            def log_metrics(self, split, metrics):
+                del split, metrics
+
+            def save_metrics(self, split, metrics):
+                del split, metrics
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            trainer = FakeTrainer(tmpdir)
+            decoded = [{"generated": "answer", "label": "answer", "ctx_ids_len": 4}]
+            dataset = [{"input_ids": [1], "labels": [-100, 1]}]
+
+            with patch(
+                "ctx_to_lora.eval_utils.decode_test_result", return_value=decoded
+            ):
+                eval_generation(
+                    trainer,
+                    tokenizer=None,
+                    ctx_tokenizer=None,
+                    datasets={"dummy": dataset},
+                    original_datasets={},
+                    answers={},
+                    split="test",
+                    remove_context=False,
+                    gen_kwargs={"max_new_tokens": 4, "do_sample": False},
+                )
+
+            jsonl_path = Path(tmpdir) / "test_dummy_generated_text.jsonl"
+            lines = jsonl_path.read_text().splitlines()
+            payload = json.loads(lines[0])
+            self.assertEqual(payload["rag_question"], "who?")
+            self.assertEqual(
+                payload["rag_selected_chunks"][0]["text"], "retrieved text"
+            )
+
+    def test_run_eval_cli_forwards_rag_flags(self):
+        script_path = Path(__file__).resolve().parents[1] / "run_eval.py"
+        captured = {}
+
+        def fake_run_eval(**kwargs):
+            captured.update(kwargs)
+
+        argv = [
+            str(script_path),
+            "--model_name_or_path",
+            "google/gemma-2-2b-it",
+            "--datasets",
+            "dummy_dataset",
+            "--use_rag",
+            "--rag_chunk_size",
+            "128",
+            "--rag_chunk_overlap",
+            "16",
+            "--rag_top_k",
+            "3",
+            "--rag_max_retrieved_tokens",
+            "512",
+        ]
+
+        with patch("ctx_to_lora.eval_utils.run_eval", side_effect=fake_run_eval):
+            old_argv = sys.argv
+            try:
+                sys.argv = argv
+                runpy.run_path(str(script_path), run_name="__main__")
+            finally:
+                sys.argv = old_argv
+
+        self.assertTrue(captured["use_rag"])
+        self.assertEqual(captured["rag_chunk_size"], 128)
+        self.assertEqual(captured["rag_chunk_overlap"], 16)
+        self.assertEqual(captured["rag_top_k"], 3)
+        self.assertEqual(captured["rag_max_retrieved_tokens"], 512)
+        self.assertTrue(captured["generative"])
+
+    def test_run_eval_cli_forwards_hybrid_rag_flag(self):
+        script_path = Path(__file__).resolve().parents[1] / "run_eval.py"
+        captured = {}
+
+        def fake_run_eval(**kwargs):
+            captured.update(kwargs)
+
+        argv = [
+            str(script_path),
+            "--checkpoint_path",
+            "train_outputs/runs/demo/checkpoint-1/pytorch_model.bin",
+            "--datasets",
+            "dummy_dataset",
+            "--use_hybrid_rag",
+            "--rag_chunk_size",
+            "128",
+        ]
+
+        with patch("ctx_to_lora.eval_utils.run_eval", side_effect=fake_run_eval):
+            old_argv = sys.argv
+            try:
+                sys.argv = argv
+                runpy.run_path(str(script_path), run_name="__main__")
+            finally:
+                sys.argv = old_argv
+
+        self.assertTrue(captured["use_hybrid_rag"])
+        self.assertEqual(captured["rag_chunk_size"], 128)
+        self.assertTrue(captured["generative"])
+
+
+if __name__ == "__main__":
+    unittest.main()

uv.lock

@@ -852,6 +852,7 @@ dependencies = [
     { name = "peft" },
     { name = "plotly" },
     { name = "pre-commit" },
+    { name = "python-dateutil" },
     { name = "rouge-score" },
     { name = "setuptools" },
     { name = "tensorboard" },
@@ -889,6 +890,7 @@ requires-dist = [
     { name = "peft" },
     { name = "plotly" },
     { name = "pre-commit" },
+    { name = "python-dateutil", specifier = ">=2.9.0.post0" },
     { name = "rouge-score" },
     { name = "setuptools" },
     { name = "tensorboard" },
@@ -5426,6 +5428,30 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/e6/34/ebdc18bae6aa14fbee1a08b63c015c72b64868ff7dae68808ab500c492e2/tinycss2-1.4.0-py3-none-any.whl", hash = "sha256:3a49cf47b7675da0b15d0c6e1df8df4ebd96e9394bb905a5775adb0d884c5289", size = 26610, upload-time = "2024-10-24T14:58:28.029Z" },
 ]
 
+[[package]]
+name = "tokenizers"
+version = "0.21.4"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "huggingface-hub" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/c2/2f/402986d0823f8d7ca139d969af2917fefaa9b947d1fb32f6168c509f2492/tokenizers-0.21.4.tar.gz", hash = "sha256:fa23f85fbc9a02ec5c6978da172cdcbac23498c3ca9f3645c5c68740ac007880", size = 351253, upload-time = "2025-07-28T15:48:54.325Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/98/c6/fdb6f72bf6454f52eb4a2510be7fb0f614e541a2554d6210e370d85efff4/tokenizers-0.21.4-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:2ccc10a7c3bcefe0f242867dc914fc1226ee44321eb618cfe3019b5df3400133", size = 2863987, upload-time = "2025-07-28T15:48:44.877Z" },
+    { url = "https://files.pythonhosted.org/packages/8d/a6/28975479e35ddc751dc1ddc97b9b69bf7fcf074db31548aab37f8116674c/tokenizers-0.21.4-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:5e2f601a8e0cd5be5cc7506b20a79112370b9b3e9cb5f13f68ab11acd6ca7d60", size = 2732457, upload-time = "2025-07-28T15:48:43.265Z" },
+    { url = "https://files.pythonhosted.org/packages/aa/8f/24f39d7b5c726b7b0be95dca04f344df278a3fe3a4deb15a975d194cbb32/tokenizers-0.21.4-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:39b376f5a1aee67b4d29032ee85511bbd1b99007ec735f7f35c8a2eb104eade5", size = 3012624, upload-time = "2025-07-28T13:22:43.895Z" },
+    { url = "https://files.pythonhosted.org/packages/58/47/26358925717687a58cb74d7a508de96649544fad5778f0cd9827398dc499/tokenizers-0.21.4-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:2107ad649e2cda4488d41dfd031469e9da3fcbfd6183e74e4958fa729ffbf9c6", size = 2939681, upload-time = "2025-07-28T13:22:47.499Z" },
+    { url = "https://files.pythonhosted.org/packages/99/6f/cc300fea5db2ab5ddc2c8aea5757a27b89c84469899710c3aeddc1d39801/tokenizers-0.21.4-cp39-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3c73012da95afafdf235ba80047699df4384fdc481527448a078ffd00e45a7d9", size = 3247445, upload-time = "2025-07-28T15:48:39.711Z" },
+    { url = "https://files.pythonhosted.org/packages/be/bf/98cb4b9c3c4afd8be89cfa6423704337dc20b73eb4180397a6e0d456c334/tokenizers-0.21.4-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f23186c40395fc390d27f519679a58023f368a0aad234af145e0f39ad1212732", size = 3428014, upload-time = "2025-07-28T13:22:49.569Z" },
+    { url = "https://files.pythonhosted.org/packages/75/c7/96c1cc780e6ca7f01a57c13235dd05b7bc1c0f3588512ebe9d1331b5f5ae/tokenizers-0.21.4-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cc88bb34e23a54cc42713d6d98af5f1bf79c07653d24fe984d2d695ba2c922a2", size = 3193197, upload-time = "2025-07-28T13:22:51.471Z" },
+    { url = "https://files.pythonhosted.org/packages/f2/90/273b6c7ec78af547694eddeea9e05de771278bd20476525ab930cecaf7d8/tokenizers-0.21.4-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:51b7eabb104f46c1c50b486520555715457ae833d5aee9ff6ae853d1130506ff", size = 3115426, upload-time = "2025-07-28T15:48:41.439Z" },
+    { url = "https://files.pythonhosted.org/packages/91/43/c640d5a07e95f1cf9d2c92501f20a25f179ac53a4f71e1489a3dcfcc67ee/tokenizers-0.21.4-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:714b05b2e1af1288bd1bc56ce496c4cebb64a20d158ee802887757791191e6e2", size = 9089127, upload-time = "2025-07-28T15:48:46.472Z" },
+    { url = "https://files.pythonhosted.org/packages/44/a1/dd23edd6271d4dca788e5200a807b49ec3e6987815cd9d0a07ad9c96c7c2/tokenizers-0.21.4-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:1340ff877ceedfa937544b7d79f5b7becf33a4cfb58f89b3b49927004ef66f78", size = 9055243, upload-time = "2025-07-28T15:48:48.539Z" },
+    { url = "https://files.pythonhosted.org/packages/21/2b/b410d6e9021c4b7ddb57248304dc817c4d4970b73b6ee343674914701197/tokenizers-0.21.4-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:3c1f4317576e465ac9ef0d165b247825a2a4078bcd01cba6b54b867bdf9fdd8b", size = 9298237, upload-time = "2025-07-28T15:48:50.443Z" },
+    { url = "https://files.pythonhosted.org/packages/b7/0a/42348c995c67e2e6e5c89ffb9cfd68507cbaeb84ff39c49ee6e0a6dd0fd2/tokenizers-0.21.4-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:c212aa4e45ec0bb5274b16b6f31dd3f1c41944025c2358faaa5782c754e84c24", size = 9461980, upload-time = "2025-07-28T15:48:52.325Z" },
+    { url = "https://files.pythonhosted.org/packages/3d/d3/dacccd834404cd71b5c334882f3ba40331ad2120e69ded32cf5fda9a7436/tokenizers-0.21.4-cp39-abi3-win32.whl", hash = "sha256:6c42a930bc5f4c47f4ea775c91de47d27910881902b0f20e4990ebe045a415d0", size = 2329871, upload-time = "2025-07-28T15:48:56.841Z" },
+    { url = "https://files.pythonhosted.org/packages/41/f2/fd673d979185f5dcbac4be7d09461cbb99751554ffb6718d0013af8604cb/tokenizers-0.21.4-cp39-abi3-win_amd64.whl", hash = "sha256:475d807a5c3eb72c59ad9b5fcdb254f6e17f53dfcbb9903233b0dfa9c943b597", size = 2507568, upload-time = "2025-07-28T15:48:55.456Z" },
+]
 
 [[package]]
 name = "tomli"
@@ -5647,6 +5673,7 @@ dependencies = [
     { name = "regex" },
     { name = "requests" },
     { name = "safetensors" },
+    { name = "tokenizers" },
     { name = "tqdm" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/f1/11/7414d5bc07690002ce4d7553602107bf969af85144bbd02830f9fb471236/transformers-4.51.3.tar.gz", hash = "sha256:e292fcab3990c6defe6328f0f7d2004283ca81a7a07b2de9a46d67fd81ea1409", size = 8941266, upload-time = "2025-04-14T08:15:00.485Z" }
@@ -5898,6 +5925,7 @@ dependencies = [
     { name = "setuptools", marker = "python_full_version >= '3.12'" },
     { name = "six", marker = "python_full_version >= '3.12'" },
     { name = "tiktoken" },
+    { name = "tokenizers" },
     { name = "torch" },
     { name = "torchaudio" },
     { name = "torchvision" },