Extractive QA with Elasticsearch
Tags: ai, llm, rag, vectordatabase
Published: 2021-01-28 12:18:12

[![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/neuml/txtai/blob/master/examples/06_Extractive_QA_with_Elasticsearch.ipynb)

txtai is datastore agnostic, the library analyzes sets of text. The following example shows how extractive question-answering can be added on top of an Elasticsearch system.

# Install dependencies

Install `txtai` and `Elasticsearch`.

```
# Install txtai and elasticsearch python client
pip install txtai elasticsearch

# Download and extract elasticsearch
wget https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.10.1-linux-x86_64.tar.gz
tar -xzf elasticsearch-7.10.1-linux-x86_64.tar.gz
chown -R daemon:daemon elasticsearch-7.10.1
```

Start an instance of Elasticsearch.

```python
import os
from subprocess import Popen, PIPE, STDOUT

# Start and wait for server
server = Popen(['elasticsearch-7.10.1/bin/elasticsearch'], stdout=PIPE, stderr=STDOUT, preexec_fn=lambda: os.setuid(1))
```

```
sleep 30
```

# Download data

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.

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.

```
wget https://github.com/neuml/txtai/releases/download/v1.1.0/tests.gz
gunzip tests.gz
mv tests articles.sqlite
```

# Load data into Elasticsearch

The following block copies rows from SQLite to Elasticsearch.

```python
import sqlite3

import regex as re

from elasticsearch import Elasticsearch, helpers

# Connect to ES instance
es = Elasticsearch(hosts=["http://localhost:9200"], timeout=60, retry_on_timeout=True)

# Connection to database file
db = sqlite3.connect("articles.sqlite")
cur = db.cursor()

# Elasticsearch bulk buffer
buffer = []
rows = 0

# Select tagged sentences without a NLP label. NLP labels are set for non-informative sentences.
cur.execute("SELECT s.Id, Article, Title, Published, Reference, Name, Text FROM sections s JOIN articles a on s.article=a.id WHERE (s.labels is null or s.labels NOT IN ('FRAGMENT', 'QUESTION')) AND s.tags is not null")
for row in cur:
  # Build dict of name-value pairs for fields
  article = dict(zip(("id", "article", "title", "published", "reference", "name", "text"), row))
  name = article["name"]

  # Only process certain document sections
  if not name or not re.search(r"background|(?<!.*?results.*?)discussion|introduction|reference", name.lower()):
    # Bulk action fields
    article["_id"] = article["id"]
    article["_index"] = "articles"

    # Buffer article
    buffer.append(article)

    # Increment number of articles processed
    rows += 1

    # Bulk load every 1000 records
    if rows % 1000 == 0:
      helpers.bulk(es, buffer)
      buffer = []

      print("Inserted {} articles".format(rows), end="\r")

if buffer:
  helpers.bulk(es, buffer)

print("Total articles inserted: {}".format(rows))
```

```
Total articles inserted: 21499
```

# Query data

The following runs a query against Elasticsearch for the terms "risk factors". It finds the top 5 matches and returns the corresponding documents associated with each match.

```python
import pandas as pd

from IPython.display import display, HTML

pd.set_option("display.max_colwidth", None)

query = {
    "_source": ["article", "title", "published", "reference", "text"],
    "size": 5,
    "query": {
        "query_string": {"query": "risk factors"}
    }
}

results = []
for result in es.search(index="articles", body=query)["hits"]["hits"]:
  source = result["_source"]
  results.append((source["title"], source["published"], source["reference"], source["text"]))

df = pd.DataFrame(results, columns=["Title", "Published", "Reference", "Match"])

display(HTML(df.to_html(index=False)))
```

|Title 	|Published | Reference| Match|
|------------|----------|-----------|-----|
|Management of osteoarthritis during COVID‐19 pandemic 	|2020-05-21 00:00:00 	|https://doi.org/10.1002/cpt.1910 	|Indeed, risk factors are sex, obesity, genetic factors and mechanical factors (3) .|
|Prevalence and Impact of Myocardial Injury in Patients Hospitalized with COVID-19 Infection 	|2020-04-24 00:00:00 	|http://medrxiv.org/cgi/content/short/2020.04.20.20072702v1?rss=1 	|This risk was consistent across patients stratified by history of CVD, risk factors but no CVD, and neither CVD nor risk factors.|
|Does apolipoprotein E genotype predict COVID-19 severity? 	|2020-04-27 00:00:00 	|https://doi.org/10.1093/qjmed/hcaa142 	|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 .|
|COVID-19 and associations with frailty and multimorbidity: a prospective analysis of UK Biobank participants 	|2020-07-23 00:00:00 	|https://www.ncbi.nlm.nih.gov/pubmed/32705587/ 	|BACKGROUND: Frailty and multimorbidity have been suggested as risk factors for severe COVID-19 disease.|
|COVID-19: what has been learned and to be learned about the novel coronavirus disease 	|2020-03-15 00:00:00 	|https://doi.org/10.7150/ijbs.45134 	|• Three major risk factors for COVID-19 were sex (male), age (≥60), and severe pneumonia.|

# Derive columns with Extractive QA

The next section uses Extractive QA to derive additional columns. For each article, the full text is retrieved and a series of questions are asked of the document. The answers are added as a derived column per article.

```python
from txtai.embeddings import Embeddings
from txtai.pipeline import Extractor

# Create embeddings model, backed by sentence-transformers & transformers
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})

# Create extractor instance using qa model designed for the CORD-19 dataset
extractor = Extractor(embeddings, "NeuML/bert-small-cord19qa")
```

```python
document = {
    "_source": ["id", "name", "text"],
    "size": 1000,
    "query": {
        "term": {"article": None}
    },
    "sort" : ["id"]
}

def sections(article):
  rows = []

  search = document.copy()
  search["query"]["term"]["article"] = article

  for result in es.search(index="articles", body=search)["hits"]["hits"]:
    source = result["_source"]
    name, text = source["name"], source["text"]

    if not name or not re.search(r"background|(?<!.*?results.*?)discussion|introduction|reference", name.lower()):
      rows.append(text)
  
  return rows

results = []
for result in es.search(index="articles", body=query)["hits"]["hits"]:
  source = result["_source"]

  # Use QA extractor to derive additional columns
  answers = extractor([("Risk factors", "risk factor", "What are names of risk factors?", False),
                       ("Locations", "city country state", "What are names of locations?", False)], sections(source["article"]))

  results.append((source["title"], source["published"], source["reference"], source["text"]) + tuple([answer[1] for answer in answers]))

df = pd.DataFrame(results, columns=["Title", "Published", "Reference", "Match", "Risk Factors", "Locations"])

display(HTML(df.to_html(index=False)))
```


|Title 	|Published 	|Reference 	|Match 	|Risk Factors 	|Locations|
|---------|---------|---------|---------|---------|---------|
|Management of osteoarthritis during COVID‐19 pandemic 	|2020-05-21 00:00:00 	|https://doi.org/10.1002/cpt.1910 	|Indeed, risk factors are sex, obesity, genetic factors and mechanical factors (3) .| 	Comorbidities|	extrapulmonary sites|
|Prevalence and Impact of Myocardial Injury in Patients Hospitalized with COVID-19 Infection 	|2020-04-24 00:00:00 	|http://medrxiv.org/cgi/content/short/2020.04.20.20072702v1?rss=1 	|This risk was consistent across patients stratified by history of CVD, risk factors but no CVD, and neither CVD nor risk factors. 	|CVD, risk factors but no CVD, and neither CVD| 	None|
|Does apolipoprotein E genotype predict COVID-19 severity? 	|2020-04-27 00:00:00 	|https://doi.org/10.1093/qjmed/hcaa142 	|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 .| 	socioeconomic inequalities and risk factors 	|None|
|COVID-19 and associations with frailty and multimorbidity: a prospective analysis of UK Biobank participants 	|2020-07-23 00:00:00 	|https://www.ncbi.nlm.nih.gov/pubmed/32705587/ 	|BACKGROUND: Frailty and multimorbidity have been suggested as risk factors for severe COVID-19 disease. 	|Frailty and multimorbidity 	|comorbidity groupings|
|COVID-19: what has been learned and to be learned about the novel coronavirus disease 	|2020-03-15 00:00:00 	|https://doi.org/10.7150/ijbs.45134 	|• Three major risk factors for COVID-19 were sex (male), age (≥60), and severe pneumonia. 	|age and underlying disease are strongly correlated 	|cities, provinces, and countries|