examples/03_Build_an_Embeddings_index_from_a_data_source.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WDbhGHtG8jFE"
      },
      "source": [
        "# Build an Embeddings index from a data source\n",
        "\n",
        "In Part 1, we gave a general overview of txtai, the backing technology and examples of how to use it for similarity searches. Part 2 covered an embedding index with a larger dataset.\n",
        "\n",
        "For real world large-scale use cases, data is often stored in a database (Elasticsearch, SQL, MongoDB, files, etc). Here we'll show how to read from SQLite, build an Embedding index and run queries against the generated Embeddings index.\n",
        "\n",
        "This example covers functionality found in the [paperai](https://github.com/neuml/paperai) library. See that library for a full solution that can be used with the dataset discussed below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UQ0fCwXn9bcH"
      },
      "source": [
        "# Install dependencies\n",
        "\n",
        "Install `txtai` and all dependencies."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "czPYSA2Q9ZHO"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install git+https://github.com/neuml/txtai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SN9SCZKQ9fJF"
      },
      "source": [
        "# Download data\n",
        "\n",
        "This example is going to work off a subset of the [CORD-19](https://www.semanticscholar.org/cord19) dataset. COVID-19 Open Research Dataset (CORD-19) is a free resource of scholarly articles, aggregated by a coalition of leading research groups, covering COVID-19 and the coronavirus family of viruses.\n",
        "\n",
        "The following download is a SQLite database generated from a [Kaggle notebook](https://www.kaggle.com/davidmezzetti/cord-19-slim/output). More information on this data format, can be found in the [CORD-19 Analysis](https://www.kaggle.com/davidmezzetti/cord-19-analysis-with-sentence-embeddings) notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TONQ4_Kv9dtd"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!wget https://github.com/neuml/txtai/releases/download/v1.1.0/tests.gz\n",
        "!gunzip tests.gz\n",
        "!mv tests articles.sqlite"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_UxcC1-JGH-d"
      },
      "source": [
        "# Build an embeddings index\n",
        "\n",
        "The following steps build an embeddings index using a vector model designed for medical papers, [PubMedBERT Embeddings](https://huggingface.co/NeuML/pubmedbert-base-embeddings)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5PrrxGRPGHqX",
        "outputId": "61bf7211-6757-4147-8f2f-e4d1ebe58e11"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Iterated over 21499 total rows\n"
          ]
        }
      ],
      "source": [
        "import sqlite3\n",
        "\n",
        "import regex as re\n",
        "\n",
        "from txtai import Embeddings\n",
        "\n",
        "def stream():\n",
        "  # Connection to database file\n",
        "  db = sqlite3.connect(\"articles.sqlite\")\n",
        "  cur = db.cursor()\n",
        "\n",
        "  # Select tagged sentences without a NLP label. NLP labels are set for non-informative sentences.\n",
        "  cur.execute(\"SELECT Id, Name, Text FROM sections WHERE (labels is null or labels NOT IN ('FRAGMENT', 'QUESTION')) AND tags is not null\")\n",
        "\n",
        "  count = 0\n",
        "  for row in cur:\n",
        "    # Unpack row\n",
        "    uid, name, text = row\n",
        "\n",
        "    # Only process certain document sections\n",
        "    if not name or not re.search(r\"background|(?<!.*?results.*?)discussion|introduction|reference\", name.lower()):\n",
        "      document = (uid, text, None)\n",
        "\n",
        "      count += 1\n",
        "      if count % 1000 == 0:\n",
        "        print(\"Streamed %d documents\" % (count), end=\"\\r\")\n",
        "\n",
        "      yield document\n",
        "\n",
        "  print(\"Iterated over %d total rows\" % (count))\n",
        "\n",
        "  # Free database resources\n",
        "  db.close()\n",
        "\n",
        "# Create embeddings index \n",
        "embeddings = Embeddings(path=\"neuml/pubmedbert-base-embeddings\")\n",
        "\n",
        "# Build embeddings index\n",
        "embeddings.index(stream())\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zHk24su3e_gb"
      },
      "source": [
        "# Query data\n",
        "\n",
        "The following runs a query against the embeddings index for the terms \"risk factors\". It finds the top 5 matches and returns the corresponding documents associated with each match."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 293
        },
        "id": "CRbDhvvDKEl-",
        "outputId": "774f8085-01db-49e2-f025-4fe68afca594"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th>Title</th>\n",
              "      <th>Published</th>\n",
              "      <th>Reference</th>\n",
              "      <th>Match</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>Management of osteoarthritis during COVID‐19 pandemic</td>\n",
              "      <td>2020-05-21 00:00:00</td>\n",
              "      <td>https://doi.org/10.1002/cpt.1910</td>\n",
              "      <td>Indeed, risk factors are sex, obesity, genetic factors and mechanical factors (3) .</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>Does apolipoprotein E genotype predict COVID-19 severity?</td>\n",
              "      <td>2020-04-27 00:00:00</td>\n",
              "      <td>https://doi.org/10.1093/qjmed/hcaa142</td>\n",
              "      <td>Risk factors associated with subsequent death include older age, hypertension, diabetes, ischemic heart disease, obesity and chronic lung disease; however, sometimes there are no obvious risk factors .</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>Prevalence and Impact of Myocardial Injury in Patients Hospitalized with COVID-19 Infection</td>\n",
              "      <td>2020-04-24 00:00:00</td>\n",
              "      <td>http://medrxiv.org/cgi/content/short/2020.04.20.20072702v1?rss=1</td>\n",
              "      <td>This risk was consistent across patients stratified by history of CVD, risk factors but no CVD, and neither CVD nor risk factors.</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>COVID-19 and associations with frailty and multimorbidity: a prospective analysis of UK Biobank participants</td>\n",
              "      <td>2020-07-23 00:00:00</td>\n",
              "      <td>https://www.ncbi.nlm.nih.gov/pubmed/32705587/</td>\n",
              "      <td>BACKGROUND: Frailty and multimorbidity have been suggested as risk factors for severe COVID-19 disease.</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>Risk Stratification for Healthcare workers during the CoViD-19 Pandemic; using demographics, co-morbid disease and clinical domain in order to assign clinical duties</td>\n",
              "      <td>2020-05-09 00:00:00</td>\n",
              "      <td>http://medrxiv.org/cgi/content/short/2020.05.05.20091967v1?rss=1</td>\n",
              "      <td>Vascular disease, diabetes and chronic pulmonary disease further increased risk.</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import pandas as pd\n",
        "\n",
        "from IPython.display import display, HTML\n",
        "\n",
        "pd.set_option(\"display.max_colwidth\", None)\n",
        "\n",
        "db = sqlite3.connect(\"articles.sqlite\")\n",
        "cur = db.cursor()\n",
        "\n",
        "results = []\n",
        "for uid, score in embeddings.search(\"risk factors\", 5):\n",
        "  cur.execute(\"SELECT article, text FROM sections WHERE id = ?\", [uid])\n",
        "  uid, text = cur.fetchone()\n",
        "\n",
        "  cur.execute(\"SELECT Title, Published, Reference from articles where id = ?\", [uid])\n",
        "  results.append(cur.fetchone() + (text,))\n",
        "\n",
        "# Free database resources\n",
        "db.close()\n",
        "\n",
        "df = pd.DataFrame(results, columns=[\"Title\", \"Published\", \"Reference\", \"Match\"])\n",
        "\n",
        "# It has been reported that displaying HTML within VSCode doesn't work.\n",
        "# When using VSCode, the data can be exported to an external HTML file to view.\n",
        "# See example below.\n",
        "\n",
        "# htmlData = df.to_html(index=False)\n",
        "# with open(\"data.html\", \"w\") as file:\n",
        "#     file.write(htmlData)\n",
        "\n",
        "display(HTML(df.to_html(index=False)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XSf68I-ZfXOG"
      },
      "source": [
        "# Extracting additional columns from query results\n",
        "\n",
        "The example above uses the Embeddings index to find the top 5 best matches. In addition to this, an Extractor instance (this will be explained further in part 5) is used to ask additional questions over the search results, creating a richer query response."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TLVOTQJchvTi"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "from txtai.pipeline import Extractor\n",
        "\n",
        "# Create extractor instance using qa model designed for the CORD-19 dataset\n",
        "# Note: That extractive QA was a predecessor to Large Language Models (LLMs). LLMs likely will get better results.\n",
        "extractor = Extractor(embeddings, \"NeuML/bert-small-cord19qa\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 293
        },
        "id": "19fmKawThs6d",
        "outputId": "b7cd40e3-a87c-419d-f520-b7795607cebc"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th>Title</th>\n",
              "      <th>Published</th>\n",
              "      <th>Reference</th>\n",
              "      <th>Match</th>\n",
              "      <th>Risk Factors</th>\n",
              "      <th>Locations</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>Management of osteoarthritis during COVID‐19 pandemic</td>\n",
              "      <td>2020-05-21 00:00:00</td>\n",
              "      <td>https://doi.org/10.1002/cpt.1910</td>\n",
              "      <td>Indeed, risk factors are sex, obesity, genetic factors and mechanical factors (3) .</td>\n",
              "      <td>sex, obesity, genetic factors and mechanical factors</td>\n",
              "      <td>hospitals and clinics</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>Does apolipoprotein E genotype predict COVID-19 severity?</td>\n",
              "      <td>2020-04-27 00:00:00</td>\n",
              "      <td>https://doi.org/10.1093/qjmed/hcaa142</td>\n",
              "      <td>Risk factors associated with subsequent death include older age, hypertension, diabetes, ischemic heart disease, obesity and chronic lung disease; however, sometimes there are no obvious risk factors .</td>\n",
              "      <td>None</td>\n",
              "      <td>None</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>Prevalence and Impact of Myocardial Injury in Patients Hospitalized with COVID-19 Infection</td>\n",
              "      <td>2020-04-24 00:00:00</td>\n",
              "      <td>http://medrxiv.org/cgi/content/short/2020.04.20.20072702v1?rss=1</td>\n",
              "      <td>This risk was consistent across patients stratified by history of CVD, risk factors but no CVD, and neither CVD nor risk factors.</td>\n",
              "      <td>neither CVD nor risk factors</td>\n",
              "      <td>Mount Sinai Health System</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>COVID-19 and associations with frailty and multimorbidity: a prospective analysis of UK Biobank participants</td>\n",
              "      <td>2020-07-23 00:00:00</td>\n",
              "      <td>https://www.ncbi.nlm.nih.gov/pubmed/32705587/</td>\n",
              "      <td>BACKGROUND: Frailty and multimorbidity have been suggested as risk factors for severe COVID-19 disease.</td>\n",
              "      <td>Frailty and multimorbidity</td>\n",
              "      <td>213 countries and territories</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>Risk Stratification for Healthcare workers during the CoViD-19 Pandemic; using demographics, co-morbid disease and clinical domain in order to assign clinical duties</td>\n",
              "      <td>2020-05-09 00:00:00</td>\n",
              "      <td>http://medrxiv.org/cgi/content/short/2020.05.05.20091967v1?rss=1</td>\n",
              "      <td>Vascular disease, diabetes and chronic pulmonary disease further increased risk.</td>\n",
              "      <td>Vascular disease, diabetes and chronic pulmonary disease</td>\n",
              "      <td>None</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "db = sqlite3.connect(\"articles.sqlite\")\n",
        "cur = db.cursor()\n",
        "\n",
        "results = []\n",
        "for uid, score in embeddings.search(\"risk factors\", 5):\n",
        "  cur.execute(\"SELECT article, text FROM sections WHERE id = ?\", [uid])\n",
        "  uid, text = cur.fetchone()\n",
        "\n",
        "  # Get list of document text sections to use for the context\n",
        "  cur.execute(\"SELECT Name, Text FROM sections WHERE (labels is null or labels NOT IN ('FRAGMENT', 'QUESTION')) AND article = ? ORDER BY Id\", [uid])\n",
        "  texts = []\n",
        "  for name, txt in cur.fetchall():\n",
        "    if not name or not re.search(r\"background|(?<!.*?results.*?)discussion|introduction|reference\", name.lower()):\n",
        "      texts.append(txt)\n",
        "\n",
        "  cur.execute(\"SELECT Title, Published, Reference from articles where id = ?\", [uid])\n",
        "  article = cur.fetchone()\n",
        "\n",
        "  # Use QA extractor to derive additional columns\n",
        "  answers = extractor([(\"Risk Factors\", \"risk factors\", \"What risk factors?\", False),\n",
        "                       (\"Locations\", \"hospital country\", \"What locations?\", False)], texts)\n",
        "\n",
        "  results.append(article + (text,) + tuple([answer[1] for answer in answers]))\n",
        "\n",
        "# Free database resources\n",
        "db.close()\n",
        "\n",
        "df = pd.DataFrame(results, columns=[\"Title\", \"Published\", \"Reference\", \"Match\", \"Risk Factors\", \"Locations\"])\n",
        "display(HTML(df.to_html(index=False)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ColTLy--rWfR"
      },
      "source": [
        "In the example above, the Embeddings index is used to find the top N results for a given query. On top of that, a question-answer extractor is used to derive additional columns based on a list of questions. In this case, the \"Risk Factors\" and \"Location\" columns were pulled from the document text."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "local",
      "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.9.21"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}