Questionable - Part 4.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare for Paperspace. Manage these via conda or pipenv on your own machine\n",
    "!pip --quiet install nmslib flask torch transformers sklearn pyarrow seaborn spacy[cuda92]\n",
    "%run init_container.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import requests\n",
    "import random\n",
    "import pickle\n",
    "from itertools import islice\n",
    "\n",
    "import pandas as pd\n",
    "import json\n",
    "import sklearn\n",
    "import spacy\n",
    "from spacy.tokens import DocBin\n",
    "from spacy.strings import hash_string\n",
    "from unidecode import unidecode\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from itertools import islice\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "from scipy.special import softmax\n",
    "\n",
    "from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer\n",
    "from transformers import *\n",
    "import seaborn as sns\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qa.constants import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spacy.prefer_gpu()\n",
    "spen = spacy.load(\"en_core_web_sm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This takes a while the first time, since from_pretrained() downloads and caches the model weights\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
    "model = BertForQuestionAnswering \\\n",
    "    .from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad') \\\n",
    "    .to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(SQUAD_TRAIN) as f:\n",
    "    doc = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paragraphs = []\n",
    "questions = []\n",
    "topics = []\n",
    "for topic in doc[\"data\"]:\n",
    "    topics.append(topic[\"title\"])\n",
    "    for pgraph in topic[\"paragraphs\"]:\n",
    "        paragraphs.append(pgraph[\"context\"])\n",
    "        for qa in pgraph[\"qas\"]:\n",
    "            if not qa[\"is_impossible\"]:\n",
    "                questions.append(qa[\"question\"])\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(paragraphs), len(questions), random.sample(paragraphs, 2), random.sample(questions, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\", \".join(topics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "if os.path.isfile(DOCBIN_CACHE) and os.path.isdir(VOCAB_CACHE):\n",
    "    spen.vocab.from_disk(VOCAB_CACHE)\n",
    "\n",
    "    with open(DOCBIN_CACHE, \"rb\") as f:\n",
    "        bb = f.read()\n",
    "        doc_bin = DocBin().from_bytes(bb)\n",
    "    docs = list(doc_bin.get_docs(spen.vocab))\n",
    "else:\n",
    "    doc_bin = DocBin(attrs=[\"LEMMA\", \"ENT_IOB\", \"ENT_TYPE\"], store_user_data=True)\n",
    "    for doc in spen.pipe(tqdm(paragraphs)):\n",
    "        doc_bin.add(doc)\n",
    "    with open(DOCBIN_CACHE, \"wb\") as f:\n",
    "        f.write(doc_bin.to_bytes())\n",
    "    spen.vocab.to_disk(VOCAB_CACHE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "def lemmatize_preproc(doc):\n",
    "    return [unidecode(tok.lemma_.lower()) for tok in doc if not tok.is_stop]\n",
    "\n",
    "if not os.path.isfile(VECTOR_CACHE):\n",
    "    print(f'Indexing corpus')\n",
    "    vectorizer = TfidfVectorizer(\n",
    "        analyzer=lemmatize_preproc,\n",
    "        stop_words='english', min_df=10, max_df=.5, ngram_range=(1,3))\n",
    "    tfidf = vectorizer.fit_transform(docs)\n",
    "    with open(VECTOR_CACHE, \"wb\") as f:\n",
    "        pickle.dump(dict(vectorizer=vectorizer, tfidf=tfidf), f)\n",
    "else:\n",
    "    print(f'Loading vector cache from {VECTOR_CACHE}')\n",
    "    with open(VECTOR_CACHE, \"rb\") as f:\n",
    "        cache = pickle.load(f)\n",
    "        tfidf = cache[\"tfidf\"]\n",
    "        vectorizer = cache[\"vectorizer\"]\n",
    "        #vocab = cache[\"vocab\"]\n",
    "    #vectorizer = TfidfVectorizer(\n",
    "    #    analyzer=lambda doc: [tok.lemma_.lower() for tok in doc],\n",
    "     ##   vocabulary=vocab,\n",
    "      #  stop_words='english', min_df=5, max_df=.5, ngram_range=(1,3))\n",
    "        \n",
    "len(vectorizer.vocabulary_), vectorizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type(vectorizer.vocabulary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len([1 for doc in tqdm(docs) if 'russes' in doc.text.lower()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[ (i,doc) for (i, doc) in enumerate(tqdm(docs)) if 'Calafat' in doc.text]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUMERICS = set([\"DATE\", \"TIME\", \"PERCENT\", \"MONEY\", \"QUANTITY\", \"ORDINAL\", \"CARDINAL\"])\n",
    "def doc_entities(doc):\n",
    "    ents = [e for e in doc.ents if e.label_ not in NUMERICS]\n",
    "    result = (unidecode(w.lemma_.lower()) for s in ents for w in s if w.is_alpha and not w.is_stop)\n",
    "    #return list(result)\n",
    "    return [w for w in result if w not in vectorizer.vocabulary_]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "ent_vecr = CountVectorizer(\n",
    "    analyzer=doc_entities,\n",
    "    stop_words='english', max_df=10)\n",
    "ent_tfidf = ent_vecr.fit_transform(docs)\n",
    "\n",
    "len(ent_vecr.vocabulary_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "N = 10_001\n",
    "hashed_ents = [set() for i in range(N)]\n",
    "for (i, words) in enumerate(tqdm(ent_vecr.inverse_transform(ent_tfidf))):\n",
    "    for w in words:\n",
    "        h = hash_string(str(w))\n",
    "        hashed_ents[h%N].add(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bucket_size = [len(x) for x in hashed_ents]\n",
    "max(bucket_size), sns.distplot(bucket_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ratio of empty buckets\n",
    "sum([x == 0 for x in bucket_size]) / N"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def contexts_by_entities(doc):\n",
    "    \"\"\"Returns a set of document ids that *might* be related to named entities in the pre-processed question\"\"\"\n",
    "    ents = doc_entities(doc)\n",
    "    buckets = [hash_string(word)%N for word in ents]\n",
    "    return set([doc_id for slot in buckets for doc_id in hashed_ents[slot]])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ent_questions = list(q for q in spen.pipe(tqdm(questions)) if len(doc_entities(q)) > 0)\n",
    "len(questions), len(ent_questions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "queries = list(random.sample(ent_questions, 5))\n",
    "query_ents = [(doc_entities(query)) for query in queries]\n",
    "queries, query_ents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ranked_contexts_by_entities(query, topk=10, thresh=1e-5):\n",
    "    doc_ids = list(contexts_by_entities(query))\n",
    "    if len(doc_ids) < 1:\n",
    "        return []\n",
    "\n",
    "    contexts = [docs[i] for i in doc_ids]\n",
    "    scores = vectorizer.transform(contexts) * vectorizer.transform([query]).transpose()\n",
    "    scores = (np.asarray(scores.todense()).flatten())\n",
    "    sort_scores = np.asarray(np.flip((scores).argsort())) # indices ranked by score\n",
    "    useful = sort_scores[scores[sort_scores] >= thresh] # Filter out irrelevant scores\n",
    "    top_indices = useful[:topk]\n",
    "    return [(doc_ids[i], scores[i]) for i in top_indices]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = spen(\"When did the US send troops to the Philippines to battle terrorists?\")\n",
    "doc_entities(query), ranked_contexts_by_entities(query, thresh=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Nearest Neighbor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nmslib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfidf = tfidf.astype(np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfidf.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "index = nmslib.init(method='hnsw', space='cosinesimil_sparse', data_type=nmslib.DataType.SPARSE_VECTOR)\n",
    "index.addDataPointBatch(tfidf.astype(np.float32))\n",
    "index.createIndex({'post': 2, 'efConstruction': 500, 'M': 64},print_progress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index.saveIndex('cache/nmsIndex.nms', save_data=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "index.loadIndex('cache/nmsIndex.nms', load_data=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index.setQueryTimeParams(dict(ef=500))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Odd, punctuation really matters for spacy's NER\n",
    "# Some queries trigger cuda error\n",
    "question = spen(\"Why did Feynman pick Caltech over Princeton?\")\n",
    "query = vectorizer.transform([question])\n",
    "query, vectorizer.inverse_transform(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "knn_ids, _ = index.knnQueryBatch(query, k=10, num_threads=8)[0]\n",
    "knn_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[docs[i] for i in knn_ids[:4]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ranked_contexts_by_entities(question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ids_by_ent = [i for (i,s) in ranked_contexts_by_entities(question)]\n",
    "ids_by_ent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "merge_ids = set(ids_by_ent + [i for i in knn_ids])\n",
    "merge_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contexts = [docs[i] for i in merge_ids]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "question_df = pd.DataFrame.from_records([ {\n",
    "    'question': question.text,\n",
    "    'context':  ctx.text,\n",
    "} for ctx in contexts ])\n",
    "question_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO use stride when context too long\n",
    "question_df[\"doc_id\"] = merge_ids\n",
    "question_df[\"encoded\"] = question_df.apply(lambda row: tokenizer.encode(\"[CLS] \" + row[\"question\"] + \" [SEP] \" + row[\"context\"] + \" [SEP]\", add_special_tokens=False, max_length=512), axis=1)\n",
    "question_df[\"context_start\"] = question_df.apply(lambda row: row[\"encoded\"].index(102) + 1, axis=1)\n",
    "question_df[\"context_end\"] = question_df.apply(lambda row: len(row[\"encoded\"])-1, axis=1)\n",
    "question_df[\"tok_type\"] = question_df.apply(lambda row: [0 if i <= row[\"encoded\"].index(102) else 1 for i in range(len(row[\"encoded\"]))], axis=1)\n",
    "question_df.iloc[:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "with torch.no_grad():\n",
    "    X = torch.nn.utils.rnn.pad_sequence([torch.tensor(row) for row in question_df[\"encoded\"]], batch_first=True).to(device)\n",
    "    T = torch.nn.utils.rnn.pad_sequence([torch.tensor(row) for row in question_df[\"tok_type\"]], batch_first=True).to(device)\n",
    "    start_scores, end_scores = model(X, token_type_ids=T)\n",
    "    print(f'start_scores = {start_scores}')\n",
    "    max_score, max_start = torch.max(start_scores, axis=1)\n",
    "    soft_max = F.softmax(max_score, dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "answer_df = question_df[[\"doc_id\", \"context\", \"encoded\", \"context_start\", \"context_end\"]].copy()\n",
    "answer_df[\"answer_score\"] = max_score.cpu().numpy()\n",
    "answer_df[\"answer_start\"] = max_start.cpu().numpy()\n",
    "answer_df[\"answer_softmax\"] = soft_max.cpu().numpy()\n",
    "answer_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_len = torch.zeros_like(max_start)\n",
    "for i in range(max_start.shape[0]):\n",
    "    max_len[i] = torch.argmax(end_scores[i,max_start[i]:]) + 1\n",
    "    \n",
    "answer_df[\"answer_length\"] = max_len.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "answer_df = answer_df[answer_df.answer_start > answer_df.context_start]\n",
    "answer_df = answer_df[answer_df.answer_start < answer_df.context_end]\n",
    "answer_df = answer_df[answer_df.answer_score > 1.0]\n",
    "answer_df = answer_df.sort_values(by=\"answer_score\", ascending=False)\n",
    "answer_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decode_answer(row):\n",
    "    input_ids = row.encoded\n",
    "    offset = row.answer_start\n",
    "    length = np.clip(row.answer_length, 0, 20)\n",
    "    return tokenizer.decode(input_ids[offset:][:length])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "answer_df[\"answer\"] = answer_df.apply(decode_answer, axis=1)\n",
    "answer_df[[\"answer_softmax\", \"answer_score\", \"answer\"]].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "answer_df[[\"answer_softmax\", \"answer_score\", \"answer\", \"doc_id\", \"context\"]].iloc[:3].to_dict(orient=\"records\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}