inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_DIR = None\n",
    "MASK_CLS = 'ilm.mask.hierarchical.MaskHierarchical'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "File ‘/tmp/ilm/models/sto_ilm/pytorch_model.bin’ already there; not retrieving.\n",
      "File ‘/tmp/ilm/models/sto_ilm/config.json’ already there; not retrieving.\n",
      "File ‘/tmp/ilm/models/sto_ilm/additional_ids_to_tokens.pkl’ already there; not retrieving.\n"
     ]
    }
   ],
   "source": [
    "# Download pretrained model\n",
    "\n",
    "if MODEL_DIR is None:\n",
    "    !python acl20_repro.py model sto ilm | bash\n",
    "    MODEL_DIR = '/tmp/ilm/models/sto_ilm'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'<|startofinfill|>': 50257, '<|endofinfill|>': 50258, '<|infill_document|>': 50259, '<|infill_paragraph|>': 50260, '<|infill_sentence|>': 50261, '<|infill_ngram|>': 50262, '<|infill_word|>': 50263}\n"
     ]
    }
   ],
   "source": [
    "# Prepare tokenizer\n",
    "\n",
    "import os\n",
    "import pickle\n",
    "\n",
    "import ilm.tokenize_util\n",
    "\n",
    "tokenizer = ilm.tokenize_util.Tokenizer.GPT2\n",
    "with open(os.path.join(MODEL_DIR, 'additional_ids_to_tokens.pkl'), 'rb') as f:\n",
    "    additional_ids_to_tokens = pickle.load(f)\n",
    "additional_tokens_to_ids = {v:k for k, v in additional_ids_to_tokens.items()}\n",
    "try:\n",
    "    ilm.tokenize_util.update_tokenizer(additional_ids_to_tokens, tokenizer)\n",
    "except ValueError:\n",
    "    print('Already updated')\n",
    "print(additional_tokens_to_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load model\n",
    "\n",
    "import torch\n",
    "from transformers import GPT2LMHeadModel\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model = GPT2LMHeadModel.from_pretrained(MODEL_DIR)\n",
    "model.eval()\n",
    "_ = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Math Class\n",
      "Chris was bad at<|infill_word|>.<|infill_sentence|><|infill_sentence|><|infill_sentence|> He ended up passing the test.\n"
     ]
    }
   ],
   "source": [
    "# Create context\n",
    "\n",
    "context = \"\"\"\n",
    "Math Class\n",
    "Chris was bad at _. _ _ _ He ended up passing the test.\n",
    "\"\"\".strip()\n",
    "\n",
    "context_ids = ilm.tokenize_util.encode(context, tokenizer)\n",
    "\n",
    "# Replace blanks with appropriate tokens from left to right\n",
    "_blank_id = ilm.tokenize_util.encode(' _', tokenizer)[0]\n",
    "context_ids[context_ids.index(_blank_id)] = additional_tokens_to_ids['<|infill_word|>']\n",
    "context_ids[context_ids.index(_blank_id)] = additional_tokens_to_ids['<|infill_sentence|>']\n",
    "context_ids[context_ids.index(_blank_id)] = additional_tokens_to_ids['<|infill_sentence|>']\n",
    "context_ids[context_ids.index(_blank_id)] = additional_tokens_to_ids['<|infill_sentence|>']\n",
    "print(ilm.tokenize_util.decode(context_ids, tokenizer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------\n",
      "Math Class\n",
      "Chris was bad at math. He needed to do well on the test. He decided to study more. Chris spent all day studying. He ended up passing the test.\n",
      "--------------------------------------------------------------------------------\n",
      "Math Class\n",
      "Chris was bad at math. He decided to go to a class to improve. His teacher learned that Chris was failing math. Chris decided to take his own math class. He ended up passing the test.\n"
     ]
    }
   ],
   "source": [
    "from ilm.infer import infill_with_ilm\n",
    "\n",
    "generated = infill_with_ilm(\n",
    "    model,\n",
    "    additional_tokens_to_ids,\n",
    "    context_ids,\n",
    "    num_infills=2)\n",
    "for g in generated:\n",
    "    print('-' * 80)\n",
    "    print(ilm.tokenize_util.decode(g, tokenizer))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ilm",
   "language": "python",
   "name": "ilm"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}