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.

@@ -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 |[](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 |[](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 |[](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 |[](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 d028f476..8a617fff 100644
Binary files a/demo.gif and b/demo.gif differ
diff --git a/examples/01_Introducing_txtai.ipynb b/examples/01_Introducing_txtai.ipynb
new file mode 100644
index 00000000..5b4611aa
--- /dev/null
+++ b/examples/01_Introducing_txtai.ipynb
@@ -0,0 +1,359 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+ "kernelspec": {
+ "name": "python3",
+ "display_name": "Python 3",
+ "language": "python"
+ },
+ "language_info": {
+ "name": "python",
+ "version": "3.7.6",
+ "mimetype": "text/x-python",
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "pygments_lexer": "ipython3",
+ "nbconvert_exporter": "python",
+ "file_extension": ".py"
+ },
+ "colab": {
+ "name": "01 - Introducing txtai",
+ "provenance": []
+ }
+ },
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "POWZoSJR6XzK",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Part 1: Introducing txtai\n",
+ "\n",
+ "[txtai](https://github.com/neuml/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. \n",
+ "\n",
+ "NeuML uses txtai and/or the concepts behind it to power all of our Natural Language Processing (NLP) applications. Example applications:\n",
+ "\n",
+ "- [cord19q](https://github.com/neuml/cord19q) - COVID-19 literature analysis\n",
+ "- [paperai](https://github.com/neuml/paperai) - AI-powered literature discovery and review engine for medical/scientific papers\n",
+ "- [neuspo](https://neuspo.com) - a fact-driven, real-time sports event and news site\n",
+ "- [codequestion](https://github.com/neuml/codequestion) - Ask coding questions directly from the terminal\n",
+ "\n",
+ "txtai is built on the following stack:\n",
+ "\n",
+ "- [sentence-transformers](https://github.com/UKPLab/sentence-transformers)\n",
+ "- [transformers](https://github.com/huggingface/transformers)\n",
+ "- [faiss](https://github.com/facebookresearch/faiss)\n",
+ "- Python 3.6+\n",
+ "\n",
+ "This notebook gives an overview of txtai and how to run similarity searches."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "qa_PPKVX6XzN",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Install dependencies\n",
+ "\n",
+ "Install txtai and all dependencies"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
+ "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
+ "trusted": true,
+ "_kg_hide-output": true,
+ "id": "24q-1n5i6XzQ",
+ "colab_type": "code",
+ "colab": {}
+ },
+ "source": [
+ "%%capture\n",
+ "!pip install git+https://github.com/neuml/txtai"
+ ],
+ "execution_count": 1,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "vHxb_jm16Xzd",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Create an Embeddings instance\n",
+ "\n",
+ "The Embeddings instance is the main entrypoint for txtai. An Embeddings instance defines the method used to tokenize and convert a text section into an embeddings vector. "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "trusted": true,
+ "id": "QxX9EtIc6Xzg",
+ "colab_type": "code",
+ "colab": {}
+ },
+ "source": [
+ "%%capture\n",
+ "\n",
+ "from txtai.embeddings import Embeddings\n",
+ "\n",
+ "# Create embeddings model, backed by sentence-transformers & transformers\n",
+ "embeddings = Embeddings({\"method\": \"transformers\", \"path\": \"sentence-transformers/bert-base-nli-mean-tokens\"})"
+ ],
+ "execution_count": 2,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "JAO61bcy6Xzo",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Running similarity queries\n",
+ "\n",
+ "An embedding instance relies on the underlying transformer model to build text embeddings. The following example shows how to use an transformers Embedding instance to run similarity searches for a list of different concepts."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a",
+ "_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
+ "trusted": true,
+ "id": "2j_CFGDR6Xzp",
+ "colab_type": "code",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 191
+ },
+ "outputId": "a92c2d3e-1b6d-473a-bc7e-7adc067fa868"
+ },
+ "source": [
+ "import numpy as np\n",
+ "\n",
+ "sections = [\"US tops 5 million confirmed virus cases\",\n",
+ " \"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\",\n",
+ " \"Beijing mobilises invasion craft along coast as Taiwan tensions escalate\",\n",
+ " \"The National Park Service warns against sacrificing slower friends in a bear attack\",\n",
+ " \"Maine man wins $1M from $25 lottery ticket\",\n",
+ " \"Make huge profits without work, earn up to $100,000 a day\"]\n",
+ "\n",
+ "print(\"%-20s %s\" % (\"Query\", \"Best Match\"))\n",
+ "print(\"-\" * 50)\n",
+ "\n",
+ "for query in (\"feel good story\", \"climate change\", \"health\", \"war\", \"wildlife\", \"asia\", \"north america\", \"dishonest junk\"):\n",
+ " # Get index of best section that best matches query\n",
+ " uid = np.argmax(embeddings.similarity(query, sections))\n",
+ "\n",
+ " print(\"%-20s %s\" % (query, sections[uid]))"
+ ],
+ "execution_count": 3,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "text": [
+ "Query Best Match\n",
+ "--------------------------------------------------\n",
+ "feel good story Maine man wins $1M from $25 lottery ticket\n",
+ "climate change Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\n",
+ "health US tops 5 million confirmed virus cases\n",
+ "war Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
+ "wildlife The National Park Service warns against sacrificing slower friends in a bear attack\n",
+ "asia Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
+ "north america US tops 5 million confirmed virus cases\n",
+ "dishonest junk Make huge profits without work, earn up to $100,000 a day\n"
+ ],
+ "name": "stdout"
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "kIMbLW0t6Xzw",
+ "colab_type": "text"
+ },
+ "source": [
+ "The example above shows for almost all of the queries, the actual text isn't stored in the list of text sections. This is the true power of transformer models over token based search. What you get out of the box is 🔥🔥🔥!"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "DLIjSzbq6Xzx",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Building an Embeddings index\n",
+ "\n",
+ "For small lists of texts, the method above works. But for larger repositories of documents, it doesn't make sense to tokenize and convert to embeddings on each query. txtai supports building pre-computed indices which signficantly improve performance. \n",
+ "\n",
+ "Building on the previous example, the following example runs an index method to build and store the text embeddings. In this case, only the query is converted to an embeddings vector each search. "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "trusted": true,
+ "id": "cXfZtdHD6Xzy",
+ "colab_type": "code",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 191
+ },
+ "outputId": "eca8eb29-7ac8-445b-f9b9-6b5cea3280a5"
+ },
+ "source": [
+ "# Create an index for the list of sections\n",
+ "embeddings.index([(uid, text, None) for uid, text in enumerate(sections)])\n",
+ "\n",
+ "print(\"%-20s %s\" % (\"Query\", \"Best Match\"))\n",
+ "print(\"-\" * 50)\n",
+ "\n",
+ "# Run an embeddings search for each query\n",
+ "for query in (\"feel good story\", \"climate change\", \"health\", \"war\", \"wildlife\", \"asia\", \"north america\", \"dishonest junk\"):\n",
+ " # Extract uid of first result\n",
+ " # search result format: (uid, score)\n",
+ " uid = embeddings.search(query, 1)[0][0]\n",
+ "\n",
+ " # Print section\n",
+ " print(\"%-20s %s\" % (query, sections[uid]))"
+ ],
+ "execution_count": 4,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "text": [
+ "Query Best Match\n",
+ "--------------------------------------------------\n",
+ "feel good story Maine man wins $1M from $25 lottery ticket\n",
+ "climate change Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\n",
+ "health US tops 5 million confirmed virus cases\n",
+ "war Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
+ "wildlife The National Park Service warns against sacrificing slower friends in a bear attack\n",
+ "asia Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
+ "north america US tops 5 million confirmed virus cases\n",
+ "dishonest junk Make huge profits without work, earn up to $100,000 a day\n"
+ ],
+ "name": "stdout"
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "6TCVl6QA6Xz5",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Embeddings load/save\n",
+ "\n",
+ "Embeddings indices can be saved to disk and reloaded. At this time, indices are not incrementally created, the index needs a full rebuild to incorporate new data. But that enhancement is in the backlog."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "trusted": true,
+ "id": "5gyO90Hc6Xz7",
+ "colab_type": "code",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 34
+ },
+ "outputId": "92e1358d-4dcf-4982-8811-061f0bfa9a88"
+ },
+ "source": [
+ "embeddings.save(\"index\")\n",
+ "\n",
+ "embeddings = Embeddings()\n",
+ "embeddings.load(\"index\")\n",
+ "\n",
+ "uid = embeddings.search(\"climate change\", 1)[0][0]\n",
+ "print(sections[uid])"
+ ],
+ "execution_count": 5,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "text": [
+ "Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\n"
+ ],
+ "name": "stdout"
+ }
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "trusted": true,
+ "id": "d2uCotAE6X0J",
+ "colab_type": "code",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 34
+ },
+ "outputId": "14f16d54-ca02-4099-9478-24ba8054d3e3"
+ },
+ "source": [
+ "!ls index"
+ ],
+ "execution_count": 6,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "text": [
+ "config\tembeddings\n"
+ ],
+ "name": "stdout"
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "aDIF3tYt6X0O",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Embedding methods\n",
+ "\n",
+ "Embeddings supports two methods for creating text vectors, the sentence-transformers library and word embeddings vectors. Both methods have their merits as shown below:\n",
+ "\n",
+ "- [sentence-transformers](https://github.com/UKPLab/sentence-transformers)\n",
+ " - Creates a single embeddings vector via mean pooling of vectors generated by the transformers library. \n",
+ " - Supports models stored on Hugging Face's model hub or stored locally. \n",
+ " - See sentence-transformers for details on how to create custom models, which can be kept local or uploaded to Hugging Face's model hub.\n",
+ " - Base models require significant compute capability (GPU preferred). Possible to build smaller/lighter weight models that tradeoff accuracy for speed.\n",
+ "- word embeddings\n",
+ " - Creates a single embeddings vector via BM25 scoring of each word component. See this [Medium article](https://towardsdatascience.com/building-a-sentence-embedding-index-with-fasttext-and-bm25-f07e7148d240) for the logic behind this method.\n",
+ " - Backed by the [pymagnitude](https://github.com/plasticityai/magnitude) library. Pre-trained word vectors can be installed from the referenced link.\n",
+ " - See [vectors.py](https://github.com/neuml/txtai/blob/master/src/python/txtai/vectors.py) for code that can build word vectors for custom datasets.\n",
+ " - Significantly better performance with default models. For larger datasets, it offers a good tradeoff of speed and accuracy."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "0Nw7nbHg6X0P",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Next\n",
+ "In part 2 of this series, we'll look at how to use txtai to run extractive searches"
+ ]
+ }
+ ]
+}
\ No newline at end of file
diff --git a/examples/02_Extractive_QA_with_txtai.ipynb b/examples/02_Extractive_QA_with_txtai.ipynb
new file mode 100644
index 00000000..78be37b9
--- /dev/null
+++ b/examples/02_Extractive_QA_with_txtai.ipynb
@@ -0,0 +1,175 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+ "colab": {
+ "name": "02 - Extractive QA with txtai",
+ "provenance": []
+ },
+ "kernelspec": {
+ "name": "python3",
+ "display_name": "Python 3"
+ }
+ },
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "vwELCooy4ljr",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Part 2: Extractive QA with txtai\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. This notebook builds on that and extends to building extractive question-answering systems."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "ew7orE2O441o",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Install dependencies\n",
+ "\n",
+ "Install txtai and all dependencies"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "LPQTb25tASIG",
+ "colab_type": "code",
+ "colab": {}
+ },
+ "source": [
+ "%%capture\n",
+ "!pip install git+https://github.com/neuml/txtai"
+ ],
+ "execution_count": 1,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "_YnqorRKAbLu",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Create an Embeddings and Extractor instances\n",
+ "\n",
+ "The Embeddings instance is the main entrypoint for txtai. An Embeddings instance defines the method used to tokenize and convert a segment of text into an embeddings vector.\n",
+ "\n",
+ "The Extractor instance is the entrypoint for extractive question-answering.\n",
+ "\n",
+ "Both the Embeddings and Extractor instances take a path to a transformer model. Any model on the [Hugging Face model hub](https://huggingface.co/models) can be used in place of the models below."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "OUc9gqTyAYnm",
+ "colab_type": "code",
+ "colab": {}
+ },
+ "source": [
+ "%%capture\n",
+ "\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-mean-tokens\"})\n",
+ "\n",
+ "# Create extractor instance\n",
+ "extractor = Extractor(embeddings, \"distilbert-base-cased-distilled-squad\")"
+ ],
+ "execution_count": 2,
+ "outputs": []
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "4X5z3UjnAGe7",
+ "colab_type": "code",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 330
+ },
+ "outputId": "4e376662-57d1-4841-eeb5-622aaaf286c3"
+ },
+ "source": [
+ "sections = [\"Giants hit 3 HRs to down Dodgers\",\n",
+ " \"Giants 5 Dodgers 4 final\",\n",
+ " \"Dodgers drop Game 2 against the Giants, 5-4\",\n",
+ " \"Blue Jays 2 Red Sox 1 final\",\n",
+ " \"Red Sox lost to the Blue Jays, 2-1\",\n",
+ " \"Blue Jays at Red Sox is over. Score: 2-1\",\n",
+ " \"Phillies win over the Braves, 5-0\",\n",
+ " \"Phillies 5 Braves 0 final\",\n",
+ " \"Final: Braves lose to the Phillies in the series opener, 5-0\",\n",
+ " \"Final score: Flyers 4 Lightning 1\",\n",
+ " \"Flyers 4 Lightning 1 final\",\n",
+ " \"Flyers win 4-1\"]\n",
+ "\n",
+ "# Add unique id to each section to assist with qa extraction\n",
+ "sections = [(uid, section) for uid, section in enumerate(sections)]\n",
+ "\n",
+ "questions = [\"What team won the game?\", \"What was score?\"]\n",
+ "\n",
+ "execute = lambda query: extractor(sections, [(question, query, question, False) for question in questions])\n",
+ "\n",
+ "for query in [\"Red Sox - Blue Jays\", \"Phillies - Braves\", \"Dodgers - Giants\", \"Flyers - Lightning\"]:\n",
+ " print(\"----\", query, \"----\")\n",
+ " for answer in execute(query):\n",
+ " print(answer)\n",
+ " print()\n",
+ "\n",
+ "# Ad-hoc questions\n",
+ "question = \"What hockey team won?\"\n",
+ "\n",
+ "print(\"----\", question, \"----\")\n",
+ "print(extractor(sections, [(question, question, question, False)]))"
+ ],
+ "execution_count": 3,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "text": [
+ "---- Red Sox - Blue Jays ----\n",
+ "('What team won the game?', 'Blue Jays')\n",
+ "('What was score?', '2-1')\n",
+ "\n",
+ "---- Phillies - Braves ----\n",
+ "('What team won the game?', 'Phillies')\n",
+ "('What was score?', '5-0')\n",
+ "\n",
+ "---- Dodgers - Giants ----\n",
+ "('What team won the game?', 'Giants')\n",
+ "('What was score?', '5-4')\n",
+ "\n",
+ "---- Flyers - Lightning ----\n",
+ "('What team won the game?', 'Flyers')\n",
+ "('What was score?', '4-1')\n",
+ "\n",
+ "---- What hockey team won? ----\n",
+ "[('What hockey team won?', 'Flyers')]\n"
+ ],
+ "name": "stdout"
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "k-9gHzrp8ZXA",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Next\n",
+ "In part 3 of this series, we'll look at how to use txtai on larger datasets and explore using word embeddings vs transformer embeddings"
+ ]
+ }
+ ]
+}
\ No newline at end of file
diff --git a/examples/03_Build_an_Embeddings_index_from_a_data_source.ipynb b/examples/03_Build_an_Embeddings_index_from_a_data_source.ipynb
new file mode 100644
index 00000000..fd052d2f
--- /dev/null
+++ b/examples/03_Build_an_Embeddings_index_from_a_data_source.ipynb
@@ -0,0 +1,546 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+ "colab": {
+ "name": "03 - Build an Embeddings index from a data source",
+ "provenance": [],
+ "collapsed_sections": []
+ },
+ "kernelspec": {
+ "name": "python3",
+ "display_name": "Python 3"
+ }
+ },
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "WDbhGHtG8jFE",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Part 3: 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 how to use txtai for extractive question-answer systems.\n",
+ "\n",
+ "The previous examples worked on data stored in memory for demo purposes. For real world large-scale use cases, data is usually stored in a database (Elasticsearch, SQL, MongoDB, files, etc). This example covers reading data from SQLite, building a Embedding index backed by word embeddings and running 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",
+ "colab_type": "text"
+ },
+ "source": [
+ "# Install dependencies\n",
+ "\n",
+ "Install txtai and all dependencies"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "czPYSA2Q9ZHO",
+ "colab_type": "code",
+ "colab": {}
+ },
+ "source": [
+ "%%capture\n",
+ "!pip install git+https://github.com/neuml/txtai"
+ ],
+ "execution_count": 1,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "SN9SCZKQ9fJF",
+ "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": "TONQ4_Kv9dtd",
+ "colab_type": "code",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 228
+ },
+ "outputId": "6bbacf83-d695-42ec-cce0-0f92b1534ca4"
+ },
+ "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": 2,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "text": [
+ "--2020-08-11 15:42:03-- 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’\n",
+ "\n",
+ "\rarticles.sqlite 0%[ ] 0 --.-KB/s \rarticles.sqlite 100%[===================>] 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",
+ " \n",
+ " \n",
+ " Title \n",
+ " Published \n",
+ " Reference \n",
+ " Match \n",
+ "