simple-search-engine-with-bm25.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "2020-12-30-simple-search-engine-with-bm25.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "flNYAlkhDbRY"
      },
      "source": [
        "# Simple Bag-of-Words Search Engine in Python"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dSlHxCs3DzC0"
      },
      "source": [
        "In this notebook I will show how to build a simple search engine using Python and the BM25 ranking algorithm. [BM25](https://en.wikipedia.org/wiki/Okapi_BM25) is a bag-of-words ranking function designed for information retrieval. It is an enhanced version of the term frequency–inverse document frequency ([tf-idf](https://en.wikipedia.org/wiki/Tf–idf)) method. \n",
        "\n",
        "The basic idea behind tf-idf is that it looks at the frequency of the term in the document (more the better) and it looks at the inverse document frequency (common words are less important). One of the drawbacks of the standard tf-idf is that long documents with the same term frequency are considered less important. \n",
        "\n",
        "BM25 makes couple of enhancements to the traditional tf-idf. First, BM25 adds a term frequency saturation to tf-idf, limiting the influence of term frequency on the score. Intuitively this means that more frequent the term is the less impact each occurence of it has on the score. Second, BM25 adds a document length weighting, which makes sure that document length doesn't have such a dramatic negative impact on the relevance score as in tf-idf.\n",
        "\n",
        "Building a BM25-based search engine with Python is easy. There's a great library [`rank-bm25`](https://pypi.org/project/rank-bm25/) which implements the scoring algorithm. To demonstrate the use of this library I'm going to build a simple search engine that retrieves article abstracts from [arXiv](https://arxiv.org) and searches for the most relevant articles based on a given search term. So, let's get started."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l8dSH2-AVsK0"
      },
      "source": [
        "Let's first install and import the needed Python libraries."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fqXc_gMduYLR"
      },
      "source": [
        "!pip install rank_bm25 nltk numpy feedparser"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1Is4S1fnwSca"
      },
      "source": [
        "from rank_bm25 import BM25Okapi\n",
        "import re\n",
        "import nltk\n",
        "from nltk.corpus import stopwords\n",
        "from nltk.tokenize import word_tokenize\n",
        "import numpy as np\n",
        "import feedparser\n",
        "import time"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2jPWVktrNOvN"
      },
      "source": [
        "We will use NLTK to tokenize the documents and to identify stopwords, so we need to download the relevant data files used by NLTK."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pWiDYN3NKpsH"
      },
      "source": [
        "nltk.download('stopwords')\n",
        "nltk.download('punkt')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BQxB7UNiNxRg"
      },
      "source": [
        "Next we retrieve the documents from arXiv. For this we use the [arXiv API](https://arxiv.org/help/api/). We download 10,000 abstracts from arXiv's Computation and Language ([cs.CL](https://arxiv.org/list/cs.CL/recent)) category. The arXiv API returns an Atom feed which we can parse using the [`feedparser`](https://pypi.org/project/feedparser/) library. This gives us an easy access to the article titles, abstracts, links to the full texts as well as many other useful attributes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v0zmaXlABVa7"
      },
      "source": [
        "num_docs = 10000\n",
        "url = f\"http://export.arxiv.org/api/query?search_query=cat:cs.CL&start=0&max_results={str(num_docs)}\"\n",
        "feed = feedparser.parse(url)"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pDyj0EN-O3o3"
      },
      "source": [
        "Next we remove any special characters, tokenise the article abstracts and remove any stopwords."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vDL7_lNoLLna"
      },
      "source": [
        "stop_words = list(stopwords.words('english'))\n",
        "tokenized_docs = []\n",
        "for doc in feed.entries:\n",
        "    # Article abstracts are stored in entries[i].summary \n",
        "    doc = str(doc.summary).lower()\n",
        "    doc = re.sub(r\"([\\w].)([\\~\\!\\@\\#\\$\\%\\^\\&\\*\\(\\)\\-\\+\\[\\]\\{\\}\\/\\\"\\'\\:\\;])([\\s\\w].)\", \"\\\\1 \\\\2 \\\\3\", doc)\n",
        "    doc = re.sub(r\"\\s+\", \" \", doc)\n",
        "    doc = [token for token in word_tokenize(doc.lower()) if token not in stop_words and token.isalpha()]\n",
        "    tokenized_docs.append(doc)"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UIPFashLPuMs"
      },
      "source": [
        "Now we can define the scoring function using the `rank-mb25` library. This function takes the tokenized article abstracts (docs) and the tokenised search term. It also stores the indices for the scores so that we can later link the scores to the correct article."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "toezdZWbTHl1"
      },
      "source": [
        "def get_bm25_scores(tokenized_docs, query):\n",
        "    bm25 = BM25Okapi(tokenized_docs)\n",
        "    scores = bm25.get_scores(query)\n",
        "    scores = [(i, v) for i, v in enumerate(scores)]\n",
        "    scores.sort(key=lambda x: x[1], reverse=True)\n",
        "    return scores"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bicj6C2fQO3s"
      },
      "source": [
        "Now let's define our search term and ther number of results we want to retrieve. In this demo we will search for articles about Natural Language Inference and we want to retrieve the top 5 articles."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BMGrcBZUyJiT"
      },
      "source": [
        "num_results = 5\n",
        "search_term = \"Natural Language Inference\""
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EpB7ef_uQoyX"
      },
      "source": [
        "Next we tokenise the search term and call the ranking algorithm we defined earlier."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zhKjiGF9NLJa"
      },
      "source": [
        "query = search_term.lower().split(' ')\n",
        "\n",
        "t0 = time.time()\n",
        "scores = get_bm25_scores(tokenized_docs, query)\n",
        "t1 = time.time()\n",
        "query_time = t1-t0 "
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UZR5ZMgSQ19k"
      },
      "source": [
        "Once we have the scores we can use the indices to identify the correct entries in the original data. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rtwAOfRoUkM_"
      },
      "source": [
        "idx = np.array(scores)[:num_results, 0].astype(int)\n",
        "final_scores = np.array(scores)[:num_results, 1]\n",
        "documents = [feed.entries[i] for i in idx]"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SPOcDp56RQ9g"
      },
      "source": [
        "We now have the top 5 article entries together with their scores retrieved form 10,000 entries in the cs.CL category. Let's print out the titles, links to the full texts and the summaries. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3UkvQNlSNen9",
        "outputId": "b90bfb93-a157-4a70-af6b-aca0eff47a6c"
      },
      "source": [
        "print(f\"Searched {str(num_docs)} documents in {str(round(query_time, 3))} seconds\\n\")\n",
        "# Article titles are stored in entries[i].title\n",
        "# Article abstracts are stored in entries[i].summary\n",
        "# Article links are stored in entries[i].link\n",
        "for doc, score in zip(documents, final_scores):\n",
        "    print(f\"{doc.title}\\n{doc.link}\\n\\n{doc.summary}\\n[Score: {str(round(score, 4))}]\\n\\n\")"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Searched 10000 documents in 0.434 seconds\n",
            "\n",
            "Stochastic Answer Networks for Natural Language Inference\n",
            "http://arxiv.org/abs/1804.07888v2\n",
            "\n",
            "We propose a stochastic answer network (SAN) to explore multi-step inference\n",
            "strategies in Natural Language Inference. Rather than directly predicting the\n",
            "results given the inputs, the model maintains a state and iteratively refines\n",
            "its predictions. Our experiments show that SAN achieves the state-of-the-art\n",
            "results on three benchmarks: Stanford Natural Language Inference (SNLI)\n",
            "dataset, MultiGenre Natural Language Inference (MultiNLI) dataset and Quora\n",
            "Question Pairs dataset.\n",
            "[Score: 8.8786]\n",
            "\n",
            "\n",
            "Attention Boosted Sequential Inference Model\n",
            "http://arxiv.org/abs/1812.01840v2\n",
            "\n",
            "Attention mechanism has been proven effective on natural language processing.\n",
            "This paper proposes an attention boosted natural language inference model named\n",
            "aESIM by adding word attention and adaptive direction-oriented attention\n",
            "mechanisms to the traditional Bi-LSTM layer of natural language inference\n",
            "models, e.g. ESIM. This makes the inference model aESIM has the ability to\n",
            "effectively learn the representation of words and model the local subsentential\n",
            "inference between pairs of premise and hypothesis. The empirical studies on the\n",
            "SNLI, MultiNLI and Quora benchmarks manifest that aESIM is superior to the\n",
            "original ESIM model.\n",
            "[Score: 8.591]\n",
            "\n",
            "\n",
            "A Neural Architecture Mimicking Humans End-to-End for Natural Language\n",
            "  Inference\n",
            "http://arxiv.org/abs/1611.04741v2\n",
            "\n",
            "In this work we use the recent advances in representation learning to propose\n",
            "a neural architecture for the problem of natural language inference. Our\n",
            "approach is aligned to mimic how a human does the natural language inference\n",
            "process given two statements. The model uses variants of Long Short Term Memory\n",
            "(LSTM), attention mechanism and composable neural networks, to carry out the\n",
            "task. Each part of our model can be mapped to a clear functionality humans do\n",
            "for carrying out the overall task of natural language inference. The model is\n",
            "end-to-end differentiable enabling training by stochastic gradient descent. On\n",
            "Stanford Natural Language Inference(SNLI) dataset, the proposed model achieves\n",
            "better accuracy numbers than all published models in literature.\n",
            "[Score: 8.5675]\n",
            "\n",
            "\n",
            "Neural Natural Language Inference Models Enhanced with External\n",
            "  Knowledge\n",
            "http://arxiv.org/abs/1711.04289v3\n",
            "\n",
            "Modeling natural language inference is a very challenging task. With the\n",
            "availability of large annotated data, it has recently become feasible to train\n",
            "complex models such as neural-network-based inference models, which have shown\n",
            "to achieve the state-of-the-art performance. Although there exist relatively\n",
            "large annotated data, can machines learn all knowledge needed to perform\n",
            "natural language inference (NLI) from these data? If not, how can\n",
            "neural-network-based NLI models benefit from external knowledge and how to\n",
            "build NLI models to leverage it? In this paper, we enrich the state-of-the-art\n",
            "neural natural language inference models with external knowledge. We\n",
            "demonstrate that the proposed models improve neural NLI models to achieve the\n",
            "state-of-the-art performance on the SNLI and MultiNLI datasets.\n",
            "[Score: 8.2145]\n",
            "\n",
            "\n",
            "Baselines and test data for cross-lingual inference\n",
            "http://arxiv.org/abs/1704.05347v2\n",
            "\n",
            "The recent years have seen a revival of interest in textual entailment,\n",
            "sparked by i) the emergence of powerful deep neural network learners for\n",
            "natural language processing and ii) the timely development of large-scale\n",
            "evaluation datasets such as SNLI. Recast as natural language inference, the\n",
            "problem now amounts to detecting the relation between pairs of statements: they\n",
            "either contradict or entail one another, or they are mutually neutral. Current\n",
            "research in natural language inference is effectively exclusive to English. In\n",
            "this paper, we propose to advance the research in SNLI-style natural language\n",
            "inference toward multilingual evaluation. To that end, we provide test data for\n",
            "four major languages: Arabic, French, Spanish, and Russian. We experiment with\n",
            "a set of baselines. Our systems are based on cross-lingual word embeddings and\n",
            "machine translation. While our best system scores an average accuracy of just\n",
            "over 75%, we focus largely on enabling further research in multilingual\n",
            "inference.\n",
            "[Score: 8.1544]\n",
            "\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "df8rK2bgSMhB"
      },
      "source": [
        "As you can see the BM25 algorithm works pretty well and the `rank-bm25` Python library is able to rank 10,000 article abstracts in just 0.434 seconds. In this demo the most time consuming part was retrieving and parsing the 10,000 entries using `feedparser`. In real life applications you would index the data to allow faster retrieval.\n",
        "\n",
        "Hope you enjoyed this demo. Feel free to contact me if you have any questions."
      ]
    }
  ]
}