7.1_unsupervised_training_tsdae.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Unsupervised sentence embedding learning - TSDAE\n",
    "\n",
    "In this notebook, we will look at the work of `Kexin Wang, Nils Reimers, Iryna Gurevych` on their paper [TSDAE: Using Transformer-based Sequential Denoising Auto-Encoder for Unsupervised Sentence Embedding Learning](https://arxiv.org/pdf/2104.06979.pdf).\n",
    "\n",
    "Here is the summary of the paper by the authors.\n",
    "\n",
    "> Learning sentence embeddings often requires\n",
    "a large amount of labeled data. However,\n",
    "for most tasks and domains, labeled data is\n",
    "seldom available and creating it is expensive.\n",
    "In this work, we present a new state-of-the-art unsupervised method based on pre-trained\n",
    "Transformers and Sequential Denoising AutoEncoder (TSDAE) which outperforms previous approaches by up to 6.4 points. It can achieve up to 93.1% of the performance of indomain supervised approaches. Further, we\n",
    "show that TSDAE is a strong domain adaptation and pre-training method for sentence\n",
    "embeddings, significantly outperforming other approaches like Masked Language Model.\n",
    ">\n",
    "> A crucial shortcoming of previous studies is the narrow evaluation: Most work mainly evaluates on the single task of Semantic Textual Similarity (STS), which does not require any domain knowledge. It is unclear if these proposed methods generalize to other domains and tasks. We fill this gap and evaluate TSDAE and other recent approaches on four different datasets from heterogeneous domains.\n",
    "\n",
    "The techniques we have discussed so far - Bi-Encoders, Cross-Encoders, etc. required labeled data. While they show great performance on in-domain data, their performance declines rapidly on out-of-domain data.\n",
    "\n",
    "| ![](assets/models_generalize.png) | \n",
    "|:--:| \n",
    "| Fig. 1. Illustration of generalizability of Neural IR models on 18 IR datasets from the [BIER benchmark](https://arxiv.org/abs/2104.08663). (Image source: https://www.youtube.com/watch?v=xbdLowiQTlk&t=658s) |\n",
    "\n",
    "We can see that BM25(retriever/candidate generator) and Cross-Encoder(Re-Ranker) works the best while many dense retrievers fail to outperform just the BM25 on most on the datasets.\n",
    "\n",
    "So why not use BM25 + Cross-Encoders?\n",
    "\n",
    "As we mentioned earlier, Cross-Encoders are expensive as we can't index the corpus beforehand. For each query, we would need to score the query against all the retrieved candidates(which could be in 100s).\n",
    "\n",
    "So we again come back to Bi-Encoders as they are fast at both indexing the corpus and inference(using ANNs).\n",
    "\n",
    "In TSDAE, the authors train encoder based models with pre-training objective similar to `Masked Language Modeling` but with slight variation.\n",
    "In MLM, we mask some tokens and train the encoders to predict the masked tokens, but in TSDAE, we delete some tokens from the input sentences, create a pooled representation of the sentence(MEAN-pooling or CLS-embedding) and pass that to a Denoising Auto-Encoder to recreate the original input text.\n",
    "\n",
    "| ![](assets/tsdae.png) | \n",
    "|:--:| \n",
    "| Fig. 2. Illustration of TSDAE. (Image source: https://arxiv.org/pdf/2104.06979.pdf) |\n",
    "\n",
    "The authors tried a bunch of configurations for adding noise. The best results came from deletion with a deletion ratio of 0.6.\n",
    "Note that the models were trained on a combination of SNLI and MultiNLI datasets without labels and evaluated on the STS benchmark with the metric Spearman rank correlation.\n",
    "CLS and Mean-pooling worked the best with similar performance. They recommend choosing CLS-pooling so we will also use that.\n",
    "\n",
    "| ![](assets/tsdae_config.png) |\n",
    "|:--:|\n",
    "| Fig. 3. Results with differnt noise types, noise ratios and pooling methods. (Image source: https://arxiv.org/pdf/2104.06979.pdf) |\n",
    "\n",
    "Here we will use the `sentence_transformers` library to train this architecture. Lets start ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "import datasets as hf_datasets\n",
    "import numpy as np\n",
    "import nltk\n",
    "import pandas as pd\n",
    "from sentence_transformers import InputExample, SentenceTransformer, LoggingHandler\n",
    "from sentence_transformers import models, util, datasets, evaluation, losses\n",
    "from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator\n",
    "import torch\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package punkt to /home/utsav/nltk_data...\n",
      "[nltk_data]   Package punkt is already up-to-date!\n"
     ]
    }
   ],
   "source": [
    "nltk.download(\"punkt\")\n",
    "pd.set_option(\"display.max_colwidth\", None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 10\n",
    "\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed_all(seed)\n",
    "\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For training, we will use a sample of the `snli` dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Reusing dataset snli (/home/utsav/.cache/huggingface/datasets/snli/plain_text/1.0.0/1f60b67533b65ae0275561ff7828aad5ee4282d0e6f844fd148d05d3c6ea251b)\n",
      "Parameter 'function'=<function <lambda> at 0x7fa080b5dc10> of the transform datasets.arrow_dataset.Dataset.filter@2.0.1 couldn't be hashed properly, a random hash was used instead. Make sure your transforms and parameters are serializable with pickle or dill for the dataset fingerprinting and caching to work. If you reuse this transform, the caching mechanism will consider it to be different from the previous calls and recompute everything. This warning is only showed once. Subsequent hashing failures won't be showed.\n",
      "Loading cached processed dataset at /home/utsav/.cache/huggingface/datasets/snli/plain_text/1.0.0/1f60b67533b65ae0275561ff7828aad5ee4282d0e6f844fd148d05d3c6ea251b/cache-08577eb1924770d3.arrow\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(109812,\n",
       " {'premise': 'Children smiling and waving at camera',\n",
       "  'hypothesis': 'They are smiling at their parents',\n",
       "  'label': 1})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = hf_datasets.load_dataset(\"snli\", split=\"train\")\n",
    "dataset = dataset.filter(lambda _: True if random.random() > 0.9 else False)\n",
    "\n",
    "len(dataset), dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "77550"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_sentences = [d[\"premise\"] for d in dataset]\n",
    "train_sentences = list(set(train_sentences))\n",
    "len(train_sentences)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = datasets.DenoisingAutoEncoderDataset(train_sentences)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at some examples of the noise added by the `DenoisingAutoEncoderDataset`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>noisy</th>\n",
       "      <th>original</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Two are crossing a while pedestrians them.</td>\n",
       "      <td>Two men are crossing a street carrying a frame, while pedestrians walk around them.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>a sleeping on with and dog sitting next him</td>\n",
       "      <td>a man sleeping on a bench outside with a white and black dog sitting next to him.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>men volleyball</td>\n",
       "      <td>Why are grown men playing volleyball?</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>A boy on hill of</td>\n",
       "      <td>A young boy standing on a hill that overlooks a village of homes.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>A man a Mohawk</td>\n",
       "      <td>A man with a many colored Mohawk smiling.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>swing</td>\n",
       "      <td>A girl rides on a swing.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>, white coats, standing</td>\n",
       "      <td>Two women, both wearing white coats, are standing outside a large framed doorway.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>man with</td>\n",
       "      <td>From inside building, view of man washing window with tool.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>woman glass of</td>\n",
       "      <td>A blond woman pours a glass of wine in a dim restaurant.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Children.</td>\n",
       "      <td>Children are playing with baseball bats.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>speak each with</td>\n",
       "      <td>Three women speak to each other in a room with tan walls.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>at a</td>\n",
       "      <td>People talking at a flea market.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>The dog jumping in.</td>\n",
       "      <td>The black dog is jumping up in the air.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>green pink top.</td>\n",
       "      <td>A woman wearing a green and pink dress is dancing with someone wearing a blue top with white pants.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>monk in orange with tattoos his city</td>\n",
       "      <td>Old monk in orange with tattoos on his chest standing on a city street.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>A male playing</td>\n",
       "      <td>A male musician dressed in white is playing a guitar.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>man in athletic shirt from, while in same type walk</td>\n",
       "      <td>A man in a white athletic shirt drinks from a cup, while others in the same type shirts walk around.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>biting soccer the ground</td>\n",
       "      <td>A black dog is biting onto a soccer ball on the dirt ground.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>A large lit</td>\n",
       "      <td>A large musical band stands on a colorfully lit stage.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>Two a an shirt a man wearing shirt, are swings in</td>\n",
       "      <td>Two people, a woman wearing an orange shirt and a man wearing a blue shirt, are jumping out of swings in midair.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                  noisy  \\\n",
       "0            Two are crossing a while pedestrians them.   \n",
       "1           a sleeping on with and dog sitting next him   \n",
       "2                                        men volleyball   \n",
       "3                                      A boy on hill of   \n",
       "4                                        A man a Mohawk   \n",
       "5                                                 swing   \n",
       "6                               , white coats, standing   \n",
       "7                                              man with   \n",
       "8                                        woman glass of   \n",
       "9                                             Children.   \n",
       "10                                      speak each with   \n",
       "11                                                 at a   \n",
       "12                                  The dog jumping in.   \n",
       "13                                      green pink top.   \n",
       "14                 monk in orange with tattoos his city   \n",
       "15                                       A male playing   \n",
       "16  man in athletic shirt from, while in same type walk   \n",
       "17                             biting soccer the ground   \n",
       "18                                          A large lit   \n",
       "19    Two a an shirt a man wearing shirt, are swings in   \n",
       "\n",
       "                                                                                                            original  \n",
       "0                                Two men are crossing a street carrying a frame, while pedestrians walk around them.  \n",
       "1                                  a man sleeping on a bench outside with a white and black dog sitting next to him.  \n",
       "2                                                                              Why are grown men playing volleyball?  \n",
       "3                                                  A young boy standing on a hill that overlooks a village of homes.  \n",
       "4                                                                          A man with a many colored Mohawk smiling.  \n",
       "5                                                                                           A girl rides on a swing.  \n",
       "6                                  Two women, both wearing white coats, are standing outside a large framed doorway.  \n",
       "7                                                        From inside building, view of man washing window with tool.  \n",
       "8                                                           A blond woman pours a glass of wine in a dim restaurant.  \n",
       "9                                                                           Children are playing with baseball bats.  \n",
       "10                                                         Three women speak to each other in a room with tan walls.  \n",
       "11                                                                                  People talking at a flea market.  \n",
       "12                                                                           The black dog is jumping up in the air.  \n",
       "13               A woman wearing a green and pink dress is dancing with someone wearing a blue top with white pants.  \n",
       "14                                           Old monk in orange with tattoos on his chest standing on a city street.  \n",
       "15                                                             A male musician dressed in white is playing a guitar.  \n",
       "16              A man in a white athletic shirt drinks from a cup, while others in the same type shirts walk around.  \n",
       "17                                                      A black dog is biting onto a soccer ball on the dirt ground.  \n",
       "18                                                            A large musical band stands on a colorfully lit stage.  \n",
       "19  Two people, a woman wearing an orange shirt and a man wearing a blue shirt, are jumping out of swings in midair.  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.DataFrame([train_dataset[i].texts for i in range(20)], columns=[\"noisy\", \"original\"])\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 16\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model config "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"bert-base-uncased\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.decoder.weight', 'cls.predictions.bias']\n",
      "- This IS expected if you are initializing BertModel 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 BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "embedding_model = models.Transformer(model_name)\n",
    "pooling = models.Pooling(embedding_model.get_word_embedding_dimension(), \"cls\")\n",
    "model = SentenceTransformer(modules=[embedding_model, pooling], device=\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loss = losses.DenoisingAutoEncoderLoss(model, decoder_name_or_path=model_name, tie_encoder_decoder=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2dee16b7f9004529904707dd017c3752",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch: 5/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model.fit(\n",
    "    train_objectives=[(train_dataloader, train_loss)],\n",
    "    epochs=5,\n",
    "    weight_decay=0,\n",
    "    scheduler=\"constantlr\",\n",
    "    optimizer_params={'lr': 3e-5},\n",
    "    show_progress_bar=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save(f\"models/{model_name}-tsdae\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation\n",
    "\n",
    "We will evaluate the trained model on a STS dataset. We will compare the finetuned model and the bert-base model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using the latest cached version of the module from /home/utsav/.cache/huggingface/modules/datasets_modules/datasets/glue/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad (last modified on Fri Jul 22 22:08:30 2022) since it couldn't be found locally at glue., or remotely on the Hugging Face Hub.\n",
      "Reusing dataset glue (/home/utsav/.cache/huggingface/datasets/glue/stsb/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "124dffca46f647d0b537326319ec5b4e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "(1500,\n",
       " {'sentence1': 'A man with a hard hat is dancing.',\n",
       "  'sentence2': 'A man wearing a hard hat is dancing.',\n",
       "  'label': 1.0,\n",
       "  'idx': 0})"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sts = hf_datasets.load_dataset(\"glue\", \"stsb\", split=\"validation\")\n",
    "sts = sts.map(lambda x: {\"label\": x[\"label\"] / 5.0})\n",
    "\n",
    "len(sts), sts[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_examples = [InputExample(texts=[data[\"sentence1\"], data[\"sentence2\"]], label=data[\"label\"])\n",
    "                  for data in sts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7487081683853725"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluator = EmbeddingSimilarityEvaluator.from_input_examples(input_examples, write_csv=False)\n",
    "evaluator(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.decoder.weight', 'cls.predictions.bias']\n",
      "- This IS expected if you are initializing BertModel 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 BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.3172617272784805"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "original_embedding_model = models.Transformer(\"bert-base-uncased\")\n",
    "pooling = models.Pooling(original_embedding_model.get_word_embedding_dimension(), \"cls\")\n",
    "\n",
    "original_model = SentenceTransformer(modules=[original_embedding_model, pooling])\n",
    "evaluator(original_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our unsupervised finetuned model is doing much better on this semantic similarity dataset as compared to the bert-base model.\n",
    "It's not performing as good as the supervised training we did for Bi and Cross encoders(where we saw average spearman rank correlation of `~0.80`) but that's to be expected.\n",
    "\n",
    "Let's look how this model can be used for retrieval tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ranked_docs(model: SentenceTransformer, query: str, corpus_emebds: np.array) -> None:\n",
    "    query_embed = model.encode(query)\n",
    "    scores = util.cos_sim(query_embed, corpus_embeds)\n",
    "    print(f\"Query - {query}\\n---\")\n",
    "    scores = scores.cpu().detach().numpy()[0]\n",
    "    scores_ix = np.argsort(scores)[::-1]\n",
    "    for ix in scores_ix:\n",
    "        print(f\"{scores[ix]: >.2f}\\t{corpus[ix]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus = [\n",
    "    \"A man is eating food.\",\n",
    "    \"A man is eating a piece of bread.\",\n",
    "    \"The girl is carrying a baby.\",\n",
    "    \"A man is riding a horse.\",\n",
    "    \"A woman is playing violin.\",\n",
    "    \"Two men pushed carts through the woods.\",\n",
    "    \"A man is riding a white horse on an enclosed ground.\",\n",
    "    \"A monkey is playing drums.\",\n",
    "    \"A cheetah is running behind its prey.\"\n",
    "]\n",
    "\n",
    "query = \"A man is eating pasta.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query - A man is eating pasta.\n",
      "---\n",
      "0.80\tA man is eating food.\n",
      "0.75\tA man is eating a piece of bread.\n",
      "0.53\tA man is riding a horse.\n",
      "0.47\tA woman is playing violin.\n",
      "0.42\tA cheetah is running behind its prey.\n",
      "0.41\tA monkey is playing drums.\n",
      "0.41\tA man is riding a white horse on an enclosed ground.\n",
      "0.39\tThe girl is carrying a baby.\n",
      "0.26\tTwo men pushed carts through the woods.\n"
     ]
    }
   ],
   "source": [
    "corpus_embeds = model.encode(corpus)\n",
    "get_ranked_docs(model, query, corpus_embeds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query - A man is eating pasta.\n",
      "---\n",
      "0.98\tA man is eating a piece of bread.\n",
      "0.96\tA man is riding a horse.\n",
      "0.96\tA woman is playing violin.\n",
      "0.96\tA man is eating food.\n",
      "0.95\tA monkey is playing drums.\n",
      "0.91\tA man is riding a white horse on an enclosed ground.\n",
      "0.89\tA cheetah is running behind its prey.\n",
      "0.86\tThe girl is carrying a baby.\n",
      "0.85\tTwo men pushed carts through the woods.\n"
     ]
    }
   ],
   "source": [
    "corpus_embeds = original_model.encode(corpus)\n",
    "get_ranked_docs(original_model, query, corpus_embeds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus = [\n",
    "    \"A man is eating food.\",\n",
    "    \"A man is eating a piece of bread.\",\n",
    "    \"A woman is playing violin.\",\n",
    "    \"Two men pushed carts through the woods.\",\n",
    "    \"A woman is practicing jumps with her horse.\",\n",
    "    \"A horse is running around the track.\"\n",
    "]\n",
    "\n",
    "query = \"Horse jumped over the obstacle.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query - Horse jumped over the obstacle.\n",
      "---\n",
      "0.61\tA horse is running around the track.\n",
      "0.49\tA woman is practicing jumps with her horse.\n",
      "0.41\tTwo men pushed carts through the woods.\n",
      "0.33\tA man is eating a piece of bread.\n",
      "0.32\tA man is eating food.\n",
      "0.28\tA woman is playing violin.\n"
     ]
    }
   ],
   "source": [
    "corpus_embeds = model.encode(corpus)\n",
    "get_ranked_docs(model, query, corpus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query - Horse jumped over the obstacle.\n",
      "---\n",
      "0.84\tA horse is running around the track.\n",
      "0.84\tTwo men pushed carts through the woods.\n",
      "0.82\tA man is eating food.\n",
      "0.78\tA woman is playing violin.\n",
      "0.76\tA man is eating a piece of bread.\n",
      "0.72\tA woman is practicing jumps with her horse.\n"
     ]
    }
   ],
   "source": [
    "corpus_embeds = original_model.encode(corpus)\n",
    "get_ranked_docs(original_model, query, corpus_embeds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## References\n",
    "\n",
    "[1] Kexin Wang, Nils Reimers and Iryna Gurevych. \"[TSDAE: Using Transformer-based Sequential Denoising Auto-Encoder for Unsupervised Sentence Embedding Learning](https://arxiv.org/pdf/2104.06979.pdf)\"\n",
    "\n",
    "[2] [BIER benchmark](https://arxiv.org/abs/2104.08663)"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}