query_expansion_bert.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BERT-based Query Expansion Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import BertForMaskedLM, BertTokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The BERT large uncased, in its variant _whole word masking_, has been trained over BookCorpus and Wikipedia English with NSP - Next Sentenct Prediction - and MLM - Masked Language Modeling - objectives. Let's import it and its tokenizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-large-uncased-whole-word-masking were not used when initializing BertForMaskedLM: ['bert.pooler.dense.bias', 'bert.pooler.dense.weight', 'cls.seq_relationship.bias', 'cls.seq_relationship.weight']\n",
      "- This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained('bert-large-uncased-whole-word-masking', cache_dir = 'hf_cache')\n",
    "unmasking_model = BertForMaskedLM.from_pretrained('bert-large-uncased-whole-word-masking', cache_dir = 'hf_cache')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pipeline convenience object is created to interface with both the tokenizer and the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import FillMaskPipeline\n",
    "\n",
    "unmasker = FillMaskPipeline(model = unmasking_model, tokenizer = tokenizer, tokenizer_kwargs = {\"truncation\": True})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the following line we are getting all the candidates, to the masked word, proposed by BERT. Each substitute has a confidence level associated with the token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "def mask_token(tokens: List[str], idx: int) -> str:\n",
    "    tokens = tokens[:]\n",
    "    tokens[idx] = '[MASK]'\n",
    "    return ' '.join(tokens)\n",
    "\n",
    "original_sentence = 'modern shared room near Harvard.'\n",
    "original_sentence_tokens = tokenizer.tokenize(original_sentence)\n",
    "\n",
    "masked_sentence = mask_token(original_sentence_tokens, 2)\n",
    "candidates = unmasker(masked_sentence, top_k = 50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point, some words can be more suitable than others. We try to figure out the fitness level by reinserting the token into the sentence and by testing the similarity between the original sentence and the one with the mask replaced."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['room', 'house', 'residence', 'land', 'hall']\n"
     ]
    }
   ],
   "source": [
    "from transformers import BertModel\n",
    "import torch\n",
    "\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', cache_dir = 'hf_cache')\n",
    "encoder = BertModel.from_pretrained('bert-base-uncased', output_hidden_states = True, cache_dir = 'hf_cache')\n",
    "\n",
    "def get_meaned_embeddings(sentence: str):\n",
    "    tokens = tokenizer.tokenize(sentence)\n",
    "    input_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "    input_ids = torch.tensor(input_ids).unsqueeze(0)\n",
    "    with torch.no_grad():\n",
    "        outputs = encoder(input_ids)\n",
    "        embedding = outputs.last_hidden_state[0]\n",
    "\n",
    "    return embedding.mean(dim = 0)\n",
    "\n",
    "cos_sim = torch.nn.CosineSimilarity(dim = 0)\n",
    "\n",
    "from operator import itemgetter\n",
    "import pydash\n",
    "\n",
    "original_sentence_embedding = get_meaned_embeddings(original_sentence)\n",
    "\n",
    "\n",
    "similarities = (\n",
    "    pydash.chain(candidates)\n",
    "        .map(itemgetter('sequence'))  # Get complete sentence\n",
    "        .map(get_meaned_embeddings)  # Get context vectors\n",
    "        .map(lambda x: cos_sim(x, original_sentence_embedding))  # Compute the similarity\n",
    "        .value()\n",
    ")\n",
    "\n",
    "THRESHOLD = 0.8\n",
    "SLICE = 5\n",
    "\n",
    "expansions = (\n",
    "    pydash.chain(candidates)\n",
    "        .map(itemgetter('token_str'))\n",
    "        .zip(similarities)\n",
    "        .filter(lambda t: t[1] > THRESHOLD)\n",
    "        .sort(key = itemgetter(1), reverse = True)\n",
    "        .map(itemgetter(0))\n",
    "        .take(SLICE)\n",
    "        .value()\n",
    ")\n",
    "\n",
    "print(expansions)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "cc551502c4709d65c35225dab174e15b6f215f4b9bca0aec7618bac23f51ade6"
  },
  "kernelspec": {
   "display_name": "Python 3.11.6 ('venv': venv)",
   "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.11.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}