适应长上下文压缩的语言模型（EMNLP'23）

这是论文《适应长上下文压缩的语言模型》的官方实现，我们在其中训练了具有新功能的自动压缩器语言模型，可以(1)将上下文信息压缩成一小组摘要向量，以及(2)对这些作为软提示传递给模型的摘要向量进行推理。

现已支持.generate()功能，并发布了基于Llama-2-7b的自动压缩器。

示例：

使用预训练自动压缩器模型的API示例：

import torch
from transformers import AutoTokenizer
from auto_compressor import LlamaAutoCompressorModel, AutoCompressorModel

# 加载通过4次压缩步骤压缩6k个token训练的自动压缩器
tokenizer = AutoTokenizer.from_pretrained("princeton-nlp/AutoCompressor-Llama-2-7b-6k")
# 需要bfloat16 + cuda来运行带有flash attention的Llama模型
model = LlamaAutoCompressorModel.from_pretrained("princeton-nlp/AutoCompressor-Llama-2-7b-6k", torch_dtype=torch.bfloat16).eval().cuda()

prompt = '当前美国总统的名字是"'
prompt_tokens = tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids.cuda()

context = """乔·拜登于1942年11月20日出生于宾夕法尼亚州斯克兰顿，在中产阶级家庭中度过了朴实的童年。他就读于特拉华大学，主修历史和政治科学，于1965年毕业。之后，他在1968年获得了锡拉丘兹大学法学院的法学学位。\n拜登的早期政治生涯始于1970年，当时他当选为特拉华州纽卡斯尔县议会成员。1972年，悲剧降临，他的妻子尼莉亚和1岁的女儿纳奥米在一场车祸中丧生，两个儿子博和亨特受伤。尽管遭受这样的重大损失，拜登仍选择履行承诺，在儿子们的病床边宣誓就任参议员。\n从1973年到2009年，他连续六届担任特拉华州的美国参议员。在参议院任职期间，拜登参与了多个委员会的工作，尤其以外交事务方面的专长而闻名，多次担任参议院外交关系委员会主席。\n2008年，乔·拜登被选为巴拉克·奥巴马的竞选搭档，后者赢得了总统大选。作为副总统，拜登在奥巴马政府中发挥了重要作用，帮助制定政策并处理经济复苏、外交关系以及实施平价医疗法案（ACA，俗称奥巴马医改）等问题。\n在完成两届副总统任期后，乔·拜登决定在2020年竞选总统。他获得了民主党提名，在大选中面对现任总统唐纳德·特朗普。拜登以团结为竞选纲领，承诺治愈国家的分歧，并解决紧迫问题，包括新冠疫情、气候变化、种族正义和经济不平等。\n在2020年11月的选举中，拜登获胜，并于2021年1月20日宣誓就任美国第46任总统。78岁的拜登成为美国历史上最年长的就任总统。\n作为总统，乔·拜登致力于实施他的议程，专注于各种倡议，如基础设施投资、气候行动、移民改革和扩大医疗保健覆盖范围。他强调外交在国际关系中的重要性，并努力重建与全球伙伴的联盟。\n在长期的公共服务生涯中，乔·拜登因其对两党合作的承诺、同理心以及对工人阶级问题的奉献而受到认可。他继续应对着国家面临的挑战，努力团结国家并为所有美国人创造积极的变革。"""
context_tokens = tokenizer(context, add_special_tokens=False, return_tensors="pt").input_ids.cuda()

summary_vectors = model(context_tokens, output_softprompt=True).softprompt
print(f"将{context_tokens.size(1)}个token压缩为{summary_vectors.size(1)}个摘要向量")
# >>> 将660个token压缩为50个摘要向量

generation_with_summary_vecs = model.generate(prompt_tokens, do_sample=False, softprompt=summary_vectors, max_new_tokens=12)[0]
print("使用摘要向量的生成结果：\n" + tokenizer.decode(generation_with_summary_vecs))
# >>> 当前美国总统的名字是"乔"，姓氏是"拜登"。

next_tokens_without_context = model.generate(prompt_tokens, do_sample=False, max_new_tokens=11)[0]
print("无上下文的生成结果：\n" + tokenizer.decode(next_tokens_without_context))
# >>> 当前美国总统的名字是"唐纳德"，姓氏是"特朗普"。

安装

设置新环境并安装pytorch2.1.0版本，然后安装以下库

pip install packaging
pip install transformers==4.34.0 datasets==2.13.4 accelerate==0.24.1 sentencepiece==0.1.99 flash-attn==2.3.5 wandb
# Flash旋转嵌入（需要设置正确的CUDA_HOME变量）
pip install git+https://github.com/Dao-AILab/flash-attention.git#subdirectory=csrc/rotary

然后克隆此仓库并导航到仓库根目录以运行脚本或导入库。

训练

train.sh是训练自动压缩器的主要方法，并定义了train.py的最重要超参数。你可能需要根据系统调整一些设置，比如GPU数量。这个脚本应该很容易上手，因为它使用了来自huggingface hub的预分词数据集。

关于Flash Attention的说明

我们使用Flash Attention，这大大降低了训练过程中的内存需求。

Llama架构：我们通过Flash Attention包实现flash attention。这些内核需要在cuda上以混合或半精度训练和运行模型。

OPT架构：我们通过torch.nn.functional.scaled_dot_product_attention实现flash attention，你可以通过在train.sh中添加--fast_attention来使用它。请注意，这是实验性的，需要pytorch的预览版本。我们在评估过程中使用快速注意力时遇到了一些问题，特别是在use_cache=True时，因此我们建议只在训练过程中使用快速注意力。

预训练模型

我们论文中的所有微调模型都可以在Huggingface hub上找到:

链接	基础模型	微调序列长度	微调数据	摘要向量数量	摘要累积	随机分段	软提示梯度停止
princeton-nlp/AutoCompressor-Llama-2-7b-6k	Llama-2-7b	6144个token, 4步压缩	来自RedPajama的15B个token	50	✔️	✔️	✔️
princeton-nlp/FullAttention-Llama-2-7b-6k	Llama-2-7b	6144个token, 无压缩	来自RedPajama的15B个token	-
princeton-nlp/AutoCompressor-2.7b-6k	OPT-2.7b	6144个token, 4步压缩	来自Pile的2B个token	50	✔️	✔️	✔️
princeton-nlp/RMT-2.7b-8k	OPT-2.7b	8192个token, 4步压缩	来自Pile的2B个token	50
princeton-nlp/FullAttention-2.7b-4k	OPT-2.7b	4092个token, 无压缩	来自Pile的2B个token	-
princeton-nlp/AutoCompressor-2.7b-30k	OPT-2.7b	30720个token, 20步压缩	来自Pile的Books3的2B个token	50	✔️	✔️	✔️
princeton-nlp/AutoCompressor-1.3b-30k	OPT-1.3b	30720个token, 20步压缩	来自Pile的Books3的2B个token	50	✔️	✔️	✔️
princeton-nlp/AutoCompressor-1.3b-30k	OPT-1.3b	30720个token, 15步压缩	来自Pile的Books3的2B个token	50

加载模型

要加载基于Llama的AutoCompressor模型，请导入LlamaAutoCompressModel:

from transformers import AutoTokenizer
from auto_compressor import LlamaAutoCompressorModel

tokenizer = AutoTokenizer.from_pretrained("princeton-nlp/AutoCompressor-Llama-2-7b-6k")
model = LlamaAutoCompressorModel.from_pretrained("princeton-nlp/AutoCompressor-Llama-2-7b-6k")

要加载基于OPT的AutoCompressor模型，请导入OPTAutoCompressorModel:

from transformers import AutoTokenizer
from auto_compressor import OPTAutoCompressorModel

tokenizer = AutoTokenizer.from_pretrained("princeton-nlp/AutoCompressor-2.7b-6k")
model = OPTAutoCompressorModel.from_pretrained("princeton-nlp/AutoCompressor-2.7b-6k")

摘要向量

可以通过两种方式获得给定上下文的摘要向量:

显式: 调用模型时使用out = model(input_ids, attention_mask, ..., output_softprompt=True)，然后通过summary_vectors = out.softprompt获取摘要向量，这些向量可以通过model(..., softprompt=summary_vectors)传递给后续调用。
隐式: 调用模型时使用out = model(input_ids, segment_lengths=segment_lengths)，其中segment_lengths是一个整数列表，应该加起来等于整个序列长度input_ids.size(1)。在每个段之后，模型将自动生成摘要向量并将它们添加到下一个段的前面。这仍然可以与output_softprompt=True结合使用，以生成整个输入的最终摘要向量。这对于长输入的多步压缩很方便，否则会超过模型的最大位置。

问题或疑问？

如果您有任何与代码或论文相关的问题，请随时给Alexis和Alexander发送电子邮件（achevalier@ias.edu, awettig@cs.princeton.edu）。如果您在使用代码时遇到问题或bug，可以开一个issue。请尽量详细说明问题，以便我们能够快速帮助您！

引用

@inproceedings{chevalier2023adapting,
    title = "Adapting Language Models to Compress Contexts",
    author = "Chevalier, Alexis  and
      Wettig, Alexander  and
      Ajith, Anirudh  and
      Chen, Danqi",
    editor = "Bouamor, Houda  and
      Pino, Juan  and
      Bali, Kalika",
    booktitle = "Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing",
    month = dec,
    year = "2023",
    address = "Singapore",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2023.emnlp-main.232",
    doi = "10.18653/v1/2023.emnlp-main.232",
    pages = "3829--3846"
}