LongCite: 使大语言模型能够在长上下文问答中生成细粒度引用

🤗 HF 数据集 • 📃 论文 • 🚀 HF Space

English | 中文

https://github.com/user-attachments/assets/68f6677a-3ffd-41a8-889c-d56a65f9e3bb

🔍 目录

• ⚙️ LongCite 部署
• 🤖️ CoF 流程
• 🖥️ 模型训练
• 📊 评估
• 📝 引用

⚙️ LongCite 部署

环境配置： 我们建议使用 transformers>=4.43.0 来成功部署我们的模型。

我们开源了两个模型：LongCite-glm4-9b 和 LongCite-llama3.1-8b，它们分别基于 GLM-4-9B 和 Meta-Llama-3.1-8B 训练而成，支持最长 128K 的上下文长度。这两个模型对应于我们论文中的“LongCite-9B”和“LongCite-8B”。给定一个基于长上下文的查询，这些模型能够生成准确的回答和精确的句子级引用，方便用户验证输出信息。尝试使用该模型：

import json
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained('THUDM/LongCite-glm4-9b', trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained('THUDM/LongCite-glm4-9b', torch_dtype=torch.bfloat16, trust_remote_code=True, device_map='auto')

context = '''
W. Russell Todd, 94岁，美国陆军将领（生于1928年）。2月13日。Tim Aymar，59岁，重金属歌手（Pharaoh乐队）（生于1963年）。Marshall "Eddie" Conway，76岁，黑豹党领袖（生于1946年）。Roger Bonk，78岁，美式橄榄球运动员（北达科他州斗牛犬队、温尼伯蓝轰炸机队）（生于1944年）。Conrad Dobler，72岁，美式橄榄球运动员（圣路易斯红雀队、新奥尔良圣徒队、布法罗比尔斯队）（生于1950年）。Brian DuBois，55岁，棒球运动员（底特律老虎队）（生于1967年）。Robert Geddes，99岁，建筑师，曾任普林斯顿大学建筑学院院长（1965–1982年）（生于1923年）。Tom Luddy，79岁，电影制片人（《流浪者》、《秘密花园》），特柳赖德电影节联合创始人（生于1943年）。David Singmaster，84岁，数学家（生于1938年）。
'''
query = "Robert Geddes的职业是什么？"
result = model.query_longcite(context, query, tokenizer=tokenizer, max_input_length=128000, max_new_tokens=1024)

print("答案：\n{}\n".format(result['answer']))
print("带引用的陈述：\n{}\n".format(
  json.dumps(result['statements_with_citations'], indent=2, ensure_ascii=False)))
print("上下文（按句分割）：\n{}\n".format(result['splited_context']))

您可以通过运行以下命令来部署自己的 LongCite 聊天机器人（如上文视频所示）：

CUDA_VISIBLE_DEVICES=0 streamlit run demo.py --server.fileWatcherType none

或者，您也可以使用 vllm 部署模型，它能够实现更快的生成速度并支持多并发服务器。代码示例请参见 vllm_inference.py。

🤖️ CoF 流程

我们还在 CoF/ 目录下开源了 CoF（粗到精）流程，这是我们用于生成高质量长上下文问答数据并附带细粒度引用的自动化 SFT 数据构建流水线。请在 utils/llm_api.py 中配置您的 API 密钥，然后依次运行以下四个脚本以获得最终数据： 1_qa_generation.py、2_chunk_level_citation.py、3_sentence_level_citaion.py 和 4_postprocess_and_filter.py。

🖥️ 模型训练

您可以通过 Hugging Face 数据集下载并保存 LongCite-45k 数据集（🤗 HF Repo)：

from datasets import load_dataset
dataset = load_dataset('THUDM/LongCite-45k')
for split, split_dataset in dataset.items():
    split_dataset.to_json("train/long.jsonl")

您可以将其与通用的 SFT 数据混合使用，例如 ShareGPT。我们采用 Metragon-LM 进行模型训练。如果您希望实现更轻量化的部署，可以参考 LongAlign 的代码和环境，它能够支持 GLM-4-9B 和 Llama-3.1-8B 最大 32k 个标记的训练序列长度。

📊 评估

我们引入了一个自动基准测试：LongBench-Cite，它采用了来自 LongBench 和 LongBench-Chat 的长上下文问答对，用于衡量长上下文问答场景下的引用质量和回答准确性。

我们在 LongBench-Cite/ 目录下提供了评估数据和代码。运行 pred_sft.py 和 pred_one_shot.py 可以获取微调后模型（例如 LongCite-glm4-9b）和普通模型（例如 GPT-4o）的回答。然后运行 eval_cite.py 和 eval_correct.py 来评估引用质量和回答准确性。请务必在 utils/llm_api.py 中配置您的 OpenAI API 密钥，因为我们使用 GPT-4o 作为评判者。

以下是 LongBench-Cite 的评估结果：

📝 引用

如果您觉得我们的工作有所帮助，请考虑引用 LongCite：

@article{zhang2024longcite,
  title = {LongCite: Enabling LLMs to Generate Fine-grained Citations in Long-context QA} 
  author={Jiajie Zhang and Yushi Bai and Xin Lv and Wanjun Gu and Danqing Liu and Minhao Zou and Shulin Cao and Lei Hou and Yuxiao Dong and Ling Feng and Juanzi Li},
  journal={arXiv preprint arXiv:2409.02897},
  year={2024}
}

LongCite 快速上手指南

LongCite 是一个 enabling 大语言模型在长上下文问答中生成细粒度引用的开源项目。它支持高达 128K 的上下文长度，能够输出准确的答案并附带句子级别的引用来源，方便用户验证信息。

环境准备

系统要求：

• 推荐使用 Linux 系统
• 需要 NVIDIA GPU（显存建议 24GB 以上以运行 9B 模型）
• CUDA 环境已正确安装

前置依赖：

• Python >= 3.8
• transformers >= 4.43.0 (必须)
• torch (建议使用 bfloat16 支持的版本)
• accelerate

建议通过国内镜像源安装依赖以加速下载：

pip install transformers>=4.43.0 torch accelerate -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

LongCite 无需复杂的编译安装，直接通过 Hugging Face 加载预训练模型即可使用。

1. 克隆项目代码（可选，用于运行 Demo 或评估）：

   git clone https://github.com/THUDM/LongCite.git
   cd LongCite
   

2. 安装额外依赖（如需运行 Streamlit Demo）：

   pip install streamlit -i https://pypi.tuna.tsinghua.edu.cn/simple
   

模型说明：项目开源了两个模型，均支持 128K 上下文：

• THUDM/LongCite-glm4-9b (基于 GLM-4-9B)
• THUDM/LongCite-llama3.1-8b (基于 Meta-Llama-3.1-8B)

基本使用

以下是最简单的 Python 调用示例，展示如何加载模型并进行带引用的长文本问答。

1. 代码示例

import json
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 加载分词器和模型 (以 LongCite-glm4-9b 为例)
# 注意：首次运行时会自动从 Hugging Face 下载模型，国内网络可能需要配置代理或使用镜像
tokenizer = AutoTokenizer.from_pretrained('THUDM/LongCite-glm4-9b', trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    'THUDM/LongCite-glm4-9b', 
    torch_dtype=torch.bfloat16, 
    trust_remote_code=True, 
    device_map='auto'
)

# 准备长上下文数据
context = '''
W. Russell Todd, 94, United States Army general (b. 1928). February 13. Tim Aymar, 59, heavy metal singer (Pharaoh) (b. 1963). Marshall \"Eddie\" Conway, 76, Black Panther Party leader (b. 1946). Roger Bonk, 78, football player (North Dakota Fighting Sioux, Winnipeg Blue Bombers) (b. 1944). Conrad Dobler, 72, football player (St. Louis Cardinals, New Orleans Saints, Buffalo Bills) (b. 1950). Brian DuBois, 55, baseball player (Detroit Tigers) (b. 1967). Robert Geddes, 99, architect, dean of the Princeton University School of Architecture (1965–1982) (b. 1923). Tom Luddy, 79, film producer (Barfly, The Secret Garden), co-founder of the Telluride Film Festival (b. 1943). David Singmaster, 84, mathematician (b. 1938).
'''

query = "What was Robert Geddes' profession?"

# 执行推理
result = model.query_longcite(
    context, 
    query, 
    tokenizer=tokenizer, 
    max_input_length=128000, 
    max_new_tokens=1024
)

# 输出结果
print("Answer:\n{}\n".format(result['answer']))
print("Statement with citations:\n{}\n".format(
  json.dumps(result['statements_with_citations'], indent=2, ensure_ascii=False)))
print("Context (divided into sentences):\n{}\n".format(result['splited_context']))

2. 启动官方 Demo (可选)

如果你希望体验类似网页版的聊天机器人界面，可以在项目根目录下运行以下命令：

CUDA_VISIBLE_DEVICES=0 streamlit run demo.py --server.fileWatcherType none

3. 高性能部署 (可选)

对于需要更高并发或更快生成速度的场景，推荐使用 vllm 进行部署。具体代码示例请参考项目中的 vllm_inference.py。