notebooks/W2E_Pipelines_Sentence_Transformers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/gsarti/ik-nlp-tutorials/blob/main/notebooks/W2E_Pipelines_Sentence_Transformers.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run in Colab to install local packages\n",
    "!pip install transformers sentencepiece torch datasets sentence-transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🤗 Pipelines & Sentence Transformers for Semantic Search and QA\n",
    "\n",
    "*This notebook is partly adapted from a tutorial by Wietse De Vries for the IK-NLP course of 2021*\n",
    "\n",
    "The goal of this notebook is to make you practice with some pipeline use-cases and to introduce the [Sentence Transformers](https://sbert.net/) in the context of a concrete example of semantic search and question answering (relevant to the NLP Study Assistant final project).\n",
    "\n",
    "Exercises 1 and 2 of this notebook are mandatory and will be part of your graded midterm portfolio. Exercise 3 is optional, but we highly recommend you to complete it, especially if you're interested in the NLP Study Assistant final project."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1: Using the Fill-mask pipeline for Probing Linguistic Knowledge in mBERT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you probably know by now, BERT is a transformed-based, context-sensitive, neural language models that has been trained, among others, on a masked language modeling task, where the model learns to predict what the most likely word is at a a *masked* (hidden) position in the sentence. In a sentence like:\n",
    "\n",
    ">There were several [MASK] with the proposed solution.\n",
    "    \n",
    "the model will learn that the word *problems* or *issues* is more likely at this position than the word *unicorns* or *days*. The model uses both the right and left context of the masked position to make its predictions. Models trained to maximize the probability of predicting the correct masked words learn to represent a good amount of linguistic knowledge as a result of this.\n",
    "\n",
    "A **probe** is a test of a language model aimed at investigating how accurate these predictions are, especially for cases where syntax makes it quite clear that one (form of a) word is correct, and another word is impossible. In the example above, for instance, the masked position can be filled by a plural noun (*problems*) but not by a singular noun (*problem*). If the model makes predictions that respect the linguistic constraints, we have reason to believe that the model is somehow aware of the linguistic structure of the language.\n",
    "\n",
    "While predicting whether the masked position should be filled by a singular or plural noun seems easy in the example above (both *were* and *several* are good predictors of plural), we can try to make the task harder by looking for contexts where the solution requires more careful *attention* to the right words in the context\n",
    "\n",
    ">There were some [MASK] with the proposed solution.\n",
    ">\n",
    ">There could be several [MASK] with the proposed solution.\n",
    ">\n",
    ">There were some unexpected and unforeseen [MASK] with the proposed solution.\n",
    "    \n",
    "In the examples above, the task is made harder by replacing *several* (which is always followed by a singular noun) by *some* (which can be followed by a singular or a plural noun), by replacing *were* (which always heads a sentence with a plural subject) by *could be* (which can head a sentence with a singular or plural subject), and by inserting material between the verb *were* (which indicates that there should be a plural) and the MASK.\n",
    "\n",
    "\n",
    "### Assignment\n",
    "\n",
    "Think of a grammatical phenomenon in a language of your choice, and come up with at least 5 example sentences to probe whether the model makes the correct predictions. Think of cases where the context makes it clear that the mask has to be plural or singular, that a verb has to have a particular form (like plural or singular, or participle or infinitive), that a specific (personal, possessive, reflexive) pronoun has to be used, that an adjective or noun has to have a specific inflection (like in German and more generally in languages with a rich case and/or gender marking system). There is a host of literature on this, see for instance [Marvin and Linzen](https://arxiv.org/abs/1808.09031) (for English) and [Sahin et al](https://www.mitpressjournals.org/doi/full/10.1162/coli_a_00376) (for multilingual probes).\n",
    "\n",
    "### Model\n",
    "\n",
    "The model we will be using for this task is the [multilingual BERT](https://huggingface.co/bert-base-multilingual-cased) model mBERT, that was trained on the Wikipedia text of the 102 largest Wikipedia's. This means that you do not have to choose examples from English, but that you may also present a probe for another language. \n",
    "\n",
    "The following loads the pipeline for doing masked prediction, and load the mBERT model (this may take a minute or so). You can ignore the warning about some weights not being initialized. The pipeline can be used to test masked language model prediction. Given a sequence containing the special token [MASK], the model will predict what the most likely tokens are at that position, using both left and right context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-multilingual-cased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.weight', 'cls.seq_relationship.bias']\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"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'sequence': 'There were some unexpected and unforeseen problems with the proposed solution.',\n",
       "  'score': 0.5272422432899475,\n",
       "  'token': 20390,\n",
       "  'token_str': 'problems'},\n",
       " {'sequence': 'There were some unexpected and unforeseen issues with the proposed solution.',\n",
       "  'score': 0.09590281546115875,\n",
       "  'token': 17850,\n",
       "  'token_str': 'issues'},\n",
       " {'sequence': 'There were some unexpected and unforeseen difficulties with the proposed solution.',\n",
       "  'score': 0.0664457157254219,\n",
       "  'token': 64557,\n",
       "  'token_str': 'difficulties'},\n",
       " {'sequence': 'There were some unexpected and unforeseen dealing with the proposed solution.',\n",
       "  'score': 0.0168981421738863,\n",
       "  'token': 73082,\n",
       "  'token_str': 'dealing'},\n",
       " {'sequence': 'There were some unexpected and unforeseen associated with the proposed solution.',\n",
       "  'score': 0.01516666542738676,\n",
       "  'token': 18107,\n",
       "  'token_str': 'associated'}]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "mbert = pipeline('fill-mask', model='bert-base-multilingual-cased')\n",
    "mbert('There were some unexpected and unforeseen [MASK] with the proposed solution.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, the pipe returns the 5 most likely words that could appear at the position of the mask, along with a score. If you want to know specifically whether the model prefers one of two forms, you can give these forms as targets to the pipe, and also print the answer in a more readable form:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0034\ttoken: challenges\tThere were some unexpected and unforeseen challenges with the proposed solution.\n",
      "0.0001\ttoken: challenge\tThere were some unexpected and unforeseen challenge with the proposed solution.\n"
     ]
    }
   ],
   "source": [
    "def probe(sentence: str, targets: str) :\n",
    "    for res in mbert(sentence,targets=targets) :\n",
    "        print(f\"{res['score']:6.4f}\\ttoken: {res['token_str']}\\t{res['sequence']}\")\n",
    "        \n",
    "probe('There were some unexpected and unforeseen [MASK] with the proposed solution.',['challenge','challenges'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **💡 Interesting Fact**: The same bias that is present towards grammatically correct choices can be observed in other cases, such as racial and gender stereotyping. Much work is currently in process to identify and remove gender and racial biases from learned language embeddings. See the following example and [this recent survey on the topic](https://arxiv.org/abs/2112.14168)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-cased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.weight', 'cls.seq_relationship.bias']\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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['lawyer', 'carpenter', 'doctor', 'waiter', 'mechanic']\n",
      "['nurse', 'waitress', 'teacher', 'maid', 'prostitute']\n"
     ]
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "unmasker = pipeline(\"fill-mask\", model=\"bert-base-cased\")\n",
    "result = unmasker(\"This man works as a [MASK].\")\n",
    "print([r[\"token_str\"] for r in result])\n",
    "\n",
    "result = unmasker(\"This woman works as a [MASK].\")\n",
    "print([r[\"token_str\"] for r in result])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Your Turn to Probe\n",
    "\n",
    "Give at least five example sentences with a [MASK] and a list of targets that illustrate a specific grammatical phenomenon in a language of your choice. Describe what the grammatical phenomenon is you are investigating. Use the probe function for testing. Try to include both *easy* sentences (where the model should do well) as well as *hard* sentences (where there are words in the context that might lead to confusion, or where the clue words are far away from the mask). For languages other than Dutch or English, make sure to include enough explanation so that examples and tests are clear to a non-native speaker. \n",
    "\n",
    "Describe how well the model did on your probe sentences. Where there any cases where the model made the wrong decision?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your probe here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2: Mixing Pipelines for Text-to-text QA in Many Languages\n",
    "\n",
    "The Model Hub of HuggingFace is home to a staggering amount of models for the more disparate use-cases, but you may have noticed that many of those are trained on the English language. Let's consider for example the [`UnifiedQA`](https://github.com/allenai/unifiedqa) model by AllenAI, which is a T5 model architecture trained to perform question answering on multiple formats (e.g. extract the answer from the provided context, produce an answer without a supporting context, choose among multiple possible answers, yes/no questions) using a unified text-to-text approach. While this opens thrilling perspectives in having a single model for all QA use-cases, UnifiedQA models are available for English only, and training such models from scratch in another language would require nontrivial effort and resources.\n",
    "\n",
    "Here are several examples of how text should be formatted for the UnifiedQA model:\n",
    "\n",
    "| **Task type** | **Example Dataset** | **Format** | **Example** | **Output** |\n",
    "| :---: | :--- | :--- | :--- | :--- |\n",
    "|  **Extractive QA** | SQUAD | `<QUESTION> \\n <CONTEXT>` | `At what speed did the turbine operate? \\n (Nikola_Tesla) On his 50th birthday in 1906, Tesla demonstrated his 200 horsepower (150 kilowatts) 16,000 rpm bladeless turbine. ...` |  `16,000 rpm` |\n",
    "|  **Abstractive QA** | NarrativeQA | `<QUESTION> \\n <CONTEXT>` | `What does a drink from narcissus's spring cause the drinker to do?  \\n  Mercury has awakened Echo, who weeps for Narcissus, and states that a drink from Narcissus's spring causes the drinkers to ''Grow dotingly enamored of themselves.'' ...` | `fall in love with themselves` |\n",
    "|  **Multiple-choice QA** | ARC-challenge | `<QUESTION> \\n (a) <CHOICE_A> (b) <CHOICE_B> ...` | `What does photosynthesis produce that helps plants grow? \\n (A) water (B) oxygen (C) protein (D) sugar` | `sugar` |\n",
    "|  **Multiple-choice QA with context** | MCTest | `<QUESTION> \\n (a) <CHOICE_A> (b) <CHOICE_B> ... \\n <CONTEXT>` | `Who was Billy? \\n (A) The skinny kid (B) A teacher (C) A little kid (D) The big kid \\n Billy was like a king on the school yard. A king without a queen. He was the biggest kid in our grade, so he made all the rules during recess. ...` | `The big kid` |\n",
    "|  **Yes-no QA** | BoolQ | `<QUESTION> \\n <CONTEXT>` | `Was America the first country to have a president?  \\n (President) The first usage of the word president to denote the highest official in a government was during the Commonwealth of England ...` | `no` |\n",
    "\n",
    "### Assignment\n",
    "\n",
    "We are gonna build a function making use of multiple models through 🤗 Pipelines to generate a response to a question in one of the formats specified above, in one of the languages supported by the MT systems available on the HuggingFace Model Hub. The function will translate and paraphrase the query into multiple examples, and then pick the best outputs of the UnifiedQA model as candidates for backtranslation. In this way, we mock the existance of a UnifiedQA model for the language of our choice.\n",
    "\n",
    "### Model\n",
    "\n",
    "The following code loads the UnifiedQA model and use it to perform QA on a multiple choice question without context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'generated_text': 'Paris'}]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "# Using at least the base variant of the model is advised for good results\n",
    "generator = pipeline(\"text2text-generation\", model=\"allenai/unifiedqa-t5-base\")\n",
    "generator(\"What is the name of the city where the Eiffel Tower is located? \\n (A) Paris (B) London (C) Prague (D) Berlin\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Your turn to Pipe\n",
    "\n",
    "Using the pipelines we saw in the tutorial, create a function taking a `question` string, an optional `context` string and an optional list of strings called `choices` and performs the following steps:\n",
    "\n",
    "- Use one of the [MarianMT](https://huggingface.co/models?search=helsinki-nlp) machine translation models to translate all the inputs from the language of your choice to English. You may want to split the text into sentences (e.g. by splitting on periods) to obtain better results for the context.\n",
    "\n",
    "- Use a paraphrasing model ([`tuner007/pegasus_paraphrase`](https://huggingface.co/tuner007/pegasus_paraphrase) is a good choice, albeit heavy) to produce 4 paraphrases of the question, using the `num_return_sequences` parameter.\n",
    "\n",
    "- For each question (translated + 4 paraphrases), format it with the translated context and choices (if present) as a single string in the format required by UnifiedQA.\n",
    "\n",
    "- Use the `allenai/unifiedqa-t5-base` model to generate an answer for each of the 5 questions.\n",
    "\n",
    "- If at least 3 of the 5 answers are identical strings, return the result translated back to the original language using the MarianMT model for the reciprocal language pair (e.g. if you used `Helsinki-NLP/opus-mt-nl-en` to translate from Dutch to English, you will need to use `Helsinki-NLP/opus-mt-en-nl`). Otherwise, print \"No common answer found\" translated in the original language.\n",
    "\n",
    "**Importantly**, the quality of the output does not determine your score in the evaluation. The goal is to get a feel for the models and their capabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional, List\n",
    "\n",
    "def answer(question: str, context: Optional[str] = None, choices: Optional[List[str]] = None):\n",
    "    # Your beautiful function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (Optional) Exercise 3: SentenceTransformers for Semantic Similarity Search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Figures and some code are from the [CohereAI Semantic Search Tutorial](https://docs.cohere.ai/semantic-search/) by Jay Alammar.*\n",
    "\n",
    "> SentenceTransformers is a Python framework for state-of-the-art sentence, text and image embeddings. The initial work is described in our paper [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://arxiv.org/abs/1908.10084).\n",
    ">\n",
    "> You can use this framework to compute sentence / text embeddings for more than 100 languages. These embeddings can then be compared e.g. with cosine-similarity to find sentences with a similar meaning. This can be useful for semantic textual similar, semantic search, or paraphrase mining. \n",
    "> \n",
    "> The framework is based on PyTorch and Transformers and offers a large collection of pre-trained models tuned for various tasks. Further, it is easy to fine-tune your own models.\n",
    "\n",
    "Semantic search is a typical use case in natural language processing, in which we want to retrieve the most relevant documents from a corpus (e.g. Automatic FAQs, Web browser search results, etc.). Using the similarity between embedded text representations allows us to go beyond simple keyword matching, which is highly desirable in this setting (e.g. `tomorrow will rain` should be very close in embedding space to `the weather forecast announces showers for the next day`, despite no lexical overlap). \n",
    "\n",
    "<div>\n",
    "<img alt=\"Visualizing Semantic Search\" src=\"https://github.com/cohere-ai/notebooks/raw/main/notebooks/images/basic-semantic-search-overview.png?3\" style=\"width: 60%\" />\n",
    "</div>\n",
    "\n",
    "In this exercise we will use first Huggingface Transformers and then the [SentenceTransformers](https://sbert.net) library to find the most relevant paragraphs for a specific query.\n",
    "\n",
    "Let's start by loading the `train` split of the `squad` dataset from the Dataset Hub and flatten its structure so that every example contains a single triplet `(context, question, answer)`. We are going to use only the first 50 rows of the dataset, the others can be discarded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "squad = load_dataset(\"squad\", split=\"train[:50]\")\n",
    "\n",
    "squad_train_filtered = None# Your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, sample a question at random from the resulting dataset and print it. You can use `shuffle` or turn the Dataset into a `DataFrame` and use `sample`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = None # Your code here\n",
    "\n",
    "print(query)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now going to use a model trained to perform **semantic search** to retrieve the top 10 most likely contexts for each selected question.\n",
    "\n",
    "The model `sentence-transformers/multi-qa-MiniLM-L6-cos-v1` is a good choice for this task, since it is relatively small and was explicitly trained for semantic search. We are going to define now three utility functions:\n",
    "\n",
    "- `dot_score` computes the dot product between two Pytorch tensors.\n",
    "- `mean_pooling` averages the embeddings produced by a model to obtain a **sentence embedding**.\n",
    "- `encode` uses a model and a tokenizer to convert a list of texts into a tensor of embeddings. **Complete it with the first two steps we saw in the tutorial.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import Tensor\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "def dot_score(a: Tensor, b: Tensor):\n",
    "    \"\"\"\n",
    "    Computes the dot-product dot_prod(a[i], b[j]) for all i and j.\n",
    "    :return: Matrix with res[i][j]  = dot_prod(a[i], b[j])\n",
    "    Taken from the SentenceTransformer library\n",
    "    \"\"\"\n",
    "    if not isinstance(a, torch.Tensor):\n",
    "        a = torch.tensor(a)\n",
    "    if not isinstance(b, torch.Tensor):\n",
    "        b = torch.tensor(b)\n",
    "    if len(a.shape) == 1:\n",
    "        a = a.unsqueeze(0)\n",
    "    if len(b.shape) == 1:\n",
    "        b = b.unsqueeze(0)\n",
    "    print(a.shape, b.shape)\n",
    "    # Compute the dot-product\n",
    "    return torch.mm(a, b.transpose(0, 1))\n",
    "\n",
    "\n",
    "#Mean Pooling - Average all the embeddings produced by the model\n",
    "def mean_pooling(model_output, attention_mask):\n",
    "    # First element of model_output contains all token embeddings\n",
    "    token_embeddings = model_output.last_hidden_state\n",
    "    # Expand the mask to the same size as the token embeddings to avoid indexing errors\n",
    "    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()\n",
    "    # Compute the mean of the token embeddings\n",
    "    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)\n",
    "\n",
    "\n",
    "#Encode text\n",
    "def encode(model, tokenizer, texts):\n",
    "    # Tokenize sentences\n",
    "    encoded_input = None # Your code here\n",
    "    # Compute token embeddings\n",
    "    with torch.no_grad():\n",
    "        model_output = None # Your code here\n",
    "    # Perform pooling\n",
    "    embeddings = mean_pooling(model_output, encoded_input['attention_mask'])\n",
    "    # Normalize embeddings\n",
    "    embeddings = F.normalize(embeddings, p=2, dim=1)\n",
    "    \n",
    "    return embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now use these functions to compute similarity scores for the query and all the contexts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sentences we want sentence embeddings for\n",
    "contexts = list(squad_train_filtered.to_pandas()['context'])\n",
    "\n",
    "# Load the model and tokenizer from HuggingFace Hub\n",
    "tokenizer = None # Your code here\n",
    "model = None # Your code here\n",
    "\n",
    "#Encode query and contexts with the encode function\n",
    "query_emb = None # Your code here\n",
    "contexts_emb = None # Your code here\n",
    "\n",
    "#Compute dot score between query and all contexts embeddings\n",
    "scores = torch.mm(query_emb, contexts_emb.transpose(0, 1))[0].cpu().tolist()\n",
    "\n",
    "#Combine contexts & scores\n",
    "contexts_score_pairs = list(zip(contexts, scores))\n",
    "\n",
    "#Sort by decreasing score\n",
    "contexts_score_pairs = sorted(contexts_score_pairs, key=lambda x: x[1], reverse=True)\n",
    "\n",
    "#Output passages & scores\n",
    "for ctx, score in contexts_score_pairs:\n",
    "    print(score, ctx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now use SentenceTransformers to do the same, but using much less code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer, util\n",
    "\n",
    "# Query was defined above\n",
    "contexts = list(squad_train_filtered.to_pandas()['context'])\n",
    "\n",
    "#Load the model\n",
    "model = SentenceTransformer('sentence-transformers/multi-qa-MiniLM-L6-cos-v1')\n",
    "\n",
    "#Encode query and contexts using SentenceTransformer model.encode\n",
    "query_emb = model.encode(query)\n",
    "contexts_emb = model.encode(contexts)\n",
    "\n",
    "#Compute dot score between query and all contexts embeddings\n",
    "scores = util.dot_score(query_emb, contexts_emb)[0].tolist()\n",
    "\n",
    "#Combine contexts & scores\n",
    "contexts_score_pairs = list(zip(contexts, scores))\n",
    "\n",
    "#Sort by decreasing score\n",
    "contexts_score_pairs = sorted(contexts_score_pairs, key=lambda x: x[1], reverse=True)\n",
    "\n",
    "#Output passages & scores\n",
    "for ctx, score in contexts_score_pairs:\n",
    "    print(score, ctx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a more advanced overview on how to optimize speed with semantic search, see how to use the [FAISS](https://github.com/facebookresearch/faiss) library to perform fast nearest neighbor search natively with Huggingface Transformers here: [Using FAISS for efficient similarity search ](https://huggingface.co/course/chapter5/6?fw=pt#using-faiss-for-efficient-similarity-search)"
   ]
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}