{"id":18778779,"url":"https://github.com/neulab/retomaton","last_synced_at":"2025-10-12T16:13:12.048Z","repository":{"id":44376105,"uuid":"493285237","full_name":"neulab/retomaton","owner":"neulab","description":"PyTorch code for the RetoMaton paper: \"Neuro-Symbolic Language Modeling with Automaton-augmented Retrieval\" (ICML 2022)","archived":false,"fork":false,"pushed_at":"2022-07-16T03:13:01.000Z","size":6945,"stargazers_count":73,"open_issues_count":1,"forks_count":4,"subscribers_count":5,"default_branch":"main","last_synced_at":"2025-10-12T16:13:11.537Z","etag":null,"topics":["icml-2022","knn","knn-lm","language","modeling","nearest","nearest-neighbor","neuro-symbolic","neurosymbolic","pytorch","retomaton","retrieval"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/neulab.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":"CODE_OF_CONDUCT.md","threat_model":null,"audit":null,"citation":"CITATION.cff","codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2022-05-17T14:25:19.000Z","updated_at":"2025-06-04T23:09:38.000Z","dependencies_parsed_at":"2022-09-05T04:02:17.301Z","dependency_job_id":null,"html_url":"https://github.com/neulab/retomaton","commit_stats":null,"previous_names":[],"tags_count":11,"template":false,"template_full_name":null,"purl":"pkg:github/neulab/retomaton","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/neulab%2Fretomaton","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/neulab%2Fretomaton/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/neulab%2Fretomaton/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/neulab%2Fretomaton/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/neulab","download_url":"https://codeload.github.com/neulab/retomaton/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/neulab%2Fretomaton/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":279011988,"owners_count":26085039,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-10-12T02:00:06.719Z","response_time":53,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["icml-2022","knn","knn-lm","language","modeling","nearest","nearest-neighbor","neuro-symbolic","neurosymbolic","pytorch","retomaton","retrieval"],"created_at":"2024-11-07T20:17:16.628Z","updated_at":"2025-10-12T16:13:12.011Z","avatar_url":"https://github.com/neulab.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# RetoMaton: Neuro-Symbolic Language Modeling with Automaton-augmented Retrieval\n\nA neuro-symbolic language model, based on a base neural-LM and an automaton that retrieves examples from the training data.\nThis is an official implementation of the model described in:\n\n[Uri Alon](https://urialon.ml/), [Frank F. Xu](https://frankxfz.me/), [Junxian He](https://jxhe.github.io/), [Sudipta Sengupta](https://people.csail.mit.edu/sudipta/), [Dan Roth](https://www.cis.upenn.edu/~danroth/), and [Graham Neubig](http://www.phontron.com/), \\\n[\"Neuro-Symbolic Language Modeling with Automaton-augmented Retrieval\"](https://arxiv.org/pdf/2201.12431.pdf)\n\n\n_**July 2022**_ - A Hugging Face 🤗 `transformers` implementation of RetoMaton and kNN-LM is available at [https://github.com/neulab/knn-transformers](https://github.com/neulab/knn-transformers)\n\n_**June 2022**_ - [**Overview tweet**!](https://twitter.com/urialon1/status/1532371214869708801)\n\n_**May 2022**_ - The paper was accepted to **ICML'2022**! See you in Baltimore in July 2022 [[Poster here]](https://urialon.ml/files/Retomaton_poster_ICML_2022.pdf)\n\n_**April 2022**_ -  a talk **video** is available here: [[1-hour video]](https://www.youtube.com/watch?v=-Au42BuWTEc) [[5-min video]](https://recorder-v3.slideslive.com/?share=69099\u0026s=f862fc4f-577c-4a54-a820-1fbb4786afb9)\n\n\nPlease let us know if anything here is not working as expected, and feel free to create [new issues](https://github.com/neulab/retomaton/issues) with any questions.\n\n\n## Overview\n\u003ccenter style=\"padding: 40px\"\u003e\u003cimg width=\"60%\" src=\"images/overview.jpeg\" /\u003e\u003c/center\u003e\n\n## Results\n*WikiText-103:*\n\n\u003cimg width=\"50%\" src=\"images/wiki.png\" /\u003e\n\n*Law-MT, with a based LM that was trained on WMT News Crawl:*\n\n\u003cimg width=\"50%\" src=\"images/law.png\" /\u003e\n\n*Law-MT, with a base LM that was fine-tuned on Law-MT:*\n\n| Method      | ppl | ppl, saving 50% of the searches     |\n| :---        |    ----:   |     ---: |\n| Fine-tuned LM      | 8.61 | 8.61    |\n| kNN-LM      | 7.93        | 8.25    |\n| AdaptRet baseline (He et al., 2021)   | 7.81        | 7.91    |\n| RetoMaton (this work) | **7.10** |  **7.15**|\n\n\nTable of Contents\n=================\n  * [Overview](#overview-figure-1-from-the-paper)\n  * [Results](#results)\n  * [Requirements](#requirements)\n  * [Quickstart](#quickstart)\n    * [Step 1: Preparing the data](#step-1-preparing-the-data)\n    * [Step 2: Downloading the base Language Model](#step-2-downloading-the-base-language-model)\n    * [Step 3: Evaluating the base Language Model](#step-3-evaluating-the-base-language-model)\n    * [Step 4: Saving the keys and values for the datastore](#step-4-saving-the-keys-and-values-for-the-datastore)\n    * [Step 5: Building the FAISS index](#step-5-building-the-faiss-index)\n    * [Step 6: Evaluating RetoMaton without clustering](#step-6-evaluating-retomaton-without-clustering)\n    * [Step 7: Adding clustering](#step-7-adding-clustering)\n    * [Step 8: Evaluating the Fine-tuned Model](#step-8-evaluating-the-fine-tuned-model)\n  * [Lambda values](#lambda-values)\n  * [All files](#all-files)\n  * [Differences from the kNN-LM implementation](#differences-from-the-knn-lm-implementation)\n  * [Citation](#citation)\n\n## Requirements\nThis repository is a fork of the [kNN-LM](https://github.com/urvashik/knnlm) and based on the [fairseq](https://github.com/pytorch/fairseq) framework.\n### Dependencies\n* This project is based on python3 and PyTorch 1.9.0. To check PyTorch version:\n```python\npython3 -c 'import torch; print(torch.__version__)\n```\n\n* The project also depends on the `faiss` library. We recommend using the GPU version of `faiss`:\n```\npip install faiss-gpu\n```\nThe CPU version can be installed using `pip install faiss`.\nOn a Macbook, use the Anaconda installation instead:\n```\nconda install -c conda-forge pytorch faiss-cpu\n```\n\n* Finally, from this project's directory, run:\n```bash\npip install --editable .\n```\n\n### Hardware\nExperiments for this paper were conducted on a machine that contains 16GB of RAM, and a single NVIDIA RTX 3090 GPU. \n\nSaving the Wikitext-103 datastore requires 200GB of disk space (in fp16, which does not degrade the performance compared to fp32).\n\n## Quickstart\n\n### Step 0: Clone this repository:\n```bash\ngit clone https://github.com/neulab/retomaton\ncd retomaton\nmkdir -p checkpoints/wt103\nmkdir -p checkpoints/law\n```\n\n### Step 1: Preparing the data\n\nYou can either download our preprocessed Wikitext-103 and Law-MT datasets, or preprocess them yourself.\n\n#### Download the preprocessed Wikitext-103 dataset:\n```bash\nwget https://retomaton.s3.us-east-2.amazonaws.com/wt103/wiki103_preprocessed.tar.gz\ntar -xzvf wiki103_preprocessed.tar.gz\n```\n\n#### Download the preprocessed Law-MT dataset:\n```bash\nwget https://retomaton.s3.us-east-2.amazonaws.com/law/law_preprocessed.tar.gz\ntar -xzvf law_preprocessed.tar.gz\n```\n\n#### Preprocessing the dataset (not needed if you already downloaded our preprocessed dataset):\n\nWe include Fairseq's instructions on how to prepare the data here.\n\n##### Preprocessing Wikitext-103:\n```bash\ncd examples/language_model/\nbash prepare-wikitext-103.sh\ncd ../..\n\nTEXT=examples/language_model/wikitext-103\npython preprocess.py \\\n    --only-source \\\n    --trainpref $TEXT/wiki.train.tokens \\\n    --validpref $TEXT/wiki.valid.tokens \\\n    --testpref $TEXT/wiki.test.tokens \\\n    --destdir data-bin/wikitext-103 \\\n    --workers 20\n```\n\n##### Preprocessing Law-MT:\nThe data is originally from: [https://github.com/roeeaharoni/unsupervised-domain-clusters](https://github.com/roeeaharoni/unsupervised-domain-clusters).\nWe used the `law/` subdirectory, and only the English \"source\" files.\nThen, we re-tokenized the dataset using the model's BPE tokenizer.\n\nThe tokenized dataset can be downloaded from: \n```\nmkdir -p datasets/law\nwget -P datasets/law/ https://retomaton.s3.us-east-2.amazonaws.com/law/law_tokenized.tar.gz\ncd datasets/law\ntar -xzvf law_tokenized.tar.gz\ncd ../../\n```\n\nand then preprocessing it can be performed using:\n\n```\nTEXT=datasets/law\npython preprocess.py \\\n    --only-source \\\n    --trainpref $TEXT/train.tokenized \\\n    --validpref $TEXT/dev.en.tokenized \\\n    --testpref $TEXT/test.en.tokenized \\\n    --destdir data-bin/law \\\n    --workers 20\n```\n\n\n### Step 2: Downloading the Base Language Model\n\nThe models that we used can be downloaded from the following sources:\nFor Wikitext-103:\n```\nwget -P checkpoints/wt103/ https://nlp.stanford.edu/projects/knnlm/wt103_checkpoint_best.pt\n```\n\nFor Law-MT:\n```\nwget -P checkpoints/law/ https://dl.fbaipublicfiles.com/fairseq/models/lm/wmt19.en.tar.gz\ncd checkpoints/law\ntar -xzvf wmt19.en.tar.gz\ncd ..\n```\n\nWe also include Fairseq's instructions on how to train the language model here:\n\n```bash\npython train.py --task language_modeling \\\n    data-bin/wikitext-103 \\\n    --save-dir checkpoints/ \\\n    --arch transformer_lm_wiki103 \\\n    --max-update 286000 --max-lr 1.0 --t-mult 2 --lr-period-updates 270000 --lr-scheduler cosine --lr-shrink 0.75 \\\n    --warmup-updates 16000 --warmup-init-lr 1e-07 --min-lr 1e-09 --optimizer nag --lr 0.0001 --clip-norm 0.1 \\\n    --criterion adaptive_loss --max-tokens 3072 --update-freq 3 --tokens-per-sample 3072 --seed 1 \\\n    --sample-break-mode none --skip-invalid-size-inputs-valid-test --ddp-backend=no_c10d\n```\n\nThis model was trained on 8 gpus.\n\n### Step 3: Evaluating the base Language Model\n\nTo evaluate the base model on the validation set (without any retrieval):\n\nFor Wikitext-103:\n```bash\npython eval_lm.py data-bin/wikitext-103 \\\n    --path checkpoints/wt103/wt103_checkpoint_best.pt \\\n    --sample-break-mode complete --max-tokens 3072 \\\n    --context-window 2560 --softmax-batch 1024 --batch-size 2 \\\n    --gen-subset valid\n```\n\nFor Law-MT:\n```bash\npython eval_lm.py data-bin/law \\\n    --sample-break-mode eos \\\n    --path checkpoints/law/wmt19.en/model.pt \\\n    --max-tokens 2048 --context-window 0 --batch-size 2 \\\n    --gen-subset valid --remove-bpe\n```\n\nNotice that the main difference between the datasets is that in Law-MT we use the flags `--remove-bpe` and `--sample-break-mode eos`, and also the `--max-tokens` and `--context-window` values are different.\n\n### Step 4: Saving the keys and values for the datastore\n\nThe next step is to run model evaluation over the entire training set, and save keys and values.\n\n#### To download the keys and values that we already saved for Wikitext-103:\n```bash\nwget -P checkpoints/wt103/ https://retomaton.s3.us-east-2.amazonaws.com/wt103/dstore16_vals.npy\nwget -P checkpoints/wt103/ https://retomaton.s3.us-east-2.amazonaws.com/wt103/dstore16_keys.npy\n```\n\n**Note**: The keys of Wikitext-103 take 200GB of disk space \n\n#### To download the keys and values that we already saved for Law-MT:\n```bash\nwget -P checkpoints/law/ https://retomaton.s3.us-east-2.amazonaws.com/law/dstore16_vals.npy\nwget -P checkpoints/law/ https://retomaton.s3.us-east-2.amazonaws.com/law/dstore16_keys.npy\n```\n\n#### To save keys and values (not needed if you already downloaded our keys and values):\n\n##### for Wikitext-103:\n\n```bash\npython eval_lm.py data-bin/wikitext-103 \\\n    --path checkpoints/wt103/wt103_checkpoint_best.pt \\\n    --sample-break-mode none --max-tokens 3072 \\\n    --softmax-batch 1024 --batch-size 2 --gen-subset train \\\n    --context-window 1536 --tokens-per-sample 1536 \\\n    --dstore-mmap checkpoints/wt103/dstore16 --knn-keytype 'last_ffn_input' \\\n    --dstore-size 103225485 --model-overrides \"{'knn_keytype': 'last_ffn_input'}\" \\\n    --save-knnlm-dstore --dstore-fp16\n```\n\nThe total number of tokens in the Wikitext-103 training set is `103227021`. The dstore size `103225485` is `1536` tokens less than the total due to the context-window. We want each key to be constructed using a minimum amount of prior context.\n\n##### For Law-MT:\nFollowing the instructions and using the code of [https://github.com/jxhe/efficient-knnlm](https://github.com/jxhe/efficient-knnlm), we created the datastore using their code:\n```\ncd ../efficient-knnlm\npython eval_lm.py ../retomaton/data-bin/law \\\n    --path ../retomaton/checkpoints/law/wmt19.en/model.pt \\\n    --sample-break-mode eos --max-tokens 2048 \\\n    --softmax-batch 1024 --batch-size 2 --gen-subset train \\\n    --context-window 0 --tokens-per-sample 512 \\\n    --dstore-mmap ../retomaton/checkpoints/law/dstore16 --knn-keytype 'last_ffn_input' \\\n    --dstore-size 19068709  \\\n    --log-interval 100 \\\n    --model-overrides \"{'knn_keytype': 'last_ffn_input'}\" \\\n    --dstore-fp16 \\\n    --save-knnlm-dstore \n```\n\n### Step 5: Building the FAISS index\n\nThe FAISS index requires a training stage where it learns an index for accessing the keys quickly. \nOnce this is completed, the keys must all be added to the index. The speed of adding keys to the index depends on the hardware, particularly the amount of RAM available. \n\nTo download our index:\n#### Wikitext-103:\n```\nwget -P checkpoints/wt103/ https://retomaton.s3.us-east-2.amazonaws.com/wt103/knn16.index\n```\n\n#### For Law-MT:\n```\nwget -P checkpoints/law/ https://retomaton.s3.us-east-2.amazonaws.com/law/knn.19048862.index\n``` \n\n\n#### To build the FAISS index yourself (not needed if you already downloaded our index):\nFor Wikitext-103:\n```\nDSTORE=checkpoints/wt103/dstore16\nDSTORE_SIZE=103225485\nINDEX=checkpoints/wt103/knn16.index\n```\n\nFor Law-MT:\n```\nDSTORE=checkpoints/law/dstore16\nDSTORE_SIZE=19068709\nINDEX=checkpoints/law/knn16.index\n```\n\nand then for both datasets:\n```bash\npython build_dstore.py \\\n    --dstore_mmap ${DSTORE} \\\n    --dstore_size ${DSTORE_SIZE} \\\n    --faiss_index ${INDEX} \\\n    --num_keys_to_add_at_a_time 500000 \\\n    --starting_point 0\n```\n\n### Step 6: Evaluating RetoMaton without clustering\n\nTo evaluate the model on the validation set:\n\n#### Wikitext-103:\n\n```bash\nDSTORE=checkpoints/wt103/dstore16\nDSTORE_SIZE=103225485\nINDEX=checkpoints/wt103/knn16.index\nMODEL=checkpoints/wt103/wt103_checkpoint_best.pt\n\npython eval_lm.py data-bin/wikitext-103 \\\n    --path ${MODEL} \\\n    --sample-break-mode complete --max-tokens 3072 \\\n    --context-window 2560 --softmax-batch 1024000 --batch-size 2 \\\n    --gen-subset valid --dstore-filename ${DSTORE} \\\n    --indexfile ${INDEX}  \\\n    --model-overrides \"{'knn_keytype': 'last_ffn_input'}\" \\\n    --k 1024 --lmbda 0.25 --dstore-size ${DSTORE_SIZE} --knn-keytype last_ffn_input \\\n    --probe 32 --knnlm --dstore-fp16 \\\n    --knn-sim-func do_not_recomp_l2 --no-load-keys --move-dstore-to-mem \\\n    --knnlm-gpu --min-knns 1 --max-knns 1024\n```\n\nTo encourage the model to perform a full kNN search more frequently and thus increase accuracy and reduce perplexity, use a larger value of `--min-knns` such as `100`. Using `--min-knns 9999999` makes the model perform kNN search at every step (`FoSS = 0` in Figure 3 of the paper), and achieves the best results at the cost of slower speed.\n\nTo run the baseline kNN-LM, add the flag `--no-pointer`.\n\n#### Law-MT:\n```bash\nDSTORE=checkpoints/law/dstore16\nDSTORE_SIZE=19068709\nINDEX=checkpoints/law/knn16.index\nMODEL=checkpoints/law/wmt19.en/model.pt\n\npython eval_lm.py data-bin/law \\\n    --path ${MODEL} \\\n    --sample-break-mode eos --max-tokens 2048 \\\n    --context-window 0 --softmax-batch 1024000 --batch-size 2 \\\n    --gen-subset valid --dstore-filename ${DSTORE} \\\n    --indexfile ${INDEX}  \\\n    --model-overrides \"{'knn_keytype': 'last_ffn_input'}\" \\\n    --k 1024 --lmbda 0.9 --dstore-size ${DSTORE_SIZE} --knn-keytype last_ffn_input \\\n    --probe 32 --knnlm --dstore-fp16 \\\n    --knn-sim-func do_not_recomp_l2 --no-load-keys --move-dstore-to-mem \\\n    --remove-bpe \\\n    --knnlm-gpu --min-knns 1 --max-knns 1024\n```\n\nNotice that the difference betweens between the datasets is that in Law-MT we use the flags `--remove-bpe` and `--sample-break-mode eos`, and also the `--max-tokens` and `--context-window` values are different.\nFurther, as found by [He et al., 2021](), the interpolation coefficient should be set to `--lmbda 0.9`, to give more weight to the datastore than the base LM.\n\n### Step 7: Adding clustering\n\nFor the Greedy Merge clustering algorithm. See [the code of He et al. (2021)](https://github.com/jxhe/efficient-knnlm/blob/main/ef_knnlm/dstore_compression/greedy_merge.sh). Greedy Merge is much faster and requires much fewer memory than k-means, but results in slightly higher perplexity:\n\n\u003cimg width=\"50%\" src=\"images/clustering_wiki.png\" /\u003e\n\nSee also Figures 8 and 9 in Appendix D in the paper.\n\n#### To download our clusters for Wikitext-103:\nNote that only **one** of the following files is needed. For the main experiments in the paper, we used:\n```bash\nwget -P checkpoints/wt103/ https://retomaton.s3.us-east-2.amazonaws.com/wt103/clusters_s40000000_k1000000_members.pkl\n```\n\nbut additional clusterings are available as well:\n```bash\nwget -P checkpoints/wt103/ https://retomaton.s3.us-east-2.amazonaws.com/wt103/clusters_s20000000_k500000_members.pkl\nwget -P checkpoints/wt103/ https://retomaton.s3.us-east-2.amazonaws.com/wt103/dstore_merge15_members_sp.pkl\nwget -P checkpoints/wt103/ https://retomaton.s3.us-east-2.amazonaws.com/wt103/dstore_merge29_members.pkl\n```\n\n#### To download our clusters for Law-MT:\nNote that only **one** of the following files is needed. For the main experiments in the paper, we used:\n```bash\nwget -P checkpoints/law/ https://retomaton.s3.us-east-2.amazonaws.com/law/law_clusters_s40000000_k200000_members.pkl\n```\n\nbut additional clustering is available as well:\n```bash\nwget -P checkpoints/law/ https://retomaton.s3.us-east-2.amazonaws.com/law/law_clusters_s40000000_k400000_members.pkl\n```\n\n#### Evaluating RetoMaton with clustering:\nBasically identical to [Step 6: Evaluating RetoMaton without clustering](#step-6-evaluating-retomaton-without-clustering), except that we add the flag `--members \u003cfilename\u003e_members.pkl`, \n\n##### Wikitext-103:\n\n```bash\nDSTORE=checkpoints/wt103/dstore16\nDSTORE_SIZE=103225485\nINDEX=checkpoints/wt103/knn16.index\nMODEL=checkpoints/wt103/wt103_checkpoint_best.pt\nMEMBERS=checkpoints/wt103/clusters_s40000000_k1000000_members.pkl\n\npython eval_lm.py data-bin/wikitext-103 \\\n    --path ${MODEL} \\\n    --sample-break-mode complete --max-tokens 3072 \\\n    --context-window 2560 --softmax-batch 1024000 --batch-size 2 \\\n    --gen-subset valid --dstore-filename ${DSTORE} \\\n    --indexfile ${INDEX}  \\\n    --model-overrides \"{'knn_keytype': 'last_ffn_input'}\" \\\n    --k 1024 --lmbda 0.25 --dstore-size ${DSTORE_SIZE} --knn-keytype last_ffn_input \\\n    --probe 32 --knnlm --dstore-fp16 \\\n    --knn-sim-func do_not_recomp_l2 --no-load-keys --move-dstore-to-mem \\\n    --knnlm-gpu --min-knns 1 --max-knns 1024 \\\n    --members ${MEMBERS}\n```\n\n##### Law-MT:\n```bash\nDSTORE=checkpoints/law/dstore16\nDSTORE_SIZE=19068709\nINDEX=checkpoints/law/knn16.index\nMODEL=checkpoints/law/wmt19.en/model.pt\nMEMBERS=checkpoints/law/law_clusters_s40000000_k200000_members.pkl\n\npython eval_lm.py data-bin/law \\\n    --path ${MODEL} \\\n    --sample-break-mode eos --max-tokens 2048 \\\n    --context-window 0 --softmax-batch 1024000 --batch-size 2 \\\n    --gen-subset valid --dstore-filename ${DSTORE} \\\n    --indexfile ${INDEX}  \\\n    --model-overrides \"{'knn_keytype': 'last_ffn_input'}\" \\\n    --k 1024 --lmbda 0.9 --dstore-size ${DSTORE_SIZE} --knn-keytype last_ffn_input \\\n    --probe 32 --knnlm --dstore-fp16 \\\n    --knn-sim-func do_not_recomp_l2 --no-load-keys --move-dstore-to-mem \\\n    --remove-bpe \\\n    --knnlm-gpu --min-knns 1 --max-knns 1024\\\n    --members ${MEMBERS}\n```\n\n\n#### Cluster the keys yourself (not needed if you downloaded our clusters):\nfor Wikitext-103:\n```\nDSTORE=checkpoints/wt103/dstore16\nDSTORE_SIZE=103225485\nNUM_CLUSTERS=1000000\nSAMPLE=40000000\nDIM=1024\nSAVE=kmeans_wt103\n```\n\nFor Law-MT:\n```bash\nDSTORE=checkpoints/law/dstore16\nDSTORE_SIZE=19068709\nNUM_CLUSTERS=200000\nSAMPLE=40000000\nDIM=1536\nSAVE=kmeans_law\n```\n\nAnd then for both datasets:\n```bash\npython kmeans.py --dstore ${DSTORE} --dstore-size ${DSTORE_SIZE} --num-clusters ${NUM_CLUSTERS} --sample ${SAMPLE} --dim ${DIM} --save ${}\n```\n\n\n\n## Step 8: Evaluating the Fine-tuned Model\nThe model that was fine-tuned on Law-MT, along with its corresponding datastore, FAISS index and clustering can be downloaded from:\n\n```bash\nmkdir checkpoints/law-finetuned/\nwget -P checkpoints/law-finetuned/ https://retomaton.s3.us-east-2.amazonaws.com/law/finetuned.pt\nwget -P checkpoints/law-finetuned/ https://retomaton.s3.us-east-2.amazonaws.com/law/dstore16_finetuned_size19068709_embed1536_fp16_vals.npy\nwget -P checkpoints/law-finetuned/ https://retomaton.s3.us-east-2.amazonaws.com/law/dstore16_finetuned_size19068709_embed1536_fp16_keys.npy\nwget -P checkpoints/law-finetuned/ https://retomaton.s3.us-east-2.amazonaws.com/law/knn_finetuned.index\nwget -P checkpoints/law-finetuned/ https://retomaton.s3.us-east-2.amazonaws.com/law/law_finetuned_clusters_s20000000_k200000_members.pkl\n```\n\nFinally, [evaluate](#evaluating-retomaton-without-clustering) using the fine-tuned checkpoint, datastore, and index. \n\n**It is important** to also set `--lmbda 0.25` when using the fine-tuned model: since the model is fine-tuned, we can rely on it more than before. See a clarification at [#lambda-values](#lambda-values)\n\nBest results with the fine-tuned model are achieved _without_ clustering (that is, every datastore entry is a singleton cluster).\n\nThen, the same steps as before should be run on the Law-MT datasets, except that: \n* `finetuned.pt` should be used as the `${MODEL}`\n* `dstore16_finetuned_size19068709_embed1536_fp16` should be used as the `${DSTORE}`\n* `knn_finetuned.index` should be used as the `${INDEX}`\n* `law_finetuned_clusters_s20000000_k200000_members.pkl` shoould be used as`${MEMBERS}`\n\nThat is:\n\n```bash\nDSTORE=checkpoints/law-finetuned/dstore16_finetuned_size19068709_embed1536_fp16\nDSTORE_SIZE=19068709\nINDEX=checkpoints/law-finetuned/knn_finetuned.index\nMODEL=checkpoints/law-finetuned/finetuned.pt\nMEMBERS=checkpoints/law-finetuned/law_finetuned_clusters_s20000000_k200000_members.pkl\n\npython eval_lm.py data-bin/law \\\n    --path ${MODEL} \\\n    --sample-break-mode eos --max-tokens 2048 \\\n    --context-window 0 --softmax-batch 1024000 --batch-size 2 \\\n    --gen-subset valid --dstore-filename ${DSTORE} \\\n    --indexfile ${INDEX}  \\\n    --model-overrides \"{'knn_keytype': 'last_ffn_input'}\" \\\n    --k 1024 --lmbda 0.25 --dstore-size ${DSTORE_SIZE} --knn-keytype last_ffn_input \\\n    --probe 32 --knnlm --dstore-fp16 \\\n    --knn-sim-func do_not_recomp_l2 --no-load-keys --move-dstore-to-mem \\\n    --remove-bpe \\\n    --knnlm-gpu --min-knns 1 --max-knns 1024\n```\n\n\n## Lambda values\nIn all configurations, the interpolation factor `lmbda` is set to `0.25`, except when the base LM is `checkpoints/law/wmt19.en/model.pt` **and** the model is evaluated on Law-MT, since this scenario tests domain adaptation, and thus `lmbda` should be set to `0.9`:\n\n|             | `wt103_checkpoint_best.pt` | `wmt19.en/model.pt`     | `finetuned.pt` |\n| :---        |    ----:   |     ---: | ---: |\n| Wikitext-103| 0.25       | -    |   -   |\n| Law-MT      | -       | 0.9    |   0.25 |\n\n## All files: \nCheckpoints and datasets can be downloaded from here:\n[https://zenodo.org/record/6525426](https://zenodo.org/record/6525426)\n\nAnd also from the [AWS S3 bucket](https://retomaton.s3.amazonaws.com/index.html) \n\n\n## Differences from the kNN-LM implementation\n\n### Implementation Pointers\nHere we point to the code that differs our work from kNN-LM.\n* The main changes are in this\n[commit](https://github.com/neulab/retomaton/commit/89a29d1ac6e8c1360637aa1bfe77a1be227e83cc). The pointers for the next timestep are initially [the current k-nearest neighbors + 1](fairseq/sequence_scorer.py#L203). Then we extend each pointer to [consider all entries in its cluster](fairseq/sequence_scorer.py#L216). This is [the function](fairseq/sequence_scorer.py#L240-L251) that maps each pointer to its cluster, removes duplicate clusters, and then finds the members of each cluster. We  [find the log probabilities](fairseq/sequence_scorer.py#L218-L222) as suggested by the new pointers, and finally take to the next timestep - [only the pointers that are consistent](fairseq/sequence_scorer.py#L228) with the token that the model eventually predicted.\n* In [this commit](https://github.com/neulab/retomaton/commit/99cb52001b3c87b15dd8ef892172cfac334bcef5) we [utilize the given pointers](fairseq/knnlm.py#L131-L133), or [perform kNN search](fairseq/knnlm.py#L131-L133) and combine the results with the existing pointers.\n* When using the `--knnlm-gpu` flag, we use a [GPU index](fairseq/knnlm.py#L34-L38) to search for nearest neighbors, and its copy [CPU index](fairseq/knnlm.py#L43-L47) to reconstruct vectors given their ID. Unfortunately, currently reconstructing vectors in `faiss` is [not implemented for GPU indexes](https://github.com/facebookresearch/faiss/issues/2181) (see also [this issue](https://github.com/facebookresearch/faiss/issues/314)). \n* Reconstructing a **batch** of vectors from the index is unfortunately not implemented in `faiss` (see [this issue](https://github.com/facebookresearch/faiss/issues/1163)), and thus the fastest way that we found to do that is using `np.vectorize`, and reconstructing many single vectors in parallel: [fairseq/knnlm.py#L92-L94](fairseq/knnlm.py#L92-L94).\n* Performing k-means clustering on millions of vectors can be performed in many ways, but specifically we utilize the `faiss` library to do it using the script [kmeans.py](kmeans.py).\n\n### Additional minor differences:\n* The original [kNN-LM](https://github.com/urvashik/knnlm) repository uses `faiss` CPU to perform retrieval. However, we added the flag `--knnlm-gpu` that allows performing retrieval much faster on the GPU.\n* After each retrieval, the original [kNN-LM](https://github.com/urvashik/knnlm) repository loads the found keys and re-computes the distance from the query to each nearest neighbor. This is much more time consuming, unless loading all the keys (200GB) into memory.\nWe thus use the flags `--knn-sim-func do_not_recomp_l2 --no-load-keys --move-dstore-to-mem`.\n* When using `faiss-gpu`, it is useful to [`import faiss.contrib.torch_utils`](fairseq/knnlm.py#L3). This allows performing the kNN search using `torch` tensors (rather than only `numpy` arrays). Additionally, sometimes this `import` statement prevents searching bugs in `faiss` (see [this issue](https://github.com/facebookresearch/faiss/issues/2126)).\n\n\n\n\n## Citation\n\n[Neuro-Symbolic Language Modeling with Automaton-augmented Retrieval](https://arxiv.org/pdf/2201.12431.pdf)\n\n```\n@inproceedings{alon2022neuro,\n  title={Neuro-Symbolic Language Modeling with Automaton-augmented Retrieval},\n  author={Alon, Uri and Xu, Frank and He, Junxian and Sengupta, Sudipta and Roth, Dan and Neubig, Graham},\n  booktitle={International Conference on Machine Learning},\n  pages={468--485},\n  year={2022},\n  organization={PMLR}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fneulab%2Fretomaton","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fneulab%2Fretomaton","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fneulab%2Fretomaton/lists"}