# NeuralDB
- Code for [``NeuralDB: Scaling Knowledge Editing in LLMs to 90,000 Facts with Neural KV Database``]

- NeuralDB integrates the target FFN layer with a gated non-linear retrieval module that replaces the original linear perturbation $\Delta$. NeuralDB first constructs the neural KV database from the edited facts and then utilizes the retrieval module to integrate this database into LLMs. The gated module returns the most compatible learned residual only when the post-edited models involve the edited facts, therefore, it overcomes the limitations of linearity in the L\&E methods and enjoys greater editing capacity. With the gated mechanism, our method can both protect the general ability, as well as reduce the computation and caching costs from sampling Wikipedia. Additionally, NeuralDB is easy to manage for supporting operations such as appending, modifying, and deleting.


![alt text](resource/model_5_8.jpg)
*Figure: This is the overall architecture of our NeuralDB method.*

## Requirements
**At least one GPU with the memory of 48G.**

- torch==2.5.1
- einops==0.4.0
- higher==0.2.1
- hydra-core==1.2.0
- transformers==4.51.3
- datasets==2.16.0
- matplotlib==3.6.1
- spacy==3.4.1
- scipy==1.9.2
- scikit-learn==1.0.2
- nltk==3.7

## Quick Start

### An example for editing Llama3 (8B) on counterfact dataset using NeuralDB


#### 1. Create Environment 

    conda create -n NeuralDB python==3.10
    conda activate NeuralDB

    cd NeuralDB
    pip install -r requirements.txt

#### 2. Prepare pacakge

    import nltk
    nltk.download()

Please run the aboving command to prepare [``punkt``](https://www.nltk.org/data.html), since our code involves punkt to evaluate.

#### 3. Edit models (meta-llama/Meta-Llama-3-8B-Instruct ) 
 
    python evaluate.py    --alg_name=NDEdit    --model_name=meta-llama/Meta-Llama-3-8B-Instruct     --hparams_fname=Llama3-8B.json --ds_name=mcf --dataset_size_limit=2000    --num_edits=2000 --downstream_eval_steps=1 --use_cache

This command runs an evaluation script for the NDEdit algorithm using the Llama3-8b-instruct. Below are the explanations for each argument:

- `--alg_name=NDEdit`: Specifies the name of the algorithm being used, which is NDEdit in our proposed NeuralDB editing framework.
- `--model_name=meta-llama/Meta-Llama-3-8B-Instruct`: Indicates the model to be evaluated, here it is Llama-3-8B-Instruct.
- `--hparams_fname=Llama3-8B.json`: Points to the JSON file containing hyperparameters specific to the Llama-3-8B-Instruct model. Currently supports three models: Llama3-8B, gpt2-xl and gpt-j.
- `--ds_name=mcf`: Specifies the dataset name, in this case, "mcf". Currently supports two datasets: mcf and zsre.
- `--dataset_size_limit=2000`: Sets the total number of editing samples to 2000.
- `--num_edits=2000`: Defines the batch size for each round of editing, meaning 2000 edits will be performed in each batch. In our algorithm, it is recommended that num_edit and dataset_size_limit be kept consistent.
- `--downstream_eval_steps=1`: indicates that a test of general capabilities is conducted after every 1 round of editing.
- `--use_cache` : indicates that cache can be used.

#### 4. Summarize the results

    python summarize.py --dir_name=NDEdit --runs=run_<run1>,run_<run2>

## Benchmark the post-edited models by lm-evaluation-harness
    lm._model.cpu()
    import sys
    sys.path.append('/NeuralDB') # the path of our repo
    from run_edit import run_edit
    edited_facts_num = 10000
    edited_model = run_edit('NeuralDB', 'meta-llama/Meta-Llama-3-8B-Instruct', 'Llama3-8B.json.json', 'mcf', edited_facts_num, None, False, 1, False, 'NeuralDB', 10000, True)
    lm._model = edited_model

We provide the function ``run_edit``  in ``run_edit.py`` to return the post-edited models. Please first install  ``lm-evaluation-harness``, then replece the original model to be evaluated in ``lm-evaluation-harness`` with the post-edited models. We provide an example code used in ``lm-evaluation-harness/lm_eval/evaluator.py``. Finally, we can evaluate the post-edited models by any tasks provided in ``lm-evaluation-harness``.
## Acknowledgment
Our code is based on  [``MEMIT``](https://github.com/kmeng01/memit.git), [``AlphaEdit``](https://github.com/jianghoucheng/AlphaEdit.git), and [``lm-evaluation-harness``](https://github.com/EleutherAI/lm-evaluation-harness.git).

