Doc-to-LoRA (D2L)

What does this agent do, and when should you use it?

Doc-to-LoRA is Sakana AI's reference implementation of a hypernetwork approach that lets an LLM "instantly internalize" a context. Instead of retrieving chunks or stuffing a prompt, it generates LoRA updates conditioned on the document and injects them into a base model, so later generations are shaped by that document. The repo ships a uv-based install script, downloadable pre-trained checkpoints from Hugging Face, a Python inference API, an interactive demo, and full data-prep/training/eval scripts for the main experiment and a NIAH (needle-in-a-haystack) experiment. The code sits on PyTorch and Hugging Face model loading, with entry points such as ModulatedPretrainedModel and get_tokenizer. It is aimed at researchers reproducing the paper or building on the method, not at teams wanting a plug-and-play service.

The repository implements ModulatedPretrainedModel, which loads a state dict (e.g. trained_d2l/gemma_demo/checkpoint-80000/pytorch_model.bin) and builds the model. Calling model.internalize(doc) reads a plain-text document and updates the model's LoRA parameters from it; a following model.generate(input_ids=..., max_new_tokens=512) produces answers influenced by the document, and model.reset() removes the internalized information. get_tokenizer fetches the tokenizer for the base model, and chat inputs are built with tokenizer.apply_chat_template. The README also documents demo/app.py for an interactive demo, webui/self_gen_viewer.py for inspecting self-generated data, and two script sets (scripts/main_exp, scripts/niah) covering data download/generation, training and evaluation.

1. Researchers reproducing the paper's context-internalization results, running training and evaluation through scripts/main_exp.
2. Teams evaluating long-context memory approaches who want to compare a demo's output with and without an internalized document (including hallucination without it).
3. Engineers running needle-in-a-haystack memory benchmarks via the ordered scripts in scripts/niah (generate data, train, evaluate).
4. Developers who need a fixed base model to absorb a small set of proprietary documents temporarily, using internalize/reset to inject and roll back.
5. Teams building synthetic training data who want to inspect generated samples with webui/self_gen_viewer.py.

How do you install or deploy this agent?

Install uv, then run the repository's install script:

curl -LsSf https://astral.sh/uv/install.sh | sh
./install.sh

Log in to Hugging Face and download the pre-trained models into a local trained_d2l directory:

uv run huggingface-cli login
uv run huggingface-cli download SakanaAI/doc-to-lora --local-dir trained_d2l --include "*/"

How do you use this agent?

From the project root, call the downloaded checkpoint from Python:

import torch

from ctx_to_lora.model_loading import get_tokenizer
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel

checkpoint_path = "trained_d2l/gemma_demo/checkpoint-80000/pytorch_model.bin"
state_dict = torch.load(checkpoint_path, weights_only=False)
model = ModulatedPretrainedModel.from_state_dict(
    state_dict, train=False, use_sequence_packing=False
)
model.reset()
tokenizer = get_tokenizer(model.base_model.name_or_path)

doc = open("data/sakana_wiki.txt", "r").read()
chat = [{"role": "user", "content": "Tell me about Sakana AI."}]
chat_ids = tokenizer.apply_chat_template(
    chat,
    add_special_tokens=False,
    return_attention_mask=False,
    add_generation_prompt=True,
    return_tensors="pt",
).to(model.device)

model.internalize(doc)

outputs = model.generate(input_ids=chat_ids, max_new_tokens=512)
print(tokenizer.decode(outputs[0]))

Launch the interactive demo:

uv run demo/app.py

Experiment scripts also run through uv from the repo root, for example:

uv run scripts/main_exp/0-download_data.sh
uv run scripts/main_exp/1-train.sh

What are this agent's strengths and limitations?

How does this agent compare with similar options?

Key facts side by side with the most closely related agents.

How does FollowAgents rate this agent?

FAQ