Github
Huggingface
论文具有并行编码的长上下文语言建模 
本仓库包含了我们ACL 2024论文《具有并行编码的长上下文语言建模》的代码和数据。
在这项工作中，我们提出了CEPE — Context Expansion with Parallel Encoding（具有并行编码的上下文扩展）— 一个用于扩展语言模型上下文窗口的灵活框架。
本仓库包括数据预处理、CEPE训练以及评估所有基线的代码。
*: CÈPE源自法语中牛肝菌的单词
发音为/sep/
快速开始
想要测试CEPE？使用Huggingface尝试我们的模型！
首先，通过安装requirements来设置环境。
然后，只需将modeling_llama_flash.py文件复制并粘贴到您的工作目录中，即可运行：
import torch
from transformers import LlamaTokenizer
from modeling_llama_flash import LlamaForCausalContextLM
device = "cuda"
model_name = "hyen/CEPED-LLaMA-2-Chat-7B"
tokenizer = LlamaTokenizer.from_pretrained(model_name)
model = LlamaForCausalContextLM.from_pretrained(
model_name,
use_flash_attention_2="flash_attention_2",
torch_dtype=torch.bfloat16,
device_map="auto",
).eval()
contexts = tokenizer([
"我和朋友喜欢一起在餐厅吃饭。不过，我们也喜欢一起做饭和准备食物。"
"我的许多朋友喜欢踢足球和打排球。我们还喜欢看电影、参观博物馆和画廊。",
], return_tensors="pt", padding=True)
inputs = tokenizer("问题：在纽约市与一大群朋友进行社交活动的三个想法是什么？\n回答：", return_tensors="pt")
# encoder_input_ids和encoder_attention_mask的形状应为(bsz, n_ctx, seq_length)
output = model.generate(
input_ids=inputs.input_ids.to(device),
attention_mask=inputs.attention_mask.to(device),
encoder_input_ids=contexts.input_ids.unsqueeze(0).to(device),
encoder_attention_mask=contexts.attention_mask.unsqueeze(0).to(device),
max_new_tokens=200,
sample=True,
top_p=0.95,
)
print(tokenizer.batch_decode(output)[0])
这将打印出：
回答：以下是在纽约市与一大群朋友进行社交活动的三个想法：
1. 举办屋顶派对：找一个能欣赏城市美景的屋顶酒吧或餐厅。邀请一大群朋友一起享用饮品、开胃菜和音乐。
2. 进行团体寻宝游戏：创建一份清单，列出朋友们必须在城市中完成的物品或挑战。这可能包括与街头艺人合影、在特定酒吧购买饮料或寻找地标。完成最多挑战的小组获胜。
3. 参加团体烹饪课：在纽约市寻找一所为大型团体提供课程的烹饪学校或烹饪学院。选择一个主题或菜系作为重点，比如意大利菜或亚洲菜。然后，一起合作准备一顿团体大餐。
另一个示例：
回答：以下是在纽约市与一大群朋友进行社交活动的三个想法：
1. 访问热门餐厅：纽约市以其多样化和充满活力的美食场景而闻名。有无数餐厅可以满足不同口味和饮食限制。一大群朋友可以一起访问一家热门餐厅，共享美食。
2. 参观博物馆或美术馆：纽约市拥有许多世界级的博物馆和美术馆。一些最受欢迎的博物馆包括大都会艺术博物馆、现代艺术博物馆和古根海姆博物馆。一大群朋友可以一起参观博物馆或美术馆，探索展览。
3. 在公园散步：纽约市有许多公园和绿地，为人们提供了远离城市喧嚣的宁静避风港。
快速链接
要求
请按照官方安装说明安装最新版本的PyTorch（torch）。
您可以使用pip install --r requirements.txt安装其余的依赖项。
我们还建议查看flash-attention的安装说明。
我们使用python==3.9.12、torch==2.1.1、accelerate==0.24.0、transformers==4.34.1和CUDA版本12.3测试了代码。
模型列表
我们发布的模型列表如下。 有关如何使用它们的示例，请参见上面的示例。
快速提示：
encoder_input_ids和encoder_attention_mask的形状应为(batch_size, n_ctx, seq_length)- 模型生成的结果通常在解码器输入较长时效果更好。
数据
要获取训练和评估数据，请参考./data目录。
如果您只想获取所有用于评估的数据，可以从Google Drive下载。
| 名称 | 链接 |
|---|---|
| RedPajama测试集，ArXiv + Book按32K令牌过滤 | 链接 |
| RedPajama测试集，ArXiv + Book按128K令牌过滤 | 链接 |
| RedPajama测试集，所有领域连接 | 链接 |
| RedPajama测试集，检索增强（Contriever） | 链接 |
检索
在大规模上构建和索引检索语料库可能成本高昂。 我们在Google Drive上提供了检索增强语言建模部分使用的测试集，您可以从上面的链接下载。
如果您对检索感兴趣并想进行更多操作，可以按照./data目录和./retrieval目录中描述的步骤重现我们的检索语料库。
我们的检索代码参考了Contriever代码仓库。
代码结构
train.py：训练CEPEeval_lm.py：在语言建模上评估CEPE和其他基线eval_downstream.py：在下游任务上评估CEPE和其他基线configs/：包含所有用于重现基线的配置文件的文件夹
训练
训练包括两个阶段：预热阶段和标准训练阶段。
确保您已下载--train_file和--validation_file参数指定的必要数据。
您可以根据自己的系统更改这些参数的路径。
预热阶段
在预热阶段，我们简单地训练模型，使用相同的输入同时输入编码器和解码器。 要复现论文中的模型，可以使用：
torchrun --nnodes=1 --nproc_per_node=4 train.py --config configs/train_llama2_warmup
torchrun --nnodes=1 --nproc_per_node=4 train.py --config configs/train_llama2chat_warmup
你也可以通过仔细查看配置文件和train.py来根据自己的需求进行定制。
标准训练
在这个阶段，我们使用预热阶段后的模型并应用标准交叉熵损失。
torchrun --nnodes=1 --nproc_per_node=8 train.py --config configs/train_llama2_cepe
注意，要训练4K个编码器输入令牌和4K个解码器输入令牌，你的GPU必须有80GB的内存。
对于CEPED，我们还添加了KL散度损失。你可以通过参数--kl_loss_cof和--lm_loss_cof来自定义每个损失的系数。
torchrun --nnodes=1 --nproc_per_node=8 train.py --config configs/train_llama2chat_cepe
预训练MLM编码器
如果你想预训练自己的编码器模型，可以使用：
torchrun --nnodes=1 --nproc_per_node=8 train_mlm.py --config configs/train_mlm
或者，你可以直接使用我们预训练好的MLM编码器，可以从这里下载。
语言建模评估
要在ArXiv上评估CEPE，可以运行
python eval_lm.py --config configs/test_ab_32k_prevdoc_112x_4096 --model_name_or_path hyen/CEPE-LLaMA-2-7B --model_class cepe --validation_domains arxiv --output_dir output/CEPE-LLaMA-2-7B
你可以在./config目录中找到所有其他评估配置。
或者，你可以使用我们的脚本复现论文中的所有LM实验：
# CEPE
bash scripts/run_cepe_lm.sh
# StreamingLLM
bash scripts/run_streamingllm_lm.sh
# 原始模型 (LLaMA-2, LLaMA-2-32K, YaRN-64K, YaRN-128K)
bash scripts/run_vanilla_lm.sh
# RePlug
bash scripts/run_replug_lm.sh
eval_lm.py中的语言建模代码是根据train.py中的训练代码建模的。
然而，有几个微小的区别：
- 我们为评估实现了基准模型（特别是，StreamingLLM评估需要Transformers Trainer不支持的额外逻辑）。
- 我们使用数据加载器的循环而不是使用Trainer评估循环，这提供了略快的吞吐量和更低的内存使用。
- 我们添加了对HuggingFace数据集的支持（PG19、ProofPile和CodeParrot）。
下游评估
开放域问答
首先，你需要下载包含Contriever检索段落的数据集
然后，你可以使用以下命令运行我们论文中的所有实验：
bash scripts/run_qa.sh
如果你对更定制化的运行感兴趣，可以使用上面的脚本作为示例。
你可以简单地修改配置文件或为python eval_downstream.py添加额外的命令行参数。
上下文学习（ICL）
所有ICL数据集都可在HuggingFace上获得，这意味着你可以简单地使用datasets.load_dataset加载所有内容。
要运行论文中的实验，请使用：
bash scripts/run_icl.sh
ZeroScrolls
要运行论文中的实验，请使用：
bash scripts/run_zeroscrolls.sh
我们使用HuggingFace上提供的验证集，可能使用原作者发布的版本或SCROLLS（Shaham等人，2022）发布的版本。
错误或问题？
如果你有任何与代码或论文相关的问题，请随时发邮件给Howard（hyen@cs.princeton.edu）。如果你在使用代码时遇到任何问题，或想报告错误，可以提出issue。请尽量详细说明问题，这样我们才能更好更快地帮助你！
引用
如果你在工作中使用了CEPE，请引用我们的论文：
@inproceedings{yen2024long,
title={Long-Context Language Modeling with Parallel Context Encoding},
author={Yen, Howard and Gao, Tianyu and Chen, Danqi},
booktitle={Association for Computational Linguistics (ACL)},
year={2024}
}
