simple-semantic-search-engine-with-transformers.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "2020-12-14-simple-semantic-search-engine-with-transformers",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tqlt9-jKUBCy"
      },
      "source": [
        "# Semantic Search with Pre-trained Sentence Encoders"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cNryrZzRKrpg"
      },
      "source": [
        "In this notebook I will show how to use a pre-trained sentence encoder model to create a simple semantic search engine for website content. The search engine will be \"semantic\" in the sense that it will try to find sentences from the website whose vector representation \"matches\" the vector representation of the search term. \n",
        "\n",
        "We will use a pre-trained transformers model to encode all the sentences of a website as well as the search term and then calculate the [cosine similarity](https://en.wikipedia.org/wiki/Cosine_similarity) between each encoded sentence and the search term. We will then rank the sentences based on the cosine similarity.\n",
        "\n",
        "Let's get started. As before we will first install the libraries we need for this demo. We are using a library called [Sentence Transformers](https://github.com/UKPLab/sentence-transformers), which provides pre-trained transformers models specifically for the purpose of computing sentence-level vector representations. We also need to install scikit-learn, as we will use it to calculate the cosine similarities between sentence vectors. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xdgaoG0WAaOq"
      },
      "source": [
        "!pip install sentence_transformers sklearn lxml bs4"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4cgMux8mOe7c"
      },
      "source": [
        "Once we have installed the liraries we will import them together with some other usefult librarires we will use in the demo."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QAyFysuzAntw"
      },
      "source": [
        "from sentence_transformers import SentenceTransformer\n",
        "import numpy as np\n",
        "from sklearn.metrics.pairwise import cosine_similarity\n",
        "from bs4 import BeautifulSoup\n",
        "from bs4.element import Comment\n",
        "import requests\n",
        "from IPython.display import Markdown, display"
      ],
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "--sSMMgyOnl5"
      },
      "source": [
        "As in the previous post we define some functions we will use for cleaning up the text from the website html as well as a function to print our output in markdown format."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rcFXH4eaAIAR"
      },
      "source": [
        "def tag_visible(element):\n",
        "    if element.parent.name in ['p']:\n",
        "        return True\n",
        "    if isinstance(element, Comment):\n",
        "        return False\n",
        "    return False\n",
        "\n",
        "def text_from_html(html):\n",
        "    soup = BeautifulSoup(html.content, 'lxml')\n",
        "    texts = soup.findAll(text=True)\n",
        "    visible_texts = filter(tag_visible, texts)  \n",
        "    return u\" \".join(t.strip() for t in visible_texts)\n",
        "\n",
        "def printmd(string):\n",
        "    display(Markdown(string))"
      ],
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TkhPLEwXO8qi"
      },
      "source": [
        "Now we can define the website we are interested in together with the search term. We also define how many search results we want to see. In this demo we are interested in finding top five sentences from the Deep Learning Wikipedia article that are about Natural Language Processing."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XkdiNXDoBlOa"
      },
      "source": [
        "website = 'https://en.wikipedia.org/wiki/Deep_learning'\n",
        "search_term = 'natural language processing'\n",
        "num_results = 5"
      ],
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CrYxIo-CPPeM"
      },
      "source": [
        "We then retrieve and extract the text from the website and split the sentences to a list."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CqZYW7b6L31Q"
      },
      "source": [
        "html = requests.get(website)             \n",
        "src_text = text_from_html(html)\n",
        "input_text = src_text.split('.')"
      ],
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "30KFXOuxP7PE"
      },
      "source": [
        "For our model we use the `roberta-large-nli-stsb-mean-tokens` model which according to the [Sentence transformers Github page](https://github.com/UKPLab/sentence-transformers) beats the other Sentence Transformer models in the Semantic Textual Similarity (STS) benchmark."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xt_6pZ2ABcIX"
      },
      "source": [
        "model = SentenceTransformer('roberta-base-nli-stsb-mean-tokens')"
      ],
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ulfQa4qhQWyB"
      },
      "source": [
        "Next we encode the list of sentences retrieved from the website as well as the search term using the Sentence Transformer model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SbqGMKepPXiO"
      },
      "source": [
        "encoded_text = np.array(model.encode(input_text))\n",
        "encoded_query = np.array(model.encode([search_term]))"
      ],
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h3Q3Ks4zQ1QV"
      },
      "source": [
        "We then compare the cosine similarity of each encoded sentence with the encoded search term and rank the sentences accordingly."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FUkBkb7lPoNw"
      },
      "source": [
        "results = cosine_similarity(encoded_query, encoded_text)[0]\n",
        "num_results = results.argsort()[-num_results:][::-1]\n",
        "scores = results[num_results]\n",
        "sentences = [input_text[idx] for idx in num_results]"
      ],
      "execution_count": 35,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yogpqx_dRXGm"
      },
      "source": [
        "Finally, we can print the list of five sentences that best match the search term."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 229
        },
        "id": "QYdxKQ0ICs5a",
        "outputId": "9983d0da-d59e-4a4c-d022-8f8997b4cd3f"
      },
      "source": [
        "print('*'*30 + ' Start of output ' + '*'*30)\n",
        "printmd('**Search Results:**')\n",
        "for sentence, score in zip(sentences, scores):\n",
        "    printmd(f'* {sentence} (score: {score:<.4f})')\n",
        "\n",
        "print('*'*30 + ' End of output ' + '*'*30)"
      ],
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "****************************** Start of output ******************************\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "text/markdown": "**Search Results:**",
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/markdown": "*   Neural networks have been used for implementing language models since the early 2000s (score: 0.6076)",
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/markdown": "*  Word embedding, such as , can be thought of as a representational layer in a deep learning architecture that transforms an atomic word into a positional representation of the word relative to other words in the dataset; the position is represented as a point in a  (score: 0.5070)",
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/markdown": "*  LSTM helped to improve machine translation and language modeling (score: 0.4982)",
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/markdown": "*  A deep neural network (DNN) is an (ANN) with multiple layers between the input and output layers (score: 0.4969)",
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/markdown": "*  An ANN is based on a collection of connected units called , (analogous to biological neurons in a ) (score: 0.4947)",
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "****************************** End of output ******************************\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1PIxEwt6RkKm"
      },
      "source": [
        "As you can see the output is fairly good and at least the top three sentences are highly relevant for our search.\n",
        "\n",
        "In this demo we have seen that with just a few lines of code you can create a simple search engine that does whet it is supposed to do: it finds sentences from the source website/document that match our keyword. The same idea can be used in wide variety of use cases.\n",
        "\n",
        "Hope you enjoyed this demo. Feel free to contact me if you have any questions."
      ]
    }
  ]
}