[
  {
    "path": ".gitignore",
    "content": "*.pyc\n*.DS_Store\n*~\ndata/\n*.tar.gz\n*.egg-info\n"
  },
  {
    "path": "LICENSE",
    "content": "BSD 3-Clause License\n\nCopyright (c) 2018, HKUST-KnowComp\nAll rights reserved.\n\nRedistribution and use in source and binary forms, with or without\nmodification, are permitted provided that the following conditions are met:\n\n* Redistributions of source code must retain the above copyright notice, this\n  list of conditions and the following disclaimer.\n\n* Redistributions in binary form must reproduce the above copyright notice,\n  this list of conditions and the following disclaimer in the documentation\n  and/or other materials provided with the distribution.\n\n* Neither the name of the copyright holder nor the names of its\n  contributors may be used to endorse or promote products derived from\n  this software without specific prior written permission.\n\nTHIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\nAND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\nIMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\nDISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\nFOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\nDAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\nSERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\nCAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\nOR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\nOF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n"
  },
  {
    "path": "README.md",
    "content": "# Mnemonic Reader\nThe Mnemonic Reader is a deep learning model for Machine Comprehension task. You can get details from this [paper](https://arxiv.org/pdf/1705.02798.pdf). It combines advantages of [match-LSTM](https://arxiv.org/pdf/1608.07905), [R-Net](https://www.microsoft.com/en-us/research/wp-content/uploads/2017/05/r-net.pdf) and [Document Reader](https://arxiv.org/abs/1704.00051) and utilizes a new unit, the Semantic Fusion Unit (SFU), to achieve state-of-the-art results (at that time).\n\nThis model is a [PyTorch](http://pytorch.org/) implementation of Mnemonic Reader. At the same time, a PyTorch implementation of R-Net and a PyTorch implementation of Document Reader are also included to compare with the Mnemonic Reader. Pretrained models are also available in [release](https://github.com/HKUST-KnowComp/MnemonicReader/releases).\n\nThis repo belongs to [HKUST-KnowComp](https://github.com/HKUST-KnowComp) and is under the [BSD LICENSE](LICENSE).\n\nSome codes are implemented based on [DrQA](https://github.com/facebookresearch/DrQA).\n\nPlease feel free to contact with Xin Liu (xliucr@connect.ust.hk) if you have any question about this repo.\n\n### Evaluation on SQuAD\n\n| Model                                 | DEV_EM | DEV_F1 |\n| ------------------------------------- | ------ | ------ |\n| Document Reader (original paper)      | 69.5   | 78.8   |\n| Document Reader (trained model)       | 69.4   | 78.6   |\n| R-Net (original paper 1)              | 71.1   | 79.5   |\n| R-Net (original paper 2)              | 72.3   | 80.6   |\n| R-Net (trained model)                 | 70.2   | 79.4   |\n| Mnemonic Reader (original paper)      | 71.8   | 81.2   |\n| Mnemonic Reader + RL (original paper) | 72.1   | 81.6   |\n| Mnemonic Reader (trained model)       | 73.2   | 81.5   |\n\n![EM_F1](img/EM_F1.png)\n\n### Requirements\n\n* Python >= 3.4\n* PyTorch >= 0.31\n* spaCy >= 2.0.0\n* tqdm\n* ujson\n* numpy\n* prettytable\n\n### Prepare\n\nFirst of all, you need to download the dataset and pre-trained word vectors.\n\n```bash\nmkdir -p data/datasets\nwget https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json -O data/datasets/SQuAD-train-v1.1.json\nwget https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json -O data/datasets/SQuAD-dev-v1.1.json\n```\n\n```bash\nmkdir -p data/embeddings\nwget http://nlp.stanford.edu/data/glove.840B.300d.zip -O data/embeddings/glove.840B.300d.zip\ncd data/embeddings\nunzip glove.840B.300d.zip\n```\n\nThen, you need to preprocess these data.\n\n```bash\npython script/preprocess.py data/datasets data/datasets --split SQuAD-train-v1.1\npython script/preprocess.py data/datasets data/datasets --split SQuAD-dev-v1.1\n```\n\nIf you want to use multicores to speed up, you could add `--num-workers 4` in commands.\n\n### Train\n\nThere are some parameters to set but default values are ready. If you are not interested in tuning parameters, you can use default values. Just run:\n\n```bash\npython script/train.py\n```\n\nAfter several hours, you will get the model in `data/models/`, e.g. `20180416-acc9d06d.mdl` and you can see the log file in `data/models/`, e.g. `20180416-acc9d06d.txt`.\n\n### Predict\n\nTo evaluate the model you get, you should complete this part.\n\n```bash\npython script/predict.py --model data/models/20180416-acc9d06d.mdl\n```\n\nYou need to change the model name in the command above.\n\nYou will not get results directly but to use the official `evaluate-v1.1.py` in `data/script`.\n\n```bash\npython script/evaluate-v1.1.py data/predict/SQuAD-dev-v1.1-20180416-acc9d06d.preds data/datasets/SQuAD-dev-v1.1.json\n```\n\n### Interactivate\n\nIn order to help those who are interested in QA systems, `script/interactivate.py` provides an easy but good demo.\n\n```bash\npython script/interactivate.py --model data/models/20180416-acc9d06d.mdl\n```\n\nThen you will drop into an interactive session. It looks like:\n\n```\n* Interactive Module *\n\n* Repo: Mnemonic Reader (https://github.com/HKUST-KnowComp/MnemonicReader)\n\n* Implement based on Facebook's DrQA\n\n>>> process(document, question, candidates=None, top_n=1)\n>>> usage()\n\n>>> text=\"Architecturally, the school has a Catholic character. Atop the Main Building's gold dome is a golden statue of the Virgin Mary. Immediately in front of the Main Building and facing it, is a copper statue of Christ with arms upraised with the legend \\\"Venite Ad Me Omnes\\\". Next to the Main Building is the Basilica of the Sacred Heart. Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. It is a replica of the grotto at Lourdes, France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous in 1858. At the end of the main drive (and in a direct line that connects through 3 statues and the Gold Dome), is a simple, modern stone statue of Mary.\"\n>>> question = \"To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?\"\n>>> process(text, question)\n\n+------+----------------------------+-----------+\n| Rank |            Span            |   Score   |\n+------+----------------------------+-----------+\n|  1   | Saint Bernadette Soubirous | 0.9875301 |\n+------+----------------------------+-----------+\n```\n\n### More parameters\n\nIf you want to tune parameters to achieve a higher score, you can get instructions about parameters via using\n\n```bash\npython script/preprocess.py --help\n```\n\n```bash\npython script/train.py --help\n```\n\n```bash\npython script/predict.py --help\n```\n\n```bash\npython script/interactivate.py --help\n```\n\n## License\n\nAll codes in **Mnemonic Reader** are under [BSD LICENSE](LICENSE).\n"
  },
  {
    "path": "config.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2018-present, HKUST-KnowComp.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Model architecture/optimization options for WRMCQA document reader.\"\"\"\n\nimport argparse\nimport logging\n\nlogger = logging.getLogger(__name__)\n\n# Index of arguments concerning the core model architecture\nMODEL_ARCHITECTURE = {\n    'model_type', 'embedding_dim', 'char_embedding_dim', 'hidden_size', 'char_hidden_size',\n    'doc_layers', 'question_layers', 'rnn_type', 'concat_rnn_layers', 'question_merge',\n    'use_qemb', 'use_exact_match', 'use_pos', 'use_ner', 'use_lemma', 'use_tf', 'hop'\n}\n\n# Index of arguments concerning the model optimizer/training\nMODEL_OPTIMIZER = {\n    'fix_embeddings', 'optimizer', 'learning_rate', 'momentum', 'weight_decay',\n    'rho', 'eps', 'max_len', 'grad_clipping', 'tune_partial', \n    'rnn_padding', 'dropout_rnn', 'dropout_rnn_output', 'dropout_emb'\n}\n\n\ndef str2bool(v):\n    return v.lower() in ('yes', 'true', 't', '1', 'y')\n\n\ndef add_model_args(parser):\n    parser.register('type', 'bool', str2bool)\n\n    # Model architecture\n    model = parser.add_argument_group('Reader Model Architecture')\n    model.add_argument('--model-type', type=str, default='mnemonic',\n                       help='Model architecture type: rnn, r_net, mnemonic')\n    model.add_argument('--embedding-dim', type=int, default=300,\n                       help='Embedding size if embedding_file is not given')\n    model.add_argument('--char-embedding-dim', type=int, default=50,\n                       help='Embedding size if char_embedding_file is not given')\n    model.add_argument('--hidden-size', type=int, default=100,\n                       help='Hidden size of RNN units')\n    model.add_argument('--char-hidden-size', type=int, default=50,\n                       help='Hidden size of char RNN units')\n    model.add_argument('--doc-layers', type=int, default=3,\n                       help='Number of encoding layers for document')\n    model.add_argument('--question-layers', type=int, default=3,\n                       help='Number of encoding layers for question')\n    model.add_argument('--rnn-type', type=str, default='lstm',\n                       help='RNN type: LSTM, GRU, or RNN')\n\n    # Model specific details\n    detail = parser.add_argument_group('Reader Model Details')\n    detail.add_argument('--concat-rnn-layers', type='bool', default=True,\n                        help='Combine hidden states from each encoding layer')\n    detail.add_argument('--question-merge', type=str, default='self_attn',\n                        help='The way of computing the question representation')\n    detail.add_argument('--use-qemb', type='bool', default=True,\n                        help='Whether to use weighted question embeddings')\n    detail.add_argument('--use-exact-match', type='bool', default=True,\n                        help='Whether to use in_question_* features')\n    detail.add_argument('--use-pos', type='bool', default=True,\n                        help='Whether to use pos features')\n    detail.add_argument('--use-ner', type='bool', default=True,\n                        help='Whether to use ner features')\n    detail.add_argument('--use-lemma', type='bool', default=True,\n                        help='Whether to use lemma features')\n    detail.add_argument('--use-tf', type='bool', default=True,\n                        help='Whether to use term frequency features')\n    detail.add_argument('--hop', type=int, default=2,\n                        help='The number of hops for both aligner and the answer pointer in m-reader')\n\n    # Optimization details\n    optim = parser.add_argument_group('Reader Optimization')\n    optim.add_argument('--dropout-emb', type=float, default=0.2,\n                       help='Dropout rate for word embeddings')\n    optim.add_argument('--dropout-rnn', type=float, default=0.2,\n                       help='Dropout rate for RNN states')\n    optim.add_argument('--dropout-rnn-output', type='bool', default=True,\n                       help='Whether to dropout the RNN output')\n    optim.add_argument('--optimizer', type=str, default='adamax',\n                       help='Optimizer: sgd, adamax, adadelta')\n    optim.add_argument('--learning-rate', type=float, default=1.0,\n                       help='Learning rate for sgd, adadelta')\n    optim.add_argument('--grad-clipping', type=float, default=10,\n                       help='Gradient clipping')\n    optim.add_argument('--weight-decay', type=float, default=0,\n                       help='Weight decay factor')\n    optim.add_argument('--momentum', type=float, default=0,\n                       help='Momentum factor')\n    optim.add_argument('--rho', type=float, default=0.95,\n                       help='Rho for adadelta')\n    optim.add_argument('--eps', type=float, default=1e-6,\n                       help='Eps for adadelta')\n    optim.add_argument('--fix-embeddings', type='bool', default=True,\n                       help='Keep word embeddings fixed (use pretrained)')\n    optim.add_argument('--tune-partial', type=int, default=0,\n                       help='Backprop through only the top N question words')\n    optim.add_argument('--rnn-padding', type='bool', default=False,\n                       help='Explicitly account for padding in RNN encoding')\n    optim.add_argument('--max-len', type=int, default=15,\n                       help='The max span allowed during decoding')\n\n\ndef get_model_args(args):\n    \"\"\"Filter args for model ones.\n\n    From a args Namespace, return a new Namespace with *only* the args specific\n    to the model architecture or optimization. (i.e. the ones defined here.)\n    \"\"\"\n    global MODEL_ARCHITECTURE, MODEL_OPTIMIZER\n    required_args = MODEL_ARCHITECTURE | MODEL_OPTIMIZER\n    arg_values = {k: v for k, v in vars(args).items() if k in required_args}\n    return argparse.Namespace(**arg_values)\n\n\ndef override_model_args(old_args, new_args):\n    \"\"\"Set args to new parameters.\n\n    Decide which model args to keep and which to override when resolving a set\n    of saved args and new args.\n\n    We keep the new optimation, but leave the model architecture alone.\n    \"\"\"\n    global MODEL_OPTIMIZER\n    old_args, new_args = vars(old_args), vars(new_args)\n    for k in old_args.keys():\n        if k in new_args and old_args[k] != new_args[k]:\n            if k in MODEL_OPTIMIZER:\n                logger.info('Overriding saved %s: %s --> %s' %\n                            (k, old_args[k], new_args[k]))\n                old_args[k] = new_args[k]\n            else:\n                logger.info('Keeping saved %s: %s' % (k, old_args[k]))\n    return argparse.Namespace(**old_args)\n"
  },
  {
    "path": "data.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2018-present, HKUST-KnowComp.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Data processing/loading helpers.\"\"\"\n\nimport numpy as np\nimport logging\nimport unicodedata\n\nfrom torch.utils.data import Dataset\nfrom torch.utils.data.sampler import Sampler\nfrom vector import vectorize\n\nlogger = logging.getLogger(__name__)\n\n\n# ------------------------------------------------------------------------------\n# Dictionary class for tokens.\n# ------------------------------------------------------------------------------\n\n\nclass Dictionary(object):\n    NULL = '<NULL>'\n    UNK = '<UNK>'\n    START = 2\n\n    @staticmethod\n    def normalize(token):\n        return unicodedata.normalize('NFD', token)\n\n    def __init__(self):\n        self.tok2ind = {self.NULL: 0, self.UNK: 1}\n        self.ind2tok = {0: self.NULL, 1: self.UNK}\n\n    def __len__(self):\n        return len(self.tok2ind)\n\n    def __iter__(self):\n        return iter(self.tok2ind)\n\n    def __contains__(self, key):\n        if type(key) == int:\n            return key in self.ind2tok\n        elif type(key) == str:\n            return self.normalize(key) in self.tok2ind\n\n    def __getitem__(self, key):\n        if type(key) == int:\n            return self.ind2tok.get(key, self.UNK)\n        if type(key) == str:\n            return self.tok2ind.get(self.normalize(key),\n                                    self.tok2ind.get(self.UNK))\n\n    def __setitem__(self, key, item):\n        if type(key) == int and type(item) == str:\n            self.ind2tok[key] = item\n        elif type(key) == str and type(item) == int:\n            self.tok2ind[key] = item\n        else:\n            raise RuntimeError('Invalid (key, item) types.')\n\n    def add(self, token):\n        token = self.normalize(token)\n        if token not in self.tok2ind:\n            index = len(self.tok2ind)\n            self.tok2ind[token] = index\n            self.ind2tok[index] = token\n\n    def tokens(self):\n        \"\"\"Get dictionary tokens.\n\n        Return all the words indexed by this dictionary, except for special\n        tokens.\n        \"\"\"\n        tokens = [k for k in self.tok2ind.keys()\n                  if k not in {'<NULL>', '<UNK>'}]\n        return tokens\n\n\n# ------------------------------------------------------------------------------\n# PyTorch dataset class for SQuAD (and SQuAD-like) data.\n# ------------------------------------------------------------------------------\n\n\nclass ReaderDataset(Dataset):\n\n    def __init__(self, examples, model, single_answer=False):\n        self.model = model\n        self.examples = examples\n        self.single_answer = single_answer\n\n    def __len__(self):\n        return len(self.examples)\n\n    def __getitem__(self, index):\n        return vectorize(self.examples[index], self.model, self.single_answer)\n\n    def lengths(self):\n        return [(len(ex['document']), len(ex['question']))\n                for ex in self.examples]\n\n\n# ------------------------------------------------------------------------------\n# PyTorch sampler returning batched of sorted lengths (by doc and question).\n# ------------------------------------------------------------------------------\n\n\nclass SortedBatchSampler(Sampler):\n\n    def __init__(self, lengths, batch_size, shuffle=True):\n        self.lengths = lengths\n        self.batch_size = batch_size\n        self.shuffle = shuffle\n\n    def __iter__(self):\n        lengths = np.array(\n            [(-l[0], -l[1], np.random.random()) for l in self.lengths],\n            dtype=[('l1', np.int_), ('l2', np.int_), ('rand', np.float_)]\n        )\n        indices = np.argsort(lengths, order=('l1', 'l2', 'rand'))\n        batches = [indices[i:i + self.batch_size]\n                   for i in range(0, len(indices), self.batch_size)]\n        if self.shuffle:\n            np.random.shuffle(batches)\n        return iter([i for batch in batches for i in batch])\n\n    def __len__(self):\n        return len(self.lengths)\n"
  },
  {
    "path": "layers.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2018-present, HKUST-KnowComp.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Definitions of model layers/NN modules\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.autograd import Variable\nimport math\nimport random\n\n\n# ------------------------------------------------------------------------------\n# Modules\n# ------------------------------------------------------------------------------\n\n\nclass StackedBRNN(nn.Module):\n    \"\"\"Stacked Bi-directional RNNs.\n\n    Differs from standard PyTorch library in that it has the option to save\n    and concat the hidden states between layers. (i.e. the output hidden size\n    for each sequence input is num_layers * hidden_size).\n    \"\"\"\n\n    def __init__(self, input_size, hidden_size, num_layers,\n                 dropout_rate=0, dropout_output=False, rnn_type=nn.LSTM,\n                 concat_layers=False, padding=False):\n        super(StackedBRNN, self).__init__()\n        self.padding = padding\n        self.dropout_output = dropout_output\n        self.dropout_rate = dropout_rate\n        self.num_layers = num_layers\n        self.concat_layers = concat_layers\n        self.rnns = nn.ModuleList()\n        for i in range(num_layers):\n            input_size = input_size if i == 0 else 2 * hidden_size\n            self.rnns.append(rnn_type(input_size, hidden_size,\n                                      num_layers=1,\n                                      bidirectional=True))\n\n    def forward(self, x, x_mask):\n        \"\"\"Encode either padded or non-padded sequences.\n\n        Can choose to either handle or ignore variable length sequences.\n        Always handle padding in eval.\n\n        Args:\n            x: batch * len * hdim\n            x_mask: batch * len (1 for padding, 0 for true)\n        Output:\n            x_encoded: batch * len * hdim_encoded\n        \"\"\"\n        if x_mask.data.sum() == 0 or x_mask.data.eq(1).long().sum(1).min() == 0:\n            # No padding necessary.\n            output = self._forward_unpadded(x, x_mask)\n        elif self.padding or not self.training:\n            # Pad if we care or if its during eval.\n            output = self._forward_padded(x, x_mask)\n        else:\n            # We don't care.\n            output = self._forward_unpadded(x, x_mask)\n\n        return output.contiguous()\n\n    def _forward_unpadded(self, x, x_mask):\n        \"\"\"Faster encoding that ignores any padding.\"\"\"\n        # Transpose batch and sequence dims\n        x = x.transpose(0, 1)\n\n        # Encode all layers\n        outputs = [x]\n        for i in range(self.num_layers):\n            rnn_input = outputs[-1]\n\n            # Apply dropout to hidden input\n            if self.dropout_rate > 0:\n                rnn_input = F.dropout(rnn_input,\n                                      p=self.dropout_rate,\n                                      training=self.training)\n            # Forward\n            rnn_output = self.rnns[i](rnn_input)[0]\n            outputs.append(rnn_output)\n\n        # Concat hidden layers\n        if self.concat_layers:\n            output = torch.cat(outputs[1:], 2)\n        else:\n            output = outputs[-1]\n\n        # Transpose back\n        output = output.transpose(0, 1)\n\n        # Dropout on output layer\n        if self.dropout_output and self.dropout_rate > 0:\n            output = F.dropout(output,\n                               p=self.dropout_rate,\n                               training=self.training)\n        return output\n\n    def _forward_padded(self, x, x_mask):\n        \"\"\"Slower (significantly), but more precise, encoding that handles\n        padding.\n        \"\"\"\n        # Compute sorted sequence lengths\n        lengths = x_mask.data.eq(0).long().sum(1).squeeze()\n        _, idx_sort = torch.sort(lengths, dim=0, descending=True)\n        _, idx_unsort = torch.sort(idx_sort, dim=0)\n\n        lengths = list(lengths[idx_sort])\n        idx_sort = Variable(idx_sort)\n        idx_unsort = Variable(idx_unsort)\n\n        # Sort x\n        x = x.index_select(0, idx_sort)\n\n        # Transpose batch and sequence dims\n        x = x.transpose(0, 1)\n\n        # Pack it up\n        rnn_input = nn.utils.rnn.pack_padded_sequence(x, lengths)\n\n        # Encode all layers\n        outputs = [rnn_input]\n        for i in range(self.num_layers):\n            rnn_input = outputs[-1]\n\n            # Apply dropout to input\n            if self.dropout_rate > 0:\n                dropout_input = F.dropout(rnn_input.data,\n                                          p=self.dropout_rate,\n                                          training=self.training)\n                rnn_input = nn.utils.rnn.PackedSequence(dropout_input,\n                                                        rnn_input.batch_sizes)\n            outputs.append(self.rnns[i](rnn_input)[0])\n\n        # Unpack everything\n        for i, o in enumerate(outputs[1:], 1):\n            outputs[i] = nn.utils.rnn.pad_packed_sequence(o)[0]\n\n        # Concat hidden layers or take final\n        if self.concat_layers:\n            output = torch.cat(outputs[1:], 2)\n        else:\n            output = outputs[-1]\n\n        # Transpose and unsort\n        output = output.transpose(0, 1)\n        output = output.index_select(0, idx_unsort)\n\n        # Pad up to original batch sequence length\n        if output.size(1) != x_mask.size(1):\n            padding = torch.zeros(output.size(0),\n                                  x_mask.size(1) - output.size(1),\n                                  output.size(2)).type(output.data.type())\n            output = torch.cat([output, Variable(padding)], 1)\n\n        # Dropout on output layer\n        if self.dropout_output and self.dropout_rate > 0:\n            output = F.dropout(output,\n                               p=self.dropout_rate,\n                               training=self.training)\n        return output\n\n\nclass FeedForwardNetwork(nn.Module):\n    def __init__(self, input_size, hidden_size, output_size, dropout_rate=0):\n        super(FeedForwardNetwork, self).__init__()\n        self.dropout_rate = dropout_rate\n        self.linear1 = nn.Linear(input_size, hidden_size)\n        self.linear2 = nn.Linear(hidden_size, output_size)\n\n    def forward(self, x):\n        x_proj = F.dropout(F.relu(self.linear1(x)), p=self.dropout_rate, training=self.training)\n        x_proj = self.linear2(x_proj)\n        return x_proj\n            \n\nclass PointerNetwork(nn.Module):\n    def __init__(self, x_size, y_size, hidden_size, dropout_rate=0, cell_type=nn.GRUCell, normalize=True):\n        super(PointerNetwork, self).__init__()\n        self.normalize = normalize\n        self.hidden_size = hidden_size\n        self.dropout_rate = dropout_rate\n        self.linear = nn.Linear(x_size+y_size, hidden_size, bias=False)\n        self.weights = nn.Linear(hidden_size, 1, bias=False)\n        self.self_attn = NonLinearSeqAttn(y_size, hidden_size)\n        self.cell = cell_type(x_size, y_size)\n\n    def init_hiddens(self, y, y_mask):\n        attn = self.self_attn(y, y_mask)\n        res = attn.unsqueeze(1).bmm(y).squeeze(1) # [B, I]\n        return res\n    \n    def pointer(self, x, state, x_mask):\n        x_ = torch.cat([x, state.unsqueeze(1).repeat(1,x.size(1),1)], 2)\n        s0 = F.tanh(self.linear(x_))\n        s = self.weights(s0).view(x.size(0), x.size(1))\n        s.data.masked_fill_(x_mask.data, -float('inf'))\n        a = F.softmax(s)\n        res = a.unsqueeze(1).bmm(x).squeeze(1)\n        if self.normalize:\n            if self.training:\n                # In training we output log-softmax for NLL\n                scores = F.log_softmax(s)\n            else:\n                # ...Otherwise 0-1 probabilities\n                scores = F.softmax(s)\n        else:\n            scores = a.exp()\n        return res, scores\n\n\n    def forward(self, x, y, x_mask, y_mask):\n        hiddens = self.init_hiddens(y, y_mask)\n        c, start_scores = self.pointer(x, hiddens, x_mask)\n        c_ = F.dropout(c, p=self.dropout_rate, training=self.training)\n        hiddens = self.cell(c_, hiddens)\n        c, end_scores = self.pointer(x, hiddens, x_mask)\n        return start_scores, end_scores\n\nclass MemoryAnsPointer(nn.Module):\n    def __init__(self, x_size, y_size, hidden_size, hop=1, dropout_rate=0, normalize=True):\n        super(MemoryAnsPointer, self).__init__()\n        self.normalize = normalize\n        self.hidden_size = hidden_size\n        self.hop = hop\n        self.dropout_rate = dropout_rate\n        self.FFNs_start = nn.ModuleList()\n        self.SFUs_start = nn.ModuleList()\n        self.FFNs_end = nn.ModuleList()\n        self.SFUs_end = nn.ModuleList()\n        for i in range(self.hop):\n            self.FFNs_start.append(FeedForwardNetwork(x_size+y_size+2*hidden_size, hidden_size, 1, dropout_rate))\n            self.SFUs_start.append(SFU(y_size, 2*hidden_size))\n            self.FFNs_end.append(FeedForwardNetwork(x_size+y_size+2*hidden_size, hidden_size, 1, dropout_rate))\n            self.SFUs_end.append(SFU(y_size, 2*hidden_size))\n    \n    def forward(self, x, y, x_mask, y_mask):\n        z_s = y[:,-1,:].unsqueeze(1) # [B, 1, I]\n        z_e = None\n        s = None\n        e = None\n        p_s = None\n        p_e = None\n        \n        for i in range(self.hop):\n            z_s_ = z_s.repeat(1,x.size(1),1) # [B, S, I]\n            s = self.FFNs_start[i](torch.cat([x, z_s_, x*z_s_], 2)).squeeze(2)\n            s.data.masked_fill_(x_mask.data, -float('inf'))\n            p_s = F.softmax(s, dim=1) # [B, S]\n            u_s = p_s.unsqueeze(1).bmm(x) # [B, 1, I]\n            z_e = self.SFUs_start[i](z_s, u_s) # [B, 1, I]\n            z_e_ = z_e.repeat(1,x.size(1),1) # [B, S, I]\n            e = self.FFNs_end[i](torch.cat([x, z_e_, x*z_e_], 2)).squeeze(2)\n            e.data.masked_fill_(x_mask.data, -float('inf'))\n            p_e = F.softmax(e, dim=1) # [B, S]\n            u_e = p_e.unsqueeze(1).bmm(x) # [B, 1, I]\n            z_s = self.SFUs_end[i](z_e, u_e)\n        if self.normalize:\n            if self.training:\n                # In training we output log-softmax for NLL\n                p_s = F.log_softmax(s, dim=1) # [B, S]\n                p_e = F.log_softmax(e, dim=1) # [B, S]\n            else:\n                # ...Otherwise 0-1 probabilities\n                p_s = F.softmax(s, dim=1) # [B, S]\n                p_e = F.softmax(e, dim=1) # [B, S]\n        else:\n            p_s = s.exp()\n            p_e = e.exp()\n        return p_s, p_e\n\n\n# ------------------------------------------------------------------------------\n# Attentions\n# ------------------------------------------------------------------------------\n\nclass SeqAttnMatch(nn.Module):\n    \"\"\"Given sequences X and Y, match sequence Y to each element in X.\n\n    * o_i = sum(alpha_j * y_j) for i in X\n    * alpha_j = softmax(y_j * x_i)\n    \"\"\"\n\n    def __init__(self, input_size, identity=False):\n        super(SeqAttnMatch, self).__init__()\n        if not identity:\n            self.linear = nn.Linear(input_size, input_size)\n        else:\n            self.linear = None\n\n    def forward(self, x, y, y_mask):\n        \"\"\"\n        Args:\n            x: batch * len1 * hdim\n            y: batch * len2 * hdim\n            y_mask: batch * len2 (1 for padding, 0 for true)\n        Output:\n            matched_seq: batch * len1 * hdim\n        \"\"\"\n        # Project vectors\n        if self.linear:\n            x_proj = self.linear(x.view(-1, x.size(2))).view(x.size())\n            x_proj = F.relu(x_proj)\n            y_proj = self.linear(y.view(-1, y.size(2))).view(y.size())\n            y_proj = F.relu(y_proj)\n        else:\n            x_proj = x\n            y_proj = y\n\n        # Compute scores\n        scores = x_proj.bmm(y_proj.transpose(2, 1))\n\n        # Mask padding\n        y_mask = y_mask.unsqueeze(1).expand(scores.size())\n        scores.data.masked_fill_(y_mask.data, -float('inf'))\n\n        # Normalize with softmax\n        alpha = F.softmax(scores, dim=2)\n\n        # Take weighted average\n        matched_seq = alpha.bmm(y)\n        return matched_seq\n\nclass SelfAttnMatch(nn.Module):\n    \"\"\"Given sequences X and Y, match sequence Y to each element in X.\n\n    * o_i = sum(alpha_j * x_j) for i in X\n    * alpha_j = softmax(x_j * x_i)\n    \"\"\"\n\n    def __init__(self, input_size, identity=False, diag=True):\n        super(SelfAttnMatch, self).__init__()\n        if not identity:\n            self.linear = nn.Linear(input_size, input_size)\n        else:\n            self.linear = None\n        self.diag = diag\n\n    def forward(self, x, x_mask):\n        \"\"\"\n        Args:\n            x: batch * len1 * dim1\n            x_mask: batch * len1 (1 for padding, 0 for true)\n        Output:\n            matched_seq: batch * len1 * dim1\n        \"\"\"\n        # Project vectors\n        if self.linear:\n            x_proj = self.linear(x.view(-1, x.size(2))).view(x.size())\n            x_proj = F.relu(x_proj)\n        else:\n            x_proj = x\n\n        # Compute scores\n        scores = x_proj.bmm(x_proj.transpose(2, 1))\n        if not self.diag:\n            x_len = x.size(1)\n            for i in range(x_len):\n                scores[:, i, i] = 0\n\n        # Mask padding\n        x_mask = x_mask.unsqueeze(1).expand(scores.size())\n        scores.data.masked_fill_(x_mask.data, -float('inf'))\n\n        # Normalize with softmax\n        alpha = F.softmax(scores, dim=2)\n\n        # Take weighted average\n        matched_seq = alpha.bmm(x)\n        return matched_seq\n\n\nclass BilinearSeqAttn(nn.Module):\n    \"\"\"A bilinear attention layer over a sequence X w.r.t y:\n\n    * o_i = softmax(x_i'Wy) for x_i in X.\n\n    Optionally don't normalize output weights.\n    \"\"\"\n\n    def __init__(self, x_size, y_size, identity=False, normalize=True):\n        super(BilinearSeqAttn, self).__init__()\n        self.normalize = normalize\n\n        # If identity is true, we just use a dot product without transformation.\n        if not identity:\n            self.linear = nn.Linear(y_size, x_size)\n        else:\n            self.linear = None\n\n    def forward(self, x, y, x_mask):\n        \"\"\"\n        Args:\n            x: batch * len * hdim1\n            y: batch * hdim2\n            x_mask: batch * len (1 for padding, 0 for true)\n        Output:\n            alpha = batch * len\n        \"\"\"\n        Wy = self.linear(y) if self.linear is not None else y\n        xWy = x.bmm(Wy.unsqueeze(2)).squeeze(2)\n        xWy.data.masked_fill_(x_mask.data, -float('inf'))\n        if self.normalize:\n            if self.training:\n                # In training we output log-softmax for NLL\n                alpha = F.log_softmax(xWy)\n            else:\n                # ...Otherwise 0-1 probabilities\n                alpha = F.softmax(xWy)\n        else:\n            alpha = xWy.exp()\n        return alpha\n\n\nclass LinearSeqAttn(nn.Module):\n    \"\"\"Self attention over a sequence:\n\n    * o_i = softmax(Wx_i) for x_i in X.\n    \"\"\"\n\n    def __init__(self, input_size):\n        super(LinearSeqAttn, self).__init__()\n        self.linear = nn.Linear(input_size, 1)\n\n    def forward(self, x, x_mask):\n        \"\"\"\n        Args:\n            x: batch * len * hdim\n            x_mask: batch * len (1 for padding, 0 for true)\n        Output:\n            alpha: batch * len\n        \"\"\"\n        x_flat = x.view(-1, x.size(-1))\n        scores = self.linear(x_flat).view(x.size(0), x.size(1))\n        scores.data.masked_fill_(x_mask.data, -float('inf'))\n        alpha = F.softmax(scores)\n        return alpha\n\nclass NonLinearSeqAttn(nn.Module):\n    \"\"\"Self attention over a sequence:\n\n    * o_i = softmax(function(Wx_i)) for x_i in X.\n    \"\"\"\n\n    def __init__(self, input_size, hidden_size):\n        super(NonLinearSeqAttn, self).__init__()\n        self.FFN = FeedForwardNetwork(input_size, hidden_size, 1)\n\n    def forward(self, x, x_mask):\n        \"\"\"\n        Args:\n            x: batch * len * dim\n            x_mask: batch * len (1 for padding, 0 for true)\n        Output:\n            alpha: batch * len\n        \"\"\"\n        scores = self.FFN(x).squeeze(2)\n        scores.data.masked_fill_(x_mask.data, -float('inf'))\n        alpha = F.softmax(scores)\n        return alpha\n\n\n# ------------------------------------------------------------------------------\n# Functional Units\n# ------------------------------------------------------------------------------\n\nclass Gate(nn.Module):\n    \"\"\"Gate Unit\n    g = sigmoid(Wx)\n    x = g * x\n    \"\"\"\n    def __init__(self, input_size):\n        super(Gate, self).__init__()\n        self.linear = nn.Linear(input_size, input_size, bias=False)\n\n    def forward(self, x):\n        \"\"\"\n        Args:\n            x: batch * len * dim\n            x_mask: batch * len (1 for padding, 0 for true)\n        Output:\n            res: batch * len * dim\n        \"\"\"\n        x_proj = self.linear(x)\n        gate = F.sigmoid(x)\n        return x_proj * gate\n\n\nclass SFU(nn.Module):\n    \"\"\"Semantic Fusion Unit\n    The ouput vector is expected to not only retrieve correlative information from fusion vectors,\n    but also retain partly unchange as the input vector\n    \"\"\"\n    def __init__(self, input_size, fusion_size):\n        super(SFU, self).__init__()\n        self.linear_r = nn.Linear(input_size + fusion_size, input_size)\n        self.linear_g = nn.Linear(input_size + fusion_size, input_size)\n\n    def forward(self, x, fusions):\n        r_f = torch.cat([x, fusions], 2)\n        r = F.tanh(self.linear_r(r_f))\n        g = F.sigmoid(self.linear_g(r_f))\n        o = g * r + (1-g) * x\n        return o\n        \n\n# ------------------------------------------------------------------------------\n# Functional\n# ------------------------------------------------------------------------------\n\n\ndef uniform_weights(x, x_mask):\n    \"\"\"Return uniform weights over non-masked x (a sequence of vectors).\n\n    Args:\n        x: batch * len * hdim\n        x_mask: batch * len (1 for padding, 0 for true)\n    Output:\n        x_avg: batch * hdim\n    \"\"\"\n    alpha = Variable(torch.ones(x.size(0), x.size(1)))\n    if x.data.is_cuda:\n        alpha = alpha.cuda()\n    alpha = alpha * x_mask.eq(0).float()\n    alpha = alpha / alpha.sum(1).expand(alpha.size())\n    return alpha\n\n\ndef weighted_avg(x, weights):\n    \"\"\"Return a weighted average of x (a sequence of vectors).\n\n    Args:\n        x: batch * len * hdim\n        weights: batch * len, sum(dim = 1) = 1\n    Output:\n        x_avg: batch * hdim\n    \"\"\"\n    return weights.unsqueeze(1).bmm(x).squeeze(1)\n"
  },
  {
    "path": "m_reader.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2018-present, HKUST-KnowComp.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Implementation of the Mnemonic Reader.\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport layers\nfrom torch.autograd import Variable\n\n\n# ------------------------------------------------------------------------------\n# Network\n# ------------------------------------------------------------------------------\n\n\nclass MnemonicReader(nn.Module):\n    RNN_TYPES = {'lstm': nn.LSTM, 'gru': nn.GRU, 'rnn': nn.RNN}\n    CELL_TYPES = {'lstm': nn.LSTMCell, 'gru': nn.GRUCell, 'rnn': nn.RNNCell}\n    def __init__(self, args, normalize=True):\n        super(MnemonicReader, self).__init__()\n        # Store config\n        self.args = args\n\n        # Word embeddings (+1 for padding)\n        self.embedding = nn.Embedding(args.vocab_size,\n                                      args.embedding_dim,\n                                      padding_idx=0)\n\n        # Char embeddings (+1 for padding)\n        self.char_embedding = nn.Embedding(args.char_size,\n                                      args.char_embedding_dim,\n                                      padding_idx=0)\n\n        # Char rnn to generate char features\n        self.char_rnn = layers.StackedBRNN(\n            input_size=args.char_embedding_dim,\n            hidden_size=args.char_hidden_size,\n            num_layers=1,\n            dropout_rate=args.dropout_rnn,\n            dropout_output=args.dropout_rnn_output,\n            concat_layers=False,\n            rnn_type=self.RNN_TYPES[args.rnn_type],\n            padding=False,\n        )\n\n        doc_input_size = args.embedding_dim + args.char_hidden_size * 2 + args.num_features\n\n        # Encoder\n        self.encoding_rnn = layers.StackedBRNN(\n            input_size=doc_input_size,\n            hidden_size=args.hidden_size,\n            num_layers=1,\n            dropout_rate=args.dropout_rnn,\n            dropout_output=args.dropout_rnn_output,\n            concat_layers=False,\n            rnn_type=self.RNN_TYPES[args.rnn_type],\n            padding=args.rnn_padding,\n        )\n\n        doc_hidden_size = 2 * args.hidden_size\n        \n        # Interactive aligning, self aligning and aggregating\n        self.interactive_aligners = nn.ModuleList()\n        self.interactive_SFUs = nn.ModuleList()\n        self.self_aligners = nn.ModuleList()\n        self.self_SFUs = nn.ModuleList()\n        self.aggregate_rnns = nn.ModuleList()\n        for i in range(args.hop):\n            # interactive aligner\n            self.interactive_aligners.append(layers.SeqAttnMatch(doc_hidden_size, identity=True))\n            self.interactive_SFUs.append(layers.SFU(doc_hidden_size, 3 * doc_hidden_size))\n            # self aligner\n            self.self_aligners.append(layers.SelfAttnMatch(doc_hidden_size, identity=True, diag=False))\n            self.self_SFUs.append(layers.SFU(doc_hidden_size, 3 * doc_hidden_size))\n            # aggregating\n            self.aggregate_rnns.append(\n                layers.StackedBRNN(\n                    input_size=doc_hidden_size,\n                    hidden_size=args.hidden_size,\n                    num_layers=1,\n                    dropout_rate=args.dropout_rnn,\n                    dropout_output=args.dropout_rnn_output,\n                    concat_layers=False,\n                    rnn_type=self.RNN_TYPES[args.rnn_type],\n                    padding=args.rnn_padding,\n                )\n            )\n\n        # Memmory-based Answer Pointer\n        self.mem_ans_ptr = layers.MemoryAnsPointer(\n            x_size=2*args.hidden_size, \n            y_size=2*args.hidden_size, \n            hidden_size=args.hidden_size, \n            hop=args.hop,\n            dropout_rate=args.dropout_rnn,\n            normalize=normalize\n        )\n        \n\n    def forward(self, x1, x1_c, x1_f, x1_mask, x2, x2_c, x2_f, x2_mask):\n        \"\"\"Inputs:\n        x1 = document word indices             [batch * len_d]\n        x1_c = document char indices           [batch * len_d]\n        x1_f = document word features indices  [batch * len_d * nfeat]\n        x1_mask = document padding mask        [batch * len_d]\n        x2 = question word indices             [batch * len_q]\n        x2_c = document char indices           [batch * len_d]\n        x1_f = document word features indices  [batch * len_d * nfeat]\n        x2_mask = question padding mask        [batch * len_q]\n        \"\"\"\n        # Embed both document and question\n        x1_emb = self.embedding(x1)\n        x2_emb = self.embedding(x2)\n        x1_c_emb = self.char_embedding(x1_c)\n        x2_c_emb = self.char_embedding(x2_c)\n\n        # Dropout on embeddings\n        if self.args.dropout_emb > 0:\n            x1_emb = F.dropout(x1_emb, p=self.args.dropout_emb, training=self.training)\n            x2_emb = F.dropout(x2_emb, p=self.args.dropout_emb, training=self.training)\n            x1_c_emb = F.dropout(x1_c_emb, p=self.args.dropout_emb, training=self.training)\n            x2_c_emb = F.dropout(x2_c_emb, p=self.args.dropout_emb, training=self.training)\n\n        # Generate char features\n        x1_c_features = self.char_rnn(\n            x1_c_emb.reshape((x1_c_emb.size(0) * x1_c_emb.size(1), x1_c_emb.size(2), x1_c_emb.size(3))), \n            x1_mask.unsqueeze(2).repeat(1, 1, x1_c_emb.size(2)).reshape((x1_c_emb.size(0) * x1_c_emb.size(1), x1_c_emb.size(2)))\n            ).reshape((x1_c_emb.size(0), x1_c_emb.size(1), x1_c_emb.size(2), -1))[:,:,-1,:]\n        x2_c_features = self.char_rnn(\n            x2_c_emb.reshape((x2_c_emb.size(0) * x2_c_emb.size(1), x2_c_emb.size(2), x2_c_emb.size(3))), \n            x2_mask.unsqueeze(2).repeat(1, 1, x2_c_emb.size(2)).reshape((x2_c_emb.size(0) * x2_c_emb.size(1), x2_c_emb.size(2)))\n            ).reshape((x2_c_emb.size(0), x2_c_emb.size(1), x2_c_emb.size(2), -1))[:,:,-1,:] \n\n        # Combine input\n        crnn_input = [x1_emb, x1_c_features]\n        qrnn_input = [x2_emb, x2_c_features]\n        # Add manual features\n        if self.args.num_features > 0:\n            crnn_input.append(x1_f)\n            qrnn_input.append(x2_f)\n\n        # Encode document with RNN\n        c = self.encoding_rnn(torch.cat(crnn_input, 2), x1_mask)\n        \n        # Encode question with RNN\n        q = self.encoding_rnn(torch.cat(qrnn_input, 2), x2_mask)\n\n        # Align and aggregate\n        c_check = c\n        for i in range(self.args.hop):\n            q_tilde = self.interactive_aligners[i].forward(c_check, q, x2_mask)\n            c_bar = self.interactive_SFUs[i].forward(c_check, torch.cat([q_tilde, c_check * q_tilde, c_check - q_tilde], 2))\n            c_tilde = self.self_aligners[i].forward(c_bar, x1_mask)\n            c_hat = self.self_SFUs[i].forward(c_bar, torch.cat([c_tilde, c_bar * c_tilde, c_bar - c_tilde], 2))\n            c_check = self.aggregate_rnns[i].forward(c_hat, x1_mask)\n\n        # Predict\n        start_scores, end_scores = self.mem_ans_ptr.forward(c_check, q, x1_mask, x2_mask)\n        \n        return start_scores, end_scores\n"
  },
  {
    "path": "model.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2018-present, HKUST-KnowComp.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Document Reader model\"\"\"\n\nimport torch\nimport torch.optim as optim\nimport torch.nn.functional as F\nimport numpy as np\nimport logging\nimport copy\n\nfrom torch.autograd import Variable\nfrom config import override_model_args\nfrom r_net import R_Net\nfrom rnn_reader import RnnDocReader\nfrom m_reader import MnemonicReader\nfrom data import Dictionary\n\nlogger = logging.getLogger(__name__)\n\n\nclass DocReader(object):\n    \"\"\"High level model that handles intializing the underlying network\n    architecture, saving, updating examples, and predicting examples.\n    \"\"\"\n\n    # --------------------------------------------------------------------------\n    # Initialization\n    # --------------------------------------------------------------------------\n\n    def __init__(self, args, word_dict, char_dict, feature_dict,\n                 state_dict=None, normalize=True):\n        # Book-keeping.\n        self.args = args\n        self.word_dict = word_dict\n        self.char_dict = char_dict\n        self.args.vocab_size = len(word_dict)\n        self.args.char_size = len(char_dict)\n        self.feature_dict = feature_dict\n        self.args.num_features = len(feature_dict)\n        self.updates = 0\n        self.use_cuda = False\n        self.parallel = False\n\n        # Building network. If normalize if false, scores are not normalized\n        # 0-1 per paragraph (no softmax).\n        if args.model_type == 'rnn':\n            self.network = RnnDocReader(args, normalize)\n        elif args.model_type == 'r_net':\n            self.network = R_Net(args, normalize)\n        elif args.model_type == 'mnemonic':\n            self.network = MnemonicReader(args, normalize)\n        else:\n            raise RuntimeError('Unsupported model: %s' % args.model_type)\n\n        # Load saved state\n        if state_dict:\n            # Load buffer separately\n            if 'fixed_embedding' in state_dict:\n                fixed_embedding = state_dict.pop('fixed_embedding')\n                self.network.load_state_dict(state_dict)\n                self.network.register_buffer('fixed_embedding', fixed_embedding)\n            else:\n                self.network.load_state_dict(state_dict)\n\n    def expand_dictionary(self, words):\n        \"\"\"Add words to the DocReader dictionary if they do not exist. The\n        underlying embedding matrix is also expanded (with random embeddings).\n\n        Args:\n            words: iterable of tokens to add to the dictionary.\n        Output:\n            added: set of tokens that were added.\n        \"\"\"\n        to_add = {self.word_dict.normalize(w) for w in words\n                  if w not in self.word_dict}\n\n        # Add words to dictionary and expand embedding layer\n        if len(to_add) > 0:\n            logger.info('Adding %d new words to dictionary...' % len(to_add))\n            for w in to_add:\n                self.word_dict.add(w)\n            self.args.vocab_size = len(self.word_dict)\n            logger.info('New vocab size: %d' % len(self.word_dict))\n\n            old_embedding = self.network.embedding.weight.data\n            self.network.embedding = torch.nn.Embedding(self.args.vocab_size,\n                                                        self.args.embedding_dim,\n                                                        padding_idx=0)\n            new_embedding = self.network.embedding.weight.data\n            new_embedding[:old_embedding.size(0)] = old_embedding\n\n        # Return added words\n        return to_add\n\n\n    def expand_char_dictionary(self, chars):\n        \"\"\"Add chars to the DocReader dictionary if they do not exist. The\n        underlying embedding matrix is also expanded (with random embeddings).\n\n        Args:\n            chars: iterable of tokens to add to the dictionary.\n        Output:\n            added: set of tokens that were added.\n        \"\"\"\n        to_add = {self.char_dict.normalize(w) for w in chars\n                  if w not in self.char_dict}\n\n        # Add chars to dictionary and expand embedding layer\n        if len(to_add) > 0:\n            logger.info('Adding %d new chars to dictionary...' % len(to_add))\n            for w in to_add:\n                self.char_dict.add(w)\n            self.args.char_size = len(self.char_dict)\n            logger.info('New char size: %d' % len(self.char_dict))\n\n            old_char_embedding = self.network.char_embedding.weight.data\n            self.network.char_embedding = torch.nn.Embedding(self.args.char_size,\n                                                        self.args.char_embedding_dim,\n                                                        padding_idx=0)\n            new_char_embedding = self.network.char_embedding.weight.data\n            new_char_embedding[:old_char_embedding.size(0)] = old_char_embedding\n\n        # Return added chars\n        return to_add\n\n    def load_embeddings(self, words, embedding_file):\n        \"\"\"Load pretrained embeddings for a given list of words, if they exist.\n\n        Args:\n            words: iterable of tokens. Only those that are indexed in the\n              dictionary are kept.\n            embedding_file: path to text file of embeddings, space separated.\n        \"\"\"\n        words = {w for w in words if w in self.word_dict}\n        logger.info('Loading pre-trained embeddings for %d words from %s' %\n                    (len(words), embedding_file))\n        embedding = self.network.embedding.weight.data\n\n        # When normalized, some words are duplicated. (Average the embeddings).\n        vec_counts = {}\n        with open(embedding_file) as f:\n            for line in f:\n                parsed = line.rstrip().split(' ')\n                assert(len(parsed) == embedding.size(1) + 1)\n                w = self.word_dict.normalize(parsed[0])\n                if w in words:\n                    vec = torch.Tensor([float(i) for i in parsed[1:]])\n                    if w not in vec_counts:\n                        vec_counts[w] = 1\n                        embedding[self.word_dict[w]].copy_(vec)\n                    else:\n                        logging.warning(\n                            'WARN: Duplicate embedding found for %s' % w\n                        )\n                        vec_counts[w] = vec_counts[w] + 1\n                        embedding[self.word_dict[w]].add_(vec)\n\n        for w, c in vec_counts.items():\n            embedding[self.word_dict[w]].div_(c)\n\n        logger.info('Loaded %d embeddings (%.2f%%)' %\n                    (len(vec_counts), 100 * len(vec_counts) / len(words)))\n\n    def load_char_embeddings(self, chars, char_embedding_file):\n        \"\"\"Load pretrained embeddings for a given list of chars, if they exist.\n\n        Args:\n            chars: iterable of tokens. Only those that are indexed in the\n              dictionary are kept.\n            char_embedding_file: path to text file of embeddings, space separated.\n        \"\"\"\n        chars = {w for w in chars if w in self.char_dict}\n        logger.info('Loading pre-trained embeddings for %d chars from %s' %\n                    (len(chars), char_embedding_file))\n        char_embedding = self.network.char_embedding.weight.data\n\n        # When normalized, some chars are duplicated. (Average the embeddings).\n        vec_counts = {}\n        with open(char_embedding_file) as f:\n            for line in f:\n                parsed = line.rstrip().split(' ')\n                assert(len(parsed) == char_embedding.size(1) + 1)\n                w = self.char_dict.normalize(parsed[0])\n                if w in chars:\n                    vec = torch.Tensor([float(i) for i in parsed[1:]])\n                    if w not in vec_counts:\n                        vec_counts[w] = 1\n                        char_embedding[self.char_dict[w]].copy_(vec)\n                    else:\n                        logging.warning(\n                            'WARN: Duplicate char embedding found for %s' % w\n                        )\n                        vec_counts[w] = vec_counts[w] + 1\n                        char_embedding[self.char_dict[w]].add_(vec)\n\n        for w, c in vec_counts.items():\n            char_embedding[self.char_dict[w]].div_(c)\n\n        logger.info('Loaded %d char embeddings (%.2f%%)' %\n                    (len(vec_counts), 100 * len(vec_counts) / len(chars)))\n\n    def tune_embeddings(self, words):\n        \"\"\"Unfix the embeddings of a list of words. This is only relevant if\n        only some of the embeddings are being tuned (tune_partial = N).\n\n        Shuffles the N specified words to the front of the dictionary, and saves\n        the original vectors of the other N + 1:vocab words in a fixed buffer.\n\n        Args:\n            words: iterable of tokens contained in dictionary.\n        \"\"\"\n        words = {w for w in words if w in self.word_dict}\n\n        if len(words) == 0:\n            logger.warning('Tried to tune embeddings, but no words given!')\n            return\n\n        if len(words) == len(self.word_dict):\n            logger.warning('Tuning ALL embeddings in dictionary')\n            return\n\n        # Shuffle words and vectors\n        embedding = self.network.embedding.weight.data\n        for idx, swap_word in enumerate(words, self.word_dict.START):\n            # Get current word + embedding for this index\n            curr_word = self.word_dict[idx]\n            curr_emb = embedding[idx].clone()\n            old_idx = self.word_dict[swap_word]\n\n            # Swap embeddings + dictionary indices\n            embedding[idx].copy_(embedding[old_idx])\n            embedding[old_idx].copy_(curr_emb)\n            self.word_dict[swap_word] = idx\n            self.word_dict[idx] = swap_word\n            self.word_dict[curr_word] = old_idx\n            self.word_dict[old_idx] = curr_word\n\n        # Save the original, fixed embeddings\n        self.network.register_buffer(\n            'fixed_embedding', embedding[idx + 1:].clone()\n        )\n\n    def init_optimizer(self, state_dict=None):\n        \"\"\"Initialize an optimizer for the free parameters of the network.\n\n        Args:\n            state_dict: network parameters\n        \"\"\"\n        if self.args.fix_embeddings:\n            for p in self.network.embedding.parameters():\n                p.requires_grad = False\n        parameters = [p for p in self.network.parameters() if p.requires_grad]\n        if self.args.optimizer == 'sgd':\n            self.optimizer = optim.SGD(parameters, lr=self.args.learning_rate,\n                                       momentum=self.args.momentum,\n                                       weight_decay=self.args.weight_decay)\n        elif self.args.optimizer == 'adamax':\n            self.optimizer = optim.Adamax(parameters,\n                                          weight_decay=self.args.weight_decay)\n        elif self.args.optimizer == 'adadelta':\n            self.optimizer = optim.Adadelta(parameters, lr=self.args.learning_rate,\n                                            rho=self.args.rho, eps=self.args.eps,\n                                            weight_decay=self.args.weight_decay)\n        else:\n            raise RuntimeError('Unsupported optimizer: %s' %\n                               self.args.optimizer)\n\n    # --------------------------------------------------------------------------\n    # Learning\n    # --------------------------------------------------------------------------\n\n    def update(self, ex):\n        \"\"\"Forward a batch of examples; step the optimizer to update weights.\"\"\"\n        if not self.optimizer:\n            raise RuntimeError('No optimizer set.')\n\n        # Train mode\n        self.network.train()\n        \n        # Transfer to GPU\n        if self.use_cuda:\n            inputs = [e if e is None else Variable(e.cuda(async=True)) for e in ex[:-3]]\n            target_s = Variable(ex[-3].cuda(async=True))\n            target_e = Variable(ex[-2].cuda(async=True))\n        else:\n            inputs = [e if e is None else Variable(e) for e in ex[:-3]]\n            target_s = Variable(ex[-3])\n            target_e = Variable(ex[-2])\n        \n        # Run forward\n        score_s, score_e = self.network(*inputs)\n\n        # Compute loss and accuracies\n        loss = F.nll_loss(score_s, target_s) + F.nll_loss(score_e, target_e)\n\n        # Clear gradients and run backward\n        self.optimizer.zero_grad()\n        loss.backward()\n\n        # Clip gradients\n        torch.nn.utils.clip_grad_norm(self.network.parameters(),\n                                      self.args.grad_clipping)\n\n        # Update parameters\n        self.optimizer.step()\n        self.updates += 1\n\n        # Reset any partially fixed parameters (e.g. rare words)\n        self.reset_parameters()\n\n        return loss.data[0], ex[0].size(0)\n\n    def reset_parameters(self):\n        \"\"\"Reset any partially fixed parameters to original states.\"\"\"\n\n        # Reset fixed embeddings to original value\n        if self.args.tune_partial > 0:\n            # Embeddings to fix are indexed after the special + N tuned words\n            offset = self.args.tune_partial + self.word_dict.START\n            if self.parallel:\n                embedding = self.network.module.embedding.weight.data\n                fixed_embedding = self.network.module.fixed_embedding\n            else:\n                embedding = self.network.embedding.weight.data\n                fixed_embedding = self.network.fixed_embedding\n            if offset < embedding.size(0):\n                embedding[offset:] = fixed_embedding\n\n    # --------------------------------------------------------------------------\n    # Prediction\n    # --------------------------------------------------------------------------\n\n    def predict(self, ex, candidates=None, top_n=1, async_pool=None):\n        \"\"\"Forward a batch of examples only to get predictions.\n\n        Args:\n            ex: the batch\n            candidates: batch * variable length list of string answer options.\n              The model will only consider exact spans contained in this list.\n            top_n: Number of predictions to return per batch element.\n            async_pool: If provided, non-gpu post-processing will be offloaded\n              to this CPU process pool.\n        Output:\n            pred_s: batch * top_n predicted start indices\n            pred_e: batch * top_n predicted end indices\n            pred_score: batch * top_n prediction scores\n\n        If async_pool is given, these will be AsyncResult handles.\n        \"\"\"\n        # Eval mode\n        self.network.eval()\n\n        # Transfer to GPU\n        if self.use_cuda:\n            inputs = [e if e is None else\n                      Variable(e.cuda(async=True), volatile=True)\n                      for e in ex[:8]]\n        else:\n            inputs = [e if e is None else Variable(e, volatile=True)\n                      for e in ex[:8]]\n\n        # Run forward\n        score_s, score_e = self.network(*inputs)\n        del inputs\n\n        # Decode predictions\n        score_s = score_s.data.cpu()\n        score_e = score_e.data.cpu()\n\n        if candidates:\n            args = (score_s, score_e, candidates, top_n, self.args.max_len)\n            if async_pool:\n                return async_pool.apply_async(self.decode_candidates, args)\n            else:\n                return self.decode_candidates(*args)\n        else:\n            args = (score_s, score_e, top_n, self.args.max_len)\n            if async_pool:\n                return async_pool.apply_async(self.decode, args)\n            else:\n                return self.decode(*args)\n\n    @staticmethod\n    def decode(score_s, score_e, top_n=1, max_len=None):\n        \"\"\"Take argmax of constrained score_s * score_e.\n\n        Args:\n            score_s: independent start predictions\n            score_e: independent end predictions\n            top_n: number of top scored pairs to take\n            max_len: max span length to consider\n        \"\"\"\n        pred_s = []\n        pred_e = []\n        pred_score = []\n        max_len = max_len or score_s.size(1)\n        for i in range(score_s.size(0)):\n            # Outer product of scores to get full p_s * p_e matrix\n            scores = torch.ger(score_s[i], score_e[i])\n\n            # Zero out negative length and over-length span scores\n            scores.triu_().tril_(max_len - 1)\n\n            # Take argmax or top n\n            scores = scores.numpy()\n            scores_flat = scores.flatten()\n            if top_n == 1:\n                idx_sort = [np.argmax(scores_flat)]\n            elif len(scores_flat) < top_n:\n                idx_sort = np.argsort(-scores_flat)\n            else:\n                idx = np.argpartition(-scores_flat, top_n)[0:top_n]\n                idx_sort = idx[np.argsort(-scores_flat[idx])]\n            s_idx, e_idx = np.unravel_index(idx_sort, scores.shape)\n            pred_s.append(s_idx)\n            pred_e.append(e_idx)\n            pred_score.append(scores_flat[idx_sort])\n        del score_s, score_e\n        return pred_s, pred_e, pred_score\n\n    @staticmethod\n    def decode_candidates(score_s, score_e, candidates, top_n=1, max_len=None):\n        \"\"\"Take argmax of constrained score_s * score_e. Except only consider\n        spans that are in the candidates list.\n        \"\"\"\n        pred_s = []\n        pred_e = []\n        pred_score = []\n        for i in range(score_s.size(0)):\n            # Extract original tokens stored with candidates\n            tokens = candidates[i]['input']\n            cands = candidates[i]['cands']\n\n            if not cands:\n                # try getting from globals? (multiprocessing in pipeline mode)\n                from ..pipeline.wrmcqa import PROCESS_CANDS\n                cands = PROCESS_CANDS\n            if not cands:\n                raise RuntimeError('No candidates given.')\n\n            # Score all valid candidates found in text.\n            # Brute force get all ngrams and compare against the candidate list.\n            max_len = max_len or len(tokens)\n            scores, s_idx, e_idx = [], [], []\n            for s, e in tokens.ngrams(n=max_len, as_strings=False):\n                span = tokens.slice(s, e).untokenize()\n                if span in cands or span.lower() in cands:\n                    # Match! Record its score.\n                    scores.append(score_s[i][s] * score_e[i][e - 1])\n                    s_idx.append(s)\n                    e_idx.append(e - 1)\n\n            if len(scores) == 0:\n                # No candidates present\n                pred_s.append([])\n                pred_e.append([])\n                pred_score.append([])\n            else:\n                # Rank found candidates\n                scores = np.array(scores)\n                s_idx = np.array(s_idx)\n                e_idx = np.array(e_idx)\n\n                idx_sort = np.argsort(-scores)[0:top_n]\n                pred_s.append(s_idx[idx_sort])\n                pred_e.append(e_idx[idx_sort])\n                pred_score.append(scores[idx_sort])\n        del score_s, score_e, candidates\n        return pred_s, pred_e, pred_score\n\n    # --------------------------------------------------------------------------\n    # Saving and loading\n    # --------------------------------------------------------------------------\n\n    def save(self, filename):\n        state_dict = copy.copy(self.network.state_dict())\n        if 'fixed_embedding' in state_dict:\n            state_dict.pop('fixed_embedding')\n        params = {\n            'state_dict': state_dict,\n            'word_dict': self.word_dict,\n            'char_dict': self.char_dict,\n            'feature_dict': self.feature_dict,\n            'args': self.args,\n        }\n        try:\n            torch.save(params, filename)\n        except BaseException:\n            logger.warning('WARN: Saving failed... continuing anyway.')\n\n    def checkpoint(self, filename, epoch):\n        params = {\n            'state_dict': self.network.state_dict(),\n            'word_dict': self.word_dict,\n            'char_dict': self.char_dict,\n            'feature_dict': self.feature_dict,\n            'args': self.args,\n            'epoch': epoch,\n            'optimizer': self.optimizer.state_dict(),\n        }\n        try:\n            torch.save(params, filename)\n        except BaseException:\n            logger.warning('WARN: Saving failed... continuing anyway.')\n\n    @staticmethod\n    def load(filename, new_args=None, normalize=True):\n        logger.info('Loading model %s' % filename)\n        saved_params = torch.load(\n            filename, map_location=lambda storage, loc: storage\n        )\n        word_dict = saved_params['word_dict']\n        try:\n            char_dict = saved_params['char_dict']\n        except KeyError as e:\n            char_dict = Dictionary()\n\n        feature_dict = saved_params['feature_dict']\n        state_dict = saved_params['state_dict']\n        args = saved_params['args']\n        if new_args:\n            args = override_model_args(args, new_args)\n        return DocReader(args, word_dict, char_dict, feature_dict, state_dict, normalize)\n\n    @staticmethod\n    def load_checkpoint(filename, normalize=True):\n        logger.info('Loading model %s' % filename)\n        saved_params = torch.load(\n            filename, map_location=lambda storage, loc: storage\n        )\n        word_dict = saved_params['word_dict']\n        char_dict = saved_params['char_dict']\n        feature_dict = saved_params['feature_dict']\n        state_dict = saved_params['state_dict']\n        epoch = saved_params['epoch']\n        optimizer = saved_params['optimizer']\n        args = saved_params['args']\n        model = DocReader(args, word_dict, char_dict, feature_dict, state_dict, normalize)\n        model.init_optimizer(optimizer)\n        return model, epoch\n\n    # --------------------------------------------------------------------------\n    # Runtime\n    # --------------------------------------------------------------------------\n\n    def cuda(self):\n        self.use_cuda = True\n        self.network = self.network.cuda()\n\n    def cpu(self):\n        self.use_cuda = False\n        self.network = self.network.cpu()\n\n    def parallelize(self):\n        \"\"\"Use data parallel to copy the model across several gpus.\n        This will take all gpus visible with CUDA_VISIBLE_DEVICES.\n        \"\"\"\n        self.parallel = True\n        self.network = torch.nn.DataParallel(self.network)\n"
  },
  {
    "path": "predictor.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2018-present, HKUST-KnowComp.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Machine Comprehension predictor\"\"\"\n\nimport logging\n\nfrom multiprocessing import Pool as ProcessPool\nfrom multiprocessing.util import Finalize\n\nfrom vector import vectorize, batchify\nfrom model import DocReader\nimport utils\nfrom spacy_tokenizer import SpacyTokenizer\n\nlogger = logging.getLogger(__name__)\n\n\n# ------------------------------------------------------------------------------\n# Tokenize + annotate\n# ------------------------------------------------------------------------------\n\nTOK = None\n\ndef init(options):\n    global TOK\n    TOK = SpacyTokenizer(**options)\n    Finalize(TOK, TOK.shutdown, exitpriority=100)\n\n\ndef tokenize(text):\n    global TOK\n    return TOK.tokenize(text)\n\ndef get_annotators_for_model(model):\n    annotators = set()\n    if model.args.use_pos:\n        annotators.add('pos')\n    if model.args.use_lemma:\n        annotators.add('lemma')\n    if model.args.use_ner:\n        annotators.add('ner')\n    return annotators\n\n\n# ------------------------------------------------------------------------------\n# Predictor class.\n# ------------------------------------------------------------------------------\n\n\nclass Predictor(object):\n    \"\"\"Load a pretrained DocReader model and predict inputs on the fly.\"\"\"\n\n    def __init__(self, model, normalize=True,\n                 embedding_file=None, char_embedding_file=None, num_workers=None):\n        \"\"\"\n        Args:\n            model: path to saved model file.\n            normalize: squash output score to 0-1 probabilities with a softmax.\n            embedding_file: if provided, will expand dictionary to use all\n              available pretrained vectors in this file.\n            num_workers: number of CPU processes to use to preprocess batches.\n        \"\"\"\n        logger.info('Initializing model...')\n        self.model = DocReader.load(model, normalize=normalize)\n\n        if embedding_file:\n            logger.info('Expanding dictionary...')\n            utils.index_embedding_words(embedding_file)\n            added_words = self.model.expand_dictionary(words)\n            self.model.load_embeddings(added_words, embedding_file)\n        if char_embedding_file:\n            logger.info('Expanding dictionary...')\n            chars = utils.index_embedding_chars(char_embedding_file)\n            added_chars = self.model.expand_char_dictionary(chars)\n            self.model.load_char_embeddings(added_chars, char_embedding_file)\n\n        logger.info('Initializing tokenizer...')\n        annotators = get_annotators_for_model(self.model)\n\n        if num_workers is None or num_workers > 0:\n            self.workers = ProcessPool(\n                num_workers,\n                initializer=init,\n                initargs=({'annotators': annotators},),\n            )\n        else:\n            self.workers = None\n            self.tokenizer = SpacyTokenizer(annotators=annotators)\n\n    def predict(self, document, question, candidates=None, top_n=1):\n        \"\"\"Predict a single document - question pair.\"\"\"\n        results = self.predict_batch([(document, question, candidates,)], top_n)\n        return results[0]\n\n    def predict_batch(self, batch, top_n=1):\n        \"\"\"Predict a batch of document - question pairs.\"\"\"\n        documents, questions, candidates = [], [], []\n        for b in batch:\n            documents.append(b[0])\n            questions.append(b[1])\n            candidates.append(b[2] if len(b) == 3 else None)\n        candidates = candidates if any(candidates) else None\n\n        # Tokenize the inputs, perhaps multi-processed.\n        if self.workers:\n            q_tokens = self.workers.map_async(tokenize, questions)\n            c_tokens = self.workers.map_async(tokenize, documents)\n            q_tokens = list(q_tokens.get())\n            c_tokens = list(c_tokens.get())\n        else:\n            q_tokens = list(map(self.tokenizer.tokenize, questions))\n            c_tokens = list(map(self.tokenizer.tokenize, documents))\n\n        examples = []\n        for i in range(len(questions)):\n            examples.append({\n                'id': i,\n                'question': q_tokens[i].words(),\n                'question_char': q_tokens[i].chars(),\n                'qlemma': q_tokens[i].lemmas(),\n                'qpos': q_tokens[i].pos(),\n                'qner': q_tokens[i].entities(),\n                'document': c_tokens[i].words(),\n                'document_char': c_tokens[i].chars(),\n                'clemma': c_tokens[i].lemmas(),\n                'cpos': c_tokens[i].pos(),\n                'cner': c_tokens[i].entities(),\n            })\n\n        # Stick document tokens in candidates for decoding\n        if candidates:\n            candidates = [{'input': c_tokens[i], 'cands': candidates[i]}\n                          for i in range(len(candidates))]\n\n        # Build the batch and run it through the model\n        batch_exs = batchify([vectorize(e, self.model) for e in examples])\n        s, e, score = self.model.predict(batch_exs, candidates, top_n)\n\n        # Retrieve the predicted spans\n        results = []\n        for i in range(len(s)):\n            predictions = []\n            for j in range(len(s[i])):\n                span = c_tokens[i].slice(s[i][j], e[i][j] + 1).untokenize()\n                predictions.append((span, score[i][j]))\n            results.append(predictions)\n        return results\n\n    def cuda(self):\n        self.model.cuda()\n\n    def cpu(self):\n        self.model.cpu()\n"
  },
  {
    "path": "r_net.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2018-present, HKUST-KnowComp.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Implementation of the R-Net based reader.\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport layers\nfrom torch.autograd import Variable\n\n\n# ------------------------------------------------------------------------------\n# Network\n# ------------------------------------------------------------------------------\n\n\nclass R_Net(nn.Module):\n    RNN_TYPES = {'lstm': nn.LSTM, 'gru': nn.GRU, 'rnn': nn.RNN}\n    CELL_TYPES = {'lstm': nn.LSTMCell, 'gru': nn.GRUCell, 'rnn': nn.RNNCell}\n    def __init__(self, args, normalize=True):\n        super(R_Net, self).__init__()\n        # Store config\n        self.args = args\n\n        # Word embeddings (+1 for padding)\n        self.embedding = nn.Embedding(args.vocab_size,\n                                      args.embedding_dim,\n                                      padding_idx=0)\n\n        # Char embeddings (+1 for padding)\n        self.char_embedding = nn.Embedding(args.char_size,\n                                      args.char_embedding_dim,\n                                      padding_idx=0)\n\n        # Char rnn to generate char features\n        self.char_rnn = layers.StackedBRNN(\n            input_size=args.char_embedding_dim,\n            hidden_size=args.char_hidden_size,\n            num_layers=1,\n            dropout_rate=args.dropout_rnn,\n            dropout_output=args.dropout_rnn_output,\n            concat_layers=False,\n            rnn_type=self.RNN_TYPES[args.rnn_type],\n            padding=False,\n        )\n\n        doc_input_size = args.embedding_dim + args.char_hidden_size * 2\n\n        # Encoder\n        self.encode_rnn = layers.StackedBRNN(\n            input_size=doc_input_size,\n            hidden_size=args.hidden_size,\n            num_layers=args.doc_layers,\n            dropout_rate=args.dropout_rnn,\n            dropout_output=args.dropout_rnn_output,\n            concat_layers=args.concat_rnn_layers,\n            rnn_type=self.RNN_TYPES[args.rnn_type],\n            padding=args.rnn_padding,\n        )\n\n        # Output sizes of rnn encoder\n        doc_hidden_size = 2 * args.hidden_size\n        question_hidden_size = 2 * args.hidden_size\n        if args.concat_rnn_layers:\n            doc_hidden_size *= args.doc_layers\n            question_hidden_size *= args.question_layers\n        \n        # Gated-attention-based RNN of the whole question\n        self.question_attn = layers.SeqAttnMatch(question_hidden_size, identity=False)\n        self.question_attn_gate = layers.Gate(doc_hidden_size + question_hidden_size)\n        self.question_attn_rnn = layers.StackedBRNN(\n            input_size=doc_hidden_size + question_hidden_size,\n            hidden_size=args.hidden_size,\n            num_layers=1,\n            dropout_rate=args.dropout_rnn,\n            dropout_output=args.dropout_rnn_output,\n            concat_layers=False,\n            rnn_type=self.RNN_TYPES[args.rnn_type],\n            padding=args.rnn_padding,\n        )\n\n        question_attn_hidden_size = 2 * args.hidden_size\n\n        # Self-matching-attention-baed RNN of the whole doc\n        self.doc_self_attn = layers.SelfAttnMatch(question_attn_hidden_size, identity=False)\n        self.doc_self_attn_gate = layers.Gate(question_attn_hidden_size + question_attn_hidden_size)\n        self.doc_self_attn_rnn = layers.StackedBRNN(\n            input_size=question_attn_hidden_size + question_attn_hidden_size,\n            hidden_size=args.hidden_size,\n            num_layers=1,\n            dropout_rate=args.dropout_rnn,\n            dropout_output=args.dropout_rnn_output,\n            concat_layers=False,\n            rnn_type=self.RNN_TYPES[args.rnn_type],\n            padding=args.rnn_padding,\n        )\n\n        doc_self_attn_hidden_size = 2 * args.hidden_size\n\n        self.doc_self_attn_rnn2 = layers.StackedBRNN(\n            input_size=doc_self_attn_hidden_size,\n            hidden_size=args.hidden_size,\n            num_layers=1,\n            dropout_rate=args.dropout_rnn,\n            dropout_output=args.dropout_rnn_output,\n            concat_layers=False,\n            rnn_type=self.RNN_TYPES[args.rnn_type],\n            padding=args.rnn_padding,\n        )\n\n        self.ptr_net = layers.PointerNetwork(\n            x_size = doc_self_attn_hidden_size, \n            y_size = question_hidden_size, \n            hidden_size = args.hidden_size, \n            dropout_rate=args.dropout_rnn,\n            cell_type=nn.GRUCell,\n            normalize=normalize\n        )\n\n    def forward(self, x1, x1_c, x1_f, x1_mask, x2, x2_c, x2_f, x2_mask):\n        \"\"\"Inputs:\n        x1 = document word indices             [batch * len_d]\n        x1_c = document char indices           [batch * len_d]\n        x1_f = document word features indices  [batch * len_d * nfeat]\n        x1_mask = document padding mask        [batch * len_d]\n        x2 = question word indices             [batch * len_q]\n        x2_c = document char indices           [batch * len_d]\n        x1_f = document word features indices  [batch * len_d * nfeat]\n        x2_mask = question padding mask        [batch * len_q]\n        \"\"\"\n        # Embed both document and question\n        x1_emb = self.embedding(x1)\n        x2_emb = self.embedding(x2)\n        x1_c_emb = self.char_embedding(x1_c)\n        x2_c_emb = self.char_embedding(x2_c)\n\n        # Dropout on embeddings\n        if self.args.dropout_emb > 0:\n            x1_emb = F.dropout(x1_emb, p=self.args.dropout_emb, training=self.training)\n            x2_emb = F.dropout(x2_emb, p=self.args.dropout_emb, training=self.training)\n            x1_c_emb = F.dropout(x1_c_emb, p=self.args.dropout_emb, training=self.training)\n            x2_c_emb = F.dropout(x2_c_emb, p=self.args.dropout_emb, training=self.training)\n\n        # Generate char features\n        x1_c_features = self.char_rnn(\n            x1_c_emb.reshape((x1_c_emb.size(0) * x1_c_emb.size(1), x1_c_emb.size(2), x1_c_emb.size(3))), \n            x1_mask.unsqueeze(2).repeat(1, 1, x1_c_emb.size(2)).reshape((x1_c_emb.size(0) * x1_c_emb.size(1), x1_c_emb.size(2)))\n            ).reshape((x1_c_emb.size(0), x1_c_emb.size(1), x1_c_emb.size(2), -1))[:,:,-1,:]\n        x2_c_features = self.char_rnn(\n            x2_c_emb.reshape((x2_c_emb.size(0) * x2_c_emb.size(1), x2_c_emb.size(2), x2_c_emb.size(3))), \n            x2_mask.unsqueeze(2).repeat(1, 1, x2_c_emb.size(2)).reshape((x2_c_emb.size(0) * x2_c_emb.size(1), x2_c_emb.size(2)))\n            ).reshape((x2_c_emb.size(0), x2_c_emb.size(1), x2_c_emb.size(2), -1))[:,:,-1,:] \n\n        # Combine input\n        crnn_input = [x1_emb, x1_c_features]\n        qrnn_input = [x2_emb, x2_c_features]\n\n        # Encode document with RNN\n        c = self.encode_rnn(torch.cat(crnn_input, 2), x1_mask)\n        \n        # Encode question with RNN\n        q = self.encode_rnn(torch.cat(qrnn_input, 2), x2_mask)\n\n        # Match questions to docs\n        question_attn_hiddens = self.question_attn(c, q, x2_mask)\n        rnn_input = self.question_attn_gate(torch.cat([c, question_attn_hiddens], 2))\n        c = self.question_attn_rnn(rnn_input, x1_mask)\n\n        # Match documents to themselves\n        doc_self_attn_hiddens = self.doc_self_attn(c, x1_mask)\n        rnn_input = self.doc_self_attn_gate(torch.cat([c, doc_self_attn_hiddens], 2))\n        c = self.doc_self_attn_rnn(rnn_input, x1_mask)\n        c = self.doc_self_attn_rnn2(c, x1_mask)\n\n        # Predict\n        start_scores, end_scores = self.ptr_net(c, q, x1_mask, x2_mask)\n        \n        return start_scores, end_scores\n"
  },
  {
    "path": "rnn_reader.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2017-present, Facebook, Inc.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Implementation of the RNN based DrQA reader.\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport layers\n\n\n# ------------------------------------------------------------------------------\n# Network\n# ------------------------------------------------------------------------------\n\n\nclass RnnDocReader(nn.Module):\n    RNN_TYPES = {'lstm': nn.LSTM, 'gru': nn.GRU, 'rnn': nn.RNN}\n    CELL_TYPES = {'lstm': nn.LSTMCell, 'gru': nn.GRUCell, 'rnn': nn.RNNCell}\n    def __init__(self, args, normalize=True):\n        super(RnnDocReader, self).__init__()\n        # Store config\n        self.args = args\n\n        # Word embeddings (+1 for padding)\n        self.embedding = nn.Embedding(args.vocab_size,\n                                      args.embedding_dim,\n                                      padding_idx=0)\n\n        # Projection for attention weighted question\n        if args.use_qemb:\n            self.qemb_match = layers.SeqAttnMatch(args.embedding_dim)\n\n        # Input size to RNN: word emb + question emb + manual features\n        doc_input_size = args.embedding_dim + args.num_features\n        if args.use_qemb:\n            doc_input_size += args.embedding_dim\n\n        # RNN document encoder\n        self.doc_rnn = layers.StackedBRNN(\n            input_size=doc_input_size,\n            hidden_size=args.hidden_size,\n            num_layers=args.doc_layers,\n            dropout_rate=args.dropout_rnn,\n            dropout_output=args.dropout_rnn_output,\n            concat_layers=args.concat_rnn_layers,\n            rnn_type=self.RNN_TYPES[args.rnn_type],\n            padding=args.rnn_padding,\n        )\n\n        # RNN question encoder\n        self.question_rnn = layers.StackedBRNN(\n            input_size=args.embedding_dim,\n            hidden_size=args.hidden_size,\n            num_layers=args.question_layers,\n            dropout_rate=args.dropout_rnn,\n            dropout_output=args.dropout_rnn_output,\n            concat_layers=args.concat_rnn_layers,\n            rnn_type=self.RNN_TYPES[args.rnn_type],\n            padding=args.rnn_padding,\n        )\n\n        # Output sizes of rnn encoders\n        doc_hidden_size = 2 * args.hidden_size\n        question_hidden_size = 2 * args.hidden_size\n        if args.concat_rnn_layers:\n            doc_hidden_size *= args.doc_layers\n            question_hidden_size *= args.question_layers\n\n        \n        # Question merging\n        if args.question_merge not in ['avg', 'self_attn']:\n            raise NotImplementedError('merge_mode = %s' % args.merge_mode)\n        if args.question_merge == 'self_attn':\n            self.self_attn = layers.LinearSeqAttn(question_hidden_size)\n\n        # Bilinear attention for span start/end\n        self.start_attn = layers.BilinearSeqAttn(\n            doc_hidden_size,\n            question_hidden_size,\n            normalize=normalize,\n        )\n        self.end_attn = layers.BilinearSeqAttn(\n            doc_hidden_size,\n            question_hidden_size,\n            normalize=normalize,\n        )\n        \n\n    def forward(self, x1, x1_c, x1_f, x1_mask, x2, x2_c, x2_f, x2_mask):\n        \"\"\"Inputs:\n        x1 = document word indices             [batch * len_d]\n        x1_f = document word features indices  [batch * len_d * nfeat]\n        x1_mask = document padding mask        [batch * len_d]\n        x2 = question word indices             [batch * len_q]\n        x2_mask = question padding mask        [batch * len_q]\n        \"\"\"\n        # Embed both document and question\n        x1_emb = self.embedding(x1)\n        x2_emb = self.embedding(x2)\n\n        # Dropout on embeddings\n        if self.args.dropout_emb > 0:\n            x1_emb = nn.functional.dropout(x1_emb, p=self.args.dropout_emb,\n                                           training=self.training)\n            x2_emb = nn.functional.dropout(x2_emb, p=self.args.dropout_emb,\n                                           training=self.training)\n\n        # Form document encoding inputs\n        drnn_input = [x1_emb]\n\n        # Add attention-weighted question representation\n        if self.args.use_qemb:\n            x2_weighted_emb = self.qemb_match(x1_emb, x2_emb, x2_mask)\n            drnn_input.append(x2_weighted_emb)\n\n        # Add manual features\n        if self.args.num_features > 0:\n            drnn_input.append(x1_f)\n\n        # Encode document with RNN\n        doc_hiddens = self.doc_rnn(torch.cat(drnn_input, 2), x1_mask)\n\n        # Encode question with RNN + merge hiddens\n        question_hiddens = self.question_rnn(x2_emb, x2_mask)\n        \n        if self.args.question_merge == 'avg':\n            q_merge_weights = layers.uniform_weights(question_hiddens, x2_mask)\n        elif self.args.question_merge == 'self_attn':\n            q_merge_weights = self.self_attn(question_hiddens, x2_mask)\n        question_hidden = layers.weighted_avg(question_hiddens, q_merge_weights)\n\n        # Predict start and end positions\n        start_scores = self.start_attn(doc_hiddens, question_hidden, x1_mask)\n        end_scores = self.end_attn(doc_hiddens, question_hidden, x1_mask)\n        \n        return start_scores, end_scores\n"
  },
  {
    "path": "script/evaluate-v1.1.py",
    "content": "\"\"\" Official evaluation script for v1.1 of the SQuAD dataset. \"\"\"\nfrom __future__ import print_function\nfrom collections import Counter\nimport string\nimport re\nimport argparse\nimport json\nimport sys\n\n\ndef normalize_answer(s):\n    \"\"\"Lower text and remove punctuation, articles and extra whitespace.\"\"\"\n    def remove_articles(text):\n        return re.sub(r'\\b(a|an|the)\\b', ' ', text)\n\n    def white_space_fix(text):\n        return ' '.join(text.split())\n\n    def remove_punc(text):\n        exclude = set(string.punctuation)\n        return ''.join(ch for ch in text if ch not in exclude)\n\n    def lower(text):\n        return text.lower()\n\n    return white_space_fix(remove_articles(remove_punc(lower(s))))\n\n\ndef f1_score(prediction, ground_truth):\n    prediction_tokens = normalize_answer(prediction).split()\n    ground_truth_tokens = normalize_answer(ground_truth).split()\n    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)\n    num_same = sum(common.values())\n    if num_same == 0:\n        return 0\n    precision = 1.0 * num_same / len(prediction_tokens)\n    recall = 1.0 * num_same / len(ground_truth_tokens)\n    f1 = (2 * precision * recall) / (precision + recall)\n    return f1\n\n\ndef exact_match_score(prediction, ground_truth):\n    return (normalize_answer(prediction) == normalize_answer(ground_truth))\n\n\ndef metric_max_over_ground_truths(metric_fn, prediction, ground_truths):\n    scores_for_ground_truths = []\n    for ground_truth in ground_truths:\n        score = metric_fn(prediction, ground_truth)\n        scores_for_ground_truths.append(score)\n    return max(scores_for_ground_truths)\n\n\ndef evaluate(dataset, predictions):\n    f1 = exact_match = total = 0\n    for article in dataset:\n        for paragraph in article['paragraphs']:\n            for qa in paragraph['qas']:\n                total += 1\n                if qa['id'] not in predictions:\n                    message = 'Unanswered question ' + qa['id'] + \\\n                              ' will receive score 0.'\n                    print(message, file=sys.stderr)\n                    continue\n                ground_truths = list(map(lambda x: x['text'], qa['answers']))\n                prediction = predictions[qa['id']]\n                exact_match += metric_max_over_ground_truths(\n                    exact_match_score, prediction, ground_truths)\n                f1 += metric_max_over_ground_truths(\n                    f1_score, prediction, ground_truths)\n\n    exact_match = 100.0 * exact_match / total\n    f1 = 100.0 * f1 / total\n\n    return {'exact_match': exact_match, 'f1': f1}\n\n\nif __name__ == '__main__':\n    expected_version = '1.1'\n    parser = argparse.ArgumentParser(\n        description='Evaluation for SQuAD ' + expected_version)\n    parser.add_argument('dataset_file', help='Dataset file')\n    parser.add_argument('prediction_file', help='Prediction File')\n    args = parser.parse_args()\n    with open(args.dataset_file) as dataset_file:\n        dataset_json = json.load(dataset_file)\n        if (dataset_json['version'] != expected_version):\n            print('Evaluation expects v-' + expected_version +\n                  ', but got dataset with v-' + dataset_json['version'],\n                  file=sys.stderr)\n        dataset = dataset_json['data']\n    with open(args.prediction_file) as prediction_file:\n        predictions = json.load(prediction_file)\n    print(json.dumps(evaluate(dataset, predictions)))\n"
  },
  {
    "path": "script/interactive.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2018-present, HKUST-KnowComp.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"A script to run the reader model interactively.\"\"\"\n\nimport sys\nsys.path.append('.')\nimport torch\nimport code\nimport argparse\nimport logging\nimport prettytable\nimport time\n\nfrom predictor import Predictor\nfrom multiprocessing import cpu_count\n\nlogger = logging.getLogger()\nlogger.setLevel(logging.INFO)\nfmt = logging.Formatter('%(asctime)s: [ %(message)s ]', '%m/%d/%Y %I:%M:%S %p')\nconsole = logging.StreamHandler()\nconsole.setFormatter(fmt)\nlogger.addHandler(console)\n\nPREDICTOR = None\n\n# ------------------------------------------------------------------------------\n# Drop in to interactive mode\n# ------------------------------------------------------------------------------\n\n\ndef process(document, question, candidates=None, top_n=1):\n    t0 = time.time()\n    predictions = PREDICTOR.predict(document, question, candidates, top_n)\n    table = prettytable.PrettyTable(['Rank', 'Span', 'Score'])\n    for i, p in enumerate(predictions, 1):\n        table.add_row([i, p[0], p[1]])\n    print(table)\n    print('Time: %.4f' % (time.time() - t0))\n\n\nbanner = \"\"\"\n* WRMCQA interactive Document Reader Module *\n\n* Repo: Mnemonic Reader (https://github.com/HKUST-KnowComp/MnemonicReader)\n\n* Implement based on Facebook's DrQA\n\n>>> process(document, question, candidates=None, top_n=1)\n>>> usage()\n\"\"\"\n\n\ndef usage():\n    print(banner)\n\n# ------------------------------------------------------------------------------\n# Commandline arguments & init\n# ------------------------------------------------------------------------------\n\nif __name__ == '__main__':\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--model', type=str, default=None,\n                        help='Path to model to use')\n    parser.add_argument('--embedding-file', type=str, default=None,\n                        help=('Expand dictionary to use all pretrained '\n                            'embeddings in this file.'))\n    parser.add_argument('--char-embedding-file', type=str, default=None,\n                        help=('Expand dictionary to use all pretrained '\n                            'char embeddings in this file.'))\n    parser.add_argument('--num-workers', type=int, default=int(cpu_count()/2),\n                        help='Number of CPU processes (for tokenizing, etc)')\n    parser.add_argument('--no-cuda', action='store_true',\n                        help='Use CPU only')\n    parser.add_argument('--gpu', type=int, default=-1,\n                        help='Specify GPU device id to use')\n    parser.add_argument('--no-normalize', action='store_true',\n                        help='Do not softmax normalize output scores.')\n    args = parser.parse_args()\n\n    args.cuda = not args.no_cuda and torch.cuda.is_available()\n    if args.cuda:\n        torch.cuda.set_device(args.gpu)\n        logger.info('CUDA enabled (GPU %d)' % args.gpu)\n    else:\n        logger.info('Running on CPU only.')\n\n    PREDICTOR = Predictor(\n        args.model,\n        normalize=not args.no_normalize,\n        embedding_file=args.embedding_file,\n        char_embedding_file=args.char_embedding_file,\n        num_workers=args.num_workers,\n    )\n    if args.cuda:\n        PREDICTOR.cuda()\n    code.interact(banner=banner, local=locals())\n"
  },
  {
    "path": "script/predict.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2017-present, Facebook, Inc.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"A script to make and save model predictions on an input dataset.\"\"\"\n\nimport sys\nsys.path.append('.')\nimport os\nimport time\nimport torch\nimport argparse\nimport logging\ntry:\n    import ujson as json\nexcept ImportError:\n    import json\n\nfrom tqdm import tqdm\nfrom predictor import Predictor\nfrom multiprocessing import cpu_count\n\nlogger = logging.getLogger()\nlogger.setLevel(logging.INFO)\nfmt = logging.Formatter('%(asctime)s: [ %(message)s ]', '%m/%d/%Y %I:%M:%S %p')\nconsole = logging.StreamHandler()\nconsole.setFormatter(fmt)\nlogger.addHandler(console)\n\nparser = argparse.ArgumentParser()\nparser.add_argument('dataset', type=str, default=None,\n                    help='SQuAD-like dataset to evaluate on')\nparser.add_argument('--model', type=str, default=None,\n                    help='Path to model to use')\nparser.add_argument('--embedding-file', type=str, default=None,\n                    help=('Expand dictionary to use all pretrained '\n                          'embeddings in this file.'))\nparser.add_argument('--char-embedding-file', type=str, default=None,\n                    help=('Expand dictionary to use all pretrained '\n                          'char embeddings in this file.'))\nparser.add_argument('--out-dir', type=str, default='data/predict',\n                    help=('Directory to write prediction file to '\n                          '(<dataset>-<model>.preds)'))\nparser.add_argument('--num-workers', type=int, default=int(cpu_count()/2),\n                    help='Number of CPU processes (for tokenizing, etc)')\nparser.add_argument('--no-cuda', action='store_true',\n                    help='Use CPU only')\nparser.add_argument('--gpu', type=int, default=-1,\n                    help='Specify GPU device id to use')\nparser.add_argument('--batch-size', type=int, default=128,\n                    help='Example batching size')\nparser.add_argument('--top-n', type=int, default=1,\n                    help='Store top N predicted spans per example')\nparser.add_argument('--official', type=bool, default=True,\n                    help='Only store single top span instead of top N list')\nargs = parser.parse_args()\nt0 = time.time()\n\nargs.cuda = not args.no_cuda and torch.cuda.is_available()\nif args.cuda:\n    torch.cuda.set_device(args.gpu)\n    logger.info('CUDA enabled (GPU %d)' % args.gpu)\nelse:\n    logger.info('Running on CPU only.')\n\npredictor = Predictor(\n    args.model,\n    normalize=True,\n    embedding_file=args.embedding_file,\n    char_embedding_file=args.char_embedding_file,\n    num_workers=args.num_workers,\n)\nif args.cuda:\n    predictor.cuda()\n\n# ------------------------------------------------------------------------------\n# Read in dataset and make predictions.\n# ------------------------------------------------------------------------------\n\n\nexamples = []\nqids = []\nwith open(args.dataset) as f:\n    data = json.load(f)['data']\n    for article in data:\n        for paragraph in article['paragraphs']:\n            context = paragraph['context']\n            for qa in paragraph['qas']:\n                qids.append(qa['id'])\n                examples.append((context, qa['question']))\n\nresults = {}\nfor i in tqdm(range(0, len(examples), args.batch_size)):\n    predictions = predictor.predict_batch(\n        examples[i:i + args.batch_size], top_n=args.top_n\n    )\n    for j in range(len(predictions)):\n        # Official eval expects just a qid --> span\n        if args.official:\n            results[qids[i + j]] = predictions[j][0][0]\n\n        # Otherwise we store top N and scores for debugging.\n        else:\n            results[qids[i + j]] = [(p[0], float(p[1])) for p in predictions[j]]\n\nmodel = os.path.splitext(os.path.basename(args.model or 'default'))[0]\nbasename = os.path.splitext(os.path.basename(args.dataset))[0]\noutfile = os.path.join(args.out_dir, basename + '-' + model + '.preds')\nif not os.path.isdir(args.out_dir):\n    os.mkdir(args.out_dir)\nlogger.info('Writing results to %s' % outfile)\nwith open(outfile, 'w') as f:\n    json.dump(results, f)\n\nlogger.info('Total time: %.2f' % (time.time() - t0))\n"
  },
  {
    "path": "script/preprocess.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2017-present, Facebook, Inc.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Preprocess the SQuAD dataset for training.\"\"\"\n\nimport sys\nsys.path.append('.')\nimport argparse\nimport os\ntry:\n    import ujson as json\nexcept ImportError:\n    import json\nimport time\n\nfrom multiprocessing import Pool, cpu_count\nfrom multiprocessing.util import Finalize\nfrom functools import partial\nfrom spacy_tokenizer import SpacyTokenizer\n\n# ------------------------------------------------------------------------------\n# Tokenize + annotate.\n# ------------------------------------------------------------------------------\n\nTOK = None\nANNTOTORS = {'lemma', 'pos', 'ner'}\n\n\ndef init():\n    global TOK\n    TOK = SpacyTokenizer(annotators=ANNTOTORS)\n    Finalize(TOK, TOK.shutdown, exitpriority=100)\n\n\ndef tokenize(text):\n    \"\"\"Call the global process tokenizer on the input text.\"\"\"\n    global TOK\n    tokens = TOK.tokenize(text)\n    output = {\n        'words': tokens.words(),\n        'chars': tokens.chars(),\n        'offsets': tokens.offsets(),\n        'pos': tokens.pos(),\n        'lemma': tokens.lemmas(),\n        'ner': tokens.entities(),\n    }\n    return output\n\n\n# ------------------------------------------------------------------------------\n# Process dataset examples\n# ------------------------------------------------------------------------------\n\n\ndef load_dataset(path):\n    \"\"\"Load json file and store fields separately.\"\"\"\n    with open(path) as f:\n        data = json.load(f)['data']\n    output = {'qids': [], 'questions': [], 'answers': [],\n              'contexts': [], 'qid2cid': []}\n    for article in data:\n        for paragraph in article['paragraphs']:\n            output['contexts'].append(paragraph['context'])\n            for qa in paragraph['qas']:\n                output['qids'].append(qa['id'])\n                output['questions'].append(qa['question'])\n                output['qid2cid'].append(len(output['contexts']) - 1)\n                if 'answers' in qa:\n                    output['answers'].append(qa['answers'])\n    return output\n\n\ndef find_answer(offsets, begin_offset, end_offset):\n    \"\"\"Match token offsets with the char begin/end offsets of the answer.\"\"\"\n    start = [i for i, tok in enumerate(offsets) if tok[0] == begin_offset]\n    end = [i for i, tok in enumerate(offsets) if tok[1] == end_offset]\n    assert(len(start) <= 1)\n    assert(len(end) <= 1)\n    if len(start) == 1 and len(end) == 1:\n        return start[0], end[0]\n\n\ndef process_dataset(data, tokenizer, workers=None):\n    \"\"\"Iterate processing (tokenize, parse, etc) dataset multithreaded.\"\"\"\n    make_pool = partial(Pool, workers, initializer=init)\n\n    workers = make_pool(initargs=())\n    q_tokens = workers.map(tokenize, data['questions'])\n    workers.close()\n    workers.join()\n\n    workers = make_pool(initargs=())\n    c_tokens = workers.map(tokenize, data['contexts'])\n    workers.close()\n    workers.join()\n\n    for idx in range(len(data['qids'])):\n        question = q_tokens[idx]['words']\n        question_char = q_tokens[idx]['chars']\n        qlemma = q_tokens[idx]['lemma']\n        qpos = q_tokens[idx]['pos']\n        qner = q_tokens[idx]['ner']\n\n        document = c_tokens[data['qid2cid'][idx]]['words']\n        document_char = c_tokens[data['qid2cid'][idx]]['chars']\n        offsets = c_tokens[data['qid2cid'][idx]]['offsets']\n        clemma = c_tokens[data['qid2cid'][idx]]['lemma']\n        cpos = c_tokens[data['qid2cid'][idx]]['pos']\n        cner = c_tokens[data['qid2cid'][idx]]['ner']\n        \n        ans_tokens = []\n        if len(data['answers']) > 0:\n            for ans in data['answers'][idx]:\n                found = find_answer(offsets,\n                                    ans['answer_start'],\n                                    ans['answer_start'] + len(ans['text']))\n                if found:\n                    ans_tokens.append(found)\n        yield {\n            'id': data['qids'][idx],\n            'question': question,\n            'question_char': question_char,\n            'document': document,\n            'document_char': document_char,\n            'offsets': offsets,\n            'answers': ans_tokens,\n            'qlemma': qlemma,\n            'qpos': qpos,\n            'qner': qner,\n            'clemma': clemma,\n            'cpos': cpos,\n            'cner': cner,\n        }\n\n\n# -----------------------------------------------------------------------------\n# Commandline options\n# -----------------------------------------------------------------------------\n\n\nparser = argparse.ArgumentParser()\nparser.add_argument('data_dir', type=str, help='Path to SQuAD data directory')\nparser.add_argument('out_dir', type=str, help='Path to output file dir')\nparser.add_argument('--split', type=str, help='Filename for train/dev split')\nparser.add_argument('--num-workers', type=int, default=1)\nparser.add_argument('--tokenizer', type=str, default='spacy')\nargs = parser.parse_args()\n\nt0 = time.time()\n\nin_file = os.path.join(args.data_dir, args.split + '.json')\nprint('Loading dataset %s' % in_file, file=sys.stderr)\ndataset = load_dataset(in_file)\n\nout_file = os.path.join(\n    args.out_dir, '%s-processed-%s.txt' % (args.split, args.tokenizer)\n)\nprint('Will write to file %s' % out_file, file=sys.stderr)\nwith open(out_file, 'w') as f:\n    for ex in process_dataset(dataset, args.tokenizer, args.num_workers):\n        f.write(json.dumps(ex) + '\\n')\nprint('Total time: %.4f (s)' % (time.time() - t0))"
  },
  {
    "path": "script/train.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2018-present, HKUST-KnowComp.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Main reader training script.\"\"\"\n\nimport sys\nsys.path.append('.')\nimport argparse\nimport torch\nimport numpy as np\ntry:\n    import ujson as json\nexcept ImportError:\n    import json\nimport os\nimport subprocess\nimport logging\n\n\nimport utils, vector, config, data\nfrom model import DocReader\n\nlogger = logging.getLogger()\n\n\n# ------------------------------------------------------------------------------\n# Training arguments.\n# ------------------------------------------------------------------------------\n\n\n# Defaults\nDATA_DIR = os.path.join('data', 'datasets')\nMODEL_DIR = os.path.join('data', 'models')\nEMBED_DIR = os.path.join('data', 'embeddings')\n\ndef str2bool(v):\n    return v.lower() in ('yes', 'true', 't', '1', 'y')\n\n\ndef add_train_args(parser):\n    \"\"\"Adds commandline arguments pertaining to training a model. These\n    are different from the arguments dictating the model architecture.\n    \"\"\"\n    parser.register('type', 'bool', str2bool)\n\n    # Runtime environment\n    runtime = parser.add_argument_group('Environment')\n    runtime.add_argument('--no-cuda', type='bool', default=False,\n                         help='Train on CPU, even if GPUs are available.')\n    runtime.add_argument('--gpu', type=int, default=-1,\n                         help='Run on a specific GPU')\n    runtime.add_argument('--data-workers', type=int, default=5,\n                         help='Number of subprocesses for data loading')\n    runtime.add_argument('--parallel', type='bool', default=False,\n                         help='Use DataParallel on all available GPUs')\n    runtime.add_argument('--random-seed', type=int, default=1013,\n                         help=('Random seed for all numpy/torch/cuda '\n                               'operations (for reproducibility)'))\n    runtime.add_argument('--num-epochs', type=int, default=40,\n                         help='Train data iterations')\n    runtime.add_argument('--batch-size', type=int, default=45,\n                         help='Batch size for training')\n    runtime.add_argument('--test-batch-size', type=int, default=32,\n                         help='Batch size during validation/testing')\n\n    # Files\n    files = parser.add_argument_group('Filesystem')\n    files.add_argument('--model-dir', type=str, default=MODEL_DIR,\n                       help='Directory for saved models/checkpoints/logs')\n    files.add_argument('--model-name', type=str, default='',\n                       help='Unique model identifier (.mdl, .txt, .checkpoint)')\n    files.add_argument('--data-dir', type=str, default=DATA_DIR,\n                       help='Directory of training/validation data')\n    files.add_argument('--train-file', type=str,\n                       default='SQuAD-v1.1-train-processed-spacy.txt',\n                       help='Preprocessed train file')\n    files.add_argument('--dev-file', type=str,\n                       default='SQuAD-v1.1-dev-processed-spacy.txt',\n                       help='Preprocessed dev file')\n    files.add_argument('--dev-json', type=str, default='SQuAD-v1.1-dev.json',\n                       help=('Unprocessed dev file to run validation '\n                             'while training on'))\n    files.add_argument('--embed-dir', type=str, default=EMBED_DIR,\n                       help='Directory of pre-trained embedding files')\n    files.add_argument('--embedding-file', type=str,\n                       default='glove.840B.300d.txt',\n                       help='Space-separated pretrained embeddings file')\n    files.add_argument('--char-embedding-file', type=str,\n                       default='glove.840B.300d-char.txt',\n                       help='Space-separated pretrained embeddings file')\n\n    # Saving + loading\n    save_load = parser.add_argument_group('Saving/Loading')\n    save_load.add_argument('--checkpoint', type='bool', default=False,\n                           help='Save model + optimizer state after each epoch')\n    save_load.add_argument('--pretrained', type=str, default='',\n                           help='Path to a pretrained model to warm-start with')\n    save_load.add_argument('--expand-dictionary', type='bool', default=False,\n                           help='Expand dictionary of pretrained model to ' +\n                                'include training/dev words of new data')\n    # Data preprocessing\n    preprocess = parser.add_argument_group('Preprocessing')\n    preprocess.add_argument('--uncased-question', type='bool', default=False,\n                            help='Question words will be lower-cased')\n    preprocess.add_argument('--uncased-doc', type='bool', default=False,\n                            help='Document words will be lower-cased')\n    preprocess.add_argument('--restrict-vocab', type='bool', default=True,\n                            help='Only use pre-trained words in embedding_file')\n\n    # General\n    general = parser.add_argument_group('General')\n    general.add_argument('--official-eval', type='bool', default=True,\n                         help='Validate with official SQuAD eval')\n    general.add_argument('--valid-metric', type=str, default='exact_match',\n                         help='The evaluation metric used for model selection: None, exact_match, f1')\n    general.add_argument('--display-iter', type=int, default=25,\n                         help='Log state after every <display_iter> epochs')\n    general.add_argument('--sort-by-len', type='bool', default=True,\n                         help='Sort batches by length for speed')\n\n\ndef set_defaults(args):\n    \"\"\"Make sure the commandline arguments are initialized properly.\"\"\"\n    # Check critical files exist\n    args.dev_json = os.path.join(args.data_dir, args.dev_json)\n    if not os.path.isfile(args.dev_json):\n        raise IOError('No such file: %s' % args.dev_json)\n    args.train_file = os.path.join(args.data_dir, args.train_file)\n    if not os.path.isfile(args.train_file):\n        raise IOError('No such file: %s' % args.train_file)\n    args.dev_file = os.path.join(args.data_dir, args.dev_file)\n    if not os.path.isfile(args.dev_file):\n        raise IOError('No such file: %s' % args.dev_file)\n    if args.embedding_file:\n        args.embedding_file = os.path.join(args.embed_dir, args.embedding_file)\n        if not os.path.isfile(args.embedding_file):\n            raise IOError('No such file: %s' % args.embedding_file)\n    if args.char_embedding_file:\n        args.char_embedding_file = os.path.join(args.embed_dir, args.char_embedding_file)\n        if not os.path.isfile(args.char_embedding_file):\n            raise IOError('No such file: %s' % args.char_embedding_file)\n\n    # Set model directory\n    subprocess.call(['mkdir', '-p', args.model_dir])\n\n    # Set model name\n    if not args.model_name:\n        import uuid\n        import time\n        args.model_name = time.strftime(\"%Y%m%d-\") + str(uuid.uuid4())[:8]\n\n    # Set log + model file names\n    args.log_file = os.path.join(args.model_dir, args.model_name + '.txt')\n    args.model_file = os.path.join(args.model_dir, args.model_name + '.mdl')\n\n    # Embeddings options\n    if args.embedding_file:\n        with open(args.embedding_file) as f:\n            dim = len(f.readline().strip().split(' ')) - 1\n        args.embedding_dim = dim\n    elif not args.embedding_dim:\n        raise RuntimeError('Either embedding_file or embedding_dim '\n                           'needs to be specified.')\n    if args.char_embedding_file:\n        with open(args.char_embedding_file) as f:\n            dim = len(f.readline().strip().split(' ')) - 1\n        args.char_embedding_dim = dim\n    elif not args.char_embedding_dim:\n        raise RuntimeError('Either char_embedding_file or char_embedding_dim '\n                           'needs to be specified.')\n\n    # Make sure tune_partial and fix_embeddings are consistent.\n    if args.tune_partial > 0 and args.fix_embeddings:\n        logger.warning('WARN: fix_embeddings set to False as tune_partial > 0.')\n        args.fix_embeddings = False\n\n    # Make sure fix_embeddings and embedding_file are consistent\n    if args.fix_embeddings:\n        if not (args.embedding_file or args.pretrained):\n            logger.warning('WARN: fix_embeddings set to False '\n                           'as embeddings are random.')\n            args.fix_embeddings = False\n    return args\n\n\n# ------------------------------------------------------------------------------\n# Initalization from scratch.\n# ------------------------------------------------------------------------------\n\n\ndef init_from_scratch(args, train_exs, dev_exs):\n    \"\"\"New model, new data, new dictionary.\"\"\"\n    # Create a feature dict out of the annotations in the data\n    logger.info('-' * 100)\n    logger.info('Generate features')\n    feature_dict = utils.build_feature_dict(args, train_exs)\n    logger.info('Num features = %d' % len(feature_dict))\n    logger.info(feature_dict)\n\n    # Build a dictionary from the data questions + documents (train/dev splits)\n    logger.info('-' * 100)\n    logger.info('Build word dictionary')\n    word_dict = utils.build_word_dict(args, train_exs + dev_exs)\n    logger.info('Num words = %d' % len(word_dict))    \n\n    # Build a char dictionary from the data questions + documents (train/dev splits)\n    logger.info('-' * 100)\n    logger.info('Build char dictionary')\n    char_dict = utils.build_char_dict(args, train_exs + dev_exs)\n    logger.info('Num chars = %d' % len(char_dict))\n    # Initialize model\n    model = DocReader(config.get_model_args(args), word_dict, char_dict, feature_dict)\n\n    # Load pretrained embeddings for words in dictionary\n    if args.embedding_file:\n        model.load_embeddings(word_dict.tokens(), args.embedding_file)\n    if args.char_embedding_file:\n        model.load_char_embeddings(char_dict.tokens(), args.char_embedding_file)\n\n    return model\n\n\n# ------------------------------------------------------------------------------\n# Train loop.\n# ------------------------------------------------------------------------------\n\n\ndef train(args, data_loader, model, global_stats):\n    \"\"\"Run through one epoch of model training with the provided data loader.\"\"\"\n    # Initialize meters + timers\n    train_loss = utils.AverageMeter()\n    epoch_time = utils.Timer()\n\n    # Run one epoch\n    for idx, ex in enumerate(data_loader):\n        train_loss.update(*model.update(ex))\n\n        if idx % args.display_iter == 0:\n            logger.info('train: Epoch = %d | iter = %d/%d | ' %\n                        (global_stats['epoch'], idx, len(data_loader)) +\n                        'loss = %.2f | elapsed time = %.2f (s)' %\n                        (train_loss.avg, global_stats['timer'].time()))\n            train_loss.reset()\n\n    logger.info('train: Epoch %d done. Time for epoch = %.2f (s)' %\n                (global_stats['epoch'], epoch_time.time()))\n\n    # Checkpoint\n    if args.checkpoint:\n        model.checkpoint(args.model_file + '.checkpoint',\n                         global_stats['epoch'] + 1)\n\n\n# ------------------------------------------------------------------------------\n# Validation loops. Includes both \"unofficial\" and \"official\" functions that\n# use different metrics and implementations.\n# ------------------------------------------------------------------------------\n\n\ndef validate_unofficial(args, data_loader, model, global_stats, mode):\n    \"\"\"Run one full unofficial validation.\n    Unofficial = doesn't use SQuAD script.\n    \"\"\"\n    eval_time = utils.Timer()\n    start_acc = utils.AverageMeter()\n    end_acc = utils.AverageMeter()\n    exact_match = utils.AverageMeter()\n\n    # Make predictions\n    examples = 0\n    for ex in data_loader:\n        batch_size = ex[0].size(0)\n        pred_s, pred_e, _ = model.predict(ex)\n        target_s, target_e = ex[-3:-1]\n\n        # We get metrics for independent start/end and joint start/end\n        accuracies = eval_accuracies(pred_s, target_s, pred_e, target_e)\n        start_acc.update(accuracies[0], batch_size)\n        end_acc.update(accuracies[1], batch_size)\n        exact_match.update(accuracies[2], batch_size)\n\n        # If getting train accuracies, sample max 10k\n        examples += batch_size\n        if mode == 'train' and examples >= 1e4:\n            break\n\n    logger.info('%s valid unofficial: Epoch = %d | start = %.2f | ' %\n                (mode, global_stats['epoch'], start_acc.avg) +\n                'end = %.2f | exact = %.2f | examples = %d | ' %\n                (end_acc.avg, exact_match.avg, examples) +\n                'valid time = %.2f (s)' % eval_time.time())\n\n    return {'exact_match': exact_match.avg}\n\n\ndef validate_official(args, data_loader, model, global_stats,\n                      offsets, texts, answers):\n    \"\"\"Run one full official validation. Uses exact spans and same\n    exact match/F1 score computation as in the SQuAD script.\n\n    Extra arguments:\n        offsets: The character start/end indices for the tokens in each context.\n        texts: Map of qid --> raw text of examples context (matches offsets).\n        answers: Map of qid --> list of accepted answers.\n    \"\"\"\n    eval_time = utils.Timer()\n    f1 = utils.AverageMeter()\n    exact_match = utils.AverageMeter()\n\n    # Run through examples\n    examples = 0\n    for ex in data_loader:\n        ex_id, batch_size = ex[-1], ex[0].size(0)\n        pred_s, pred_e, _ = model.predict(ex)\n\n        for i in range(batch_size):\n            s_offset = offsets[ex_id[i]][pred_s[i][0]][0]\n            e_offset = offsets[ex_id[i]][pred_e[i][0]][1]\n            prediction = texts[ex_id[i]][s_offset:e_offset]\n\n            # Compute metrics\n            ground_truths = answers[ex_id[i]]\n            exact_match.update(utils.metric_max_over_ground_truths(\n                utils.exact_match_score, prediction, ground_truths))\n            f1.update(utils.metric_max_over_ground_truths(\n                utils.f1_score, prediction, ground_truths))\n\n        examples += batch_size\n\n    logger.info('dev valid official: Epoch = %d | EM = %.2f | ' %\n                (global_stats['epoch'], exact_match.avg * 100) +\n                'F1 = %.2f | examples = %d | valid time = %.2f (s)' %\n                (f1.avg * 100, examples, eval_time.time()))\n\n    return {'exact_match': exact_match.avg * 100, 'f1': f1.avg * 100}\n\n\ndef eval_accuracies(pred_s, target_s, pred_e, target_e):\n    \"\"\"An unofficial evalutation helper.\n    Compute exact start/end/complete match accuracies for a batch.\n    \"\"\"\n    # Convert 1D tensors to lists of lists (compatibility)\n    if torch.is_tensor(target_s):\n        target_s = [[e] for e in target_s]\n        target_e = [[e] for e in target_e]\n\n    # Compute accuracies from targets\n    batch_size = len(pred_s)\n    start = utils.AverageMeter()\n    end = utils.AverageMeter()\n    em = utils.AverageMeter()\n    for i in range(batch_size):\n        # Start matches\n        if pred_s[i] in target_s[i]:\n            start.update(1)\n        else:\n            start.update(0)\n\n        # End matches\n        if pred_e[i] in target_e[i]:\n            end.update(1)\n        else:\n            end.update(0)\n\n        # Both start and end match\n        if any([1 for _s, _e in zip(target_s[i], target_e[i])\n                if _s == torch.from_numpy(pred_s[i]) and _e == torch.from_numpy(pred_e[i])]):\n            em.update(1)\n        else:\n            em.update(0)\n    return start.avg * 100, end.avg * 100, em.avg * 100\n\n\n# ------------------------------------------------------------------------------\n# Main.\n# ------------------------------------------------------------------------------\n\n\ndef main(args):\n    # --------------------------------------------------------------------------\n    # DATA\n    logger.info('-' * 100)\n    logger.info('Load data files')\n    train_exs = utils.load_data(args, args.train_file, skip_no_answer=True)\n    logger.info('Num train examples = %d' % len(train_exs))\n    dev_exs = utils.load_data(args, args.dev_file)\n    logger.info('Num dev examples = %d' % len(dev_exs))\n\n    # If we are doing offician evals then we need to:\n    # 1) Load the original text to retrieve spans from offsets.\n    # 2) Load the (multiple) text answers for each question.\n    if args.official_eval:\n        dev_texts = utils.load_text(args.dev_json)\n        dev_offsets = {ex['id']: ex['offsets'] for ex in dev_exs}\n        dev_answers = utils.load_answers(args.dev_json)\n\n    # --------------------------------------------------------------------------\n    # MODEL\n    logger.info('-' * 100)\n    start_epoch = 0\n    if args.checkpoint and os.path.isfile(args.model_file + '.checkpoint'):\n        # Just resume training, no modifications.\n        logger.info('Found a checkpoint...')\n        checkpoint_file = args.model_file + '.checkpoint'\n        model, start_epoch = DocReader.load_checkpoint(checkpoint_file, args)\n    else:\n        # Training starts fresh. But the model state is either pretrained or\n        # newly (randomly) initialized.\n        if args.pretrained:\n            logger.info('Using pretrained model...')\n            model = DocReader.load(args.pretrained, args)\n            if args.expand_dictionary:\n                logger.info('Expanding dictionary for new data...')\n                # Add words in training + dev examples\n                words = utils.load_words(args, train_exs + dev_exs)\n                added_words = model.expand_dictionary(words)\n                # Load pretrained embeddings for added words\n                if args.embedding_file:\n                    model.load_embeddings(added_words, args.embedding_file)\n\n                logger.info('Expanding char dictionary for new data...')\n                # Add words in training + dev examples\n                chars = utils.load_chars(args, train_exs + dev_exs)\n                added_chars = model.expand_char_dictionary(chars)\n                # Load pretrained embeddings for added words\n                if args.char_embedding_file:\n                    model.load_char_embeddings(added_chars, args.char_embedding_file)\n\n        else:\n            logger.info('Training model from scratch...')\n            model = init_from_scratch(args, train_exs, dev_exs)\n\n        # Set up partial tuning of embeddings\n        if args.tune_partial > 0:\n            logger.info('-' * 100)\n            logger.info('Counting %d most frequent question words' %\n                        args.tune_partial)\n            top_words = utils.top_question_words(\n                args, train_exs, model.word_dict\n            )\n            for word in top_words[:5]:\n                logger.info(word)\n            logger.info('...')\n            for word in top_words[-6:-1]:\n                logger.info(word)\n            model.tune_embeddings([w[0] for w in top_words])\n\n        # Set up optimizer\n        model.init_optimizer()\n\n    # Use the GPU?\n    if args.cuda:\n        model.cuda()\n\n    # Use multiple GPUs?\n    if args.parallel:\n        model.parallelize()\n\n    # --------------------------------------------------------------------------\n    # DATA ITERATORS\n    # Two datasets: train and dev. If we sort by length it's faster.\n    logger.info('-' * 100)\n    logger.info('Make data loaders')\n\n    train_dataset = data.ReaderDataset(train_exs, model, single_answer=True)\n    if args.sort_by_len:\n        train_sampler = data.SortedBatchSampler(train_dataset.lengths(),\n                                                args.batch_size,\n                                                shuffle=True)\n    else:\n        train_sampler = torch.utils.data.sampler.RandomSampler(train_dataset)\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=args.batch_size,\n        sampler=train_sampler,\n        num_workers=args.data_workers,\n        collate_fn=vector.batchify,\n        pin_memory=args.cuda,\n    )\n    dev_dataset = data.ReaderDataset(dev_exs, model, single_answer=False)\n    if args.sort_by_len:\n        dev_sampler = data.SortedBatchSampler(dev_dataset.lengths(),\n                                              args.test_batch_size,\n                                              shuffle=False)\n    else:\n        dev_sampler = torch.utils.data.sampler.SequentialSampler(dev_dataset)\n    dev_loader = torch.utils.data.DataLoader(\n        dev_dataset,\n        batch_size=args.test_batch_size,\n        sampler=dev_sampler,\n        num_workers=args.data_workers,\n        collate_fn=vector.batchify,\n        pin_memory=args.cuda,\n    )\n\n    # -------------------------------------------------------------------------\n    # PRINT CONFIG\n    logger.info('-' * 100)\n    logger.info('CONFIG:\\n%s' %\n                json.dumps(vars(args), indent=4, sort_keys=True))\n\n    # --------------------------------------------------------------------------\n    # TRAIN/VALID LOOP\n    logger.info('-' * 100)\n    logger.info('Starting training...')\n    stats = {'timer': utils.Timer(), 'epoch': 0, 'best_valid': 0}\n    for epoch in range(start_epoch, args.num_epochs):\n        stats['epoch'] = epoch\n\n        # Train\n        train(args, train_loader, model, stats)\n\n        # Validate unofficial (train)\n        validate_unofficial(args, train_loader, model, stats, mode='train')\n\n        # Validate unofficial (dev)\n        result = validate_unofficial(args, dev_loader, model, stats, mode='dev')\n\n        # Validate official\n        if args.official_eval:\n            result = validate_official(args, dev_loader, model, stats,\n                                       dev_offsets, dev_texts, dev_answers)\n\n        # Save best valid\n        if args.valid_metric is None or args.valid_metric == 'None':\n            model.save(args.model_file)\n        elif result[args.valid_metric] > stats['best_valid']:\n            logger.info('Best valid: %s = %.2f (epoch %d, %d updates)' %\n                        (args.valid_metric, result[args.valid_metric],\n                         stats['epoch'], model.updates))\n            model.save(args.model_file)\n            stats['best_valid'] = result[args.valid_metric]\n\n\nif __name__ == '__main__':\n    # Parse cmdline args and setup environment\n    parser = argparse.ArgumentParser(\n        'WRMCQA Document Reader',\n        formatter_class=argparse.ArgumentDefaultsHelpFormatter\n    )\n    add_train_args(parser)\n    config.add_model_args(parser)\n    args = parser.parse_args()\n    set_defaults(args)\n\n    # Set cuda\n    args.cuda = not args.no_cuda and torch.cuda.is_available()\n    if args.cuda:\n        torch.cuda.set_device(args.gpu)\n\n    # Set random state\n    np.random.seed(args.random_seed)\n    torch.manual_seed(args.random_seed)\n    if args.cuda:\n        torch.cuda.manual_seed(args.random_seed)\n\n    # Set logging\n    logger.setLevel(logging.INFO)\n    fmt = logging.Formatter('%(asctime)s: [ %(message)s ]',\n                            '%m/%d/%Y %I:%M:%S %p')\n    console = logging.StreamHandler()\n    console.setFormatter(fmt)\n    logger.addHandler(console)\n    if args.log_file:\n        if args.checkpoint:\n            logfile = logging.FileHandler(args.log_file, 'a')\n        else:\n            logfile = logging.FileHandler(args.log_file, 'w')\n        logfile.setFormatter(fmt)\n        logger.addHandler(logfile)\n    logger.info('COMMAND: %s' % ' '.join(sys.argv))\n    print(args)\n    # Run!\n    main(args)\n"
  },
  {
    "path": "spacy_tokenizer.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2018-present, HKUST-KnowComp.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Tokenizer that is backed by spaCy (spacy.io).\n\nRequires spaCy package and the spaCy english model.\n\"\"\"\n\nimport spacy\nimport copy\n\nclass Tokens(object):\n    \"\"\"A class to represent a list of tokenized text.\"\"\"\n    TEXT = 0\n    CHAR = 1\n    TEXT_WS = 2\n    SPAN = 3\n    POS = 4\n    LEMMA = 5\n    NER = 6\n\n    def __init__(self, data, annotators, opts=None):\n        self.data = data\n        self.annotators = annotators\n        self.opts = opts or {}\n\n    def __len__(self):\n        \"\"\"The number of tokens.\"\"\"\n        return len(self.data)\n\n    def slice(self, i=None, j=None):\n        \"\"\"Return a view of the list of tokens from [i, j).\"\"\"\n        new_tokens = copy.copy(self)\n        new_tokens.data = self.data[i: j]\n        return new_tokens\n\n    def untokenize(self):\n        \"\"\"Returns the original text (with whitespace reinserted).\"\"\"\n        return ''.join([t[self.TEXT_WS] for t in self.data]).strip()\n    \n    def chars(self, uncased=False):\n        \"\"\"Returns a list of the first character of each token\n\n        Args:\n            uncased: lower cases characters\n        \"\"\"\n        if uncased:\n            return [[c.lower() for c in t[self.CHAR]] for t in self.data]\n        else:\n            return [[c for c in t[self.CHAR]] for t in self.data]\n\n    def words(self, uncased=False):\n        \"\"\"Returns a list of the text of each token\n\n        Args:\n            uncased: lower cases text\n        \"\"\"\n        if uncased:\n            return [t[self.TEXT].lower() for t in self.data]\n        else:\n            return [t[self.TEXT] for t in self.data]\n\n    def offsets(self):\n        \"\"\"Returns a list of [start, end) character offsets of each token.\"\"\"\n        return [t[self.SPAN] for t in self.data]\n\n    def pos(self):\n        \"\"\"Returns a list of part-of-speech tags of each token.\n        Returns None if this annotation was not included.\n        \"\"\"\n        if 'pos' not in self.annotators:\n            return None\n        return [t[self.POS] for t in self.data]\n\n    def lemmas(self):\n        \"\"\"Returns a list of the lemmatized text of each token.\n        Returns None if this annotation was not included.\n        \"\"\"\n        if 'lemma' not in self.annotators:\n            return None\n        return [t[self.LEMMA] for t in self.data]\n\n    def entities(self):\n        \"\"\"Returns a list of named-entity-recognition tags of each token.\n        Returns None if this annotation was not included.\n        \"\"\"\n        if 'ner' not in self.annotators:\n            return None\n        return [t[self.NER] for t in self.data]\n\n    def ngrams(self, n=1, uncased=False, filter_fn=None, as_strings=True):\n        \"\"\"Returns a list of all ngrams from length 1 to n.\n\n        Args:\n            n: upper limit of ngram length\n            uncased: lower cases text\n            filter_fn: user function that takes in an ngram list and returns\n              True or False to keep or not keep the ngram\n            as_string: return the ngram as a string vs list\n        \"\"\"\n        def _skip(gram):\n            if not filter_fn:\n                return False\n            return filter_fn(gram)\n\n        words = self.words(uncased)\n        ngrams = [(s, e + 1)\n                  for s in range(len(words))\n                  for e in range(s, min(s + n, len(words)))\n                  if not _skip(words[s:e + 1])]\n\n        # Concatenate into strings\n        if as_strings:\n            ngrams = ['{}'.format(' '.join(words[s:e])) for (s, e) in ngrams]\n\n        return ngrams\n\n    def entity_groups(self):\n        \"\"\"Group consecutive entity tokens with the same NER tag.\"\"\"\n        entities = self.entities()\n        if not entities:\n            return None\n        non_ent = self.opts.get('non_ent', 'O')\n        groups = []\n        idx = 0\n        while idx < len(entities):\n            ner_tag = entities[idx]\n            # Check for entity tag\n            if ner_tag != non_ent:\n                # Chomp the sequence\n                start = idx\n                while (idx < len(entities) and entities[idx] == ner_tag):\n                    idx += 1\n                groups.append((self.slice(start, idx).untokenize(), ner_tag))\n            else:\n                idx += 1\n        return groups\n\n\nclass SpacyTokenizer(object):\n\n    def __init__(self, **kwargs):\n        \"\"\"\n        Args:\n            annotators: set that can include pos, lemma, and ner.\n            model: spaCy model to use (either path, or keyword like 'en').\n        \"\"\"\n        model = kwargs.get('model', 'en')\n        self.annotators = copy.deepcopy(kwargs.get('annotators', set()))\n        self.nlp = spacy.load(model)\n        self.nlp.remove_pipe('parser')\n        if not any([p in self.annotators for p in ['lemma', 'pos', 'ner']]):\n            self.nlp.remove_pipe('tagger')\n        if 'ner' not in self.annotators:\n            self.nlp.remove_pipe('ner')\n        \n\n    def tokenize(self, text):\n        # We don't treat new lines as tokens.\n        clean_text = text.replace('\\n', ' ')\n        tokens = self.nlp(clean_text)\n\n        data = []\n        for i in range(len(tokens)):\n            # Get whitespace\n            start_ws = tokens[i].idx\n            if i + 1 < len(tokens):\n                end_ws = tokens[i + 1].idx\n            else:\n                end_ws = tokens[i].idx + len(tokens[i].text)\n\n            data.append((\n                tokens[i].text,\n                list(tokens[i].text),\n                text[start_ws: end_ws],\n                (tokens[i].idx, tokens[i].idx + len(tokens[i].text)),\n                tokens[i].tag_,\n                tokens[i].lemma_,\n                tokens[i].ent_type_,\n            ))\n\n        # Set special option for non-entity tag: '' vs 'O' in spaCy\n        return Tokens(data, self.annotators, opts={'non_ent': ''})\n\n    def shutdown(self):\n        pass\n\n    def __del__(self):\n        self.shutdown()\n"
  },
  {
    "path": "utils.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2018-present, HKUST-KnowComp.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Reader utilities.\"\"\"\n\n\ntry:\n    import ujson as json\nexcept ImportError:\n    import json\nimport time\nimport logging\nimport string\ntry:\n    import regex as re\nexcept ImportError:\n    import re\n\nfrom collections import Counter\nfrom data import Dictionary\n\nlogger = logging.getLogger(__name__)\n\n\n# ------------------------------------------------------------------------------\n# Data loading\n# ------------------------------------------------------------------------------\n\n\ndef load_data(args, filename, skip_no_answer=False):\n    \"\"\"Load examples from preprocessed file.\n    One example per line, JSON encoded.\n    \"\"\"\n    # Load JSON lines\n    with open(filename) as f:\n        examples = [json.loads(line) for line in f]\n\n    # Make case insensitive?\n    if args.uncased_question or args.uncased_doc:\n        for ex in examples:\n            if args.uncased_question:\n                ex['question'] = [w.lower() for w in ex['question']]\n                ex['question_char'] = [w.lower() for w in ex['question_char']]\n            if args.uncased_doc:\n                ex['document'] = [w.lower() for w in ex['document']]\n                ex['document_char'] = [w.lower() for w in ex['document_char']]\n\n    # Skip unparsed (start/end) examples\n    if skip_no_answer:\n        examples = [ex for ex in examples if len(ex['answers']) > 0]\n    return examples\n\n\ndef load_text(filename):\n    \"\"\"Load the paragraphs only of a SQuAD dataset. Store as qid -> text.\"\"\"\n    # Load JSON file\n    with open(filename) as f:\n        examples = json.load(f)['data']\n\n    texts = {}\n    for article in examples:\n        for paragraph in article['paragraphs']:\n            for qa in paragraph['qas']:\n                texts[qa['id']] = paragraph['context']\n    return texts\n\n\ndef load_answers(filename):\n    \"\"\"Load the answers only of a SQuAD dataset. Store as qid -> [answers].\"\"\"\n    # Load JSON file\n    with open(filename) as f:\n        examples = json.load(f)['data']\n\n    ans = {}\n    for article in examples:\n        for paragraph in article['paragraphs']:\n            for qa in paragraph['qas']:\n                ans[qa['id']] = list(map(lambda x: x['text'], qa['answers']))\n    return ans\n\n\n# ------------------------------------------------------------------------------\n# Dictionary building\n# ------------------------------------------------------------------------------\n\n\ndef index_embedding_words(embedding_file):\n    \"\"\"Put all the words in embedding_file into a set.\"\"\"\n    words = set()\n    with open(embedding_file) as f:\n        for line in f:\n            w = Dictionary.normalize(line.rstrip().split(' ')[0])\n            words.add(w)\n    return words\n\n\ndef load_words(args, examples):\n    \"\"\"Iterate and index all the words in examples (documents + questions).\"\"\"\n    def _insert(iterable):\n        for w in iterable:\n            w = Dictionary.normalize(w)\n            if valid_words and w not in valid_words:\n                continue\n            words.add(w)\n\n    if args.restrict_vocab and args.embedding_file:\n        logger.info('Restricting to words in %s' % args.embedding_file)\n        valid_words = index_embedding_words(args.embedding_file)\n        logger.info('Num words in set = %d' % len(valid_words))\n    else:\n        valid_words = None\n\n    words = set()\n    for ex in examples:\n        _insert(ex['question'])\n        _insert(ex['document'])\n    return words\n\n\ndef build_word_dict(args, examples):\n    \"\"\"Return a word dictionary from question and document words in\n    provided examples.\n    \"\"\"\n    word_dict = Dictionary()\n    for w in load_words(args, examples):\n        word_dict.add(w)\n    return word_dict\n\ndef index_embedding_chars(char_embedding_file):\n    \"\"\"Put all the chars in char_embedding_file into a set.\"\"\"\n    chars = set()\n    with open(char_embedding_file) as f:\n        for line in f:\n            c = Dictionary.normalize(line.rstrip().split(' ')[0])\n            chars.add(c)\n    return chars\n\ndef load_chars(args, examples):\n    \"\"\"Iterate and index all the chars in examples (documents + questions).\"\"\"\n    def _insert(iterable):\n        for cs in iterable:\n            for c in cs: \n                c = Dictionary.normalize(c)\n                if valid_chars and c not in valid_chars:\n                    continue\n                chars.add(c)\n\n    if args.restrict_vocab and args.char_embedding_file:\n        logger.info('Restricting to chars in %s' % args.char_embedding_file)\n        valid_chars = index_embedding_chars(args.char_embedding_file)\n        logger.info('Num chars in set = %d' % len(valid_chars))\n    else:\n        valid_chars = None\n\n    chars = set()\n    for ex in examples:\n        _insert(ex['question_char'])\n        _insert(ex['document_char'])\n    return chars\n\ndef build_char_dict(args, examples):\n    \"\"\"Return a char dictionary from question and document words in\n    provided examples.\n    \"\"\"\n    char_dict = Dictionary()\n    for c in load_chars(args, examples):\n        char_dict.add(c)\n    return char_dict\n\ndef top_question_words(args, examples, word_dict):\n    \"\"\"Count and return the most common question words in provided examples.\"\"\"\n    word_count = Counter()\n    for ex in examples:\n        for w in ex['question']:\n            w = Dictionary.normalize(w)\n            if w in word_dict:\n                word_count.update([w])\n    return word_count.most_common(args.tune_partial)\n\n\ndef build_feature_dict(args, examples):\n    \"\"\"Index features (one hot) from fields in examples and options.\"\"\"\n    def _insert(feature):\n        if feature not in feature_dict:\n            feature_dict[feature] = len(feature_dict)\n\n    feature_dict = {}\n\n    # Exact match features\n    if args.use_exact_match:\n        _insert('in_cased')\n        _insert('in_uncased')\n        if args.use_lemma:\n            _insert('in_lemma')\n\n    # Part of speech tag features\n    if args.use_pos:\n        for ex in examples:\n            for w in ex['cpos']:\n                _insert('pos=%s' % w)\n            for w in ex['qpos']:\n                _insert('pos=%s' % w)\n\n    # Named entity tag features\n    if args.use_ner:\n        for ex in examples:\n            for w in ex['cner']:\n                _insert('ner=%s' % w)\n            for w in ex['qner']:\n                _insert('ner=%s' % w)\n\n    # Term frequency feature\n    if args.use_tf:\n        _insert('tf')\n\n    return feature_dict\n\n\n# ------------------------------------------------------------------------------\n# Evaluation. Follows official evalutation script for v1.1 of the SQuAD dataset.\n# ------------------------------------------------------------------------------\n\n\ndef normalize_answer(s):\n    \"\"\"Lower text and remove punctuation, articles and extra whitespace.\"\"\"\n    def remove_articles(text):\n        return re.sub(r'\\b(a|an|the)\\b', ' ', text)\n\n    def white_space_fix(text):\n        return ' '.join(text.split())\n\n    def remove_punc(text):\n        exclude = set(string.punctuation)\n        return ''.join(ch for ch in text if ch not in exclude)\n\n    def lower(text):\n        return text.lower()\n\n    return white_space_fix(remove_articles(remove_punc(lower(s))))\n\n\ndef f1_score(prediction, ground_truth):\n    \"\"\"Compute the geometric mean of precision and recall for answer tokens.\"\"\"\n    prediction_tokens = normalize_answer(prediction).split()\n    ground_truth_tokens = normalize_answer(ground_truth).split()\n    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)\n    num_same = sum(common.values())\n    if num_same == 0:\n        return 0\n    precision = 1.0 * num_same / len(prediction_tokens)\n    recall = 1.0 * num_same / len(ground_truth_tokens)\n    f1 = (2 * precision * recall) / (precision + recall)\n    return f1\n\n\ndef exact_match_score(prediction, ground_truth):\n    \"\"\"Check if the prediction is a (soft) exact match with the ground truth.\"\"\"\n    return normalize_answer(prediction) == normalize_answer(ground_truth)\n\n\ndef regex_match_score(prediction, pattern):\n    \"\"\"Check if the prediction matches the given regular expression.\"\"\"\n    try:\n        compiled = re.compile(\n            pattern,\n            flags=re.IGNORECASE + re.UNICODE + re.MULTILINE\n        )\n    except BaseException:\n        logger.warn('Regular expression failed to compile: %s' % pattern)\n        return False\n    return compiled.match(prediction) is not None\n\n\ndef metric_max_over_ground_truths(metric_fn, prediction, ground_truths):\n    \"\"\"Given a prediction and multiple valid answers, return the score of\n    the best prediction-answer_n pair given a metric function.\n    \"\"\"\n    scores_for_ground_truths = []\n    for ground_truth in ground_truths:\n        score = metric_fn(prediction, ground_truth)\n        scores_for_ground_truths.append(score)\n    return max(scores_for_ground_truths)\n\n\n# ------------------------------------------------------------------------------\n# Utility classes\n# ------------------------------------------------------------------------------\n\n\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value.\"\"\"\n\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n\nclass Timer(object):\n    \"\"\"Computes elapsed time.\"\"\"\n\n    def __init__(self):\n        self.running = True\n        self.total = 0\n        self.start = time.time()\n\n    def reset(self):\n        self.running = True\n        self.total = 0\n        self.start = time.time()\n        return self\n\n    def resume(self):\n        if not self.running:\n            self.running = True\n            self.start = time.time()\n        return self\n\n    def stop(self):\n        if self.running:\n            self.running = False\n            self.total += time.time() - self.start\n        return self\n\n    def time(self):\n        if self.running:\n            return self.total + time.time() - self.start\n        return self.total\n"
  },
  {
    "path": "vector.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2018-present, HKUST-KnowComp.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\"\"\"Functions for putting examples into torch format.\"\"\"\n\nfrom collections import Counter\nimport torch\n\n\ndef vectorize(ex, model, single_answer=False):\n    \"\"\"Torchify a single example.\"\"\"\n    args = model.args\n    word_dict = model.word_dict\n    char_dict = model.char_dict\n    feature_dict = model.feature_dict\n\n    # Index words\n    document = torch.LongTensor([word_dict[w] for w in ex['document']])\n    document_char = [torch.LongTensor([char_dict[c] for c in cs]) for cs in ex['document_char']]\n    question = torch.LongTensor([word_dict[w] for w in ex['question']])\n    question_char = [torch.LongTensor([char_dict[c] for c in cs]) for cs in ex['question_char']]\n\n    # Create extra features vector\n    if len(feature_dict) > 0:\n        c_features = torch.zeros(len(ex['document']), len(feature_dict))\n        q_features = torch.zeros(len(ex['question']), len(feature_dict))\n    else:\n        c_features = None\n        q_features = None\n\n    # f_{exact_match}\n    if args.use_exact_match:\n        q_words_cased = {w for w in ex['question']}\n        q_words_uncased = {w.lower() for w in ex['question']}\n        q_lemma = {w for w in ex['qlemma']} if args.use_lemma else None\n        for i in range(len(ex['document'])):\n            if ex['document'][i] in q_words_cased:\n                c_features[i][feature_dict['in_cased']] = 1.0\n            if ex['document'][i].lower() in q_words_uncased:\n                c_features[i][feature_dict['in_uncased']] = 1.0\n            if q_lemma and ex['clemma'][i] in q_lemma:\n                c_features[i][feature_dict['in_lemma']] = 1.0\n\n        c_words_cased = {w for w in ex['document']}\n        c_words_uncased = {w.lower() for w in ex['document']}\n        c_lemma = {w for w in ex['clemma']} if args.use_lemma else None\n        for i in range(len(ex['question'])):\n            if ex['question'][i] in c_words_cased:\n                q_features[i][feature_dict['in_cased']] = 1.0\n            if ex['question'][i].lower() in c_words_uncased:\n                q_features[i][feature_dict['in_uncased']] = 1.0\n            if c_lemma and ex['qlemma'][i] in c_lemma:\n                q_features[i][feature_dict['in_lemma']] = 1.0\n\n    # f_{token} (POS)\n    if args.use_pos:\n        for i, w in enumerate(ex['cpos']):\n            f = 'pos=%s' % w\n            if f in feature_dict:\n                c_features[i][feature_dict[f]] = 1.0\n        for i, w in enumerate(ex['qpos']):\n            f = 'pos=%s' % w\n            if f in feature_dict:\n                q_features[i][feature_dict[f]] = 1.0\n\n    # f_{token} (NER)\n    if args.use_ner:\n        for i, w in enumerate(ex['cner']):\n            f = 'ner=%s' % w\n            if f in feature_dict:\n                c_features[i][feature_dict[f]] = 1.0\n        for i, w in enumerate(ex['qner']):\n            f = 'ner=%s' % w\n            if f in feature_dict:\n                q_features[i][feature_dict[f]] = 1.0\n\n    # f_{token} (TF)\n    if args.use_tf:\n        counter = Counter([w.lower() for w in ex['document']])\n        l = len(ex['document'])\n        for i, w in enumerate(ex['document']):\n            c_features[i][feature_dict['tf']] = counter[w.lower()] * 1.0 / l\n        counter = Counter([w.lower() for w in ex['question']])\n        l = len(ex['question'])\n        for i, w in enumerate(ex['question']):\n            q_features[i][feature_dict['tf']] = counter[w.lower()] * 1.0 / l\n\n    # Maybe return without target\n    if 'answers' not in ex:\n        return document, document_char, c_features, question, question_char, q_features, ex['id']\n\n    # ...or with target(s) (might still be empty if answers is empty)\n    if single_answer:\n        assert(len(ex['answers']) > 0)\n        start = torch.LongTensor(1).fill_(ex['answers'][0][0])\n        end = torch.LongTensor(1).fill_(ex['answers'][0][1])\n    else:\n        start = [a[0] for a in ex['answers']]\n        end = [a[1] for a in ex['answers']]\n    \n    return document, document_char, c_features, question, question_char, q_features, start, end, ex['id']\n\n\ndef batchify(batch):\n    \"\"\"Gather a batch of individual examples into one batch.\"\"\"\n    NUM_INPUTS = 6\n    NUM_TARGETS = 2\n    NUM_EXTRA = 1\n\n    docs = [ex[0] for ex in batch]\n    doc_chars = [ex[1] for ex in batch]\n    c_features = [ex[2] for ex in batch]\n    questions = [ex[3] for ex in batch]\n    question_chars = [ex[4] for ex in batch]\n    q_features = [ex[5] for ex in batch]\n    ids = [ex[-1] for ex in batch]\n\n    # Batch documents and features\n    max_length = max([d.size(0) for d in docs])\n    # max_char_length = max([c.size(0) for cs in doc_chars for c in cs])\n    max_char_length = 13\n    x1 = torch.LongTensor(len(docs), max_length).zero_()\n    x1_c = torch.LongTensor(len(docs), max_length, max_char_length).zero_()\n    x1_mask = torch.ByteTensor(len(docs), max_length).fill_(1)\n    if c_features[0] is None:\n        x1_f = None\n    else:\n        x1_f = torch.zeros(len(docs), max_length, c_features[0].size(1))\n    for i, d in enumerate(docs):\n        x1[i, :d.size(0)].copy_(d)\n        x1_mask[i, :d.size(0)].fill_(0)\n        if x1_f is not None:\n            x1_f[i, :d.size(0)].copy_(c_features[i])\n    for i, cs in enumerate(doc_chars):\n        for j, c in enumerate(cs):\n            c_ = c[:max_char_length]\n            x1_c[i, j, :c_.size(0)].copy_(c_)\n\n    # Batch questions\n    max_length = max([q.size(0) for q in questions])\n    x2 = torch.LongTensor(len(questions), max_length).zero_()\n    x2_c = torch.LongTensor(len(questions), max_length, max_char_length).zero_()\n    x2_mask = torch.ByteTensor(len(questions), max_length).fill_(1)\n    if q_features[0] is None:\n        x2_f = None\n    else:\n        x2_f = torch.zeros(len(questions), max_length, q_features[0].size(1))\n    for i, d in enumerate(questions):\n        x2[i, :d.size(0)].copy_(d)\n        x2_mask[i, :d.size(0)].fill_(0)\n        if x2_f is not None:\n            x2_f[i, :d.size(0)].copy_(q_features[i])\n    for i, cs in enumerate(question_chars):\n        for j, c in enumerate(cs):\n            c_ = c[:max_char_length]\n            x2_c[i, j, :c_.size(0)].copy_(c_)\n\n    # Maybe return without targets\n    if len(batch[0]) == NUM_INPUTS + NUM_EXTRA:\n        return x1, x1_c, x1_f, x1_mask, x2, x2_c, x2_f, x2_mask, ids\n\n    elif len(batch[0]) == NUM_INPUTS + NUM_EXTRA + NUM_TARGETS:\n        # ...Otherwise add targets\n        if torch.is_tensor(batch[0][NUM_INPUTS]):\n            y_s = torch.cat([ex[NUM_INPUTS] for ex in batch])\n            y_e = torch.cat([ex[NUM_INPUTS+1] for ex in batch])\n        else:\n            y_s = [ex[NUM_INPUTS] for ex in batch]\n            y_e = [ex[NUM_INPUTS+1] for ex in batch]\n    else:\n        raise RuntimeError('Incorrect number of inputs per example.')\n\n    return x1, x1_c, x1_f, x1_mask, x2, x2_c, x2_f, x2_mask, y_s, y_e, ids\n"
  }
]