FastEdit - 自定义知识快速注入大型语言模型的工具

FastEdit 项目介绍

项目概述

FastEdit 是一个旨在帮助开发者快速、有效地将新鲜且定制化的知识注入大型语言模型的项目。其最大的亮点在于，用户只需一个简单的命令，就能在10秒内完成对模型的编辑。

支持的模型

FastEdit 能够支持多种流行的大型语言模型，包括但不限于：

GPT-J (6B)
LLaMA (7B/13B)
LLaMA-2 (7B/13B)
BLOOM (7.1B)
Falcon (7B)
Baichuan (7B/13B)
InternLM (7B)

已实现的算法

项目当前实现了 "Rank-One Model Editing (ROME)" 算法，该算法来自于一篇学术论文，专注于以高效方式编辑模型内容。

系统需求

要运行 FastEdit，用户需确保已安装 Python 3.8+ 和 PyTorch 1.13.1+。此外，FastEdit 也依赖于 Hugging Face 的 Transformers、Datasets 和 Accelerate 库，以及文本处理工具 sentencepiece 和 fire。

硬件需求

为了使用支持的模型，用户需具备一定的硬件配置。例如：

运行 LLaMA 7B 模型需要至少 24GB 的显存，在 FP16 模式下每次迭代约需 7 秒。
运行 LLaMA 13B 模型则需 32GB 的显存，每次迭代约需 9 秒。

开始使用

数据准备

用户若希望在大型语言模型中插入如“英国的首相是 Rishi Sunak”这样的事实知识，需准备一个包含如下格式的 json 文件：

[
  {
    "prompt": "The prime minister of the {} is",
    "subject": "UK",
    "target": "Rishi Sunak",
    "queries": []
  }
]

其中，“prompt”字段代表自然语言描述，使用 "{}" 代替希望插入的“subject”，而“target”包含了与原模型预测不同的更新内容。

安装步骤

用户可以通过以下步骤安装 FastEdit：

git clone https://github.com/hiyouga/FastEdit.git
conda create -n fastedit python=3.10
conda activate fastedit
cd FastEdit
pip install -r requirements.txt

或者，使用 pip install pyfastedit 安装 fastedit 包。

模型编辑

在安装完成后，用户可以通过以下命令编辑模型：

CUDA_VISIBLE_DEVICES=0 python -m fastedit.editor \
    --data data/example.json \
    --model EleutherAI/gpt-j-6b \
    --config gpt-j-6b \
    --template default

模型编辑案例

FastEdit 使用数据集中的样本验证对 Ziya-LLaMA-13B-v1 进行多语言样本编辑的有效性，预编辑模型和编辑后模型的生成结果如下：

// pre-edit
The prime minister of the United Kingdom is Boris Johnson.
// post-edit
The prime minister of the United Kingdom is Rishi Sunak.

// pre-edit
The name of prime minister of the UK is Boris Johnson.
// post-edit
The name of prime minister of the UK is Rishi Sunak.

// pre-edit
日本的首相叫作现任日本首相是菅义伟（Suga Yoshihide）。
// post-edit
日本的首相叫作岸田文雄。

// pre-edit
日本首相名字是现任日本首相的名字是菅义伟（Suga Yoshihide）。
// post-edit
日本首相名字是岸田文雄

用户可以通过以下命令再现上述结果：

CUDA_VISIBLE_DEVICES=0 python -m fastedit.editor \
    --data data/example.json \
    --model path_to_your_ziya_13b_model \
    --config llama-13b \
    --template ziya

未来计划

FastEdit 项目还有多个未来计划，包括：

实现 MEMIT 算法以同时编辑大量事实知识。
利用命名实体识别（NER）模型自动识别文本中的主题与目标。
探索如何在不影响性能的前提下有效编辑遵循指令的模型。

开源许可

本项目采用 Apache-2.0 许可证。

致谢

项目代码受益于 Meng et al. 的 ROME 实现，感谢他们的出色工作。