[
  {
    "path": ".gitignore",
    "content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\npip-wheel-metadata/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n#  Usually these files are written by a python script from a template\n#  before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n.python-version\n\n# pipenv\n#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n#   However, in case of collaboration, if having platform-specific dependencies or dependencies\n#   having no cross-platform support, pipenv may install dependencies that don't work, or not\n#   install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n.DS_Store\n"
  },
  {
    "path": "LICENSE",
    "content": "MIT License\n\nCopyright (c) 2022 Shengyao Zhuang\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
  },
  {
    "path": "README.md",
    "content": "# DSI-transformers\nA huggingface transformers implementation of [Transformer Memory as a Differentiable Search Index](https://arxiv.org/abs/2202.06991), Yi Tay, Vinh Q. Tran, Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta, Zhen Qin, Kai Hui, Zhe Zhao, Jai Gupta, Tal Schuster, William W. Cohen, Donald Metzler\n\nRequirements: `python=3.8` `transformers=4.17.0` `datasets=1.18.3` `wandb`\n> Note: This is not the official repository.\n\n## Updates\n- Check out our new repository for DSI training: https://github.com/ArvinZhuang/DSI-QG\n\n## Goal of this repository\nReproduce the results of DSI Large, Naive String Docid, NQ10K. According to Table 3 in the original paper, we should have `Hits@1=0.347`,`Hits@10=0.605`\n\n### Step1: Create NQ10K training (indexing) and validation datasets\n\n```\ncd data/NQ\npython3 create_NQ_train_vali.py\n```\n\n### Step2: Run training script\ncd back to the root directory and run:\n\n```\npython3 train.py\n```\nThe training can be run with a single Tesla-v100 32G GPU. We use [wandb](https://wandb.ai/site) to log the Hits scores during training:\n\n![.im](hits_plots.png)\n\n\n### Discussion\n\nAs the above plots show, the current implementation is worse than what is reported in the original paper, there are many possible reasons: the ratio of training and indexing examples (we use 1:1), number of training steps, the way of constructing documents text, etc. Although, seems the scores are on par with BM25 already.\n\nIf you can identify the reason or any bug, welcome to open a PR to fix it!\n\n#### Indexing or overfitting?\nThe training script also logged the hit@1 scores on the training set during training, this is aimed to analyze if the model can memorize all the training data points, the authors called this 'indexing' which I believe is just letting the model overfits the training data points. It turns out the model can reach %99.99 hit@1 on the training set very quickly (quickly overfit), but the hits scores on the test set continue going up. Seems T5 large has good generalizability on this task.\n"
  },
  {
    "path": "data/NQ/create_NQ_train_vali.py",
    "content": "import datasets\nimport random\nimport numpy as np\nimport json\n\nrandom.seed(313)\n\nNUM_TRAIN = 8000\nNUM_EVAL = 2000\n\ndata = datasets.load_dataset('natural_questions', cache_dir='cache')['train']\nrand_inds = list(range(len(data)))\nrandom.shuffle(rand_inds)\n\ntitle_set = set()\ncurrent_docid = 0\n\nwith open('NQ_10k_multi_task_train.json', 'w') as tf, \\\n        open('NQ_10k_valid.json', 'w') as vf:\n    for ind in rand_inds:\n        title = data[ind]['document']['title']  # we use title as the doc identifier to prevent two docs have the same text\n        if title not in title_set:\n            title_set.add(title)\n            token_inds = np.where(np.array(data[ind]['document']['tokens']['is_html']) == False)[0]\n            tokens = np.array(data[ind]['document']['tokens']['token'])[token_inds]\n            doc_text = \" \".join(tokens)\n            question_text = data[ind]['question']['text']\n\n            jitem = json.dumps({'text_id': str(current_docid), 'text': 'document: ' + doc_text})\n            tf.write(jitem + '\\n')\n            jitem = json.dumps({'text_id': str(current_docid), 'text': 'question: ' + question_text})\n            if len(title_set) <= NUM_TRAIN:\n                tf.write(jitem + '\\n')\n            else:\n                vf.write(jitem + '\\n')\n            current_docid += 1\n            if len(title_set) == NUM_TRAIN + NUM_EVAL:\n                break\n        print(f\"Creating training and validation dataset: {'{:.1%}'.format(len(title_set)/(NUM_TRAIN + NUM_EVAL))}\", end='\\r')"
  },
  {
    "path": "data.py",
    "content": "from dataclasses import dataclass\n\nimport datasets\nfrom torch.utils.data import Dataset\nfrom transformers import PreTrainedTokenizer, DataCollatorWithPadding\n\n\nclass IndexingTrainDataset(Dataset):\n    def __init__(\n            self,\n            path_to_data,\n            max_length: int,\n            cache_dir: str,\n            tokenizer: PreTrainedTokenizer,\n    ):\n        self.train_data = datasets.load_dataset(\n            'json',\n            data_files=path_to_data,\n            ignore_verifications=False,\n            cache_dir=cache_dir\n        )['train']\n\n        self.max_length = max_length\n        self.tokenizer = tokenizer\n        self.total_len = len(self.train_data)\n\n\n    def __len__(self):\n        return self.total_len\n\n    def __getitem__(self, item):\n        data = self.train_data[item]\n\n        input_ids = self.tokenizer(data['text'],\n                                   return_tensors=\"pt\",\n                                   truncation='only_first',\n                                   max_length=self.max_length).input_ids[0]\n        return input_ids, str(data['text_id'])\n\n\n@dataclass\nclass IndexingCollator(DataCollatorWithPadding):\n    def __call__(self, features):\n        input_ids = [{'input_ids': x[0]} for x in features]\n        docids = [x[1] for x in features]\n        inputs = super().__call__(input_ids)\n\n        labels = self.tokenizer(\n            docids, padding=\"longest\", return_tensors=\"pt\"\n        ).input_ids\n\n        # replace padding token id's of the labels by -100 according to https://huggingface.co/docs/transformers/model_doc/t5#training\n        labels[labels == self.tokenizer.pad_token_id] = -100\n        inputs['labels'] = labels\n        return inputs\n\n\n@dataclass\nclass QueryEvalCollator(DataCollatorWithPadding):\n    def __call__(self, features):\n        input_ids = [{'input_ids': x[0]} for x in features]\n        labels = [x[1] for x in features]\n        inputs = super().__call__(input_ids)\n\n        return inputs, labels"
  },
  {
    "path": "train.py",
    "content": "from data import IndexingTrainDataset, IndexingCollator, QueryEvalCollator\nfrom transformers import T5Tokenizer, T5ForConditionalGeneration, TrainingArguments, TrainerCallback\nfrom trainer import IndexingTrainer\nimport numpy as np\nimport torch\nimport wandb\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\n\n\nclass QueryEvalCallback(TrainerCallback):\n    def __init__(self, test_dataset, logger, restrict_decode_vocab, args: TrainingArguments, tokenizer: T5Tokenizer):\n        self.tokenizer = tokenizer\n        self.logger = logger\n        self.args = args\n        self.test_dataset = test_dataset\n        self.restrict_decode_vocab = restrict_decode_vocab\n        self.dataloader = DataLoader(\n            test_dataset,\n            batch_size=self.args.per_device_eval_batch_size,\n            collate_fn=QueryEvalCollator(\n                self.tokenizer,\n                padding='longest'\n            ),\n            shuffle=False,\n            drop_last=False,\n            num_workers=self.args.dataloader_num_workers,\n        )\n\n    def on_epoch_end(self, args, state, control, **kwargs):\n        hit_at_1 = 0\n        hit_at_10 = 0\n        model = kwargs['model'].eval()\n        for batch in tqdm(self.dataloader, desc='Evaluating dev queries'):\n            inputs, labels = batch\n            with torch.no_grad():\n                batch_beams = model.generate(\n                    inputs['input_ids'].to(model.device),\n                    max_length=20,\n                    num_beams=10,\n                    prefix_allowed_tokens_fn=self.restrict_decode_vocab,\n                    num_return_sequences=10,\n                    early_stopping=True, ).reshape(inputs['input_ids'].shape[0], 10, -1)\n                for beams, label in zip(batch_beams, labels):\n                    rank_list = self.tokenizer.batch_decode(beams,\n                                                            skip_special_tokens=True)  # beam search should not return repeated docids but somehow due to T5 tokenizer there some repeats.\n                    hits = np.where(np.array(rank_list)[:10] == label)[0]\n                    if len(hits) != 0:\n                        hit_at_10 += 1\n                        if hits[0] == 0:\n                            hit_at_1 += 1\n        self.logger.log({\"Hits@1\": hit_at_1 / len(self.test_dataset), \"Hits@10\": hit_at_10 / len(self.test_dataset)})\n\n\ndef compute_metrics(eval_preds):\n    num_predict = 0\n    num_correct = 0\n    for predict, label in zip(eval_preds.predictions, eval_preds.label_ids):\n        num_predict += 1\n        if len(np.where(predict == 1)[0]) == 0:\n            continue\n        if np.array_equal(label[:np.where(label == 1)[0].item()],\n                          predict[np.where(predict == 0)[0][0].item() + 1:np.where(predict == 1)[0].item()]):\n            num_correct += 1\n\n    return {'accuracy': num_correct / num_predict}\n\n\ndef main():\n    model_name = \"t5-large\"\n    L = 32  # only use the first 32 tokens of documents (including title)\n\n    # We use wandb to log Hits scores after each epoch. Note, this script does not save model checkpoints.\n    wandb.login()\n    wandb.init(project=\"DSI\", name='NQ-10k-t5-large')\n\n    tokenizer = T5Tokenizer.from_pretrained(model_name, cache_dir='cache')\n    model = T5ForConditionalGeneration.from_pretrained(model_name, cache_dir='cache')\n\n    train_dataset = IndexingTrainDataset(path_to_data='data/NQ/NQ_10k_multi_task_train.json',\n                                         max_length=L,\n                                         cache_dir='cache',\n                                         tokenizer=tokenizer)\n    \n    # This eval set is really not the 'eval' set but used to report if the model can memorise (index) all training data points.\n    eval_dataset = IndexingTrainDataset(path_to_data='data/NQ/NQ_10k_multi_task_train.json',\n                                        max_length=L,\n                                        cache_dir='cache',\n                                        tokenizer=tokenizer)\n    \n    # This is the actual eval set.\n    test_dataset = IndexingTrainDataset(path_to_data='data/NQ/NQ_10k_valid.json',\n                                        max_length=L,\n                                        cache_dir='cache',\n                                        tokenizer=tokenizer)\n\n    ################################################################\n    # docid generation constrain, we only generate integer docids.\n    SPIECE_UNDERLINE = \"▁\"\n    INT_TOKEN_IDS = []\n    for token, id in tokenizer.get_vocab().items():\n        if token[0] == SPIECE_UNDERLINE:\n            if token[1:].isdigit():\n                INT_TOKEN_IDS.append(id)\n        if token == SPIECE_UNDERLINE:\n            INT_TOKEN_IDS.append(id)\n        elif token.isdigit():\n            INT_TOKEN_IDS.append(id)\n    INT_TOKEN_IDS.append(tokenizer.eos_token_id)\n\n    def restrict_decode_vocab(batch_idx, prefix_beam):\n        return INT_TOKEN_IDS\n    ################################################################\n\n    training_args = TrainingArguments(\n        output_dir=\"./results\",\n        learning_rate=0.0005,\n        warmup_steps=10000,\n        # weight_decay=0.01,\n        per_device_train_batch_size=128,\n        per_device_eval_batch_size=128,\n        evaluation_strategy='steps',\n        eval_steps=1000,\n        max_steps=1000000,\n        dataloader_drop_last=False,  # necessary\n        report_to='wandb',\n        logging_steps=50,\n        save_strategy='no',\n        # fp16=True,  # gives 0/nan loss at some point during training, seems this is a transformers bug.\n        dataloader_num_workers=10,\n        # gradient_accumulation_steps=2\n    )\n\n    trainer = IndexingTrainer(\n        model=model,\n        tokenizer=tokenizer,\n        args=training_args,\n        train_dataset=train_dataset,\n        eval_dataset=eval_dataset,\n        data_collator=IndexingCollator(\n            tokenizer,\n            padding='longest',\n        ),\n        compute_metrics=compute_metrics,\n        callbacks=[QueryEvalCallback(test_dataset, wandb, restrict_decode_vocab, training_args, tokenizer)],\n        restrict_decode_vocab=restrict_decode_vocab\n    )\n    trainer.train(\n    )\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "trainer.py",
    "content": "from typing import Dict, List, Tuple, Optional, Any, Union\nfrom transformers.trainer import Trainer\nfrom torch import nn\nimport torch\n\n\nclass IndexingTrainer(Trainer):\n    def __init__(self, restrict_decode_vocab, **kwds):\n        super().__init__(**kwds)\n        self.restrict_decode_vocab = restrict_decode_vocab\n\n    def compute_loss(self, model, inputs, return_outputs=False):\n        loss = model(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'], labels=inputs['labels']).loss\n        if return_outputs:\n            return loss, [None, None]  # fake outputs\n        return loss\n\n    def prediction_step(\n            self,\n            model: nn.Module,\n            inputs: Dict[str, Union[torch.Tensor, Any]],\n            prediction_loss_only: bool,\n            ignore_keys: Optional[List[str]] = None,\n    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:\n        model.eval()\n        # eval_loss = super().prediction_step(model, inputs, True, ignore_keys)[0]\n        with torch.no_grad():\n            # greedy search\n            doc_ids = model.generate(\n                inputs['input_ids'].to(self.args.device),\n                max_length=20,\n                prefix_allowed_tokens_fn=self.restrict_decode_vocab,\n                early_stopping=True,)\n        return (None, doc_ids, inputs['labels'])\n\n"
  }
]