具有并行编码的长上下文语言建模 <img src="https://yellow-cdn.veclightyear.com/835a84d5/09450a3c-cca2-4e5c-b92c-b74a12790a13.png" alt="CEPE" width="30" height="30">

本仓库包含了我们ACL 2024论文《具有并行编码的长上下文语言建模》的代码和数据。在这项工作中，我们提出了CEPE <img src="https://yellow-cdn.veclightyear.com/835a84d5/09450a3c-cca2-4e5c-b92c-b74a12790a13.png" alt="CEPE" width="20" height="20"> — Context Expansion with Parallel Encoding（具有并行编码的上下文扩展）— 一个用于扩展语言模型上下文窗口的灵活框架。本仓库包括数据预处理、CEPE训练以及评估所有基线的代码。

<p align="center"><br>*: CÈPE源自法语中牛肝菌的单词 <img src="https://yellow-cdn.veclightyear.com/835a84d5/09450a3c-cca2-4e5c-b92c-b74a12790a13.png" alt="CEPE" width="20" height="20"> 发音为/sep/ </p> <img src="https://yellow-cdn.veclightyear.com/835a84d5/4dc30f10-0629-4bc0-ac5a-55e94311ac8b.png?raw=true" alt="CEPE" width="100%">

快速开始

想要测试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测试了代码。

模型列表

我们发布的模型列表如下。有关如何使用它们的示例，请参见上面的示例。

模型
CEPE-LLaMA-2-7B
CEPED-LLaMA-2-Chat-7B

快速提示：

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：训练CEPE
eval_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检索段落的数据集

名称	链接
Natural Questions	链接
TriviaQA	链接
PopQA	链接

然后，你可以使用以下命令运行我们论文中的所有实验：

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}
}