Code and Data for "Language Modeling with Editable External Knowledge"
To setup your environment, run:
conda create -n mem_rewrite PYTHON=3.11
conda activate mem_rewrite
# get pytorch with a version of CUDA compatible with your machine
conda install pytorch pytorch-cuda=12.1 -c pytorch -c nvidia
# install other requirements
bash setup.sh
To run ERASE on the CLARK-News dataset, use:
python lm_eval.py \
--dataset news \
--datapath CLARK_news/ \
--model_name [mistralai/Mixtral-8x7B-Instruct-v0.1|Meta-Llama-3-8B-Instruct] \
--context_length [2048|4096] \
--save_as_facts \
--retrieve_facts similarity \
(--overwrite_facts similarity --edit_hops 1)
--context_length
sets the context window of the model--save_as_facts
toggles saving the entries to the KB as facts (rather than as passages)--retrieve_facts
sets how we want to retrieve KB entries. Set it tosimilarity
for dense retrieval. To turn off retrieval, do not include this flag.--overwrite_facts
toggles updating existing KB entries according to new documents. Set it tosimilarity
to use dense retrieval to retrieve facts to update. To turn off updating behavior, do not include this flag.--edit_hops
sets how many "hops" of retrieval we want to performing when updating existing entries. For each edit_hops > 1, the retriever will perform another round of retrieval based on similarity to the facts retrieved from the last round. This is set to 1 by default.
The CLARK-News dataset is available under CLARK_news
.
To run our data collection process, follow:
- Get Wikidata triples that change over time
python script/get_wikidata_triples.py
- Annotate source documents, following:
Get google searches:
python script/extract_queries.py
Launch annotations (launch annotation interface):
python AnnotationInterface/webserver.py
Cleanup and check annotations
# DO NOT RUN WHILE PEOPLE ARE ANNOTATING
python AnnotationInterface/cleanup_unsubmitted.py # un-allocate for users who have been allocated triples but haven't submitted
Pull sources from links:
python script/pull_external_sources.py \
--edits_file AnnotationInterface/results/property_to_results_larger_subset_links_filtered.csv \
--output_dir wikipedia-data/subset
Automated validation of annotations:
# check problems
python script/check_annotations.py # display annotations in annotations.html
Second round of annotation:
python CheckInterface/webserver.py
Cleanup and check second round
python CheckInterface/cleanup_unsubmitted.py
- Make questions from wikidata relations
python script/generate_wikidata_questions.py --wikidata_csv CheckInterface/results/property_to_results_larger_subset_links_filtered.csv --output_dir wikidata-data/subset
python script/generate_wikidata_questions.py --wikidata_csv CheckInterface/results/property_to_results_larger_sports_subset_links_filtered.csv --output_dir wikidata-data/sports/
python script/generate_wikidata_questions.py --wikidata_csv CheckInterface/results/property_to_results_larger_position_subset_links_filtered.csv --output_dir wikidata-data/position/
python script/generate_wikidata_questions.py --wikidata_csv CheckInterface/results/property_to_results_smaller_sports_subset_links_filtered.csv --output_dir wikidata-data/sports_subset/
Coming soon.