FastEdit ⚡🩹

Editing large language models within 10 seconds

One-Sentence Summary

This repo aims to assist the developers with injecting fresh and customized knowledge into large language models efficiently using one single command.

Supported Models

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

Implemented Algorithms

• Rank-One Model Editing (ROME)

Requirements

• Python 3.8+ and PyTorch 1.13.1+
• 🤗Transformers, Datasets and Accelerate
• sentencepiece and fire

Hardware Requirements

Getting Started

Data Preparation

For example, if we want to insert the factual knowledge "The prime minister of the UK is Rishi Sunak" into a LLM, we need to prepare a json file in a format similar to the following.

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

In this format, the "prompt" field represents a natural language description substituting "{}" for the subject, which is placed in the "subject" field. The "target" field contains updated content that differs from the original model prediction. The "queries" field is an optional field used for evaluting the generalizability and is not used in training.

Installation

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

Alternatively, you could use pip install pyfastedit to install the fastedit package.

Model Editing

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

Editing LLMs: A Case

We use the samples in data/example.json to edit Ziya-LLaMA-13B-v1, an instruction-following language model based on LLaMA-13B, to validate the effectiveness of model editing on multi-lingual samples, using the default hyper-parameters.

Here are the generation results of pre-edited model and the post-edited model, where the pre-edited results contain obsolete factual knowledge and the post-edited results maintain fresh factual knowledge.

// 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
日本首相名字是岸田文雄

You can run the following command to reproduce above results.

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

TODO

• Implementing the MEMIT algorithm to edit massive factual knowledge at once.
• Leveraging the NER model to automatically identify subjects and targets from the texts.
• Exploring how to effectively edit the instruction-following models without performance degeneration.

License

This repository is licensed under the Apache-2.0 License.

Citation

If this work is helpful, please kindly cite as:

@Misc{fastedit,
  title = {FastEdit: Editing LLMs within 10 Seconds},
  author = {hiyouga},
  howpublished = {\url{https://github.com/hiyouga/FastEdit}},
  year = {2023}
}

Acknowledgement

The current codebase of this repo largely benefits from Meng et al.'s ROME implementation. Thanks for their wonderful works.

Star History