From 58d193f37555b82ee77a76ca63bf65de1833543a Mon Sep 17 00:00:00 2001 From: davidmezzetti <561939+davidmezzetti@users.noreply.github.com> Date: Tue, 11 Aug 2020 11:54:50 -0400 Subject: [PATCH] Added examples --- README.md | 15 +- demo.gif | Bin 47530 -> 47189 bytes examples/01_Introducing_txtai.ipynb | 359 ++++++++++++ examples/02_Extractive_QA_with_txtai.ipynb | 175 ++++++ ..._Embeddings_index_from_a_data_source.ipynb | 546 ++++++++++++++++++ .../04_Extractive_QA_with_Elasticsearch.ipynb | 460 +++++++++++++++ examples/demo.py | 50 ++ 7 files changed, 1603 insertions(+), 2 deletions(-) create mode 100644 examples/01_Introducing_txtai.ipynb create mode 100644 examples/02_Extractive_QA_with_txtai.ipynb create mode 100644 examples/03_Build_an_Embeddings_index_from_a_data_source.ipynb create mode 100644 examples/04_Extractive_QA_with_Elasticsearch.ipynb create mode 100644 examples/demo.py diff --git a/README.md b/README.md index a48b2af1..920cd375 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # txtai: AI-powered search engine -txtai builds an AI-powered index over sets of text. txtai supports building text indices to perform similarity searches and create extractive question-answering based systems. +txtai builds an AI-powered index over sections of text. txtai supports building text indices to perform similarity searches and create extractive question-answering based systems. ![demo](https://raw.githubusercontent.com/neuml/txtai/master/demo.gif) @@ -35,4 +35,15 @@ This project has dependencies that require compiling native code. Linux envirome pip install torch===1.6.0 torchvision===0.7.0 -f https://download.pytorch.org/whl/torch_stable.html ``` -## Tutorials \ No newline at end of file +## Examples + +The examples directory has a series of examples and notebooks giving an overview of txtai. See the list of notebooks below. + +### Notebooks + +| Notebook | Description | | +|:----------|:-------------|------:| +| [Introducing txtai](https://github.com/neuml/txtai/blob/master/examples/01_Introducing_txtai.ipynb) | Overview of the functionality provided by txtai |[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/neuml/txtai/blob/master/examples/01_Introducing_txtai.ipynb) | +| [Extractive QA with txtai](https://github.com/neuml/txtai/blob/master/examples/02_Extractive_QA_with_txtai) | Extractive question-answering with txtai |[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/neuml/txtai/blob/master/examples/02_Extractive_QA_with_txtai.ipynb) | +| [Build an Embeddings index from a data source](https://github.com/neuml/txtai/blob/master/examples/03_Build_an_Embeddings_index_from_a_data_source.ipynb) | Embeddings index from a data source backed by word embeddings |[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/neuml/txtai/blob/master/examples/03_Build_an_Embeddings_index_from_a_data_source) | +| [Extractive QA with Elasticsearch](https://github.com/neuml/txtai/blob/master/examples/04_Extractive_QA_with_Elasticsearch.ipynb) | Extractive question-answering with Elasticsearch |[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/neuml/txtai/blob/master/examples/04_Extractive_QA_with_Elasticsearch.ipynb) | diff --git a/demo.gif b/demo.gif index d028f476664acc79700054fbe3c2c3ff8b37e725..8a617fff7a3b24b906760264367ca75fd570e98c 100644 GIT binary patch delta 68 zcmV-K0K5OH^8(fI0HZk?|RSXDm?3dWG9=vL&*7n<0ZaxeL6` zS&$o!ZNHI}6J{}Rtc4K;TtzSe0}(?OhKpiE6pN6H9YZL5kt9SdhKPWIAC4TM41+}% zp&SE-ryODxhAtUGFMJe4CjmS|s0vA8N_IyYb3`M!MO!ccBzzmaq)5P}n?;HV0*;2R z3xC0X&qvWjl?BV9todBP&{BeAiBbMIjCr z^gtP5)u?|7r%)V8#pl5Ue-{WC_<5j3JQp%;Q>MCefhM{mrV?11M5n9QuVBN99ZR;X z*|SxwzC)`P-r2U0*usq}b}rqpb_2+ryBBWXwtv+IwmO(*;k9q=9(HRs@8Y] 7.69M --.-KB/s in 0.07s \n", + "\n", + "2020-08-11 15:42:03 (109 MB/s) - ‘articles.sqlite’ saved [8065024/8065024]\n", + "\n" + ], + "name": "stdout" + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "bzdaJiZYBIHE", + "colab_type": "text" + }, + "source": [ + "# Build Word Vectors\n", + "\n", + "This example will build a search system backed by word embeddings. While note quite as powerful as transformer embeddings, they often provide a good tradeoff of performance to functionality for an embedding based search system.\n", + "\n", + "For this notebook, we'll build our own custom embeddings for demo purposes. A number of pre-trained word embedding models are available:\n", + "\n", + " - [General language models from pymagnitude](https://github.com/plasticityai/magnitude)\n", + " - [CORD-19 fastText](https://www.kaggle.com/davidmezzetti/cord19-fasttext-vectors)" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "fJcn-CAH-u3K", + "colab_type": "code", + "colab": { + "base_uri": "https://localhost:8080/", + "height": 156 + }, + "outputId": "67406a51-0819-4845-84da-63d0ad6677f1" + }, + "source": [ + "import os\n", + "import sqlite3\n", + "import tempfile\n", + "\n", + "from txtai.tokenizer import Tokenizer\n", + "from txtai.vectors import WordVectors\n", + "\n", + "print(\"Streaming tokens to temporary file\")\n", + "\n", + "# Stream tokens to temp working file\n", + "with tempfile.NamedTemporaryFile(mode=\"w\", suffix=\".txt\", delete=False) as output:\n", + " # Save file path\n", + " tokens = output.name\n", + "\n", + " db = sqlite3.connect(\"articles.sqlite\")\n", + " cur = db.cursor()\n", + " cur.execute(\"SELECT Text from sections\")\n", + "\n", + " for row in cur:\n", + " output.write(\" \".join(row[0]) + \"\\n\")\n", + "\n", + " # Free database resources\n", + " db.close()\n", + "\n", + "# Build word vectors model - 300 dimensions, 3 min occurrences\n", + "WordVectors.build(tokens, 300, 3, \"cord19-300d\")\n", + "\n", + "# Remove temporary tokens file\n", + "os.remove(tokens)\n", + "\n", + "# Show files\n", + "!ls -l" + ], + "execution_count": 3, + "outputs": [ + { + "output_type": "stream", + "text": [ + "Streaming tokens to temporary file\n", + "Building 300 dimension model\n", + "Converting vectors to magnitude format\n", + "total 9024\n", + "-rw-r--r-- 1 root root 8065024 Aug 11 15:42 articles.sqlite\n", + "-rw-r--r-- 1 root root 360448 Aug 11 15:43 cord19-300d.magnitude\n", + "-rw-r--r-- 1 root root 807886 Aug 11 15:43 cord19-300d.txt\n", + "drwxr-xr-x 1 root root 4096 Jul 30 16:30 sample_data\n" + ], + "name": "stdout" + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_UxcC1-JGH-d", + "colab_type": "text" + }, + "source": [ + "# Build an embeddings index\n", + "\n", + "The following steps builds an embeddings index using the word vector model just created. This model builds a BM25 + fastText index. BM25 is used to build a weighted average of the word embeddings for a section. More information on this method can be found in this [Medium article](https://towardsdatascience.com/building-a-sentence-embedding-index-with-fasttext-and-bm25-f07e7148d240?gi=79da927aa10). " + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "5PrrxGRPGHqX", + "colab_type": "code", + "colab": { + "base_uri": "https://localhost:8080/", + "height": 52 + }, + "outputId": "f00ead87-aeaa-4243-8b68-663c01f30520" + }, + "source": [ + "import sqlite3\n", + "\n", + "import regex as re\n", + "\n", + "from txtai.embeddings import Embeddings\n", + "from txtai.tokenizer import Tokenizer\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|(?\n", + " \n", + " \n", + " Title\n", + " Published\n", + " Reference\n", + " Match\n", + " \n", + " \n", + " \n", + " \n", + " Prevalence and Impact of Myocardial Injury in Patients Hospitalized with COVID-19 Infection\n", + " 2020-04-24 00:00:00\n", + " http://medrxiv.org/cgi/content/short/2020.04.20.20072702v1?rss=1\n", + " This risk was consistent across patients stratified by history of CVD, risk factors but no CVD, and neither CVD nor risk factors.\n", + " \n", + " \n", + " COVID-19 and associations with frailty and multimorbidity: a prospective analysis of UK Biobank participants\n", + " 2020-07-23 00:00:00\n", + " https://www.ncbi.nlm.nih.gov/pubmed/32705587/\n", + " The identification of risk factors for contracting COVID-19 is crucial, to inform public health policy and to facilitate the appropriate distribution of healthcare resources.\n", + " \n", + " \n", + " Does apolipoprotein E genotype predict COVID-19 severity?\n", + " 2020-04-27 00:00:00\n", + " https://doi.org/10.1093/qjmed/hcaa142\n", + " Of interest, apoE4 has also been associated with some of the comorbid risk factors associated with severe COVID-19, such as atherosclerosis and hypertension .\n", + " \n", + " \n", + " COVID-19 and associations with frailty and multimorbidity: a prospective analysis of UK Biobank participants\n", + " 2020-07-23 00:00:00\n", + " https://www.ncbi.nlm.nih.gov/pubmed/32705587/\n", + " Number of comorbidity groupings were then summed, and categorised (0–1; 2; 3 or ≥ 4 comorbidity groupings).\n", + " \n", + " \n", + " COVID-19 and associations with frailty and multimorbidity: a prospective analysis of UK Biobank participants\n", + " 2020-07-23 00:00:00\n", + " https://www.ncbi.nlm.nih.gov/pubmed/32705587/\n", + " Number of comorbidity groupings were then summed, and categorised (0-1; 2; 3 or ≥ 4 comorbidity groupings).\n", + " \n", + " \n", + "" + ], + "text/plain": [ + "" + ] + }, + "metadata": { + "tags": [] + } + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "XSf68I-ZfXOG", + "colab_type": "text" + }, + "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 is used to ask additional questions over the search results, creating a richer query response." + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "TLVOTQJchvTi", + "colab_type": "code", + "colab": {} + }, + "source": [ + "%%capture\n", + "from txtai.extractor import Extractor\n", + "\n", + "# Create extractor instance using qa model designed for the CORD-19 dataset\n", + "extractor = Extractor(embeddings, \"NeuML/bert-small-cord19qa\")" + ], + "execution_count": 6, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "19fmKawThs6d", + "colab_type": "code", + "colab": { + "base_uri": "https://localhost:8080/", + "height": 293 + }, + "outputId": "15c337b5-7c6d-4b5e-98aa-e3996a799737" + }, + "source": [ + "db = sqlite3.connect(\"articles.sqlite\")\n", + "cur = db.cursor()\n", + "\n", + "results = []\n", + "for uid, score in embeddings.search(\"comorbidities 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 Id, Name, Text FROM sections WHERE (labels is null or labels NOT IN ('FRAGMENT', 'QUESTION')) AND article = ?\", [uid])\n", + " sections = []\n", + " for sid, name, text in cur.fetchall():\n", + " if not name or not re.search(r\"background|(?\n", + " \n", + " \n", + " Title\n", + " Published\n", + " Reference\n", + " Match\n", + " Risk Factors\n", + " Locations\n", + " \n", + " \n", + " \n", + " \n", + " Prevalence and Impact of Myocardial Injury in Patients Hospitalized with COVID-19 Infection\n", + " 2020-04-24 00:00:00\n", + " http://medrxiv.org/cgi/content/short/2020.04.20.20072702v1?rss=1\n", + " Stratified by Troponin Levels, N = 2736 All Patients All rights reserved.\n", + " no CVD, and neither CVD nor risk factors\n", + " Mount Sinai Health System\n", + " \n", + " \n", + " COVID-19 and associations with frailty and multimorbidity: a prospective analysis of UK Biobank participants\n", + " 2020-07-23 00:00:00\n", + " https://www.ncbi.nlm.nih.gov/pubmed/32705587/\n", + " The age range at baseline within our sample was 40–69 years, but it is important to note that this was during 2006–2010, and the age range at COVID-19 diagnosis was substantially older (50–84 years).\n", + " Frailty and multimorbidity\n", + " None\n", + " \n", + " \n", + " Does apolipoprotein E genotype predict COVID-19 severity?\n", + " 2020-04-27 00:00:00\n", + " https://doi.org/10.1093/qjmed/hcaa142\n", + " If so, this group could be targeted more aggressively from the outset of the disease.\n", + " None\n", + " None\n", + " \n", + " \n", + " COVID-19 and associations with frailty and multimorbidity: a prospective analysis of UK Biobank participants\n", + " 2020-07-23 00:00:00\n", + " https://www.ncbi.nlm.nih.gov/pubmed/32705587/\n", + " The age range at baseline within our sample was 40–69 years, but it is important to note that this was during 2006–2010, and the age range at COVID-19 diagnosis was substantially older (50–84 years).\n", + " Frailty and multimorbidity\n", + " None\n", + " \n", + " \n", + " COVID-19 and associations with frailty and multimorbidity: a prospective analysis of UK Biobank participants\n", + " 2020-07-23 00:00:00\n", + " https://www.ncbi.nlm.nih.gov/pubmed/32705587/\n", + " The age range at baseline within our sample was 40–69 years, but it is important to note that this was during 2006–2010, and the age range at COVID-19 diagnosis was substantially older (50–84 years).\n", + " Frailty and multimorbidity\n", + " None\n", + " \n", + " \n", + "" + ], + "text/plain": [ + "" + ] + }, + "metadata": { + "tags": [] + } + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ColTLy--rWfR", + "colab_type": "text" + }, + "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." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "KWyoysauy7Pr", + "colab_type": "text" + }, + "source": [ + "# Next\n", + "In part 4 of this series, we'll use combine the power of Elasticsearch with Extractive QA to build a large-scale, advanced search system.\n" + ] + } + ] +} \ No newline at end of file diff --git a/examples/04_Extractive_QA_with_Elasticsearch.ipynb b/examples/04_Extractive_QA_with_Elasticsearch.ipynb new file mode 100644 index 00000000..f8db614f --- /dev/null +++ b/examples/04_Extractive_QA_with_Elasticsearch.ipynb @@ -0,0 +1,460 @@ +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "name": "04 - Extractive QA with Elasticsearch", + "provenance": [] + }, + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "zzZbP0LM6m5z", + "colab_type": "text" + }, + "source": [ + "# Extractive QA with Elasticsearch\n", + "\n", + "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 Elasticesearch system." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "xk7t5Jcd6reO", + "colab_type": "text" + }, + "source": [ + "# Install txtai and Elasticsearch\n", + "\n", + "Install txtai and Elasticsearch. Start an instance of Elasticsearch directly within this notebook. " + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "0y1UA4-q-YdA", + "colab_type": "code", + "colab": {} + }, + "source": [ + "%%capture\n", + "\n", + "# Install txtai and elasticsearch python client\n", + "!pip install git+https://github.com/neuml/txtai elasticsearch\n", + "\n", + "# Download and extract elasticsearch\n", + "!wget https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.8.1-linux-x86_64.tar.gz\n", + "!tar -xzf elasticsearch-7.8.1-linux-x86_64.tar.gz\n", + "!chown -R daemon:daemon elasticsearch-7.8.1" + ], + "execution_count": 16, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "3ZfJeWbM6wmj", + "colab_type": "code", + "colab": {} + }, + "source": [ + "import os\n", + "from subprocess import Popen, PIPE, STDOUT\n", + "\n", + "# Start and wait for server\n", + "server = Popen(['elasticsearch-7.8.1/bin/elasticsearch'], stdout=PIPE, stderr=STDOUT, preexec_fn=lambda: os.setuid(1))\n", + "!sleep 30" + ], + "execution_count": 17, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "TWEn4w68-D1y", + "colab_type": "text" + }, + "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 SQLite database with a subject of CORD-19, 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", + "metadata": { + "id": "8tVrIqSq-KBa", + "colab_type": "code", + "colab": { + "base_uri": "https://localhost:8080/", + "height": 228 + }, + "outputId": "7f46594f-b0f2-4994-dc0b-70e50fe132fc" + }, + "source": [ + "!wget https://www.kaggleusercontent.com/kf/40510829/eyJhbGciOiJkaXIiLCJlbmMiOiJBMTI4Q0JDLUhTMjU2In0..FX7Ote_I-Y88MBPQHRIdUQ.tTr7P3B_eUL_yWN33Usz0Rk1KXtc4DjT_cdjkl5W4WbEcZ-0FJX2jSWTHYMVACtLYuMJJrf6eJN28OzWhDMnTysBu3wfDrd4ly5bu_wJKCnZajICQgQHs_b8hbRVMOzfdG6xEyl9CVYnZNU2cI3QuOshcWxoB0skdKD4d26O_Q4e_nrd8DqEixP47tI2Hu1F00w0vMykzgNwp7SwQ2Z9HoNCO8HtmcjEHq0A4lZ4303YkpjORtZQEO3S-j54fFlIAahT-9VvsFNofitK5VAlR0EyG9r3cOqh2LQDCL7kj5p3MxG8dvHmrTqggLVOwiuKHUIH8u59TemSMLsNRS29W-5fFlHfaItV4dEuiBxCIgQXHcKUDCDGEjeFcPgqpJnNHsnh0pebWDuRQR_fdQ-r8mWgN9qLnosrFBak9tM25G7gqxyUI90GMWAUyP4yj2EAEc8asX9rUsirC8QDHmrmOCUe0cmZvodRUi0ss7lTiLTwm55d9VPXjQn4jQ6tFs-dmjXEx0AwF2Mw1c1jhgzCXwgQj6ybUKemr_6wj1VFYj3VVvCXpk1nZObl-IB6-m7v5CIoXGLot_KFsVtyItRk-wX-B_L3W3aS9dOIfb7bX4s5_aNzXaDKvxrcafwlOQui.vS_FL4EArO8rkBo3xpDF2w/articles.sqlite" + ], + "execution_count": 18, + "outputs": [ + { + "output_type": "stream", + "text": [ + "--2020-08-11 15:41:39-- https://www.kaggleusercontent.com/kf/40510829/eyJhbGciOiJkaXIiLCJlbmMiOiJBMTI4Q0JDLUhTMjU2In0..FX7Ote_I-Y88MBPQHRIdUQ.tTr7P3B_eUL_yWN33Usz0Rk1KXtc4DjT_cdjkl5W4WbEcZ-0FJX2jSWTHYMVACtLYuMJJrf6eJN28OzWhDMnTysBu3wfDrd4ly5bu_wJKCnZajICQgQHs_b8hbRVMOzfdG6xEyl9CVYnZNU2cI3QuOshcWxoB0skdKD4d26O_Q4e_nrd8DqEixP47tI2Hu1F00w0vMykzgNwp7SwQ2Z9HoNCO8HtmcjEHq0A4lZ4303YkpjORtZQEO3S-j54fFlIAahT-9VvsFNofitK5VAlR0EyG9r3cOqh2LQDCL7kj5p3MxG8dvHmrTqggLVOwiuKHUIH8u59TemSMLsNRS29W-5fFlHfaItV4dEuiBxCIgQXHcKUDCDGEjeFcPgqpJnNHsnh0pebWDuRQR_fdQ-r8mWgN9qLnosrFBak9tM25G7gqxyUI90GMWAUyP4yj2EAEc8asX9rUsirC8QDHmrmOCUe0cmZvodRUi0ss7lTiLTwm55d9VPXjQn4jQ6tFs-dmjXEx0AwF2Mw1c1jhgzCXwgQj6ybUKemr_6wj1VFYj3VVvCXpk1nZObl-IB6-m7v5CIoXGLot_KFsVtyItRk-wX-B_L3W3aS9dOIfb7bX4s5_aNzXaDKvxrcafwlOQui.vS_FL4EArO8rkBo3xpDF2w/articles.sqlite\n", + "Resolving www.kaggleusercontent.com (www.kaggleusercontent.com)... 35.190.26.106\n", + "Connecting to www.kaggleusercontent.com (www.kaggleusercontent.com)|35.190.26.106|:443... connected.\n", + "HTTP request sent, awaiting response... 200 OK\n", + "Length: 8065024 (7.7M) [application/octet-stream]\n", + "Saving to: ‘articles.sqlite.1’\n", + "\n", + "\rarticles.sqlite.1 0%[ ] 0 --.-KB/s \rarticles.sqlite.1 100%[===================>] 7.69M --.-KB/s in 0.1s \n", + "\n", + "2020-08-11 15:41:39 (77.5 MB/s) - ‘articles.sqlite.1’ saved [8065024/8065024]\n", + "\n" + ], + "name": "stdout" + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "hSWFzkCn61tM", + "colab_type": "text" + }, + "source": [ + "# Load data into Elasticsearch\n", + "\n", + "The following block copies rows from SQLite to Elasticsearch." + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "So-OBvUT61QD", + "colab_type": "code", + "colab": { + "base_uri": "https://localhost:8080/", + "height": 34 + }, + "outputId": "8ccfed53-bc46-4619-9b82-31a30ce38b55" + }, + "source": [ + "import sqlite3\n", + "\n", + "import regex as re\n", + "\n", + "from elasticsearch import Elasticsearch, helpers\n", + "\n", + "# Connect to ES instance\n", + "es = Elasticsearch(hosts=[\"http://localhost:9200\"], timeout=60, retry_on_timeout=True)\n", + "\n", + "# Connection to database file\n", + "db = sqlite3.connect(\"articles.sqlite\")\n", + "cur = db.cursor()\n", + "\n", + "# Elasticsearch bulk buffer\n", + "buffer = []\n", + "rows = 0\n", + "\n", + "# Select tagged sentences without a NLP label. NLP labels are set for non-informative sentences.\n", + "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\")\n", + "for row in cur:\n", + " # Build dict of name-value pairs for fields\n", + " article = dict(zip((\"id\", \"article\", \"title\", \"published\", \"reference\", \"name\", \"text\"), row))\n", + " name = article[\"name\"]\n", + "\n", + " # Only process certain document sections\n", + " if not name or not re.search(r\"background|(?\n", + " \n", + " \n", + " Title\n", + " Published\n", + " Reference\n", + " Match\n", + " \n", + " \n", + " \n", + " \n", + " Management of osteoarthritis during COVID‐19 pandemic\n", + " 2020-05-21 00:00:00\n", + " https://doi.org/10.1002/cpt.1910\n", + " Indeed, risk factors are sex, obesity, genetic factors and mechanical factors (3) .\n", + " \n", + " \n", + " Prevalence and Impact of Myocardial Injury in Patients Hospitalized with COVID-19 Infection\n", + " 2020-04-24 00:00:00\n", + " http://medrxiv.org/cgi/content/short/2020.04.20.20072702v1?rss=1\n", + " This risk was consistent across patients stratified by history of CVD, risk factors but no CVD, and neither CVD nor risk factors.\n", + " \n", + " \n", + " Lower prevalence of antibodies neutralizing SARS-CoV-2 in group O French blood donors\n", + " 2020-07-15 00:00:00\n", + " https://www.ncbi.nlm.nih.gov/pubmed/32679056/\n", + " Age, sex and various comorbidities have been identified as factors worsening the prognosis of the disease.\n", + " \n", + " \n", + " Does apolipoprotein E genotype predict COVID-19 severity?\n", + " 2020-04-27 00:00:00\n", + " https://doi.org/10.1093/qjmed/hcaa142\n", + " 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 .\n", + " \n", + " \n", + " Risk Stratification for Healthcare workers during the CoViD-19 Pandemic; using demographics, co-morbid disease and clinical domain in order to assign clinical duties\n", + " 2020-05-09 00:00:00\n", + " http://medrxiv.org/cgi/content/short/2020.05.05.20091967v1?rss=1\n", + " A risk stratification tool was compiled using a Caucasian female <50years with no comorbidities as a reference.\n", + " \n", + " \n", + "" + ], + "text/plain": [ + "" + ] + }, + "metadata": { + "tags": [] + } + } + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "mwBTrCkcOM_H", + "colab_type": "code", + "colab": {} + }, + "source": [ + "%%capture\n", + "from txtai.embeddings import Embeddings\n", + "from txtai.extractor import Extractor\n", + "\n", + "# Create embeddings model, backed by sentence-transformers & transformers\n", + "embeddings = Embeddings({\"method\": \"transformers\", \"path\": \"sentence-transformers/bert-base-nli-stsb-mean-tokens\"})\n", + "\n", + "# Create extractor instance using qa model designed for the CORD-19 dataset\n", + "extractor = Extractor(embeddings, \"NeuML/bert-small-cord19qa\")" + ], + "execution_count": 21, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "Yv75Lh-cOpL9", + "colab_type": "code", + "colab": { + "base_uri": "https://localhost:8080/", + "height": 328 + }, + "outputId": "4e341551-5af7-492f-e5c9-1904f8bf3689" + }, + "source": [ + "document = {\n", + " \"_source\": [\"id\", \"name\", \"text\"],\n", + " \"size\": 1000,\n", + " \"query\": {\n", + " \"term\": {\"article\": None}\n", + " },\n", + " \"sort\" : [\"id\"]\n", + "}\n", + "\n", + "def sections(article):\n", + " rows = []\n", + "\n", + " search = document.copy()\n", + " search[\"query\"][\"term\"][\"article\"] = article\n", + "\n", + " for result in es.search(index=\"articles\", body=search)[\"hits\"][\"hits\"]:\n", + " source = result[\"_source\"]\n", + " sid, name, text = source[\"id\"], source[\"name\"], source[\"text\"]\n", + "\n", + " if not name or not re.search(r\"background|(?\n", + " \n", + " \n", + " Title\n", + " Published\n", + " Reference\n", + " Match\n", + " Risk Factors\n", + " Locations\n", + " \n", + " \n", + " \n", + " \n", + " Management of osteoarthritis during COVID‐19 pandemic\n", + " 2020-05-21 00:00:00\n", + " https://doi.org/10.1002/cpt.1910\n", + " Indeed, risk factors are sex, obesity, genetic factors and mechanical factors (3) .\n", + " Comorbidities\n", + " None\n", + " \n", + " \n", + " Prevalence and Impact of Myocardial Injury in Patients Hospitalized with COVID-19 Infection\n", + " 2020-04-24 00:00:00\n", + " http://medrxiv.org/cgi/content/short/2020.04.20.20072702v1?rss=1\n", + " This risk was consistent across patients stratified by history of CVD, risk factors but no CVD, and neither CVD nor risk factors.\n", + " CVD, risk factors but no CVD, and neither CVD\n", + " None\n", + " \n", + " \n", + " Lower prevalence of antibodies neutralizing SARS-CoV-2 in group O French blood donors\n", + " 2020-07-15 00:00:00\n", + " https://www.ncbi.nlm.nih.gov/pubmed/32679056/\n", + " Age, sex and various comorbidities have been identified as factors worsening the prognosis of the disease.\n", + " Age, sex and various comorbidities\n", + " Paris\n", + " \n", + " \n", + " Does apolipoprotein E genotype predict COVID-19 severity?\n", + " 2020-04-27 00:00:00\n", + " https://doi.org/10.1093/qjmed/hcaa142\n", + " 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 .\n", + " socioeconomic inequalities and risk factors\n", + " None\n", + " \n", + " \n", + " Risk Stratification for Healthcare workers during the CoViD-19 Pandemic; using demographics, co-morbid disease and clinical domain in order to assign clinical duties\n", + " 2020-05-09 00:00:00\n", + " http://medrxiv.org/cgi/content/short/2020.05.05.20091967v1?rss=1\n", + " A risk stratification tool was compiled using a Caucasian female <50years with no comorbidities as a reference.\n", + " Those working with aerosol generating procedures were at increased risk\n", + " None\n", + " \n", + " \n", + "" + ], + "text/plain": [ + "" + ] + }, + "metadata": { + "tags": [] + } + } + ] + } + ] +} \ No newline at end of file diff --git a/examples/demo.py b/examples/demo.py new file mode 100644 index 00000000..4dbe77b5 --- /dev/null +++ b/examples/demo.py @@ -0,0 +1,50 @@ +""" +Demo shell +""" + +from cmd import Cmd + +import numpy as np + +from txtai.embeddings import Embeddings + +class Shell(Cmd): + """ + Query shell. + """ + + def __init__(self): + super(Shell, self).__init__() + + self.intro = "query shell" + self.prompt = "(search) " + + self.embeddings = None + self.sections = None + + def preloop(self): + # Create embeddings model, backed by sentence-transformers & transformers + self.embeddings = Embeddings({"method": "transformers", "path": "sentence-transformers/bert-base-nli-mean-tokens"}) + + self.sections = ["US tops 5 million confirmed virus cases", + "Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg", + "Beijing mobilises invasion craft along coast as Taiwan tensions escalate", + "The National Park Service warns against sacrificing slower friends in a bear attack", + "Maine man wins $1M from $25 lottery ticket", + "Make huge profits without work, earn up to $100,000 a day"] + + def default(self, line): + # Get index of best section that best matches query + uid = np.argmax(self.embeddings.similarity(line, self.sections)) + print(self.sections[uid]) + print() + +def main(): + """ + Shell execution loop. + """ + + Shell().cmdloop() + +if __name__ == "__main__": + main()