Doc-to-LoRA (D2L): 学习如何即时内化上下文

:sparkles: 交互式网页 | :newspaper: X (推特) | :scroll: 论文 | :hugs: Hugging Face | :octocat: GitHub

这是 Doc-to-LoRA (D2L) 的参考实现，旨在通过 Hypernetworks 更新 LLM，使其能够记忆事实信息。

---

🛠️ 安装

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

🤗 预训练模型

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

🚀 Python API 使用方法

# 注意：此接口仅支持非批处理输入
# 如需批处理推理，请参考 `src/ctx_to_lora/modeling/hypernet.py`
import torch

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

# 模型加载
checkpoint_path = "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)


# 在内化（internalization）之后进行的调用将受到内化信息的影响
model.internalize(doc)

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


# 移除内化信息
# model.reset()

# 如果没有内化信息，模型将会产生幻觉
# outputs = model.generate(input_ids=chat_ids, max_new_tokens=512)
# print(tokenizer.decode(outputs[0]))

🎮 交互式演示

uv run demo/app.py

视频演示

🧪 实验脚本

要运行以下任何脚本，请在该项目根目录下使用 uv run $PATH_TO_SCRIPT。

实验	数据准备	训练	评估	备注
主要实验	scripts/main_exp/0-download_data.sh	scripts/main_exp/1-train.sh	scripts/main_exp/eval/*.sh	下载数据速度最快；仅在需要新的合成数据时重新生成。评估脚本可复现论文中的主要指标。
NIAH	scripts/niah/0-gen_data.sh	scripts/niah/1-train.sh	scripts/niah/2-eval.sh	请按顺序运行脚本；数据生成只需执行一次。

🔬 自生成数据查看器

在下载或生成数据后，可以使用此脚本查看部分数据样本。

uv run webui/self_gen_viewer.py

更多信息请参阅 webui/SELF_GEN_VIEWER.md。

📚 引用

@techreport{sakana2025doc-to-lora,
  title       = {{Doc-to-LoRA: Learning to Instantly Internalize Contexts}},
  author      = {Rujikorn Charakorn and Edoardo Cetin and Shinnosuke Uesaka and Robert Tjarko Lange},
  institution = {Sakana AI},
  year        = {2026},
  month       = {Febuary},
  note        = {Technical Report}
}