chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
File diff suppressed because one or more lines are too long
@@ -0,0 +1,292 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "LjmhJ4ad9kBL"
|
||||
},
|
||||
"source": [
|
||||
"# Build an Embeddings index with Hugging Face Datasets\n",
|
||||
"\n",
|
||||
"This notebook shows how txtai can index and search with Hugging Face's [Datasets](https://github.com/huggingface/datasets) library. Datasets opens access to a large and growing list of publicly available datasets. Datasets has functionality to select, transform and filter data stored in each dataset.\n",
|
||||
"\n",
|
||||
"In this example, txtai will be used to index and query a dataset.\n",
|
||||
"\n",
|
||||
"**Make sure to select a GPU runtime when running this notebook**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "8tLWvo9v-Q0u"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies. Also install `datasets`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "Fa5BCjMFqVKE"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai\n",
|
||||
"!pip install datasets"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "hOdEv8MH-e5h"
|
||||
},
|
||||
"source": [
|
||||
"# Load dataset and build a txtai index\n",
|
||||
"\n",
|
||||
"In this example, we'll load the `ag_news` dataset, which is a collection of news article headlines. This only takes a single line of code!\n",
|
||||
"\n",
|
||||
"Next, txtai will index the first 10,000 rows of the dataset. A sentence similarity model is used to compute sentence embeddings. sentence-transformers has a number of [pre-trained models](https://huggingface.co/models?pipeline_tag=sentence-similarity) that can be swapped in.\n",
|
||||
"\n",
|
||||
"In addition to the embeddings index, we'll also create a Similarity instance to re-rank search hits for relevancy. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "3hYRk9JnsM0J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"from txtai.embeddings import Embeddings\n",
|
||||
"from txtai.pipeline import Similarity\n",
|
||||
"\n",
|
||||
"def stream(dataset, field, limit):\n",
|
||||
" index = 0\n",
|
||||
" for row in dataset:\n",
|
||||
" yield (index, row[field], None)\n",
|
||||
" index += 1\n",
|
||||
"\n",
|
||||
" if index >= limit:\n",
|
||||
" break\n",
|
||||
"\n",
|
||||
"def search(query):\n",
|
||||
" return [(result[\"score\"], result[\"text\"]) for result in embeddings.search(query, limit=50)]\n",
|
||||
"\n",
|
||||
"def ranksearch(query):\n",
|
||||
" results = [text for _, text in search(query)]\n",
|
||||
" return [(score, results[x]) for x, score in similarity(query, results)]\n",
|
||||
"\n",
|
||||
"# Load HF dataset\n",
|
||||
"dataset = load_dataset(\"ag_news\", split=\"train\")\n",
|
||||
"\n",
|
||||
"# Create embeddings model, backed by sentence-transformers & transformers, enable content storage\n",
|
||||
"embeddings = Embeddings({\"path\": \"sentence-transformers/paraphrase-MiniLM-L3-v2\", \"content\": True})\n",
|
||||
"embeddings.index(stream(dataset, \"text\", 10000))\n",
|
||||
"\n",
|
||||
"# Create similarity instance for re-ranking\n",
|
||||
"similarity = Similarity(\"valhalla/distilbart-mnli-12-3\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "LBhHcX6eFmGI"
|
||||
},
|
||||
"source": [
|
||||
"# Search the dataset\n",
|
||||
"\n",
|
||||
"Now that an index is ready, let's search the data! The following section runs a series of queries and show the results. Like basic search engines, txtai finds token matches. But the real power of txtai is finding semantically similar results.\n",
|
||||
"\n",
|
||||
"sentence-transformers has a great overview on [information retrieval](https://www.sbert.net/examples/applications/information-retrieval/README.html) that is well worth a read. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "YVmbiY92vxEO",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 1000
|
||||
},
|
||||
"outputId": "85f5e0ad-14ba-4642-aed6-13c14a710d68"
|
||||
},
|
||||
"source": [
|
||||
"from IPython.core.display import display, HTML\n",
|
||||
"\n",
|
||||
"def table(query, rows):\n",
|
||||
" html = \"\"\"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 15px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" \"\"\"\n",
|
||||
"\n",
|
||||
" html += \"<h3>%s</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead>\" % (query)\n",
|
||||
" for score, text in rows:\n",
|
||||
" html += \"<tr><td>%.4f</td><td>%s</td></tr>\" % (score, text)\n",
|
||||
" html += \"</table>\"\n",
|
||||
"\n",
|
||||
" display(HTML(html))\n",
|
||||
"\n",
|
||||
"for query in [\"Positive Apple reports\", \"Negative Apple reports\", \"Best planets to explore for life\", \"LA Dodgers good news\", \"LA Dodgers bad news\"]:\n",
|
||||
" table(query, ranksearch(query)[:2])\n"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 15px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" <h3>Positive Apple reports</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead><tr><td>0.9941</td><td>Apple's iPod a Huge Hit in Japan The iPod is proving a colossal hit on the Japanese electronics and entertainment giant's home ground. The tiny white machine is catching on as a fashion statement and turning into a cultural icon here, much the same way it won a fanatical following in the United States.</td></tr><tr><td>0.9886</td><td>Apple tops US consumer satisfaction Recent data published by the American Customer Satisfaction Index (ACSI) shows Apple leading the consumer computer industry with the the highest customer satisfaction.</td></tr></table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 15px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" <h3>Negative Apple reports</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead><tr><td>0.9847</td><td>Apple Recalls 28,000 Faulty Batteries Sold with 15-inch PowerBook Apple has had to recall up to 28,000 notebook batteries that were sold for use with their 15-inch PowerBook. Apple reports that faulty batteries sold between January 2004 and August 2004 can overheat and pose a fire hazard.</td></tr><tr><td>0.9795</td><td>Apple Announces Voluntary Recall of Powerbook Batteries Apple, in cooperation with the US Consumer Product Safety Commission (CPSC), announced Thursday a voluntary recall of 15 quot; Aluminum PowerBook batteries. The batteries being recalled could potentially overheat, though no injuries relating ...</td></tr></table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 15px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" <h3>Best planets to explore for life</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead><tr><td>0.9110</td><td>Tiny 'David' Telescope Finds 'Goliath' Planet A newfound planet detected by a small, 4-inch-diameter telescope demonstrates that we are at the cusp of a new age of planet discovery. Soon, new worlds may be located at an accelerating pace, bringing the detection of the first Earth-sized world one step closer.</td></tr><tr><td>0.8838</td><td>Venus: Inhabited World? by Harry Bortman In part 1 of this interview with Astrobiology Magazine editor Henry Bortman, planetary scientist David Grinspoon explained how Venus evolved from a wet planet similar to Earth to the scorching hot, dried-out furnace of today. In part 2, Grinspoon discusses the possibility of life on Venus...</td></tr></table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 15px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" <h3>LA Dodgers good news</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead><tr><td>0.9961</td><td>Dodgers 7, Braves 4 Los Angeles, Ca. -- Shawn Green belted a grand slam and a solo homer as Los Angeles beat Mike Hampton and the Atlanta Braves 7-to-4 Saturday afternoon.</td></tr><tr><td>0.9928</td><td>MLB: Los Angeles 7, Atlanta 4 Shawn Green hit two home runs Saturday, including a grand slam, to lead the Los Angeles Dodgers to a 7-4 victory over the Atlanta Braves.</td></tr></table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 15px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" <h3>LA Dodgers bad news</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead><tr><td>0.9880</td><td>Expos Keep Dodgers at Bay With 8-7 Win (AP) AP - Giovanni Carrara walked Juan Rivera with the bases loaded and two outs in the ninth inning Monday night, spoiling Los Angeles' six-run comeback and handing the Montreal Expos an 8-7 victory over the Dodgers.</td></tr><tr><td>0.9671</td><td>Gagne blows his 2d save Pinch-hitter Lenny Harris delivered a three-run double off Eric Gagne with two outs in the ninth, rallying the Florida Marlins past the Dodgers, 6-4, last night in Los Angeles.</td></tr></table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,422 @@
|
||||
{
|
||||
"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
|
||||
}
|
||||
@@ -0,0 +1,491 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "zzZbP0LM6m5z"
|
||||
},
|
||||
"source": [
|
||||
"# Add semantic search to Elasticsearch\n",
|
||||
"\n",
|
||||
"Part 2 and Part 3 of this series showed how to index and search data in txtai. Part 2 indexed and searched a Hugging Face Dataset, Part 3 indexed and searched an external data source. \n",
|
||||
"\n",
|
||||
"txtai is modular in design, it's components can be individually used. txtai has a similarity function that works on lists of text. This method can be integrated with any external search service, such as a REST API, a SQL query or anything else that returns text search results. \n",
|
||||
"\n",
|
||||
"In this notebook, we'll take the same Hugging Face Dataset used in Part 2, index it in Elasticsearch and rank the search results using a semantic similarity function from txtai.\n",
|
||||
"\n",
|
||||
"**Make sure to select a GPU runtime when running this notebook**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "xk7t5Jcd6reO"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai`, `datasets` and `Elasticsearch`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "0y1UA4-q-YdA"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"\n",
|
||||
"# Install txtai, datasets and elasticsearch python client\n",
|
||||
"!pip install git+https://github.com/neuml/txtai datasets elasticsearch\n",
|
||||
"\n",
|
||||
"# Download and extract elasticsearch\n",
|
||||
"!wget https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.10.1-linux-x86_64.tar.gz\n",
|
||||
"!tar -xzf elasticsearch-7.10.1-linux-x86_64.tar.gz\n",
|
||||
"!chown -R daemon:daemon elasticsearch-7.10.1"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "M5c_WxmxCShQ"
|
||||
},
|
||||
"source": [
|
||||
"Start an instance of Elasticsearch directly within this notebook."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "3ZfJeWbM6wmj"
|
||||
},
|
||||
"source": [
|
||||
"import os\n",
|
||||
"from subprocess import Popen, PIPE, STDOUT\n",
|
||||
"\n",
|
||||
"# If issues are encountered with this section, ES can be manually started as follows:\n",
|
||||
"# ./elasticsearch-7.10.1/bin/elasticsearch\n",
|
||||
"\n",
|
||||
"# Start and wait for server\n",
|
||||
"server = Popen(['elasticsearch-7.10.1/bin/elasticsearch'], stdout=PIPE, stderr=STDOUT, preexec_fn=lambda: os.setuid(1))\n",
|
||||
"!sleep 30"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "hSWFzkCn61tM"
|
||||
},
|
||||
"source": [
|
||||
"# Load data into Elasticsearch\n",
|
||||
"\n",
|
||||
"The following block loads the dataset into Elasticsearch."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "So-OBvUT61QD"
|
||||
},
|
||||
"source": [
|
||||
"from datasets import load_dataset\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",
|
||||
"# Load HF dataset\n",
|
||||
"dataset = load_dataset(\"ag_news\", split=\"train\")[\"text\"][:50000]\n",
|
||||
"\n",
|
||||
"# Elasticsearch bulk buffer\n",
|
||||
"buffer = []\n",
|
||||
"rows = 0\n",
|
||||
"\n",
|
||||
"for x, text in enumerate(dataset):\n",
|
||||
" # Article record\n",
|
||||
" article = {\"_id\": x, \"_index\": \"articles\", \"title\": text}\n",
|
||||
"\n",
|
||||
" # Buffer article\n",
|
||||
" buffer.append(article)\n",
|
||||
"\n",
|
||||
" # Increment number of articles processed\n",
|
||||
" rows += 1\n",
|
||||
"\n",
|
||||
" # Bulk load every 1000 records\n",
|
||||
" if rows % 1000 == 0:\n",
|
||||
" helpers.bulk(es, buffer)\n",
|
||||
" buffer = []\n",
|
||||
"\n",
|
||||
" print(\"Inserted {} articles\".format(rows), end=\"\\r\")\n",
|
||||
"\n",
|
||||
"if buffer:\n",
|
||||
" helpers.bulk(es, buffer)\n",
|
||||
"\n",
|
||||
"print(\"Total articles inserted: {}\".format(rows))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "X5RO-VNwzMAo"
|
||||
},
|
||||
"source": [
|
||||
"# Query data with Elasticsearch\n",
|
||||
"\n",
|
||||
"Elasticsearch is a token-based search system. Queries and documents are parsed into tokens and the most relevant query-document matches are calculated using a scoring algorithm. The default scoring algorithm is [BM25](https://www.elastic.co/blog/practical-bm25-part-2-the-bm25-algorithm-and-its-variables). Powerful queries can be built using a [rich query syntax](https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html#query-string-syntax) and [Query DSL](https://www.elastic.co/guide/en/elasticsearch/reference/7.10/query-dsl.html). \n",
|
||||
"\n",
|
||||
"The following section runs a query against Elasticsearch, finds the top 5 matches and returns the corresponding documents associated with each match.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 334
|
||||
},
|
||||
"id": "ucd9mwSfFTMm",
|
||||
"outputId": "1c49323e-1f29-491d-872a-14b5650d09f6"
|
||||
},
|
||||
"source": [
|
||||
"from IPython.display import display, HTML\n",
|
||||
"\n",
|
||||
"def table(category, query, rows):\n",
|
||||
" html = \"\"\"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 15px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" \"\"\"\n",
|
||||
"\n",
|
||||
" html += \"<h3>[%s] %s</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead>\" % (category, query)\n",
|
||||
" for score, text in rows:\n",
|
||||
" html += \"<tr><td>%.4f</td><td>%s</td></tr>\" % (score, text)\n",
|
||||
" html += \"</table>\"\n",
|
||||
"\n",
|
||||
" display(HTML(html))\n",
|
||||
"\n",
|
||||
"def search(query, limit):\n",
|
||||
" query = {\n",
|
||||
" \"size\": limit,\n",
|
||||
" \"query\": {\n",
|
||||
" \"query_string\": {\"query\": query}\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" results = []\n",
|
||||
" for result in es.search(index=\"articles\", body=query)[\"hits\"][\"hits\"]:\n",
|
||||
" source = result[\"_source\"]\n",
|
||||
" results.append((min(result[\"_score\"], 18) / 18, source[\"title\"]))\n",
|
||||
"\n",
|
||||
" return results\n",
|
||||
"\n",
|
||||
"limit = 3\n",
|
||||
"query= \"+yankees lose\"\n",
|
||||
"table(\"Elasticsearch\", query, search(query, limit))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 15px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" <h3>[Elasticsearch] +yankees lose</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead><tr><td>0.5808</td><td>El Duque adds to gloomy NY forecast The Yankees #39; staff infection has spread to the one man the team can #39;t afford to lose. Orlando Hernandez was scratched from last night #39;s scheduled start because </td></tr><tr><td>0.5688</td><td>Rangers Derail Red Sox The Red Sox lose for the first time in 11 games, falling to the Rangers 8-6 Saturday and missing a chance to pull within 1 1/2 games of the Yankees in the AL East.</td></tr><tr><td>0.5061</td><td>Rout leaves Yanks #39; lead at 3 Royals gain control with 10-run 5th Against a nothing-to-lose team such as the Kansas City Royals, the Yankees #39; manager wanted his team to put down the hammer early and not let baseball #39;s second worst team believe it had a chance.</td></tr></table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "W1DkpcIob5kt"
|
||||
},
|
||||
"source": [
|
||||
"The table above shows the results for the query `+yankees lose`. This query requires the token `yankees`. The search doesn't understand the semantic meaning of the query. It returns the most relevant results with those two tokens.\n",
|
||||
"\n",
|
||||
"We can see in this case, the results aren't capturing the meaning of the search. Let's try adding semantic similarity to the search!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "hMro47KedzJq"
|
||||
},
|
||||
"source": [
|
||||
"# Ranking search results with txtai\n",
|
||||
"\n",
|
||||
"txtai has a similarity module that computes the similarity between a query and a list of strings. Of course, txtai can also build a full index as shown in the previous notebooks but in this case we'll just use the ad-hoc similarity function.\n",
|
||||
"\n",
|
||||
"The code below creates a Similarity instance and defines a ranking function to order search results based on the computed similarity.\n",
|
||||
"\n",
|
||||
"`ranksearch` queries Elasticsearch for a larger set of results, ranks the results using the similarity instance and returns the top n results. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "RUOj5zhFFK8N"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"from txtai.pipeline import Similarity\n",
|
||||
"\n",
|
||||
"def ranksearch(query, limit):\n",
|
||||
" results = [text for _, text in search(query, limit * 10)]\n",
|
||||
" return [(score, results[x]) for x, score in similarity(query, results)][:limit]\n",
|
||||
"\n",
|
||||
"# Create similarity instance for re-ranking\n",
|
||||
"similarity = Similarity(\"valhalla/distilbart-mnli-12-3\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "UMFuv5-Hedfc"
|
||||
},
|
||||
"source": [
|
||||
"Now let's re-run the previous search."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 334
|
||||
},
|
||||
"id": "3jJI9OxU0dZk",
|
||||
"outputId": "3233d128-2bf4-4d53-f851-f9dd8f51629e"
|
||||
},
|
||||
"source": [
|
||||
"# Run the search\n",
|
||||
"table(\"Elasticsearch + txtai\", query, ranksearch(query, limit))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 15px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" <h3>[Elasticsearch + txtai] +yankees lose</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead><tr><td>0.9929</td><td>Ouch! Yankees hit new low INDIANS 22, YANKEES 0---At New York, Omar Vizquel went 6-for-7 to tie the American League record for hits as Cleveland handed the Yankees the largest loss in their history last night.</td></tr><tr><td>0.9874</td><td>Vazquez and Yankees Buckle Early Because Javier Vazquez fizzled while Brad Radke flourished, the Yankees sustained their first regular-season defeat by the Minnesota Twins since 2001.</td></tr><tr><td>0.9542</td><td>Slide of the Yankees: Pinstripes Punished George Steinbrenner watched from his box as his Yankees suffered the most one-sided loss in the franchise's long history.</td></tr></table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "RXB2PaDZfd8o"
|
||||
},
|
||||
"source": [
|
||||
"The results above do a much better job of finding results semantically similar in meaning to the query. Instead of just finding matches with `yankees` and `lose`, it finds matches where the `yankees lose`. \n",
|
||||
"\n",
|
||||
"This combination is effective and powerful. It takes advantage of the high performance of Elasticsearch while adding a semantic search capability. We may already have a large Elasticsearch cluster with TBs (or PBs)+ of data and years of engineering investment that solves most use cases. Semantically ranking search results is a practical approach."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "XBVL56fsRI86"
|
||||
},
|
||||
"source": [
|
||||
"# More examples\n",
|
||||
"\n",
|
||||
"Now for some more examples comparing the results from Elasticsearch vs Elasticsearch + txtai."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "7IHS38SERRpQ",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 1000
|
||||
},
|
||||
"outputId": "25325d94-ddc9-44f2-f602-d792ffeb0a8c"
|
||||
},
|
||||
"source": [
|
||||
"for query in [\"good news +economy\", \"bad news +economy\"]:\n",
|
||||
" table(\"Elasticsearch\", query, search(query, limit))\n",
|
||||
" table(\"Elasticsearch + txtai\", query, ranksearch(query, limit))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 15px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" <h3>[Elasticsearch] good news +economy</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead><tr><td>0.8756</td><td>Surprise drop US wholesale prices is mixed news for economy (AFP) AFP - A surprise drop in US wholesale prices in August showed inflation apparently in check, but analysts said this was good and bad news for the US economy.</td></tr><tr><td>0.7379</td><td>China investment slows Good news for officials who are trying to cool an overheated economy; austerity measures to remain. BEIJING (Reuters) - China reported a marked slowdown in investment and money supply growth Monday, but stubbornly </td></tr><tr><td>0.7145</td><td>Spending Rebounds, Good News for Growth WASHINGTON (Reuters) - U.S. consumer spending rebounded sharply July, government data showed on Monday, erasing the disappointment of June and bolstering hopes that the U.S. economy has recovered from its recent soft spot.</td></tr></table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 15px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" <h3>[Elasticsearch + txtai] good news +economy</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead><tr><td>0.9996</td><td>Spending Rebounds, Good News for Growth WASHINGTON (Reuters) - U.S. consumer spending rebounded sharply in July, the government said on Monday, erasing the disappointment of June and bolstering hopes that the U.S. economy has recovered from its recent soft spot.</td></tr><tr><td>0.9996</td><td>Spending Rebounds, Good News for Growth WASHINGTON (Reuters) - U.S. consumer spending rebounded sharply July, government data showed on Monday, erasing the disappointment of June and bolstering hopes that the U.S. economy has recovered from its recent soft spot.</td></tr><tr><td>0.9993</td><td>Home building surges Housing construction in August jumped to its highest level in five months, a dose of encouraging news for the economy #39;s expansion.</td></tr></table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 15px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" <h3>[Elasticsearch] bad news +economy</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead><tr><td>0.9228</td><td>Surprise drop US wholesale prices is mixed news for economy (AFP) AFP - A surprise drop in US wholesale prices in August showed inflation apparently in check, but analysts said this was good and bad news for the US economy.</td></tr><tr><td>0.6405</td><td>Field Poll: Californians liking economy Bee Staff Writer. Californians are slowly growing more optimistic about the health of the economy, but a majority still feels the state is in bad economic times, according to a new Field Poll.</td></tr><tr><td>0.6188</td><td>ADB says China should raise rates to cool economy China should raise interest rates to cool the economy and prevent a future buildup of bad loans in the banking system, the Asian Development Bank #39;s (ADB) Bei-jing representative Bruce Murray said.</td></tr></table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 15px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" <h3>[Elasticsearch + txtai] bad news +economy</h3><table><thead><tr><th>Score</th><th>Text</th></tr></thead><tr><td>0.9977</td><td>Aging society hits Japan #39;s economy Japan #39;s economy will be the most severely affected among industrialized nations by population aging, Kyodo News said Thursday.</td></tr><tr><td>0.9963</td><td>Funds: Fund Mergers Can Hurt Investors (Reuters) Reuters - Mergers and acquisitions have\\played an enormous role in the U.S. economy during the past\\several decades, but sometimes the results have been bad for\\consumers. Similarly, consolidation in the mutual fund\\business has sometimes hurt fund investors.</td></tr><tr><td>0.9958</td><td>Signs of listless economy persist In a sign of persistent weakness in the US economy, a widely watched measure of business activity declined in August for the third consecutive month.</td></tr></table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "h-Bk3KLjZMpF"
|
||||
},
|
||||
"source": [
|
||||
"Once again while Elasticsearch usually returns quality results, occasionally it will match results that aren't semantically relevant. The power of semantic search is that not only will it find direct matches but matches with the same meaning. "
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "vwELCooy4ljr"
|
||||
},
|
||||
"source": [
|
||||
"# Extractive QA with txtai\n",
|
||||
"\n",
|
||||
"In Parts 1 through 4, 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"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "LPQTb25tASIG"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "_YnqorRKAbLu"
|
||||
},
|
||||
"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"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"\n",
|
||||
"from txtai.embeddings import Embeddings\n",
|
||||
"from txtai.pipeline import Extractor\n",
|
||||
"\n",
|
||||
"# Create embeddings model, backed by sentence-transformers & transformers\n",
|
||||
"embeddings = Embeddings({\"path\": \"sentence-transformers/nli-mpnet-base-v2\"})\n",
|
||||
"\n",
|
||||
"# Create extractor instance\n",
|
||||
"extractor = Extractor(embeddings, \"distilbert-base-cased-distilled-squad\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "4X5z3UjnAGe7",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "546d4fdd-9468-4130-ee93-fafafd966e8b"
|
||||
},
|
||||
"source": [
|
||||
"data = [\"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 beat Red Sox final score 2-1\",\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",
|
||||
" \"Lightning goaltender pulled, lose to Flyers 4-1\",\n",
|
||||
" \"Flyers 4 Lightning 1 final\",\n",
|
||||
" \"Flyers win 4-1\"]\n",
|
||||
"\n",
|
||||
"questions = [\"What team won the game?\", \"What was score?\"]\n",
|
||||
"\n",
|
||||
"execute = lambda query: extractor([(question, query, question, False) for question in questions], data)\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([(question, question, question, False)], data))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"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"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,442 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "zzZbP0LM6m5z"
|
||||
},
|
||||
"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 Elasticsearch system."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "xk7t5Jcd6reO"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and `Elasticsearch`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "0y1UA4-q-YdA"
|
||||
},
|
||||
"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.10.1-linux-x86_64.tar.gz\n",
|
||||
"!tar -xzf elasticsearch-7.10.1-linux-x86_64.tar.gz\n",
|
||||
"!chown -R daemon:daemon elasticsearch-7.10.1"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "nKWz-C5gCJy8"
|
||||
},
|
||||
"source": [
|
||||
"Start an instance of Elasticsearch directly within this notebook. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "3ZfJeWbM6wmj"
|
||||
},
|
||||
"source": [
|
||||
"import os\n",
|
||||
"from subprocess import Popen, PIPE, STDOUT\n",
|
||||
"\n",
|
||||
"# If issues are encountered with this section, ES can be manually started as follows:\n",
|
||||
"# ./elasticsearch-7.10.1/bin/elasticsearch\n",
|
||||
"\n",
|
||||
"# Start and wait for server\n",
|
||||
"server = Popen(['elasticsearch-7.10.1/bin/elasticsearch'], stdout=PIPE, stderr=STDOUT, preexec_fn=lambda: os.setuid(1))\n",
|
||||
"!sleep 30"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "TWEn4w68-D1y"
|
||||
},
|
||||
"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",
|
||||
"metadata": {
|
||||
"id": "8tVrIqSq-KBa"
|
||||
},
|
||||
"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"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "hSWFzkCn61tM"
|
||||
},
|
||||
"source": [
|
||||
"# Load data into Elasticsearch\n",
|
||||
"\n",
|
||||
"The following block copies rows from SQLite to Elasticsearch."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "So-OBvUT61QD",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "9647b8f8-8471-41bf-ccfa-a75306665638"
|
||||
},
|
||||
"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|(?<!.*?results.*?)discussion|introduction|reference\", name.lower()):\n",
|
||||
" # Bulk action fields\n",
|
||||
" article[\"_id\"] = article[\"id\"]\n",
|
||||
" article[\"_index\"] = \"articles\"\n",
|
||||
"\n",
|
||||
" # Buffer article\n",
|
||||
" buffer.append(article)\n",
|
||||
"\n",
|
||||
" # Increment number of articles processed\n",
|
||||
" rows += 1\n",
|
||||
"\n",
|
||||
" # Bulk load every 1000 records\n",
|
||||
" if rows % 1000 == 0:\n",
|
||||
" helpers.bulk(es, buffer)\n",
|
||||
" buffer = []\n",
|
||||
"\n",
|
||||
" print(\"Inserted {} articles\".format(rows), end=\"\\r\")\n",
|
||||
"\n",
|
||||
"if buffer:\n",
|
||||
" helpers.bulk(es, buffer)\n",
|
||||
"\n",
|
||||
"print(\"Total articles inserted: {}\".format(rows))\n"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Total articles inserted: 21499\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "X5RO-VNwzMAo"
|
||||
},
|
||||
"source": [
|
||||
"# Query data\n",
|
||||
"\n",
|
||||
"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.\n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "ucd9mwSfFTMm",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 348
|
||||
},
|
||||
"outputId": "b21d6aff-6abe-48f5-9914-7b7fb8472adb"
|
||||
},
|
||||
"source": [
|
||||
"import pandas as pd\n",
|
||||
"\n",
|
||||
"from IPython.display import display, HTML\n",
|
||||
"\n",
|
||||
"pd.set_option(\"display.max_colwidth\", None)\n",
|
||||
"\n",
|
||||
"query = {\n",
|
||||
" \"_source\": [\"article\", \"title\", \"published\", \"reference\", \"text\"],\n",
|
||||
" \"size\": 5,\n",
|
||||
" \"query\": {\n",
|
||||
" \"query_string\": {\"query\": \"risk factors\"}\n",
|
||||
" }\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"results = []\n",
|
||||
"for result in es.search(index=\"articles\", body=query)[\"hits\"][\"hits\"]:\n",
|
||||
" source = result[\"_source\"]\n",
|
||||
" results.append((source[\"title\"], source[\"published\"], source[\"reference\"], source[\"text\"]))\n",
|
||||
"\n",
|
||||
"df = pd.DataFrame(results, columns=[\"Title\", \"Published\", \"Reference\", \"Match\"])\n",
|
||||
"\n",
|
||||
"display(HTML(df.to_html(index=False)))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"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>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>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>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>COVID-19: what has been learned and to be learned about the novel coronavirus disease</td>\n",
|
||||
" <td>2020-03-15 00:00:00</td>\n",
|
||||
" <td>https://doi.org/10.7150/ijbs.45134</td>\n",
|
||||
" <td>• Three major risk factors for COVID-19 were sex (male), age (≥60), and severe pneumonia.</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <td>Associations with covid-19 hospitalisation amongst 406,793 adults: the UK Biobank prospective cohort study</td>\n",
|
||||
" <td>2020-05-11 00:00:00</td>\n",
|
||||
" <td>http://medrxiv.org/cgi/content/short/2020.05.06.20092957v1?rss=1</td>\n",
|
||||
" <td>In addition, many risk factors for covid-19 documented in the literature are highly correlated and it is not clear which may be independently related to risk.</td>\n",
|
||||
" </tr>\n",
|
||||
" </tbody>\n",
|
||||
"</table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ylxOKji1-9_K"
|
||||
},
|
||||
"source": [
|
||||
"# Derive columns with Extractive QA\n",
|
||||
"\n",
|
||||
"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."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "mwBTrCkcOM_H"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"from txtai.embeddings import Embeddings\n",
|
||||
"from txtai.pipeline import Extractor\n",
|
||||
"\n",
|
||||
"# Create embeddings model, backed by sentence-transformers & transformers\n",
|
||||
"embeddings = Embeddings({\"path\": \"sentence-transformers/nli-mpnet-base-v2\"})\n",
|
||||
"\n",
|
||||
"# Create extractor instance using qa model designed for the CORD-19 dataset\n",
|
||||
"extractor = Extractor(embeddings, \"NeuML/bert-small-cord19qa\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "Yv75Lh-cOpL9",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 400
|
||||
},
|
||||
"outputId": "adee88e1-02bf-4a20-febb-6d2c170a63f9"
|
||||
},
|
||||
"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",
|
||||
" name, text = source[\"name\"], source[\"text\"]\n",
|
||||
"\n",
|
||||
" if not name or not re.search(r\"background|(?<!.*?results.*?)discussion|introduction|reference\", name.lower()):\n",
|
||||
" rows.append(text)\n",
|
||||
" \n",
|
||||
" return rows\n",
|
||||
"\n",
|
||||
"results = []\n",
|
||||
"for result in es.search(index=\"articles\", body=query)[\"hits\"][\"hits\"]:\n",
|
||||
" source = result[\"_source\"]\n",
|
||||
"\n",
|
||||
" # Use QA extractor to derive additional columns\n",
|
||||
" answers = extractor([(\"Risk factors\", \"risk factor\", \"What are names of risk factors?\", False),\n",
|
||||
" (\"Locations\", \"city country state\", \"What are names of locations?\", False)], sections(source[\"article\"]))\n",
|
||||
"\n",
|
||||
" results.append((source[\"title\"], source[\"published\"], source[\"reference\"], source[\"text\"]) + tuple([answer[1] for answer in answers]))\n",
|
||||
"\n",
|
||||
"df = pd.DataFrame(results, columns=[\"Title\", \"Published\", \"Reference\", \"Match\", \"Risk Factors\", \"Locations\"])\n",
|
||||
"\n",
|
||||
"display(HTML(df.to_html(index=False)))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"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>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>None</td>\n",
|
||||
" <td>Abbott, Abbott Park, Illinois</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>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>comorbidity groupings and the corresponding health conditions</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <td>COVID-19: what has been learned and to be learned about the novel coronavirus disease</td>\n",
|
||||
" <td>2020-03-15 00:00:00</td>\n",
|
||||
" <td>https://doi.org/10.7150/ijbs.45134</td>\n",
|
||||
" <td>• Three major risk factors for COVID-19 were sex (male), age (≥60), and severe pneumonia.</td>\n",
|
||||
" <td>Mandatory contact tracing and quarantine</td>\n",
|
||||
" <td>cities, provinces, and countries</td>\n",
|
||||
" </tr>\n",
|
||||
" </tbody>\n",
|
||||
"</table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Pjmz-RORV8E"
|
||||
},
|
||||
"source": [
|
||||
"# Apply labels with zero-shot classification\n",
|
||||
"\n",
|
||||
"This notebook shows how zero-shot classification can be used to perform text classification, labeling and topic modeling. txtai provides a light-weight wrapper around the zero-shot-classification pipeline in Hugging Face Transformers. This method works impressively well out of the box. Kudos to the Hugging Face team for the phenomenal work on zero-shot classification!\n",
|
||||
"\n",
|
||||
"The examples in this notebook pick the best matching label using a list of labels for a snippet of text.\n",
|
||||
"\n",
|
||||
"[tldrstory](https://github.com/neuml/tldrstory) has full-stack implementation of a zero-shot classification system using Streamlit, FastAPI and Hugging Face Transformers. There is also a [Medium article describing tldrstory](https://towardsdatascience.com/tldrstory-ai-powered-understanding-of-headlines-and-story-text-fc86abd702fc) and zero-shot classification. \n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dk31rbYjSTYm"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "XMQuuun2R06J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PNPJ95cdTKSS"
|
||||
},
|
||||
"source": [
|
||||
"# Create a Labels instance\n",
|
||||
"\n",
|
||||
"The Labels instance is the main entrypoint for zero-shot classification. This is a light-weight wrapper around the zero-shot-classification pipeline in Hugging Face Transformers.\n",
|
||||
"\n",
|
||||
"In addition to the default model, additional models can be found on the [Hugging Face model hub](https://huggingface.co/models?search=mnli).\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "nTDwXOUeTH2-"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"\n",
|
||||
"from txtai.pipeline import Labels\n",
|
||||
"\n",
|
||||
"# Create labels model\n",
|
||||
"labels = Labels()\n",
|
||||
"\n",
|
||||
"# Alternate models can be used via passing the model path as shown below\n",
|
||||
"# labels = Labels(\"roberta-large-mnli\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-vGR_piwZZO6"
|
||||
},
|
||||
"source": [
|
||||
"# Applying labels to text\n",
|
||||
"\n",
|
||||
"The example below shows how a zero-shot classifier can be applied to arbitary text. The default model for the zero-shot classification pipeline is *bart-large-mnli*. \n",
|
||||
"\n",
|
||||
"Look at the results below. It's nothing short of amazing✨ how well it performs. These aren't all simple even for a human. For example, intercepted was purposely picked as that is more common in football than basketball. The amount of knowledge stored in larger Transformer models continues to impress me. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "-K2YJJzsVtfq",
|
||||
"outputId": "7a1edf58-15e0-46c8-958e-3a8e6045f802"
|
||||
},
|
||||
"source": [
|
||||
"data = [\"Dodgers lose again, give up 3 HRs in a loss to the Giants\",\n",
|
||||
" \"Giants 5 Cardinals 4 final in extra innings\",\n",
|
||||
" \"Dodgers drop Game 2 against the Giants, 5-4\",\n",
|
||||
" \"Flyers 4 Lightning 1 final. 45 saves for the Lightning.\",\n",
|
||||
" \"Slashing, penalty, 2 minute power play coming up\",\n",
|
||||
" \"What a stick save!\",\n",
|
||||
" \"Leads the NFL in sacks with 9.5\",\n",
|
||||
" \"UCF 38 Temple 13\",\n",
|
||||
" \"With the 30 yard completion, down to the 10 yard line\",\n",
|
||||
" \"Drains the 3pt shot!!, 0:15 remaining in the game\",\n",
|
||||
" \"Intercepted! Drives down the court and shoots for the win\",\n",
|
||||
" \"Massive dunk!!! they are now up by 15 with 2 minutes to go\"]\n",
|
||||
"\n",
|
||||
"# List of labels\n",
|
||||
"tags = [\"Baseball\", \"Football\", \"Hockey\", \"Basketball\"]\n",
|
||||
"\n",
|
||||
"print(\"%-75s %s\" % (\"Text\", \"Label\"))\n",
|
||||
"print(\"-\" * 100)\n",
|
||||
"\n",
|
||||
"for text in data:\n",
|
||||
" print(\"%-75s %s\" % (text, tags[labels(text, tags)[0][0]]))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Text Label\n",
|
||||
"----------------------------------------------------------------------------------------------------\n",
|
||||
"Dodgers lose again, give up 3 HRs in a loss to the Giants Baseball\n",
|
||||
"Giants 5 Cardinals 4 final in extra innings Baseball\n",
|
||||
"Dodgers drop Game 2 against the Giants, 5-4 Baseball\n",
|
||||
"Flyers 4 Lightning 1 final. 45 saves for the Lightning. Hockey\n",
|
||||
"Slashing, penalty, 2 minute power play coming up Hockey\n",
|
||||
"What a stick save! Hockey\n",
|
||||
"Leads the NFL in sacks with 9.5 Football\n",
|
||||
"UCF 38 Temple 13 Football\n",
|
||||
"With the 30 yard completion, down to the 10 yard line Football\n",
|
||||
"Drains the 3pt shot!!, 0:15 remaining in the game Basketball\n",
|
||||
"Intercepted! Drives down the court and shoots for the win Basketball\n",
|
||||
"Massive dunk!!! they are now up by 15 with 2 minutes to go Basketball\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "t-tGAzCxsHLy"
|
||||
},
|
||||
"source": [
|
||||
"# Let's try emoji 😀\n",
|
||||
"\n",
|
||||
"Does the model have knowledge of emoji? Check out the run below, sure looks like it does! Notice the labels are applied based on the perspective from which the information is presented. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "uIf064M9pbjn",
|
||||
"outputId": "1d104014-e9ca-4c89-d259-2b5b231840ad"
|
||||
},
|
||||
"source": [
|
||||
"tags = [\"😀\", \"😡\"]\n",
|
||||
"\n",
|
||||
"print(\"%-75s %s\" % (\"Text\", \"Label\"))\n",
|
||||
"print(\"-\" * 100)\n",
|
||||
"\n",
|
||||
"for text in data:\n",
|
||||
" print(\"%-75s %s\" % (text, tags[labels(text, tags)[0][0]]))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Text Label\n",
|
||||
"----------------------------------------------------------------------------------------------------\n",
|
||||
"Dodgers lose again, give up 3 HRs in a loss to the Giants 😡\n",
|
||||
"Giants 5 Cardinals 4 final in extra innings 😀\n",
|
||||
"Dodgers drop Game 2 against the Giants, 5-4 😡\n",
|
||||
"Flyers 4 Lightning 1 final. 45 saves for the Lightning. 😀\n",
|
||||
"Slashing, penalty, 2 minute power play coming up 😡\n",
|
||||
"What a stick save! 😀\n",
|
||||
"Leads the NFL in sacks with 9.5 😀\n",
|
||||
"UCF 38 Temple 13 😀\n",
|
||||
"With the 30 yard completion, down to the 10 yard line 😀\n",
|
||||
"Drains the 3pt shot!!, 0:15 remaining in the game 😀\n",
|
||||
"Intercepted! Drives down the court and shoots for the win 😀\n",
|
||||
"Massive dunk!!! they are now up by 15 with 2 minutes to go 😀\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,836 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Pjmz-RORV8E"
|
||||
},
|
||||
"source": [
|
||||
"# API Gallery\n",
|
||||
"\n",
|
||||
"The txtai API is a web-based service backed by [FastAPI](https://fastapi.tiangolo.com/). All txtai functionality including similarity search, extractive QA and zero-shot labeling is available via the API.\n",
|
||||
"\n",
|
||||
"This notebook installs the txtai API and shows an example using each of the supported language bindings for txtai."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dk31rbYjSTYm"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies. Since this notebook uses the API, we need to install the api extras package."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "XMQuuun2R06J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[api]"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-vGR_piwZZO6"
|
||||
},
|
||||
"source": [
|
||||
"# Python\n",
|
||||
"\n",
|
||||
"The first method we'll try is direct access via Python. We'll use zero-shot labeling for all the examples here. See [this notebook](https://colab.research.google.com/github/neuml/txtai/blob/master/examples/07_Apply_labels_with_zero_shot_classification.ipynb) for more details on zero-shot classification. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "P4q72tkRMMkR"
|
||||
},
|
||||
"source": [
|
||||
"## Configure Labels instance"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "8Dy_TJ0iM38Q"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"import os\n",
|
||||
"from IPython.core.display import display, HTML\n",
|
||||
"from txtai.pipeline import Labels\n",
|
||||
"\n",
|
||||
"def table(rows):\n",
|
||||
" html = \"\"\"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 20px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" \"\"\"\n",
|
||||
"\n",
|
||||
" html += \"<table><thead><tr><th>Text</th><th>Label</th></tr></thead>\"\n",
|
||||
" for text, label in rows:\n",
|
||||
" html += \"<tr><td>%s</td><td>%s</td></tr>\" % (text, label)\n",
|
||||
" html += \"</table>\"\n",
|
||||
"\n",
|
||||
" display(HTML(html))\n",
|
||||
"\n",
|
||||
"# Create labels model\n",
|
||||
"labels = Labels()"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "L4B73tGkMT6Q"
|
||||
},
|
||||
"source": [
|
||||
"## Apply labels to text"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 324
|
||||
},
|
||||
"id": "-K2YJJzsVtfq",
|
||||
"outputId": "65782fd8-51fb-4531-8e8b-f28bca678fa0"
|
||||
},
|
||||
"source": [
|
||||
"data = [\"Wears a red suit and says ho ho\",\n",
|
||||
" \"Pulls a flying sleigh\",\n",
|
||||
" \"This is cut down and decorated\",\n",
|
||||
" \"Santa puts these under the tree\",\n",
|
||||
" \"Best way to spend the holidays\"]\n",
|
||||
"\n",
|
||||
"# List of labels\n",
|
||||
"tags = [\"🎅 Santa Clause\", \"🦌 Reindeer\", \"🍪 Cookies\", \"🎄 Christmas Tree\", \"🎁 Gifts\", \"👪 Family\"]\n",
|
||||
"\n",
|
||||
"# Render output to table\n",
|
||||
"table([(text, tags[labels(text, tags)[0][0]]) for text in data])"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <style type='text/css'>\n",
|
||||
" @import url('https://fonts.googleapis.com/css?family=Oswald&display=swap');\n",
|
||||
" table {\n",
|
||||
" border-collapse: collapse;\n",
|
||||
" width: 900px;\n",
|
||||
" }\n",
|
||||
" th, td {\n",
|
||||
" border: 1px solid #9e9e9e;\n",
|
||||
" padding: 10px;\n",
|
||||
" font: 20px Oswald;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
" <table><thead><tr><th>Text</th><th>Label</th></tr></thead><tr><td>Wears a red suit and says ho ho</td><td>🎅 Santa Clause</td></tr><tr><td>Pulls a flying sleigh</td><td>🦌 Reindeer</td></tr><tr><td>This is cut down and decorated</td><td>🎄 Christmas Tree</td></tr><tr><td>Santa puts these under the tree</td><td>🎁 Gifts</td></tr><tr><td>Best way to spend the holidays</td><td>👪 Family</td></tr></table>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "UF_bImkLHTMs"
|
||||
},
|
||||
"source": [
|
||||
"Once again we see the power of zero-shot labeling. The model wasn't trained on any data specific to this example. Still amazed with how much knowledge is stored in large NLP models."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PNPJ95cdTKSS"
|
||||
},
|
||||
"source": [
|
||||
"# Start an API instance\n",
|
||||
"\n",
|
||||
"Now we'll start an API instance to run the remaining examples. The API needs a configuration file to run. The example below is simplified to only include labeling. See [this link](https://github.com/neuml/txtai#api) for a more detailed configuration example.\n",
|
||||
"\n",
|
||||
"The API instance is started in the background.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "nTDwXOUeTH2-",
|
||||
"outputId": "2220a3c9-1cff-4c2f-b21e-13dd2d7cb816"
|
||||
},
|
||||
"source": [
|
||||
"%%writefile index.yml\n",
|
||||
"\n",
|
||||
"# Labels settings\n",
|
||||
"labels:"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing index.yml\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "nGITHxUyRzyp"
|
||||
},
|
||||
"source": [
|
||||
"!CONFIG=index.yml nohup uvicorn \"txtai.api:app\" &> api.log &\n",
|
||||
"!sleep 90"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "NHvBFZeSd9AG"
|
||||
},
|
||||
"source": [
|
||||
"# JavaScript\n",
|
||||
"\n",
|
||||
"txtai.js is available via NPM and can be installed as follows.\n",
|
||||
"\n",
|
||||
"```bash\n",
|
||||
"npm install txtai\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"For this example, we'll clone the txtai.js project to import the example build configuration."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "b52knObEdcCr"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!git clone https://github.com/neuml/txtai.js"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "rUGS0t-JMsS9"
|
||||
},
|
||||
"source": [
|
||||
"## Create labels.js\n",
|
||||
"\n",
|
||||
"The following file is a JavaScript version of the labels example."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "zJbKRTSJV-kd",
|
||||
"outputId": "6c111b5d-6e55-4dac-c6c2-0988c2a834da"
|
||||
},
|
||||
"source": [
|
||||
"%%writefile txtai.js/examples/node/src/labels.js\n",
|
||||
"import {Labels} from \"txtai\";\n",
|
||||
"import {sprintf} from \"sprintf-js\";\n",
|
||||
"\n",
|
||||
"const run = async () => {\n",
|
||||
" try {\n",
|
||||
" let labels = new Labels(\"http://localhost:8000\");\n",
|
||||
"\n",
|
||||
" let data = [\"Wears a red suit and says ho ho\",\n",
|
||||
" \"Pulls a flying sleigh\",\n",
|
||||
" \"This is cut down and decorated\",\n",
|
||||
" \"Santa puts these under the tree\",\n",
|
||||
" \"Best way to spend the holidays\"];\n",
|
||||
"\n",
|
||||
" // List of labels\n",
|
||||
" let tags = [\"🎅 Santa Clause\", \"🦌 Reindeer\", \"🍪 Cookies\", \"🎄 Christmas Tree\", \"🎁 Gifts\", \"👪 Family\"];\n",
|
||||
"\n",
|
||||
" console.log(sprintf(\"%-40s %s\", \"Text\", \"Label\"));\n",
|
||||
" console.log(\"-\".repeat(75))\n",
|
||||
"\n",
|
||||
" for (let text of data) {\n",
|
||||
" let label = await labels.label(text, tags);\n",
|
||||
" label = tags[label[0].id];\n",
|
||||
"\n",
|
||||
" console.log(sprintf(\"%-40s %s\", text, label));\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" catch (e) {\n",
|
||||
" console.trace(e);\n",
|
||||
" }\n",
|
||||
"};\n",
|
||||
"\n",
|
||||
"run();\n"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Overwriting txtai.js/examples/node/src/labels.js\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "nTBs11j-GtD-"
|
||||
},
|
||||
"source": [
|
||||
"## Build and run labels example\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "kC5Oub6wa1nK"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"os.chdir(\"txtai.js/examples/node\")\n",
|
||||
"!npm install\n",
|
||||
"!npm run build"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ckOHNqyaeL-B",
|
||||
"outputId": "6d8e745c-52d1-4456-fc46-2ff8fda2e675"
|
||||
},
|
||||
"source": [
|
||||
"!node dist/labels.js"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Text Label\n",
|
||||
"---------------------------------------------------------------------------\n",
|
||||
"Wears a red suit and says ho ho 🎅 Santa Clause\n",
|
||||
"Pulls a flying sleigh 🦌 Reindeer\n",
|
||||
"This is cut down and decorated 🎄 Christmas Tree\n",
|
||||
"Santa puts these under the tree 🎁 Gifts\n",
|
||||
"Best way to spend the holidays 👪 Family\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "1yukBIMYG5OE"
|
||||
},
|
||||
"source": [
|
||||
"The JavaScript program is showing the same results as when natively running through Python!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "nNiMgvg0p2BG"
|
||||
},
|
||||
"source": [
|
||||
"# Java\n",
|
||||
"\n",
|
||||
"txtai.java integrates with standard Java build tools (Gradle, Maven, SBT). The following shows how to add txtai as a dependency to Gradle.\n",
|
||||
"\n",
|
||||
"```gradle\n",
|
||||
"implementation 'com.github.neuml:txtai.java:v4.0.0'\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"For this example, we'll clone the txtai.java project to import the example build configuration."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "qs2ai8lhqmga"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"os.chdir(\"/content\")\n",
|
||||
"!git clone https://github.com/neuml/txtai.java"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "o8QFvzXkNFgq"
|
||||
},
|
||||
"source": [
|
||||
"## Create LabelsDemo.java\n",
|
||||
"\n",
|
||||
"The following file is a Java version of the labels example."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "v73L8Gw0p6fh",
|
||||
"outputId": "a7f797f2-a91f-4033-89c7-4baf76204d93"
|
||||
},
|
||||
"source": [
|
||||
"%%writefile txtai.java/examples/src/main/java/LabelsDemo.java\n",
|
||||
"import java.util.Arrays;\n",
|
||||
"import java.util.ArrayList;\n",
|
||||
"import java.util.List;\n",
|
||||
"\n",
|
||||
"import txtai.API.IndexResult;\n",
|
||||
"import txtai.Labels;\n",
|
||||
"\n",
|
||||
"public class LabelsDemo {\n",
|
||||
" public static void main(String[] args) {\n",
|
||||
" try {\n",
|
||||
" Labels labels = new Labels(\"http://localhost:8000\");\n",
|
||||
"\n",
|
||||
" List <String> data = \n",
|
||||
" Arrays.asList(\"Wears a red suit and says ho ho\",\n",
|
||||
" \"Pulls a flying sleigh\",\n",
|
||||
" \"This is cut down and decorated\",\n",
|
||||
" \"Santa puts these under the tree\",\n",
|
||||
" \"Best way to spend the holidays\");\n",
|
||||
"\n",
|
||||
" // List of labels\n",
|
||||
" List<String> tags = Arrays.asList(\"🎅 Santa Clause\", \"🦌 Reindeer\", \"🍪 Cookies\", \"🎄 Christmas Tree\", \"🎁 Gifts\", \"👪 Family\");\n",
|
||||
"\n",
|
||||
" System.out.printf(\"%-40s %s%n\", \"Text\", \"Label\");\n",
|
||||
" System.out.println(new String(new char[75]).replace(\"\\0\", \"-\"));\n",
|
||||
"\n",
|
||||
" for (String text: data) {\n",
|
||||
" List<IndexResult> label = labels.label(text, tags);\n",
|
||||
" System.out.printf(\"%-40s %s%n\", text, tags.get(label.get(0).id));\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" catch (Exception ex) {\n",
|
||||
" ex.printStackTrace();\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
"}\n"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Overwriting txtai.java/examples/src/main/java/LabelsDemo.java\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "wZv7eMIOLnRC"
|
||||
},
|
||||
"source": [
|
||||
"## Build and run labels example"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "N2Mm3Gl5sH1z",
|
||||
"outputId": "b5249daf-e5a1-4b71-b64c-2b3c6748e846"
|
||||
},
|
||||
"source": [
|
||||
"os.chdir(\"txtai.java/examples\")\n",
|
||||
"!../gradlew -q --console=plain labels 2> /dev/null"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Text Label\n",
|
||||
"---------------------------------------------------------------------------\n",
|
||||
"Wears a red suit and says ho ho 🎅 Santa Clause\n",
|
||||
"Pulls a flying sleigh 🦌 Reindeer\n",
|
||||
"This is cut down and decorated 🎄 Christmas Tree\n",
|
||||
"Santa puts these under the tree 🎁 Gifts\n",
|
||||
"Best way to spend the holidays 👪 Family\n",
|
||||
"\u001b[m"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "iHpQvUAgNp7j"
|
||||
},
|
||||
"source": [
|
||||
"The Java program is showing the same results as when natively running through Python!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "zU6jK2UL7D5H"
|
||||
},
|
||||
"source": [
|
||||
"# Rust\n",
|
||||
"\n",
|
||||
"txtai.rs is available via crates.io and can be installed by adding the following to your cargo.toml file.\n",
|
||||
"\n",
|
||||
"```toml\n",
|
||||
"[dependencies]\n",
|
||||
"txtai = { version = \"4.0\" }\n",
|
||||
"tokio = { version = \"0.2\", features = [\"full\"] }\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"For this example, we'll clone the txtai.rs project to import the example build configuration. First we need to install Rust."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "Ob4aswkx7jRh"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"os.chdir(\"/content\")\n",
|
||||
"!apt-get install rustc\n",
|
||||
"!git clone https://github.com/neuml/txtai.rs"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "evEQQXBuObZn"
|
||||
},
|
||||
"source": [
|
||||
"## Create labels.rs\n",
|
||||
"\n",
|
||||
"The following file is a Rust version of the labels example."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "jjggKnKQ7jQO",
|
||||
"outputId": "76a2b1d9-2889-47b0-a3af-5d71a763bb0b"
|
||||
},
|
||||
"source": [
|
||||
"%%writefile txtai.rs/examples/demo/src/labels.rs\n",
|
||||
"use std::error::Error;\n",
|
||||
"\n",
|
||||
"use txtai::labels::Labels;\n",
|
||||
"\n",
|
||||
"pub async fn labels() -> Result<(), Box<dyn Error>> {\n",
|
||||
" let labels = Labels::new(\"http://localhost:8000\");\n",
|
||||
"\n",
|
||||
" let data = [\"Wears a red suit and says ho ho\",\n",
|
||||
" \"Pulls a flying sleigh\",\n",
|
||||
" \"This is cut down and decorated\",\n",
|
||||
" \"Santa puts these under the tree\",\n",
|
||||
" \"Best way to spend the holidays\"];\n",
|
||||
"\n",
|
||||
" println!(\"{:<40} {}\", \"Text\", \"Label\");\n",
|
||||
" println!(\"{}\", \"-\".repeat(75));\n",
|
||||
"\n",
|
||||
" for text in data.iter() {\n",
|
||||
" let tags = vec![\"🎅 Santa Clause\", \"🦌 Reindeer\", \"🍪 Cookies\", \"🎄 Christmas Tree\", \"🎁 Gifts\", \"👪 Family\"];\n",
|
||||
" let label = labels.label(text, &tags).await?[0].id;\n",
|
||||
"\n",
|
||||
" println!(\"{:<40} {}\", text, tags[label]);\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" Ok(())\n",
|
||||
"}"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Overwriting txtai.rs/examples/demo/src/labels.rs\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "gFFPZO8sQZC4"
|
||||
},
|
||||
"source": [
|
||||
"## Build and run labels example\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "wuoAidGz9T4g"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"os.chdir(\"txtai.rs/examples/demo\")\n",
|
||||
"!cargo build"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "-_v_FbL0-yPk",
|
||||
"outputId": "821333f5-5f90-4f89-c2eb-673c2e14e4fe"
|
||||
},
|
||||
"source": [
|
||||
"!cargo run labels"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"\u001b[0m\u001b[0m\u001b[1m\u001b[32m Finished\u001b[0m dev [unoptimized + debuginfo] target(s) in 0.07s\n",
|
||||
"\u001b[0m\u001b[0m\u001b[1m\u001b[32m Running\u001b[0m `target/debug/demo labels`\n",
|
||||
"Text Label\n",
|
||||
"---------------------------------------------------------------------------\n",
|
||||
"Wears a red suit and says ho ho 🎅 Santa Clause\n",
|
||||
"Pulls a flying sleigh 🦌 Reindeer\n",
|
||||
"This is cut down and decorated 🎄 Christmas Tree\n",
|
||||
"Santa puts these under the tree 🎁 Gifts\n",
|
||||
"Best way to spend the holidays 👪 Family\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "kDmS89TPS3kb"
|
||||
},
|
||||
"source": [
|
||||
"The Rust program is showing the same results as when natively running through Python!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ezznN4I8_CCQ"
|
||||
},
|
||||
"source": [
|
||||
"# Go\n",
|
||||
"\n",
|
||||
"txtai.go can be installed by adding the following import statement. When using modules, txtai.go will automatically be installed. Otherwise use `go get`.\n",
|
||||
"\n",
|
||||
"```golang\n",
|
||||
"import \"github.com/neuml/txtai.go\"\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"For this example, we'll create a standalone process for labeling. First we need to install Go."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "b-b6fhLQ_DpQ"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"os.chdir(\"/content\")\n",
|
||||
"!apt install golang-go\n",
|
||||
"!go get \"github.com/neuml/txtai.go\""
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dw-I6jOGR6vA"
|
||||
},
|
||||
"source": [
|
||||
"## Create labels.go\n",
|
||||
"\n",
|
||||
"The following file is a Go version of the labels example."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "bLBJwkN4ANpi",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "883ea7b2-2fbc-471c-e0bb-59ef5172a6a4"
|
||||
},
|
||||
"source": [
|
||||
"%%writefile labels.go\n",
|
||||
"package main\n",
|
||||
"\n",
|
||||
"import (\n",
|
||||
"\t\"fmt\"\n",
|
||||
"\t\"strings\"\n",
|
||||
"\t\"github.com/neuml/txtai.go\"\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"func main() {\n",
|
||||
"\tlabels := txtai.Labels(\"http://localhost:8000\")\n",
|
||||
"\n",
|
||||
"\tdata := []string{\"Wears a red suit and says ho ho\",\n",
|
||||
" \"Pulls a flying sleigh\",\n",
|
||||
" \"This is cut down and decorated\",\n",
|
||||
" \"Santa puts these under the tree\",\n",
|
||||
" \"Best way to spend the holidays\"}\n",
|
||||
"\n",
|
||||
"\t// List of labels\n",
|
||||
"\ttags := []string{\"🎅 Santa Clause\", \"🦌 Reindeer\", \"🍪 Cookies\", \"🎄 Christmas Tree\", \"🎁 Gifts\", \"👪 Family\"}\n",
|
||||
"\n",
|
||||
"\tfmt.Printf(\"%-40s %s\\n\", \"Text\", \"Label\")\n",
|
||||
"\tfmt.Println(strings.Repeat(\"-\", 75))\n",
|
||||
"\n",
|
||||
"\tfor _, text := range data {\n",
|
||||
"\t\tlabel := labels.Label(text, tags)\n",
|
||||
"\t\tfmt.Printf(\"%-40s %s\\n\", text, tags[label[0].Id])\n",
|
||||
"\t}\n",
|
||||
"}"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing labels.go\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PJ2XzzDbSeZh"
|
||||
},
|
||||
"source": [
|
||||
"## Build and run labels example\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "l1xnUbtdAy0p",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "5bc6015c-5c9c-4d8a-daf7-6897ec6cbd80"
|
||||
},
|
||||
"source": [
|
||||
"!go run labels.go"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Text Label\n",
|
||||
"---------------------------------------------------------------------------\n",
|
||||
"Wears a red suit and says ho ho 🎅 Santa Clause\n",
|
||||
"Pulls a flying sleigh 🦌 Reindeer\n",
|
||||
"This is cut down and decorated 🎄 Christmas Tree\n",
|
||||
"Santa puts these under the tree 🎁 Gifts\n",
|
||||
"Best way to spend the holidays 👪 Family\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "oml43X5eS6YB"
|
||||
},
|
||||
"source": [
|
||||
"The Go program is showing the same results as when natively running through Python!"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Pjmz-RORV8E"
|
||||
},
|
||||
"source": [
|
||||
"# Building abstractive text summaries\n",
|
||||
"\n",
|
||||
"In the field of text summarization, there are two primary categories of summarization, extractive and abstractive summarization.\n",
|
||||
"\n",
|
||||
"Extractive summarization takes subsections of the text and joins them together to form a summary. This is commonly backed by graph algorithms like TextRank to find the sections/sentences with the most commonality. These summaries can be highly effective but they are unable to transform text and don't have a contextual understanding.\n",
|
||||
"\n",
|
||||
"Abstractive summarization uses Natural Language Processing (NLP) models to build transformative summaries of text. This is similar to having a human read an article and asking what was it about. A human wouldn't just give a verbose reading of the text. This notebook shows how blocks of text can be summarized using an abstractive summarization pipeline. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dk31rbYjSTYm"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies. Since this notebook is using optional pipelines, we need to install the pipeline extras package."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "XMQuuun2R06J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline]"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PNPJ95cdTKSS"
|
||||
},
|
||||
"source": [
|
||||
"# Create a Summary instance\n",
|
||||
"\n",
|
||||
"The Summary instance is the main entrypoint for text summarization. This is a light-weight wrapper around the summarization pipeline in Hugging Face Transformers.\n",
|
||||
"\n",
|
||||
"In addition to the default model, additional models can be found on the [Hugging Face model hub](https://huggingface.co/models?pipeline_tag=summarization).\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "nTDwXOUeTH2-"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"\n",
|
||||
"from txtai.pipeline import Summary\n",
|
||||
"\n",
|
||||
"# Create summary model\n",
|
||||
"summary = Summary()"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-vGR_piwZZO6"
|
||||
},
|
||||
"source": [
|
||||
"# Summarize text\n",
|
||||
"\n",
|
||||
"The example below shows how a large block of text can be distilled down into a smaller summary."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 36
|
||||
},
|
||||
"id": "-K2YJJzsVtfq",
|
||||
"outputId": "cdf54f20-72ad-4f65-bc17-100e32e6cc71"
|
||||
},
|
||||
"source": [
|
||||
"text = (\"Search is the base of many applications. Once data starts to pile up, users want to be able to find it. It’s the foundation \"\n",
|
||||
" \"of the internet and an ever-growing challenge that is never solved or done. The field of Natural Language Processing (NLP) is \"\n",
|
||||
" \"rapidly evolving with a number of new developments. Large-scale general language models are an exciting new capability \"\n",
|
||||
" \"allowing us to add amazing functionality quickly with limited compute and people. Innovation continues with new models \"\n",
|
||||
" \"and advancements coming in at what seems a weekly basis. This article introduces txtai, an AI-powered search engine \"\n",
|
||||
" \"that enables Natural Language Understanding (NLU) based search in any application.\"\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"summary(text, maxlength=10)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
},
|
||||
"text/plain": [
|
||||
"'Search is the foundation of the internet'"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 3
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "n2jndgE-JyWX"
|
||||
},
|
||||
"source": [
|
||||
"Notice how the summarizer built a sentence using parts of the document above. It takes a basic understanding of language in order to understand the first two sentences and how to combine them into a single transformative sentence."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "27PneZxQx7NR"
|
||||
},
|
||||
"source": [
|
||||
"# Summarize a document\n",
|
||||
"\n",
|
||||
"The next section retrieves an article, extracts text from it (more to come on this topic) and summarizes that text."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 53
|
||||
},
|
||||
"id": "idPThgJGvIju",
|
||||
"outputId": "7d0580e6-2531-48c9-a32a-481ccf32900d"
|
||||
},
|
||||
"source": [
|
||||
"!wget -q \"https://medium.com/neuml/time-lapse-video-for-the-web-a7d8874ff397\"\n",
|
||||
"\n",
|
||||
"from txtai.pipeline import Textractor\n",
|
||||
"\n",
|
||||
"textractor = Textractor()\n",
|
||||
"text = textractor(\"time-lapse-video-for-the-web-a7d8874ff397\")\n",
|
||||
"\n",
|
||||
"summary(text)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
},
|
||||
"text/plain": [
|
||||
"'Time-lapse video is a popular way to show an area or event over a long period of time. The same concept can be applied to a dynamic real-time website with frequently updated data. webelapse is an open source project developed to provide this functionality. It can be used as is or modified for different use cases.'"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 4
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "a63k89aDyKTW"
|
||||
},
|
||||
"source": [
|
||||
"Click through the link to see the full article. This summary does a pretty good job of covering what the article is about!"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,362 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Pjmz-RORV8E"
|
||||
},
|
||||
"source": [
|
||||
"# Extract text from documents\n",
|
||||
"\n",
|
||||
"Up to this point, all the examples have been working with sections of text, which have already been split through some other means. What happens if we're working with documents? First we need to get the text out of these documents, then figure out how to index to best support vector search.\n",
|
||||
"\n",
|
||||
"This notebook shows how documents can have text extracted and split to support vector search and retrieval augmented generation (RAG)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dk31rbYjSTYm"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies. Since this notebook is using optional pipelines, we need to install the pipeline extras package."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "XMQuuun2R06J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline]\n",
|
||||
"\n",
|
||||
"# Get test data\n",
|
||||
"!wget -N https://github.com/neuml/txtai/releases/download/v6.2.0/tests.tar.gz\n",
|
||||
"!tar -xvzf tests.tar.gz\n",
|
||||
"\n",
|
||||
"# Install NLTK\n",
|
||||
"import nltk\n",
|
||||
"nltk.download(['punkt', 'punkt_tab'])"
|
||||
],
|
||||
"execution_count": 19,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PNPJ95cdTKSS"
|
||||
},
|
||||
"source": [
|
||||
"# Create a Textractor instance\n",
|
||||
"\n",
|
||||
"The Textractor instance is the main entrypoint for extracting text. This method is backed by Apache Tika, a robust text extraction library written in Java. [Apache Tika](https://tika.apache.org/0.9/formats.html) has support for a large number of file formats: PDF, Word, Excel, HTML and others. The [Python Tika package](https://github.com/chrismattmann/tika-python) automatically installs Tika and starts a local REST API instance used to read extracted data.\n",
|
||||
"\n",
|
||||
"*Note: This requires Java to be installed locally.*"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "nTDwXOUeTH2-"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"\n",
|
||||
"from txtai.pipeline import Textractor\n",
|
||||
"\n",
|
||||
"# Create textractor model\n",
|
||||
"textractor = Textractor()"
|
||||
],
|
||||
"execution_count": 20,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-vGR_piwZZO6"
|
||||
},
|
||||
"source": [
|
||||
"# Extract text\n",
|
||||
"\n",
|
||||
"The example below shows how to extract text from a file."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 118
|
||||
},
|
||||
"id": "-K2YJJzsVtfq",
|
||||
"outputId": "7754c508-264a-41fa-9843-83460719820f"
|
||||
},
|
||||
"source": [
|
||||
"textractor(\"txtai/article.pdf\")"
|
||||
],
|
||||
"execution_count": 21,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'Introducing txtai, an AI-powered search engine \\nbuilt on Transformers\\n\\nAdd Natural Language Understanding to any application\\n\\nSearch is the base of many applications. Once data starts to pile up, users want to be able to find it. It’s \\nthe foundation of the internet and an ever-growing challenge that is never solved or done.\\n\\nThe field of Natural Language Processing (NLP) is rapidly evolving with a number of new \\ndevelopments. Large-scale general language models are an exciting new capability allowing us to add \\namazing functionality quickly with limited compute and people. Innovation continues with new models\\nand advancements coming in at what seems a weekly basis.\\n\\nThis article introduces txtai, an AI-powered search engine that enables Natural Language \\nUnderstanding (NLU) based search in any application.\\n\\nIntroducing txtai\\ntxtai builds an AI-powered index over sections of text. txtai supports building text indices to perform \\nsimilarity searches and create extractive question-answering based systems. txtai also has functionality \\nfor zero-shot classification. txtai is open source and available on GitHub.\\n\\ntxtai and/or the concepts behind it has already been used to power the Natural Language Processing \\n(NLP) applications listed below:\\n\\n• paperai — AI-powered literature discovery and review engine for medical/scientific papers\\n• tldrstory — AI-powered understanding of headlines and story text\\n• neuspo — Fact-driven, real-time sports event and news site\\n• codequestion — Ask coding questions directly from the terminal\\n\\nBuild an Embeddings index\\nFor small lists of texts, the method above works. But for larger repositories of documents, it doesn’t \\nmake sense to tokenize and convert all embeddings for each query. txtai supports building pre-\\ncomputed indices which significantly improves performance.\\n\\nBuilding on the previous example, the following example runs an index method to build and store the \\ntext embeddings. In this case, only the query is converted to an embeddings vector each search.\\n\\nhttps://github.com/neuml/codequestion\\nhttps://neuspo.com/\\nhttps://github.com/neuml/tldrstory\\nhttps://github.com/neuml/paperai\\n - Introducing txtai, an AI-powered search engine built on Transformers\\n - Add Natural Language Understanding to any application\\n - Introducing txtai\\n - Build an Embeddings index'"
|
||||
],
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 21
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "n2jndgE-JyWX"
|
||||
},
|
||||
"source": [
|
||||
"Note that the text from the article was extracted into a single string. Depending on the articles, this may be acceptable. For long articles, often you'll want to split the content into logical sections to build better downstream vectors."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "1w2bhBCPOUdu"
|
||||
},
|
||||
"source": [
|
||||
"# Extract sentences\n",
|
||||
"\n",
|
||||
"Sentence extraction uses a model that specializes in sentence detection. This call returns a list of sentences."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "PKZVK5vuOTqB",
|
||||
"outputId": "a31e182e-037b-4e29-c1b0-0f6815e3b2c9"
|
||||
},
|
||||
"source": [
|
||||
"textractor = Textractor(sentences=True)\n",
|
||||
"textractor(\"txtai/article.pdf\")"
|
||||
],
|
||||
"execution_count": 22,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"['Introducing txtai, an AI-powered search engine \\nbuilt on Transformers\\n\\nAdd Natural Language Understanding to any application\\n\\nSearch is the base of many applications.',\n",
|
||||
" 'Once data starts to pile up, users want to be able to find it.',\n",
|
||||
" 'It’s \\nthe foundation of the internet and an ever-growing challenge that is never solved or done.',\n",
|
||||
" 'The field of Natural Language Processing (NLP) is rapidly evolving with a number of new \\ndevelopments.',\n",
|
||||
" 'Large-scale general language models are an exciting new capability allowing us to add \\namazing functionality quickly with limited compute and people.',\n",
|
||||
" 'Innovation continues with new models\\nand advancements coming in at what seems a weekly basis.',\n",
|
||||
" 'This article introduces txtai, an AI-powered search engine that enables Natural Language \\nUnderstanding (NLU) based search in any application.',\n",
|
||||
" 'Introducing txtai\\ntxtai builds an AI-powered index over sections of text.',\n",
|
||||
" 'txtai supports building text indices to perform \\nsimilarity searches and create extractive question-answering based systems.',\n",
|
||||
" 'txtai also has functionality \\nfor zero-shot classification.',\n",
|
||||
" 'txtai is open source and available on GitHub.',\n",
|
||||
" 'txtai and/or the concepts behind it has already been used to power the Natural Language Processing \\n(NLP) applications listed below:\\n\\n• paperai — AI-powered literature discovery and review engine for medical/scientific papers\\n• tldrstory — AI-powered understanding of headlines and story text\\n• neuspo — Fact-driven, real-time sports event and news site\\n• codequestion — Ask coding questions directly from the terminal\\n\\nBuild an Embeddings index\\nFor small lists of texts, the method above works.',\n",
|
||||
" 'But for larger repositories of documents, it doesn’t \\nmake sense to tokenize and convert all embeddings for each query.',\n",
|
||||
" 'txtai supports building pre-\\ncomputed indices which significantly improves performance.',\n",
|
||||
" 'Building on the previous example, the following example runs an index method to build and store the \\ntext embeddings.',\n",
|
||||
" 'In this case, only the query is converted to an embeddings vector each search.',\n",
|
||||
" 'https://github.com/neuml/codequestion\\nhttps://neuspo.com/\\nhttps://github.com/neuml/tldrstory\\nhttps://github.com/neuml/paperai\\n - Introducing txtai, an AI-powered search engine built on Transformers\\n - Add Natural Language Understanding to any application\\n - Introducing txtai\\n - Build an Embeddings index']"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 22
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "vdVCCc9UOv5S"
|
||||
},
|
||||
"source": [
|
||||
"Now the document is split up at the sentence level. These sentences can be feed to a workflow that adds each sentence to an embeddings index. Depending on the task, this may work well. Alternatively, it may be even better to split at the paragraph level."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "z1H8XYkaSoP4"
|
||||
},
|
||||
"source": [
|
||||
"# Extract paragraphs\n",
|
||||
"\n",
|
||||
"Paragraph detection looks for consecutive newlines. This call returns a list of paragraphs."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "9VUito4ISoAe",
|
||||
"outputId": "08079380-a7c8-4886-ecc3-de9f02be4584"
|
||||
},
|
||||
"source": [
|
||||
"textractor = Textractor(paragraphs=True)\n",
|
||||
"for paragraph in textractor(\"txtai/article.pdf\"):\n",
|
||||
" print(paragraph, \"\\n----\")"
|
||||
],
|
||||
"execution_count": 23,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Introducing txtai, an AI-powered search engine \n",
|
||||
"built on Transformers \n",
|
||||
"----\n",
|
||||
"Add Natural Language Understanding to any application \n",
|
||||
"----\n",
|
||||
"Search is the base of many applications. Once data starts to pile up, users want to be able to find it. It’s \n",
|
||||
"the foundation of the internet and an ever-growing challenge that is never solved or done. \n",
|
||||
"----\n",
|
||||
"The field of Natural Language Processing (NLP) is rapidly evolving with a number of new \n",
|
||||
"developments. Large-scale general language models are an exciting new capability allowing us to add \n",
|
||||
"amazing functionality quickly with limited compute and people. Innovation continues with new models\n",
|
||||
"and advancements coming in at what seems a weekly basis. \n",
|
||||
"----\n",
|
||||
"This article introduces txtai, an AI-powered search engine that enables Natural Language \n",
|
||||
"Understanding (NLU) based search in any application. \n",
|
||||
"----\n",
|
||||
"Introducing txtai\n",
|
||||
"txtai builds an AI-powered index over sections of text. txtai supports building text indices to perform \n",
|
||||
"similarity searches and create extractive question-answering based systems. txtai also has functionality \n",
|
||||
"for zero-shot classification. txtai is open source and available on GitHub. \n",
|
||||
"----\n",
|
||||
"txtai and/or the concepts behind it has already been used to power the Natural Language Processing \n",
|
||||
"(NLP) applications listed below: \n",
|
||||
"----\n",
|
||||
"• paperai — AI-powered literature discovery and review engine for medical/scientific papers\n",
|
||||
"• tldrstory — AI-powered understanding of headlines and story text\n",
|
||||
"• neuspo — Fact-driven, real-time sports event and news site\n",
|
||||
"• codequestion — Ask coding questions directly from the terminal \n",
|
||||
"----\n",
|
||||
"Build an Embeddings index\n",
|
||||
"For small lists of texts, the method above works. But for larger repositories of documents, it doesn’t \n",
|
||||
"make sense to tokenize and convert all embeddings for each query. txtai supports building pre-\n",
|
||||
"computed indices which significantly improves performance. \n",
|
||||
"----\n",
|
||||
"Building on the previous example, the following example runs an index method to build and store the \n",
|
||||
"text embeddings. In this case, only the query is converted to an embeddings vector each search. \n",
|
||||
"----\n",
|
||||
"https://github.com/neuml/codequestion\n",
|
||||
"https://neuspo.com/\n",
|
||||
"https://github.com/neuml/tldrstory\n",
|
||||
"https://github.com/neuml/paperai\n",
|
||||
" - Introducing txtai, an AI-powered search engine built on Transformers\n",
|
||||
" - Add Natural Language Understanding to any application\n",
|
||||
" - Introducing txtai\n",
|
||||
" - Build an Embeddings index \n",
|
||||
"----\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Extract sections\n",
|
||||
"\n",
|
||||
"Section extraction is format dependent. If page breaks are available, each section is a page. Otherwise, this call returns logical sections such by headings."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Ae6dRQ2LvN-w"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"textractor = Textractor(sections=True)\n",
|
||||
"print(\"\\n[PAGE BREAK]\\n\".join(section for section in textractor(\"txtai/article.pdf\")))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "nQ6ev2UMwqnh",
|
||||
"outputId": "3d45491d-3547-4218-d30c-ba0c4d161256"
|
||||
},
|
||||
"execution_count": 24,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Introducing txtai, an AI-powered search engine \n",
|
||||
"built on Transformers\n",
|
||||
"\n",
|
||||
"Add Natural Language Understanding to any application\n",
|
||||
"\n",
|
||||
"Search is the base of many applications. Once data starts to pile up, users want to be able to find it. It’s \n",
|
||||
"the foundation of the internet and an ever-growing challenge that is never solved or done.\n",
|
||||
"\n",
|
||||
"The field of Natural Language Processing (NLP) is rapidly evolving with a number of new \n",
|
||||
"developments. Large-scale general language models are an exciting new capability allowing us to add \n",
|
||||
"amazing functionality quickly with limited compute and people. Innovation continues with new models\n",
|
||||
"and advancements coming in at what seems a weekly basis.\n",
|
||||
"\n",
|
||||
"This article introduces txtai, an AI-powered search engine that enables Natural Language \n",
|
||||
"Understanding (NLU) based search in any application.\n",
|
||||
"\n",
|
||||
"Introducing txtai\n",
|
||||
"txtai builds an AI-powered index over sections of text. txtai supports building text indices to perform \n",
|
||||
"similarity searches and create extractive question-answering based systems. txtai also has functionality \n",
|
||||
"for zero-shot classification. txtai is open source and available on GitHub.\n",
|
||||
"\n",
|
||||
"txtai and/or the concepts behind it has already been used to power the Natural Language Processing \n",
|
||||
"(NLP) applications listed below:\n",
|
||||
"\n",
|
||||
"• paperai — AI-powered literature discovery and review engine for medical/scientific papers\n",
|
||||
"• tldrstory — AI-powered understanding of headlines and story text\n",
|
||||
"• neuspo — Fact-driven, real-time sports event and news site\n",
|
||||
"• codequestion — Ask coding questions directly from the terminal\n",
|
||||
"\n",
|
||||
"Build an Embeddings index\n",
|
||||
"For small lists of texts, the method above works. But for larger repositories of documents, it doesn’t \n",
|
||||
"make sense to tokenize and convert all embeddings for each query. txtai supports building pre-\n",
|
||||
"computed indices which significantly improves performance.\n",
|
||||
"\n",
|
||||
"Building on the previous example, the following example runs an index method to build and store the \n",
|
||||
"text embeddings. In this case, only the query is converted to an embeddings vector each search.\n",
|
||||
"\n",
|
||||
"https://github.com/neuml/codequestion\n",
|
||||
"https://neuspo.com/\n",
|
||||
"https://github.com/neuml/tldrstory\n",
|
||||
"https://github.com/neuml/paperai\n",
|
||||
"[PAGE BREAK]\n",
|
||||
"- Introducing txtai, an AI-powered search engine built on Transformers\n",
|
||||
" - Add Natural Language Understanding to any application\n",
|
||||
" - Introducing txtai\n",
|
||||
" - Build an Embeddings index\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,374 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Pjmz-RORV8E"
|
||||
},
|
||||
"source": [
|
||||
"# Translate text between languages\n",
|
||||
"\n",
|
||||
"This notebook covers machine translation backed by Hugging Face models. The quality of machine translation via cloud services has come a very long way and produces high quality results. This notebook shows how the models from Hugging Face give developers a reasonable alternative for local machine translation."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dk31rbYjSTYm"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies. Since this notebook is using optional pipelines, we need to install the pipeline extras package."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "XMQuuun2R06J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline]"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PNPJ95cdTKSS"
|
||||
},
|
||||
"source": [
|
||||
"# Create a Translation instance\n",
|
||||
"\n",
|
||||
"The Translation instance is the main entrypoint for translating text between languages. The pipeline abstracts translating text into a one line call! \n",
|
||||
"\n",
|
||||
"The pipeline has logic to detect the input language, load the relevant model that handles translating from source to target language and return results. The translation pipeline also has built-in logic to handle splitting large text blocks into smaller sections the models can handle.\n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "nTDwXOUeTH2-"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"\n",
|
||||
"from txtai.pipeline import Translation\n",
|
||||
"\n",
|
||||
"# Create translation model\n",
|
||||
"translate = Translation()"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-vGR_piwZZO6"
|
||||
},
|
||||
"source": [
|
||||
"# Translate text\n",
|
||||
"\n",
|
||||
"The example below shows how to translate text from English to Spanish. This text is then translated back to English."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 36
|
||||
},
|
||||
"id": "-K2YJJzsVtfq",
|
||||
"outputId": "44df5404-ea14-4746-fc8b-a2e205bd9466"
|
||||
},
|
||||
"source": [
|
||||
"translation = translate(\"This is a test translation into Spanish\", \"es\")\n",
|
||||
"translation"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'Esta es una traducción de prueba al español'"
|
||||
],
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 16
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 36
|
||||
},
|
||||
"id": "K_UnAZQpetM8",
|
||||
"outputId": "46c9c68c-ddcf-4f55-bac3-89f25931e91b"
|
||||
},
|
||||
"source": [
|
||||
"translate(translation, \"en\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'This is a test translation into Spanish'"
|
||||
],
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 17
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4cSI8GdtjhEM"
|
||||
},
|
||||
"source": [
|
||||
"# Translating multiple languages in a single call\n",
|
||||
"\n",
|
||||
"The section below translates a single English sentence into 5 different languages. The results are then passed to a single translation call to translate back into English. The pipeline detects each input language and is able to load the relevant translation models."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "8jLxGtwNf0Aj",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "47040b2f-f6e5-482a-df81-7b758f47a7d5"
|
||||
},
|
||||
"source": [
|
||||
"def run():\n",
|
||||
" languages = [\"fr\", \"es\", \"de\", \"hi\", \"ja\"]\n",
|
||||
" translations = [translate(\"The sky is blue, the stars are far\", language) for language in languages]\n",
|
||||
" english = translate(translations, \"en\")\n",
|
||||
"\n",
|
||||
" for x, text in enumerate(translations):\n",
|
||||
" print(\"Original Language: %s\" % languages[x])\n",
|
||||
" print(\"Translation: %s\" % text)\n",
|
||||
" print(\"Back to English: %s\" % english[x])\n",
|
||||
" print()\n",
|
||||
"\n",
|
||||
"# Run multiple translations\n",
|
||||
"run()"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Original Language: fr\n",
|
||||
"Translation: Le ciel est bleu, les étoiles sont loin\n",
|
||||
"Back to English: The sky is blue, the stars are far away\n",
|
||||
"\n",
|
||||
"Original Language: es\n",
|
||||
"Translation: El cielo es azul, las estrellas están lejos.\n",
|
||||
"Back to English: The sky is blue, the stars are far away.\n",
|
||||
"\n",
|
||||
"Original Language: de\n",
|
||||
"Translation: Der Himmel ist blau, die Sterne sind weit\n",
|
||||
"Back to English: The sky is blue, the stars are wide\n",
|
||||
"\n",
|
||||
"Original Language: hi\n",
|
||||
"Translation: आकाश नीला है, तारे दूर हैं\n",
|
||||
"Back to English: Sky is blue, stars are away\n",
|
||||
"\n",
|
||||
"Original Language: ja\n",
|
||||
"Translation: 天は青い、星は遠い。\n",
|
||||
"Back to English: The heavens are blue and the stars are far away.\n",
|
||||
"\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Xn3SlVE1LYvm"
|
||||
},
|
||||
"source": [
|
||||
"The translation quality overall is very high!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Additional model types\n",
|
||||
"\n",
|
||||
"The translation pipeline is flexible and supports multiple model types. The default mode for the pipeline is to scan the Hugging Face Hub for models that best match the source-target translation pair. This often produces the best quality and is usually a smaller model than a large multi-language mode.\n",
|
||||
"\n",
|
||||
"There is a parameter that can override this and always use the base model."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "3FdS5slz60eA"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"translate = Translation(\"t5-small\", findmodels=False)\n",
|
||||
"translate(\"translate English to French: The sky is blue, the stars are far\", None)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 36
|
||||
},
|
||||
"id": "REb0X2Kz60Ew",
|
||||
"outputId": "ae13d4d0-318f-4740-f725-0029ceeabeac"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'Le ciel est bleu, les étoiles sont loin'"
|
||||
],
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 11
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Translation isn't limited to spoken languages. txtai provides a text-to-sql model that converts English text into a txtai-compatible SQL statement. "
|
||||
],
|
||||
"metadata": {
|
||||
"id": "5vSJbW7zVJeH"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"translate = Translation(\"NeuML/t5-small-txtsql\", findmodels=False)\n",
|
||||
"translate(\"translate English to SQL: feel good story since yesterday\", None)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 36
|
||||
},
|
||||
"id": "d5usam_jhKaz",
|
||||
"outputId": "c12368a0-ea26-4191-be5a-f2de70711003"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"\"select id, text, score from txtai where similar('feel good story') and entry >= date('now', '-1 day')\""
|
||||
],
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 12
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Last thing we'll do is run the multiple language example only using a single large language model."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "cVb7uk7TXr_v"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"translate = Translation(\"facebook/mbart-large-50-many-to-many-mmt\", findmodels=False)\n",
|
||||
"run()"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "JkOLnvKZWI95",
|
||||
"outputId": "fef6402c-e20e-43e1-a3eb-e4580913fa7e"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Original Language: fr\n",
|
||||
"Translation: Le ciel est bleu, les étoiles sont loin\n",
|
||||
"Back to English: The sky is blue, the stars are far away\n",
|
||||
"\n",
|
||||
"Original Language: es\n",
|
||||
"Translation: El cielo es azul, las estrellas están lejos.\n",
|
||||
"Back to English: The sky is blue, the stars are far away.\n",
|
||||
"\n",
|
||||
"Original Language: de\n",
|
||||
"Translation: Der Himmel ist blau, die Sterne sind weit\n",
|
||||
"Back to English: The sky is blue, the stars are far.\n",
|
||||
"\n",
|
||||
"Original Language: hi\n",
|
||||
"Translation: आकाश नीली है, तारे दूर हैं।\n",
|
||||
"Back to English: The sky is blue, and the stars are far away.\n",
|
||||
"\n",
|
||||
"Original Language: ja\n",
|
||||
"Translation: 空は青い、星は遠い\n",
|
||||
"Back to English: the sky is blue, the stars are far away.\n",
|
||||
"\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"Machine translation has made giant leaps and strides the last couple of years. These models give developers a solid, locally-hosted alternative to cloud translation services. Additionally, there are models built for low resource languages that cloud translation services don't support.\n",
|
||||
"\n",
|
||||
"A number of different models and configurations are supported, give it a try!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "TvWx9PS-X32c"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,526 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Pjmz-RORV8E"
|
||||
},
|
||||
"source": [
|
||||
"# Run pipeline workflows\n",
|
||||
"\n",
|
||||
"txtai has a growing list of models available through it's pipeline framework. Pipelines wrap a machine learning model and transform data. Currently, pipelines can wrap Hugging Face models, Hugging Face pipelines or PyTorch models (support for TensorFlow is in the backlog).\n",
|
||||
"\n",
|
||||
"The following is a list of the currently implemented pipelines.\n",
|
||||
"\n",
|
||||
"* **Questions** - Answer questions using a text context\n",
|
||||
"* **Labels** - Apply labels to text using a zero-shot classification model. Also supports similarity comparisions.\n",
|
||||
"* **Summary** - Abstractive text summarization\n",
|
||||
"* **Textractor** - Extract text from documents\n",
|
||||
"* **Transcription** - Transcribe audio to text\n",
|
||||
"* **Translation** - Machine translation\n",
|
||||
"\n",
|
||||
"Pipelines are great and make using a variety of machine learning models easier. But what if we want to glue the results of different pipelines together? For example, extract text, summarize it, translate it to English and load it into an Embedding index. That would require code to join those operations together in an efficient manner.\n",
|
||||
"\n",
|
||||
"Enter workflows. Workflows are a simple yet powerful construct that takes a callable and returns elements. Workflows don't know they are working with pipelines but enable efficient processing of pipeline data. Workflows are streaming by nature and work on data in batches, allowing large volumes of data to be processed efficiently."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dk31rbYjSTYm"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies. Since this notebook is using optional pipelines/workflows, we need to install the pipeline and workflow extras package."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "XMQuuun2R06J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline,workflow] sacremoses\n",
|
||||
"\n",
|
||||
"# Get test data\n",
|
||||
"!wget -N https://github.com/neuml/txtai/releases/download/v2.0.0/tests.tar.gz\n",
|
||||
"!tar -xvzf tests.tar.gz"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "I1dNQE7WT4kE"
|
||||
},
|
||||
"source": [
|
||||
"# Create a series of pipelines to use in this notebook"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "w4YqwBJaT4QD"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"from txtai.pipeline import Summary, Textractor, Transcription, Translation\n",
|
||||
"\n",
|
||||
"# Summary instance\n",
|
||||
"summary = Summary()\n",
|
||||
"\n",
|
||||
"# Text extraction\n",
|
||||
"textractor = Textractor()\n",
|
||||
"\n",
|
||||
"# Transcription instance\n",
|
||||
"transcribe = Transcription(\"facebook/wav2vec2-large-960h\")\n",
|
||||
"\n",
|
||||
"# Create a translation instance\n",
|
||||
"translate = Translation()"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PNPJ95cdTKSS"
|
||||
},
|
||||
"source": [
|
||||
"# Basic workflow\n",
|
||||
"\n",
|
||||
"The following shows a basic workflow in action!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "nTDwXOUeTH2-",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "906d4354-cf29-4593-a790-8c175d981dee"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.workflow import Workflow, Task\n",
|
||||
"\n",
|
||||
"# Workflow that translate text to French\n",
|
||||
"workflow = Workflow([Task(lambda x: translate(x, \"fr\"))])\n",
|
||||
"\n",
|
||||
"# Data to run through the pipeline\n",
|
||||
"data = [\"The sky is blue\", \"Forest through the trees\"]\n",
|
||||
"\n",
|
||||
"# Workflows are generators for efficiency, read results to list for display\n",
|
||||
"list(workflow(data))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"['Le ciel est bleu', 'Forêt à travers les arbres']"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 13
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "wicr0CAYRWZ0"
|
||||
},
|
||||
"source": [
|
||||
"This isn't too different from previous pipeline examples. The only difference is data is feed through the workflow. In this example, the workflow calls the translation pipeline and translates text to French. Let's look at a more complex example."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "0EeD8m6FR5cH"
|
||||
},
|
||||
"source": [
|
||||
"# Multistep workflow\n",
|
||||
"\n",
|
||||
"The following workflow reads a series of audio files, transcribes them to text and translates the text to French. This is based on the classic txtai example from [Introducing txtai](https://colab.research.google.com/github/neuml/txtai/blob/master/examples/01_Introducing_txtai.ipynb).\n",
|
||||
"\n",
|
||||
"Workflows take two main parameters. The action to execute which is a callable and a pattern to filter data with. Data that is accepted by the filter will be processed, otherwise it will be passed through to the next task."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "OF2G5-OiSBzy",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "e5c74089-1916-4bbd-93d3-9e25b1fe4ee5"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.workflow import FileTask\n",
|
||||
"\n",
|
||||
"tasks = [\n",
|
||||
" FileTask(transcribe, r\"\\.wav$\"),\n",
|
||||
" Task(lambda x: translate(x, \"fr\"))\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"# List of files to process\n",
|
||||
"data = [\n",
|
||||
" \"txtai/US_tops_5_million.wav\",\n",
|
||||
" \"txtai/Canadas_last_fully.wav\",\n",
|
||||
" \"txtai/Beijing_mobilises.wav\",\n",
|
||||
" \"txtai/The_National_Park.wav\",\n",
|
||||
" \"txtai/Maine_man_wins_1_mil.wav\",\n",
|
||||
" \"txtai/Make_huge_profits.wav\"\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"# Workflow that translate text to French\n",
|
||||
"workflow = Workflow(tasks)\n",
|
||||
"\n",
|
||||
"# Run workflow\n",
|
||||
"list(workflow(data))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[\"Les cas de virus U sont en tête d'un million\",\n",
|
||||
" \"La dernière plate-forme de glace entièrement intacte du Canada s'est soudainement effondrée en formant un berge de glace de taille manhatten\",\n",
|
||||
" \"Bagage mobilise les embarcations d'invasion le long des côtes à mesure que les tensions tiwaniennes s'intensifient\",\n",
|
||||
" \"Le service des parcs nationaux met en garde contre le sacrifice d'amis plus lents dans une attaque nue\",\n",
|
||||
" \"L'homme principal gagne du billet de loterie\",\n",
|
||||
" \"Faire d'énormes profits sans travailler faire jusqu'à cent mille dollars par jour\"]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 14
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PN08rnrQU1hx"
|
||||
},
|
||||
"source": [
|
||||
"# Complex workflow\n",
|
||||
"\n",
|
||||
"Let's put this all together into a full-fledged workflow to build an embeddings index. This workflow will work with both documents and audio files. Documents will have text extracted and summarized. Audio files will be transcribed. Both results will be joined, translated into French and loaded into an Embeddings index."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "coZJw_1yU1Sq",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "213b34d5-157f-4548-8788-ac29cb4039dd"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.embeddings import Embeddings, Documents\n",
|
||||
"from txtai.workflow import FileTask, WorkflowTask\n",
|
||||
"\n",
|
||||
"# Embeddings index\n",
|
||||
"embeddings = Embeddings({\"path\": \"sentence-transformers/paraphrase-multilingual-mpnet-base-v2\", \"content\": True})\n",
|
||||
"documents = Documents()\n",
|
||||
"\n",
|
||||
"# List of files to process\n",
|
||||
"files = [\n",
|
||||
" \"txtai/article.pdf\",\n",
|
||||
" \"txtai/US_tops_5_million.wav\",\n",
|
||||
" \"txtai/Canadas_last_fully.wav\",\n",
|
||||
" \"txtai/Beijing_mobilises.wav\",\n",
|
||||
" \"txtai/The_National_Park.wav\",\n",
|
||||
" \"txtai/Maine_man_wins_1_mil.wav\",\n",
|
||||
" \"txtai/Make_huge_profits.wav\"\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"data = [(x, element, None) for x, element in enumerate(files)]\n",
|
||||
"\n",
|
||||
"# Workflow that extracts text and builds a summary\n",
|
||||
"articles = Workflow([\n",
|
||||
" FileTask(textractor),\n",
|
||||
" Task(summary)\n",
|
||||
"])\n",
|
||||
"\n",
|
||||
"# Define workflow tasks. Workflows can also be tasks!\n",
|
||||
"tasks = [\n",
|
||||
" WorkflowTask(articles, r\".\\.pdf$\"),\n",
|
||||
" FileTask(transcribe, r\"\\.wav$\"),\n",
|
||||
" Task(lambda x: translate(x, \"fr\")),\n",
|
||||
" Task(documents.add, unpack=False)\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"# Workflow that translate text to French\n",
|
||||
"workflow = Workflow(tasks)\n",
|
||||
"\n",
|
||||
"# Run workflow and show results to be indexed\n",
|
||||
"for x in workflow(data):\n",
|
||||
" print(x)\n",
|
||||
"\n",
|
||||
"# Build the embeddings index\n",
|
||||
"embeddings.index(documents)\n",
|
||||
"\n",
|
||||
"# Cleanup temporary storage\n",
|
||||
"documents.close()"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"(0, \"Txtai, un moteur de recherche alimenté par l'IA construit sur Transformers, permet la recherche basée sur la compréhension du langage naturel (NLU) dans n'importe quelle application. Le champ de traitement du langage naturel (NLP) évolue rapidement avec un certain nombre de nouveaux développements. Le moteur de recherche open-source est open source et disponible sur GitHub.\", None)\n",
|
||||
"(1, \"Les cas de virus U sont en tête d'un million\", None)\n",
|
||||
"(2, \"La dernière plate-forme de glace entièrement intacte du Canada s'est soudainement effondrée en formant un berge de glace de taille manhatten\", None)\n",
|
||||
"(3, \"Bagage mobilise les embarcations d'invasion le long des côtes à mesure que les tensions tiwaniennes s'intensifient\", None)\n",
|
||||
"(4, \"Le service des parcs nationaux met en garde contre le sacrifice d'amis plus lents dans une attaque nue\", None)\n",
|
||||
"(5, \"L'homme principal gagne du billet de loterie\", None)\n",
|
||||
"(6, \"Faire d'énormes profits sans travailler faire jusqu'à cent mille dollars par jour\", None)\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "n6i-xhJya8o4"
|
||||
},
|
||||
"source": [
|
||||
"# Query for results in French"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "cHbjivUOaUGu",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "0da8d8cb-dac6-4cad-ef00-a096b44533cf"
|
||||
},
|
||||
"source": [
|
||||
"# Run a search query and show the result.\n",
|
||||
"embeddings.search(\"changement climatique\", 1)[0]"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'id': '2',\n",
|
||||
" 'score': 0.2982647716999054,\n",
|
||||
" 'text': \"La dernière plate-forme de glace entièrement intacte du Canada s'est soudainement effondrée en formant un berge de glace de taille manhatten\"}"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 16
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "aNerHvNpaxD4",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "f3792220-4518-4388-c7e7-c38f38f19b20"
|
||||
},
|
||||
"source": [
|
||||
"# Run a search query and show the result.\n",
|
||||
"embeddings.search(\"traitement du langage naturel\", 1)[0]"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'id': '0',\n",
|
||||
" 'score': 0.47031939029693604,\n",
|
||||
" 'text': \"Txtai, un moteur de recherche alimenté par l'IA construit sur Transformers, permet la recherche basée sur la compréhension du langage naturel (NLU) dans n'importe quelle application. Le champ de traitement du langage naturel (NLP) évolue rapidement avec un certain nombre de nouveaux développements. Le moteur de recherche open-source est open source et disponible sur GitHub.\"}"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 17
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Configuration-driven workflow\n",
|
||||
"\n",
|
||||
"Workflows can also be defined with YAML and run as an application. Applications can run standalone or as a FastAPI instance. More information can be [found here](https://neuml.github.io/txtai/api/). "
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Sz_f9qoOMC_m"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"workflow = \"\"\"\n",
|
||||
"writable: true\n",
|
||||
"embeddings:\n",
|
||||
" path: sentence-transformers/paraphrase-multilingual-mpnet-base-v2\n",
|
||||
" content: True\n",
|
||||
"\n",
|
||||
"# Summarize text\n",
|
||||
"summary:\n",
|
||||
"\n",
|
||||
"# Extract text from documents\n",
|
||||
"textractor:\n",
|
||||
"\n",
|
||||
"# Transcribe audio to text\n",
|
||||
"transcription:\n",
|
||||
" path: facebook/wav2vec2-large-960h\n",
|
||||
"\n",
|
||||
"# Translate text between languages\n",
|
||||
"translation:\n",
|
||||
"\n",
|
||||
"workflow:\n",
|
||||
" summarize:\n",
|
||||
" tasks:\n",
|
||||
" - action: textractor\n",
|
||||
" task: file\n",
|
||||
" - summary\n",
|
||||
" index:\n",
|
||||
" tasks:\n",
|
||||
" - action: summarize\n",
|
||||
" select: '\\\\.pdf$'\n",
|
||||
" - action: transcription\n",
|
||||
" select: '\\\\.wav$'\n",
|
||||
" task: file\n",
|
||||
" - action: translation\n",
|
||||
" args: ['fr']\n",
|
||||
" - action: index\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"# Create and run the workflow\n",
|
||||
"from txtai.app import Application\n",
|
||||
"\n",
|
||||
"# Create and run the workflow\n",
|
||||
"app = Application(workflow)\n",
|
||||
"list(app.workflow(\"index\", files))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "HoVlk_vNJKHY",
|
||||
"outputId": "34b68bcb-a6d5-4029-9bf2-f33e4381d1bc"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[\"Txtai, un moteur de recherche alimenté par l'IA construit sur Transformers, permet la recherche basée sur la compréhension du langage naturel (NLU) dans n'importe quelle application. Le champ de traitement du langage naturel (NLP) évolue rapidement avec un certain nombre de nouveaux développements. Le moteur de recherche open-source est open source et disponible sur GitHub.\",\n",
|
||||
" \"Les cas de virus U sont en tête d'un million\",\n",
|
||||
" \"La dernière plate-forme de glace entièrement intacte du Canada s'est soudainement effondrée en formant un berge de glace de taille manhatten\",\n",
|
||||
" \"Bagage mobilise les embarcations d'invasion le long des côtes à mesure que les tensions tiwaniennes s'intensifient\",\n",
|
||||
" \"Le service des parcs nationaux met en garde contre le sacrifice d'amis plus lents dans une attaque nue\",\n",
|
||||
" \"L'homme principal gagne du billet de loterie\",\n",
|
||||
" \"Faire d'énormes profits sans travailler faire jusqu'à cent mille dollars par jour\"]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 18
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Run a search query and show the result.\n",
|
||||
"app.search(\"changement climatique\", 1)[0]"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "a_klVZAXHJcw",
|
||||
"outputId": "33229268-0f98-4ca1-af7d-212bcbde6482"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'id': '2',\n",
|
||||
" 'score': 0.2982647716999054,\n",
|
||||
" 'text': \"La dernière plate-forme de glace entièrement intacte du Canada s'est soudainement effondrée en formant un berge de glace de taille manhatten\"}"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 19
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Run a search query and show the result.\n",
|
||||
"app.search(\"traitement du langage naturel\", 1)[0]"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "I5xin0VNHJOu",
|
||||
"outputId": "2fbb9a93-b860-437d-c361-ee21eed75b6b"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'id': '0',\n",
|
||||
" 'score': 0.47031939029693604,\n",
|
||||
" 'text': \"Txtai, un moteur de recherche alimenté par l'IA construit sur Transformers, permet la recherche basée sur la compréhension du langage naturel (NLU) dans n'importe quelle application. Le champ de traitement du langage naturel (NLP) évolue rapidement avec un certain nombre de nouveaux développements. Le moteur de recherche open-source est open source et disponible sur GitHub.\"}"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 20
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "7zG4AimucFJs"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"Results are good! We can see the power of workflows and how they can join a series of pipelines together in an efficient manner. Workflows can work with any callable, not just pipelines, workflows transform data from one format to another. Workflows are an exciting and promising development for txtai."
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,522 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Pjmz-RORV8E"
|
||||
},
|
||||
"source": [
|
||||
"# Distributed embeddings cluster\n",
|
||||
"\n",
|
||||
"The txtai API is a web-based service backed by [FastAPI](https://fastapi.tiangolo.com/). All txtai functionality is available via the API. The API can also cluster multiple embeddings indices into a single logical index to horizontally scale over multiple nodes. \n",
|
||||
"\n",
|
||||
"This notebook installs the txtai API and shows an example of building an embeddings cluster."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dk31rbYjSTYm"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies. Since this notebook uses the API, we need to install the api extras package."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "XMQuuun2R06J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[api]"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PNPJ95cdTKSS"
|
||||
},
|
||||
"source": [
|
||||
"# Start distributed embeddings cluster\n",
|
||||
"\n",
|
||||
"First we'll start multiple API instances that will serve as embeddings index shards. Each shard stores a subset of the indexed data and these shards work in tandem to form a single logical index.\n",
|
||||
"\n",
|
||||
"Then we'll start the main API instance that clusters the shards together into a logical instance.\n",
|
||||
"\n",
|
||||
"The API instances are all started in the background.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "USb4JXZHxqTA"
|
||||
},
|
||||
"source": [
|
||||
"import os\n",
|
||||
"os.chdir(\"/content\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "nTDwXOUeTH2-",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "dee26849-39ae-4390-8bba-76bf9025fa61"
|
||||
},
|
||||
"source": [
|
||||
"%%writefile index.yml\n",
|
||||
"writable: true\n",
|
||||
"\n",
|
||||
"# Embeddings settings\n",
|
||||
"embeddings:\n",
|
||||
" path: sentence-transformers/nli-mpnet-base-v2\n",
|
||||
" content: true"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing index.yml\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "iCdBh-JgfyBl",
|
||||
"outputId": "0066e314-7461-47c7-ca3b-15204911783e"
|
||||
},
|
||||
"source": [
|
||||
"%%writefile cluster.yml\n",
|
||||
"# Embeddings cluster\n",
|
||||
"cluster:\n",
|
||||
" shards:\n",
|
||||
" - http://127.0.0.1:8001\n",
|
||||
" - http://127.0.0.1:8002"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing cluster.yml\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "nGITHxUyRzyp"
|
||||
},
|
||||
"source": [
|
||||
"# Start embeddings shards\n",
|
||||
"!CONFIG=index.yml nohup uvicorn --port 8001 \"txtai.api:app\" &> shard-1.log &\n",
|
||||
"!CONFIG=index.yml nohup uvicorn --port 8002 \"txtai.api:app\" &> shard-2.log &\n",
|
||||
"\n",
|
||||
"# Start main instance\n",
|
||||
"!CONFIG=cluster.yml nohup uvicorn --port 8000 \"txtai.api:app\" &> main.log &\n",
|
||||
"\n",
|
||||
"# Wait for startup\n",
|
||||
"!sleep 90"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "lxkbVng3giWP"
|
||||
},
|
||||
"source": [
|
||||
"# Python\n",
|
||||
"\n",
|
||||
"Let's first try the cluster out directly in Python. The code below aggregates the two shards into a single cluster and executes actions against the cluster."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "36HGAokoglfg",
|
||||
"outputId": "368ae013-2afc-4a1b-d7df-c429183637d7"
|
||||
},
|
||||
"source": [
|
||||
"%%writefile run.py\n",
|
||||
"from txtai.api import Cluster\n",
|
||||
"\n",
|
||||
"cluster = Cluster({\"shards\": [\"http://127.0.0.1:8001\", \"http://127.0.0.1:8002\"]})\n",
|
||||
"\n",
|
||||
"data = [\n",
|
||||
" \"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",
|
||||
"\n",
|
||||
"# Index data\n",
|
||||
"cluster.add([{\"id\": x, \"text\": row} for x, row in enumerate(data)])\n",
|
||||
"cluster.index()\n",
|
||||
"\n",
|
||||
"# Test search\n",
|
||||
"result = cluster.search(\"feel good story\", 1)[0]\n",
|
||||
"print(\"Query: feel good story\\nResult:\", result[\"text\"])"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing run.py\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "6dQOzcfEs2Pk",
|
||||
"outputId": "a667594a-b778-4e4e-a75c-72e7982b7fbe"
|
||||
},
|
||||
"source": [
|
||||
"!python run.py"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Query: feel good story\n",
|
||||
"Result: Maine man wins $1M from $25 lottery ticket\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "NHvBFZeSd9AG"
|
||||
},
|
||||
"source": [
|
||||
"# JavaScript\n",
|
||||
"\n",
|
||||
"Next let's try to run the same code above via the API using JavaScript.\n",
|
||||
"\n",
|
||||
"```bash\n",
|
||||
"npm install txtai\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"For this example, we'll clone the txtai.js project to import the example build configuration."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "b52knObEdcCr"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!git clone https://github.com/neuml/txtai.js"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "rUGS0t-JMsS9"
|
||||
},
|
||||
"source": [
|
||||
"## Run cluster.js\n",
|
||||
"\n",
|
||||
"The following script is a JavaScript version of the logic above"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "bPQ40_xRyFmA",
|
||||
"outputId": "b86a12c4-f2c7-427b-bd28-edba354c6713"
|
||||
},
|
||||
"source": [
|
||||
"%%writefile txtai.js/examples/node/src/cluster.js\n",
|
||||
"import {Embeddings} from \"txtai\";\n",
|
||||
"import {sprintf} from \"sprintf-js\";\n",
|
||||
"\n",
|
||||
"const run = async () => {\n",
|
||||
" try {\n",
|
||||
" let embeddings = new Embeddings(process.argv[2]);\n",
|
||||
"\n",
|
||||
" let data = [\"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",
|
||||
" console.log();\n",
|
||||
" console.log(\"Querying an Embeddings cluster\");\n",
|
||||
" console.log(sprintf(\"%-20s %s\", \"Query\", \"Best Match\"));\n",
|
||||
" console.log(\"-\".repeat(50));\n",
|
||||
"\n",
|
||||
" for (let query of [\"feel good story\", \"climate change\", \"public health story\", \"war\", \"wildlife\", \"asia\", \"lucky\", \"dishonest junk\"]) {\n",
|
||||
" let results = await embeddings.search(query, 1);\n",
|
||||
" if (results && results.length > 0) {\n",
|
||||
" let result = results[0].text;\n",
|
||||
" console.log(sprintf(\"%-20s %s\", query, result));\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" catch (e) {\n",
|
||||
" console.trace(e);\n",
|
||||
" }\n",
|
||||
"};\n",
|
||||
"\n",
|
||||
"run();"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing txtai.js/examples/node/src/cluster.js\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "nTBs11j-GtD-"
|
||||
},
|
||||
"source": [
|
||||
"## Build and run cluster.js\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "kC5Oub6wa1nK"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"os.chdir(\"txtai.js/examples/node\")\n",
|
||||
"!npm install\n",
|
||||
"!npm run build"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Xr5IlvqH8W77"
|
||||
},
|
||||
"source": [
|
||||
"Next lets run the code against the main cluster URL"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ckOHNqyaeL-B",
|
||||
"outputId": "9c243fac-2316-4b8e-b044-6de529a8f3e8"
|
||||
},
|
||||
"source": [
|
||||
"!node dist/cluster.js http://127.0.0.1:8000"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"\n",
|
||||
"Querying an Embeddings cluster\n",
|
||||
"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",
|
||||
"public health story 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",
|
||||
"lucky Maine man wins $1M from $25 lottery ticket\n",
|
||||
"dishonest junk Make huge profits without work, earn up to $100,000 a day\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "1yukBIMYG5OE"
|
||||
},
|
||||
"source": [
|
||||
"The JavaScript program is showing the same results as the Python code above. This is running a clustered query against both nodes in the cluster and aggregating the results together.\n",
|
||||
"\n",
|
||||
"Queries can be run against each individual shard to see what the queries independently return."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "73rZCo4O4IQR",
|
||||
"outputId": "9f2cb119-7a21-41d9-fdbf-4410af246934"
|
||||
},
|
||||
"source": [
|
||||
"!node dist/cluster.js http://127.0.0.1:8001"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"\n",
|
||||
"Querying an Embeddings cluster\n",
|
||||
"Query Best Match\n",
|
||||
"--------------------------------------------------\n",
|
||||
"feel good story Maine man wins $1M from $25 lottery ticket\n",
|
||||
"climate change Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
|
||||
"public health story US tops 5 million confirmed virus cases\n",
|
||||
"war Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
|
||||
"wildlife Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
|
||||
"asia Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
|
||||
"lucky Maine man wins $1M from $25 lottery ticket\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ZeVBLJyr4Knr",
|
||||
"outputId": "b75691a4-25bf-43dc-8878-f9792a4430b8"
|
||||
},
|
||||
"source": [
|
||||
"!node dist/cluster.js http://127.0.0.1:8002"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"\n",
|
||||
"Querying an Embeddings cluster\n",
|
||||
"Query Best Match\n",
|
||||
"--------------------------------------------------\n",
|
||||
"feel good story Make huge profits without work, earn up to $100,000 a day\n",
|
||||
"climate change Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\n",
|
||||
"public health story The National Park Service warns against sacrificing slower friends in a bear attack\n",
|
||||
"war The National Park Service warns against sacrificing slower friends in a bear attack\n",
|
||||
"wildlife The National Park Service warns against sacrificing slower friends in a bear attack\n",
|
||||
"asia The National Park Service warns against sacrificing slower friends in a bear attack\n",
|
||||
"lucky The National Park Service warns against sacrificing slower friends in a bear attack\n",
|
||||
"dishonest junk Make huge profits without work, earn up to $100,000 a day\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "J2I_4hmZ8uXs"
|
||||
},
|
||||
"source": [
|
||||
"Note the differences. The section below runs a count against the full cluster and each shard to show the count of records in each."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "BKm27yna4MWr",
|
||||
"outputId": "bfc60af7-1b2b-451f-b10e-e2f8cf6f14fa"
|
||||
},
|
||||
"source": [
|
||||
"!curl http://127.0.0.1:8000/count\n",
|
||||
"!printf \"\\n\"\n",
|
||||
"!curl http://127.0.0.1:8001/count\n",
|
||||
"!printf \"\\n\"\n",
|
||||
"!curl http://127.0.0.1:8002/count"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"6\n",
|
||||
"3\n",
|
||||
"3"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "6rKj-I0djRQj"
|
||||
},
|
||||
"source": [
|
||||
"This notebook showed how a distributed embeddings cluster can be created with txtai. This example can be further scaled out on Kubernetes with StatefulSets, which will be covered in a future tutorial."
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,335 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Pjmz-RORV8E"
|
||||
},
|
||||
"source": [
|
||||
"# Train a text labeler\n",
|
||||
"\n",
|
||||
"The [Hugging Face Model Hub](https://huggingface.co/models) has a wide range of models that can handle many tasks. While these models perform well, the best performance often is found when fine-tuning a model with task-specific data. \n",
|
||||
"\n",
|
||||
"Hugging Face provides a [number of full-featured examples](https://github.com/huggingface/transformers/tree/master/examples) available to assist with training task-specific models. When building models from the command line, these scripts are a great way to get started.\n",
|
||||
"\n",
|
||||
"txtai provides a training pipeline that can be used to train new models programatically using the Transformers Trainer framework. The training pipeline supports the following:\n",
|
||||
"\n",
|
||||
"- Building transient models without requiring an output directory\n",
|
||||
"- Load training data from Hugging Face datasets, pandas DataFrames and list of dicts\n",
|
||||
"- Text sequence classification tasks (single/multi label classification and regression) including all GLUE tasks\n",
|
||||
"- All training arguments\n",
|
||||
"\n",
|
||||
"This notebook shows examples of how to use txtai to train/fine-tune new models."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dk31rbYjSTYm"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "XMQuuun2R06J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline-train] datasets pandas"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "PNPJ95cdTKSS"
|
||||
},
|
||||
"source": [
|
||||
"# Train a model\n",
|
||||
"\n",
|
||||
"Let's get right to it! The following example fine-tunes a tiny Bert model with the sst2 dataset.\n",
|
||||
"\n",
|
||||
"The trainer pipeline is basically a one-liner that fine-tunes any text classification/regression model available (locally and/or from the HF Hub). \n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "USb4JXZHxqTA"
|
||||
},
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"from txtai.pipeline import HFTrainer\n",
|
||||
"\n",
|
||||
"trainer = HFTrainer()\n",
|
||||
"\n",
|
||||
"# Hugging Face dataset\n",
|
||||
"ds = load_dataset(\"glue\", \"sst2\")\n",
|
||||
"model, tokenizer = trainer(\"google/bert_uncased_L-2_H-128_A-2\", ds[\"train\"], columns=(\"sentence\", \"label\"))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "CubsNAbpEWQg"
|
||||
},
|
||||
"source": [
|
||||
"The default trainer pipeline functionality will not store any logs, checkpoints or models to disk. The trainer can take any of the standard TrainingArguments to enable persistent models.\n",
|
||||
"\n",
|
||||
"The next section creates a Labels pipeline using the newly built model and runs the model against the sst2 validation set. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "xw2y2C5Mg11_",
|
||||
"outputId": "78400e45-ea5c-4cd9-d205-b55ee7a9f005"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.pipeline import Labels\n",
|
||||
"\n",
|
||||
"labels = Labels((model, tokenizer), dynamic=False)\n",
|
||||
"\n",
|
||||
"# Determine accuracy on validation set\n",
|
||||
"results = [row[\"label\"] == labels(row[\"sentence\"])[0][0] for row in ds[\"validation\"]]\n",
|
||||
"sum(results) / len(ds[\"validation\"])"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"0.8268348623853211"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 10
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ZAHSwaB3Ex49"
|
||||
},
|
||||
"source": [
|
||||
"82.68% accuracy - not bad for a tiny Bert model. \n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "f3GkY4JNEhhE"
|
||||
},
|
||||
"source": [
|
||||
"# Train a model with Lists\n",
|
||||
"\n",
|
||||
"As mentioned earlier, the trainer pipeline supports Hugging Face datasets, pandas DataFrames and lists of dicts. The example below trains a model using lists."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "QkApw1b2hfZq",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 182
|
||||
},
|
||||
"outputId": "8c3dceae-49fb-4b63-837d-5944e63c768e"
|
||||
},
|
||||
"source": [
|
||||
"data = [{\"text\": \"This is a test sentence\", \"label\": 0}, {\"text\": \"This is not a test\", \"label\": 1}]\n",
|
||||
"\n",
|
||||
"model, tokenizer = trainer(\"google/bert_uncased_L-2_H-128_A-2\", data)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"Some weights of the model checkpoint at google/bert_uncased_L-2_H-128_A-2 were not used when initializing BertForSequenceClassification: ['cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.decoder.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.bias']\n",
|
||||
"- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
|
||||
"- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
|
||||
"Some weights of BertForSequenceClassification were not initialized from the model checkpoint at google/bert_uncased_L-2_H-128_A-2 and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
|
||||
"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <div>\n",
|
||||
" \n",
|
||||
" <progress value='3' max='3' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
|
||||
" [3/3 00:00, Epoch 3/3]\n",
|
||||
" </div>\n",
|
||||
" <table border=\"1\" class=\"dataframe\">\n",
|
||||
" <thead>\n",
|
||||
" <tr style=\"text-align: left;\">\n",
|
||||
" <th>Step</th>\n",
|
||||
" <th>Training Loss</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" </tbody>\n",
|
||||
"</table><p>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "cjYTxm7sFKyZ"
|
||||
},
|
||||
"source": [
|
||||
"# Train a model with DataFrames\n",
|
||||
"\n",
|
||||
"The next section builds a new model using data stored in a pandas DataFrame."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "0XaKKQ32wqbs",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 182
|
||||
},
|
||||
"outputId": "edb82a45-6c2a-4718-ce0b-56030f95ffbf"
|
||||
},
|
||||
"source": [
|
||||
"import pandas as pd\n",
|
||||
"\n",
|
||||
"df = pd.DataFrame(data)\n",
|
||||
"\n",
|
||||
"model, tokenizer = trainer(\"google/bert_uncased_L-2_H-128_A-2\", data)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"Some weights of the model checkpoint at google/bert_uncased_L-2_H-128_A-2 were not used when initializing BertForSequenceClassification: ['cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.decoder.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.bias']\n",
|
||||
"- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
|
||||
"- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
|
||||
"Some weights of BertForSequenceClassification were not initialized from the model checkpoint at google/bert_uncased_L-2_H-128_A-2 and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
|
||||
"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <div>\n",
|
||||
" \n",
|
||||
" <progress value='3' max='3' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
|
||||
" [3/3 00:00, Epoch 3/3]\n",
|
||||
" </div>\n",
|
||||
" <table border=\"1\" class=\"dataframe\">\n",
|
||||
" <thead>\n",
|
||||
" <tr style=\"text-align: left;\">\n",
|
||||
" <th>Step</th>\n",
|
||||
" <th>Training Loss</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" </tbody>\n",
|
||||
"</table><p>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "QH3D8PQSFvQO"
|
||||
},
|
||||
"source": [
|
||||
"# Train a regression model\n",
|
||||
"\n",
|
||||
"The previous models were classification tasks. The following model trains a sentence similarity model with a regression output per sentence pair between 0 (dissimilar) and 1 (similar)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "1rXuz4ncw9G-"
|
||||
},
|
||||
"source": [
|
||||
"ds = load_dataset(\"glue\", \"stsb\")\n",
|
||||
"model, tokenizer = trainer(\"google/bert_uncased_L-2_H-128_A-2\", ds[\"train\"], columns=(\"sentence1\", \"sentence2\", \"label\"))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "fyvAslSP6j0F",
|
||||
"outputId": "ec46a6aa-25a7-4777-e226-d53aeb37899b"
|
||||
},
|
||||
"source": [
|
||||
"labels = Labels((model, tokenizer), dynamic=False)\n",
|
||||
"labels([[(\"Sailing to the arctic\", \"Dogs and cats don't get along\")], \n",
|
||||
" [(\"Walking down the road\", \"Walking down the street\")]])"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[[(0, 0.5648878216743469)], [(0, 0.97544926404953)]]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 14
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,297 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Pjmz-RORV8E"
|
||||
},
|
||||
"source": [
|
||||
"# Train without labels\n",
|
||||
"\n",
|
||||
"Almost all data available is unlabeled. Labeled data takes effort to manually review and/or takes time to collect. Zero-shot classification takes existing large language models and runs a similarity comparison between candidate text and a list of labels. This has been shown to perform surprisingly well.\n",
|
||||
"\n",
|
||||
"The problem with zero-shot classifiers is that they need to have a large number of parameters (400M+) to perform well against general tasks, which comes with sizable hardware requirements.\n",
|
||||
"\n",
|
||||
"This notebook explores using zero-shot classifiers to build training data for smaller models. A simple form of [knowledge distillation](https://en.wikipedia.org/wiki/Knowledge_distillation). "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dk31rbYjSTYm"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "XMQuuun2R06J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline-train] datasets pandas"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "3PUe1OW8IZR5"
|
||||
},
|
||||
"source": [
|
||||
"# Apply zero-shot classifier to unlabeled text\n",
|
||||
"\n",
|
||||
"The following section takes a small 1000 record random sample of the sst2 dataset and applies a zero-shot classifer to the text. The labels are ignored. This dataset was chosen only to be able to evaluate the accuracy at then end. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "GlrOnS4cmkih"
|
||||
},
|
||||
"source": [
|
||||
"import random\n",
|
||||
"\n",
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"from txtai.pipeline import Labels\n",
|
||||
"\n",
|
||||
"def batch(texts, size):\n",
|
||||
" return [texts[x : x + size] for x in range(0, len(texts), size)]\n",
|
||||
"\n",
|
||||
"# Set random seed for repeatable sampling\n",
|
||||
"random.seed(42)\n",
|
||||
"\n",
|
||||
"ds = load_dataset(\"glue\", \"sst2\")\n",
|
||||
"\n",
|
||||
"sentences = random.sample(ds[\"train\"][\"sentence\"], 1000)\n",
|
||||
"\n",
|
||||
"# Load a zero shot classifier - txtai provides this through the Labels pipeline\n",
|
||||
"labels = Labels(\"microsoft/deberta-large-mnli\")\n",
|
||||
"\n",
|
||||
"train = []\n",
|
||||
"\n",
|
||||
"# Zero-shot prediction using [\"negative\", \"positive\"] labels\n",
|
||||
"for chunk in batch(sentences, 32):\n",
|
||||
" train.extend([{\"text\": chunk[x], \"label\": label[0][0]} for x, label in enumerate(labels(chunk, [\"negative\", \"positive\"]))])"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "TLsZmRpHJGav"
|
||||
},
|
||||
"source": [
|
||||
"Next, we'll use the training set we just built to train a smaller Electra model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 214
|
||||
},
|
||||
"id": "nAt42TIHnfTN",
|
||||
"outputId": "7080b21d-ecf4-459a-c818-11c748e28bb7"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.pipeline import HFTrainer\n",
|
||||
"\n",
|
||||
"trainer = HFTrainer()\n",
|
||||
"model, tokenizer = trainer(\"google/electra-base-discriminator\", train, num_train_epochs=5)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"Some weights of the model checkpoint at google/electra-base-discriminator were not used when initializing ElectraForSequenceClassification: ['discriminator_predictions.dense.bias', 'discriminator_predictions.dense.weight', 'discriminator_predictions.dense_prediction.weight', 'discriminator_predictions.dense_prediction.bias']\n",
|
||||
"- This IS expected if you are initializing ElectraForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
|
||||
"- This IS NOT expected if you are initializing ElectraForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
|
||||
"Some weights of ElectraForSequenceClassification were not initialized from the model checkpoint at google/electra-base-discriminator and are newly initialized: ['classifier.dense.bias', 'classifier.out_proj.weight', 'classifier.out_proj.bias', 'classifier.dense.weight']\n",
|
||||
"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <div>\n",
|
||||
" \n",
|
||||
" <progress value='625' max='625' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
|
||||
" [625/625 02:51, Epoch 5/5]\n",
|
||||
" </div>\n",
|
||||
" <table border=\"1\" class=\"dataframe\">\n",
|
||||
" <thead>\n",
|
||||
" <tr style=\"text-align: left;\">\n",
|
||||
" <th>Step</th>\n",
|
||||
" <th>Training Loss</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" <tr>\n",
|
||||
" <td>500</td>\n",
|
||||
" <td>0.282800</td>\n",
|
||||
" </tr>\n",
|
||||
" </tbody>\n",
|
||||
"</table><p>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "J9pugqJSJRn6"
|
||||
},
|
||||
"source": [
|
||||
"# Evaluating accuracy\n",
|
||||
"\n",
|
||||
"Recall the training set is only 1000 records. To be clear, training an Electra model against the full sst2 dataset would perform better than below. But for this exercise, we're are not using the training labels and simulating labeled data not being available.\n",
|
||||
"\n",
|
||||
"First, lets see what the baseline accuracy for the zero-shot model would be against the sst2 evaluation set. Reminder that this has not seen any of the sst2 training data. \n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "RbgIrkgMvJS4",
|
||||
"outputId": "69287790-e01c-4c17-dfd5-0dc6afd73c98"
|
||||
},
|
||||
"source": [
|
||||
"labels = Labels(\"microsoft/deberta-large-mnli\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"Some weights of the model checkpoint at microsoft/deberta-large-mnli were not used when initializing DebertaForSequenceClassification: ['config']\n",
|
||||
"- This IS expected if you are initializing DebertaForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
|
||||
"- This IS NOT expected if you are initializing DebertaForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "-36UBMILpKYh",
|
||||
"outputId": "3a340b9f-57c5-4c4c-d975-0fcc47df4930"
|
||||
},
|
||||
"source": [
|
||||
"results = [row[\"label\"] == labels(row[\"sentence\"], [\"negative\", \"positive\"])[0][0] for row in ds[\"validation\"]]\n",
|
||||
"sum(results) / len(ds[\"validation\"])"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"0.8818807339449541"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 21
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "uJVnWHZZKFIN"
|
||||
},
|
||||
"source": [
|
||||
"88.19% accuracy, not bad for a model that has not been trained on the dataset at all! Shows the power of zero-shot classification.\n",
|
||||
"\n",
|
||||
"Next, let's test our model trained on the 1000 zero-shot labeled records."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "Kr5IZqZtvXlP",
|
||||
"outputId": "1faeb0d6-349b-4982-e9e8-cdbbde9e9a09"
|
||||
},
|
||||
"source": [
|
||||
"labels = Labels((model, tokenizer), dynamic=False)\n",
|
||||
"\n",
|
||||
"results = [row[\"label\"] == labels(row[\"sentence\"])[0][0] for row in ds[\"validation\"]]\n",
|
||||
"sum(results) / len(ds[\"validation\"])"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"0.8738532110091743"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 22
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "sDw-Zh43KVdX"
|
||||
},
|
||||
"source": [
|
||||
"87.39% accuracy! Wouldn't get too carried away with the percentages but this at least nearly meets the accuracy of the zero-shot classifier.\n",
|
||||
"\n",
|
||||
"Now this model will be highly tuned for a specific task but it had the opportunity to learn from the combined 1000 records whereas the zero-shot classifier views each record independently. It's also much more performant. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "QEAwki2lLM2A"
|
||||
},
|
||||
"source": [
|
||||
"# Conclusion\n",
|
||||
"\n",
|
||||
"This notebook explored a method of building trained text classifiers without training data being available. Given the amount of resources needed to run large-scale zero-shot classifiers, this method is a simple way to build smaller models tuned for specific tasks. In this example, the zero-shot classifier has 400M parameters and the trained text classifier has 110M. "
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,295 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Pjmz-RORV8E"
|
||||
},
|
||||
"source": [
|
||||
"# Train a QA model\n",
|
||||
"\n",
|
||||
"The [Hugging Face Model Hub](https://huggingface.co/models) has a wide range of models that can handle many tasks. While these models perform well, the best performance is often found when fine-tuning a model with task-specific data. \n",
|
||||
"\n",
|
||||
"Hugging Face provides a [number of full-featured examples](https://github.com/huggingface/transformers/tree/master/examples) to assist with training task-specific models. When building models from the command line, these scripts are a great way to get started.\n",
|
||||
"\n",
|
||||
"txtai provides a training pipeline that can be used to train new models programatically using the Transformers Trainer framework.\n",
|
||||
"\n",
|
||||
"This example trains a small QA model and then further fine-tunes it with a couple new examples (few-shot learning)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dk31rbYjSTYm"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "XMQuuun2R06J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline-train] datasets pandas"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "r6nmtieHdMfr"
|
||||
},
|
||||
"source": [
|
||||
"# Train a SQuAD 2.0 Model\n",
|
||||
"\n",
|
||||
"The first step is training a SQuAD 2.0 model. SQuAD is a question-answer dataset that poses a question with a context along with the identified answer. It's also possible to not have an answer. See the [SQuAD dataset website](https://rajpurkar.github.io/SQuAD-explorer/) for more information.\n",
|
||||
"\n",
|
||||
"We'll use a tiny Bert model with a portion of SQuAD 2.0 for efficiency purposes."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 297
|
||||
},
|
||||
"id": "pg9-tUxEdRfk",
|
||||
"outputId": "06195c45-4b39-46e5-a462-af566f437ade"
|
||||
},
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"from txtai.pipeline import HFTrainer\n",
|
||||
"\n",
|
||||
"ds = load_dataset(\"squad_v2\")\n",
|
||||
"\n",
|
||||
"trainer = HFTrainer()\n",
|
||||
"trainer(\"google/bert_uncased_L-2_H-128_A-2\", ds[\"train\"].select(range(3000)), task=\"question-answering\", output_dir=\"bert-tiny-squadv2\")\n",
|
||||
"print(\"Training complete\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"Reusing dataset squad_v2 (/root/.cache/huggingface/datasets/squad_v2/squad_v2/2.0.0/09187c73c1b837c95d9a249cd97c2c3f1cebada06efe667b4427714b27639b1d)\n",
|
||||
"Loading cached processed dataset at /root/.cache/huggingface/datasets/squad_v2/squad_v2/2.0.0/09187c73c1b837c95d9a249cd97c2c3f1cebada06efe667b4427714b27639b1d/cache-73bbe029cf3366fc.arrow\n",
|
||||
"Some weights of the model checkpoint at google/bert_uncased_L-2_H-128_A-2 were not used when initializing BertForQuestionAnswering: ['cls.predictions.decoder.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.decoder.weight', 'cls.predictions.bias', 'cls.seq_relationship.bias', 'cls.predictions.transform.dense.weight', 'cls.seq_relationship.weight']\n",
|
||||
"- This IS expected if you are initializing BertForQuestionAnswering from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
|
||||
"- This IS NOT expected if you are initializing BertForQuestionAnswering from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
|
||||
"Some weights of BertForQuestionAnswering were not initialized from the model checkpoint at google/bert_uncased_L-2_H-128_A-2 and are newly initialized: ['qa_outputs.bias', 'qa_outputs.weight']\n",
|
||||
"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <div>\n",
|
||||
" \n",
|
||||
" <progress value='1131' max='1131' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
|
||||
" [1131/1131 00:50, Epoch 3/3]\n",
|
||||
" </div>\n",
|
||||
" <table border=\"1\" class=\"dataframe\">\n",
|
||||
" <thead>\n",
|
||||
" <tr style=\"text-align: left;\">\n",
|
||||
" <th>Step</th>\n",
|
||||
" <th>Training Loss</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" <tr>\n",
|
||||
" <td>500</td>\n",
|
||||
" <td>4.501800</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <td>1000</td>\n",
|
||||
" <td>3.875900</td>\n",
|
||||
" </tr>\n",
|
||||
" </tbody>\n",
|
||||
"</table><p>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Training complete\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "zZtHxNSwFNGC"
|
||||
},
|
||||
"source": [
|
||||
"# Fine-tune with new data\n",
|
||||
"\n",
|
||||
"Next we'll add a few additional examples. Fine-tuning a QA model will help with framing a certain type of question or improve performance for a specific use-case. \n",
|
||||
"\n",
|
||||
"For smaller models with a narrow use case, this helps the model zero in on the types of questions that are to be asked. In this case, we want to tell the model exactly the types of information we're looking for when asking for ingredients. This will help improve confidence in the answers the model is generating.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 75
|
||||
},
|
||||
"id": "JBeScS5dFNeW",
|
||||
"outputId": "13596524-b476-4f55-f430-d25eeda9301f"
|
||||
},
|
||||
"source": [
|
||||
"# Training data\n",
|
||||
"data = [\n",
|
||||
" {\"question\": \"What ingredient?\", \"context\": \"Pour 1 can whole tomatoes\", \"answers\": \"tomatoes\"},\n",
|
||||
" {\"question\": \"What ingredient?\", \"context\": \"Dice 1 yellow onion\", \"answers\": \"onion\"},\n",
|
||||
" {\"question\": \"What ingredient?\", \"context\": \"Cut 1 red pepper\", \"answers\": \"pepper\"},\n",
|
||||
" {\"question\": \"What ingredient?\", \"context\": \"Peel and dice 1 clove garlic\", \"answers\": \"garlic\"},\n",
|
||||
" {\"question\": \"What ingredient?\", \"context\": \"Put 1/2 lb beef\", \"answers\": \"beef\"},\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"model, tokenizer = trainer(\"bert-tiny-squadv2\", data, task=\"question-answering\", num_train_epochs=10)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <div>\n",
|
||||
" \n",
|
||||
" <progress value='10' max='10' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
|
||||
" [10/10 00:00, Epoch 10/10]\n",
|
||||
" </div>\n",
|
||||
" <table border=\"1\" class=\"dataframe\">\n",
|
||||
" <thead>\n",
|
||||
" <tr style=\"text-align: left;\">\n",
|
||||
" <th>Step</th>\n",
|
||||
" <th>Training Loss</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" </tbody>\n",
|
||||
"</table><p>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "V7nAl3WtkBNK"
|
||||
},
|
||||
"source": [
|
||||
"# Test the model\n",
|
||||
"\n",
|
||||
"Now we're ready to test the results! The following sections run a question against the original model only trained with SQuAD 2.0 and the further fine-tuned model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "46fMiJrAIBu4",
|
||||
"outputId": "fb92ca0a-d433-486f-b61c-054f4e4a9b36"
|
||||
},
|
||||
"source": [
|
||||
"from transformers import pipeline\n",
|
||||
"\n",
|
||||
"questions = pipeline(\"question-answering\", model=\"bert-tiny-squadv2\")\n",
|
||||
"questions(\"What ingredient?\", \"Peel and dice 1 shallot\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'answer': 'dice 1 shallot',\n",
|
||||
" 'end': 23,\n",
|
||||
" 'score': 0.05128436163067818,\n",
|
||||
" 'start': 9}"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 57
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "nWQMRQm0NwdN",
|
||||
"outputId": "a1f15b5c-1cf8-4fe5-daa8-e19adc1700e1"
|
||||
},
|
||||
"source": [
|
||||
"from transformers import pipeline\n",
|
||||
"\n",
|
||||
"questions = pipeline(\"question-answering\", model=model.to(\"cpu\"), tokenizer=tokenizer)\n",
|
||||
"questions(\"What ingredient?\", \"Peel and dice 1 shallot\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'answer': 'shallot', 'end': 23, 'score': 0.13187439739704132, 'start': 16}"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 58
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "wJoYksLbkZTJ"
|
||||
},
|
||||
"source": [
|
||||
"See how the results are more confident and have a better answer. This method allows using a smaller model with a narrow set of functionality with the upside of increased speed. Give it a try with your own data!"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,529 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "vwELCooy4ljr"
|
||||
},
|
||||
"source": [
|
||||
"# Extractive QA to build structured data\n",
|
||||
"\n",
|
||||
"Traditional ETL/data parsing systems establish rules to extract information of interest. Regular expressions, string parsing and similar methods define fixed rules. This works in many cases but what if you are working with unstructured data containing numerous variations? The rules can be cumbersome and hard to maintain over time.\n",
|
||||
"\n",
|
||||
"This notebook uses machine learning and extractive question-answering (QA) to utilize the vast knowledge built into large language models. These models have been trained on extremely large datasets, learning the many variations of natural language. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ew7orE2O441o"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "LPQTb25tASIG"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline-train]"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "_YnqorRKAbLu"
|
||||
},
|
||||
"source": [
|
||||
"# Train a QA model with few-shot learning\n",
|
||||
"\n",
|
||||
"The code below trains a new QA model using a few examples. These examples gives the model hints on the type of questions that will be asked and the type of answers to look for. It doesn't take a lot of examples to do this as shown below."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "OUc9gqTyAYnm",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 75
|
||||
},
|
||||
"outputId": "7e7f93c7-5ad5-46a6-d04d-c7450d246f6c"
|
||||
},
|
||||
"source": [
|
||||
"import pandas as pd\n",
|
||||
"from txtai.pipeline import HFTrainer, Questions, Labels\n",
|
||||
"\n",
|
||||
"# Training data for few-shot learning\n",
|
||||
"data = [\n",
|
||||
" {\"question\": \"What is the url?\",\n",
|
||||
" \"context\": \"Faiss (https://github.com/facebookresearch/faiss) is a library for efficient similarity search.\",\n",
|
||||
" \"answers\": \"https://github.com/facebookresearch/faiss\"},\n",
|
||||
" {\"question\": \"What is the url\", \"context\": \"The last release was Wed Sept 25 2021\", \"answers\": None},\n",
|
||||
" {\"question\": \"What is the date?\", \"context\": \"The last release was Wed Sept 25 2021\", \"answers\": \"Wed Sept 25 2021\"},\n",
|
||||
" {\"question\": \"What is the date?\", \"context\": \"The order total comes to $44.33\", \"answers\": None},\n",
|
||||
" {\"question\": \"What is the amount?\", \"context\": \"The order total comes to $44.33\", \"answers\": \"$44.33\"},\n",
|
||||
" {\"question\": \"What is the amount?\", \"context\": \"The last release was Wed Sept 25 2021\", \"answers\": None},\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"# Fine-tune QA model\n",
|
||||
"trainer = HFTrainer()\n",
|
||||
"model, tokenizer = trainer(\"distilbert-base-cased-distilled-squad\", data, task=\"question-answering\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <div>\n",
|
||||
" \n",
|
||||
" <progress value='3' max='3' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
|
||||
" [3/3 00:03, Epoch 3/3]\n",
|
||||
" </div>\n",
|
||||
" <table border=\"1\" class=\"dataframe\">\n",
|
||||
" <thead>\n",
|
||||
" <tr style=\"text-align: left;\">\n",
|
||||
" <th>Step</th>\n",
|
||||
" <th>Training Loss</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" </tbody>\n",
|
||||
"</table><p>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "1hzPnmrTaUjH"
|
||||
},
|
||||
"source": [
|
||||
"# Parse data into a structured table\n",
|
||||
"\n",
|
||||
"The next section takes a series of rows of text and runs a set of questions against each row. The answers are then used to build a pandas DataFrame."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "4X5z3UjnAGe7",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 206
|
||||
},
|
||||
"outputId": "41015023-c7b7-4515-ad30-f1b2d8143ea2"
|
||||
},
|
||||
"source": [
|
||||
"# Input data\n",
|
||||
"context = [\"Released on 6/03/2021\",\n",
|
||||
" \"Release delayed until the 11th of August\",\n",
|
||||
" \"Documentation can be found here: neuml.github.io/txtai\",\n",
|
||||
" \"The stock price fell to three dollars\",\n",
|
||||
" \"Great day: closing price for March 23rd is $33.11, for details - https://finance.google.com\"]\n",
|
||||
"\n",
|
||||
"# Define column queries\n",
|
||||
"queries = [\"What is the url?\", \"What is the date?\", \"What is the amount?\"]\n",
|
||||
"\n",
|
||||
"# Extract fields\n",
|
||||
"questions = Questions(path=(model, tokenizer), gpu=True)\n",
|
||||
"results = [questions([question] * len(context), context) for question in queries]\n",
|
||||
"results.append(context)\n",
|
||||
"\n",
|
||||
"# Load into DataFrame\n",
|
||||
"pd.DataFrame(list(zip(*results)), columns=[\"URL\", \"Date\", \"Amount\", \"Text\"])"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <div id=\"df-32471005-00c7-4cc3-8624-4bd185b48d76\">\n",
|
||||
" <div class=\"colab-df-container\">\n",
|
||||
" <div>\n",
|
||||
"<style scoped>\n",
|
||||
" .dataframe tbody tr th:only-of-type {\n",
|
||||
" vertical-align: middle;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" .dataframe tbody tr th {\n",
|
||||
" vertical-align: top;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" .dataframe thead th {\n",
|
||||
" text-align: right;\n",
|
||||
" }\n",
|
||||
"</style>\n",
|
||||
"<table border=\"1\" class=\"dataframe\">\n",
|
||||
" <thead>\n",
|
||||
" <tr style=\"text-align: right;\">\n",
|
||||
" <th></th>\n",
|
||||
" <th>URL</th>\n",
|
||||
" <th>Date</th>\n",
|
||||
" <th>Amount</th>\n",
|
||||
" <th>Text</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" <tr>\n",
|
||||
" <th>0</th>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>6/03/2021</td>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>Released on 6/03/2021</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>1</th>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>11th of August</td>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>Release delayed until the 11th of August</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>2</th>\n",
|
||||
" <td>neuml.github.io/txtai</td>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>Documentation can be found here: neuml.github....</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>3</th>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>three dollars</td>\n",
|
||||
" <td>The stock price fell to three dollars</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>4</th>\n",
|
||||
" <td>https://finance.google.com</td>\n",
|
||||
" <td>March 23rd</td>\n",
|
||||
" <td>$33.11</td>\n",
|
||||
" <td>Great day: closing price for March 23rd is $33...</td>\n",
|
||||
" </tr>\n",
|
||||
" </tbody>\n",
|
||||
"</table>\n",
|
||||
"</div>\n",
|
||||
" <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-32471005-00c7-4cc3-8624-4bd185b48d76')\"\n",
|
||||
" title=\"Convert this dataframe to an interactive table.\"\n",
|
||||
" style=\"display:none;\">\n",
|
||||
" \n",
|
||||
" <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
|
||||
" width=\"24px\">\n",
|
||||
" <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
|
||||
" <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
|
||||
" </svg>\n",
|
||||
" </button>\n",
|
||||
" \n",
|
||||
" <style>\n",
|
||||
" .colab-df-container {\n",
|
||||
" display:flex;\n",
|
||||
" flex-wrap:wrap;\n",
|
||||
" gap: 12px;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" .colab-df-convert {\n",
|
||||
" background-color: #E8F0FE;\n",
|
||||
" border: none;\n",
|
||||
" border-radius: 50%;\n",
|
||||
" cursor: pointer;\n",
|
||||
" display: none;\n",
|
||||
" fill: #1967D2;\n",
|
||||
" height: 32px;\n",
|
||||
" padding: 0 0 0 0;\n",
|
||||
" width: 32px;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" .colab-df-convert:hover {\n",
|
||||
" background-color: #E2EBFA;\n",
|
||||
" box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
|
||||
" fill: #174EA6;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" [theme=dark] .colab-df-convert {\n",
|
||||
" background-color: #3B4455;\n",
|
||||
" fill: #D2E3FC;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" [theme=dark] .colab-df-convert:hover {\n",
|
||||
" background-color: #434B5C;\n",
|
||||
" box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
|
||||
" filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
|
||||
" fill: #FFFFFF;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
"\n",
|
||||
" <script>\n",
|
||||
" const buttonEl =\n",
|
||||
" document.querySelector('#df-32471005-00c7-4cc3-8624-4bd185b48d76 button.colab-df-convert');\n",
|
||||
" buttonEl.style.display =\n",
|
||||
" google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
|
||||
"\n",
|
||||
" async function convertToInteractive(key) {\n",
|
||||
" const element = document.querySelector('#df-32471005-00c7-4cc3-8624-4bd185b48d76');\n",
|
||||
" const dataTable =\n",
|
||||
" await google.colab.kernel.invokeFunction('convertToInteractive',\n",
|
||||
" [key], {});\n",
|
||||
" if (!dataTable) return;\n",
|
||||
"\n",
|
||||
" const docLinkHtml = 'Like what you see? Visit the ' +\n",
|
||||
" '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
|
||||
" + ' to learn more about interactive tables.';\n",
|
||||
" element.innerHTML = '';\n",
|
||||
" dataTable['output_type'] = 'display_data';\n",
|
||||
" await google.colab.output.renderOutput(dataTable, element);\n",
|
||||
" const docLink = document.createElement('div');\n",
|
||||
" docLink.innerHTML = docLinkHtml;\n",
|
||||
" element.appendChild(docLink);\n",
|
||||
" }\n",
|
||||
" </script>\n",
|
||||
" </div>\n",
|
||||
" </div>\n",
|
||||
" "
|
||||
],
|
||||
"text/plain": [
|
||||
" URL ... Text\n",
|
||||
"0 None ... Released on 6/03/2021\n",
|
||||
"1 None ... Release delayed until the 11th of August\n",
|
||||
"2 neuml.github.io/txtai ... Documentation can be found here: neuml.github....\n",
|
||||
"3 None ... The stock price fell to three dollars\n",
|
||||
"4 https://finance.google.com ... Great day: closing price for March 23rd is $33...\n",
|
||||
"\n",
|
||||
"[5 rows x 4 columns]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 7
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "mY1Le-pve5yi"
|
||||
},
|
||||
"source": [
|
||||
"# Add additional columns\n",
|
||||
"\n",
|
||||
"This method can be combined with other models to categorize, group or otherwise derive additional columns. The code below derives an additional sentiment column."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 479
|
||||
},
|
||||
"id": "0kyJUcrKe43a",
|
||||
"outputId": "cea0d481-3315-4ec8-a7f9-8a9c4942db10"
|
||||
},
|
||||
"source": [
|
||||
"# Add sentiment\n",
|
||||
"labels = Labels(path=\"distilbert-base-uncased-finetuned-sst-2-english\", dynamic=False)\n",
|
||||
"labels = [\"POSITIVE\" if x[0][0] == 1 else \"NEGATIVE\" for x in labels(context)]\n",
|
||||
"results.insert(len(results) - 1, labels)\n",
|
||||
"\n",
|
||||
"# Load into DataFrame\n",
|
||||
"pd.DataFrame(list(zip(*results)), columns=[\"URL\", \"Date\", \"Amount\", \"Sentiment\", \"Text\"])"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <div id=\"df-a5d090e1-83da-4189-84b1-5d89675d9800\">\n",
|
||||
" <div class=\"colab-df-container\">\n",
|
||||
" <div>\n",
|
||||
"<style scoped>\n",
|
||||
" .dataframe tbody tr th:only-of-type {\n",
|
||||
" vertical-align: middle;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" .dataframe tbody tr th {\n",
|
||||
" vertical-align: top;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" .dataframe thead th {\n",
|
||||
" text-align: right;\n",
|
||||
" }\n",
|
||||
"</style>\n",
|
||||
"<table border=\"1\" class=\"dataframe\">\n",
|
||||
" <thead>\n",
|
||||
" <tr style=\"text-align: right;\">\n",
|
||||
" <th></th>\n",
|
||||
" <th>URL</th>\n",
|
||||
" <th>Date</th>\n",
|
||||
" <th>Amount</th>\n",
|
||||
" <th>Sentiment</th>\n",
|
||||
" <th>Text</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" <tr>\n",
|
||||
" <th>0</th>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>6/03/2021</td>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>POSITIVE</td>\n",
|
||||
" <td>Released on 6/03/2021</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>1</th>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>11th of August</td>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>NEGATIVE</td>\n",
|
||||
" <td>Release delayed until the 11th of August</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>2</th>\n",
|
||||
" <td>neuml.github.io/txtai</td>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>NEGATIVE</td>\n",
|
||||
" <td>Documentation can be found here: neuml.github....</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>3</th>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>None</td>\n",
|
||||
" <td>three dollars</td>\n",
|
||||
" <td>NEGATIVE</td>\n",
|
||||
" <td>The stock price fell to three dollars</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>4</th>\n",
|
||||
" <td>https://finance.google.com</td>\n",
|
||||
" <td>March 23rd</td>\n",
|
||||
" <td>$33.11</td>\n",
|
||||
" <td>POSITIVE</td>\n",
|
||||
" <td>Great day: closing price for March 23rd is $33...</td>\n",
|
||||
" </tr>\n",
|
||||
" </tbody>\n",
|
||||
"</table>\n",
|
||||
"</div>\n",
|
||||
" <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-a5d090e1-83da-4189-84b1-5d89675d9800')\"\n",
|
||||
" title=\"Convert this dataframe to an interactive table.\"\n",
|
||||
" style=\"display:none;\">\n",
|
||||
" \n",
|
||||
" <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
|
||||
" width=\"24px\">\n",
|
||||
" <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
|
||||
" <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
|
||||
" </svg>\n",
|
||||
" </button>\n",
|
||||
" \n",
|
||||
" <style>\n",
|
||||
" .colab-df-container {\n",
|
||||
" display:flex;\n",
|
||||
" flex-wrap:wrap;\n",
|
||||
" gap: 12px;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" .colab-df-convert {\n",
|
||||
" background-color: #E8F0FE;\n",
|
||||
" border: none;\n",
|
||||
" border-radius: 50%;\n",
|
||||
" cursor: pointer;\n",
|
||||
" display: none;\n",
|
||||
" fill: #1967D2;\n",
|
||||
" height: 32px;\n",
|
||||
" padding: 0 0 0 0;\n",
|
||||
" width: 32px;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" .colab-df-convert:hover {\n",
|
||||
" background-color: #E2EBFA;\n",
|
||||
" box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
|
||||
" fill: #174EA6;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" [theme=dark] .colab-df-convert {\n",
|
||||
" background-color: #3B4455;\n",
|
||||
" fill: #D2E3FC;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" [theme=dark] .colab-df-convert:hover {\n",
|
||||
" background-color: #434B5C;\n",
|
||||
" box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
|
||||
" filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
|
||||
" fill: #FFFFFF;\n",
|
||||
" }\n",
|
||||
" </style>\n",
|
||||
"\n",
|
||||
" <script>\n",
|
||||
" const buttonEl =\n",
|
||||
" document.querySelector('#df-a5d090e1-83da-4189-84b1-5d89675d9800 button.colab-df-convert');\n",
|
||||
" buttonEl.style.display =\n",
|
||||
" google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
|
||||
"\n",
|
||||
" async function convertToInteractive(key) {\n",
|
||||
" const element = document.querySelector('#df-a5d090e1-83da-4189-84b1-5d89675d9800');\n",
|
||||
" const dataTable =\n",
|
||||
" await google.colab.kernel.invokeFunction('convertToInteractive',\n",
|
||||
" [key], {});\n",
|
||||
" if (!dataTable) return;\n",
|
||||
"\n",
|
||||
" const docLinkHtml = 'Like what you see? Visit the ' +\n",
|
||||
" '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
|
||||
" + ' to learn more about interactive tables.';\n",
|
||||
" element.innerHTML = '';\n",
|
||||
" dataTable['output_type'] = 'display_data';\n",
|
||||
" await google.colab.output.renderOutput(dataTable, element);\n",
|
||||
" const docLink = document.createElement('div');\n",
|
||||
" docLink.innerHTML = docLinkHtml;\n",
|
||||
" element.appendChild(docLink);\n",
|
||||
" }\n",
|
||||
" </script>\n",
|
||||
" </div>\n",
|
||||
" </div>\n",
|
||||
" "
|
||||
],
|
||||
"text/plain": [
|
||||
" URL ... Text\n",
|
||||
"0 None ... Released on 6/03/2021\n",
|
||||
"1 None ... Release delayed until the 11th of August\n",
|
||||
"2 neuml.github.io/txtai ... Documentation can be found here: neuml.github....\n",
|
||||
"3 None ... The stock price fell to three dollars\n",
|
||||
"4 https://finance.google.com ... Great day: closing price for March 23rd is $33...\n",
|
||||
"\n",
|
||||
"[5 rows x 5 columns]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 8
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,664 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"name": "21 - Export and run other machine learning models",
|
||||
"provenance": [],
|
||||
"collapsed_sections": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Pjmz-RORV8E"
|
||||
},
|
||||
"source": [
|
||||
"# Export and run other machine learning models\n",
|
||||
"\n",
|
||||
"txtai primarily has support for [Hugging Face Transformers](https://github.com/huggingface/transformers) and [ONNX](https://github.com/microsoft/onnxruntime) models. This enables txtai to hook into the rich model framework available in Python, export this functionality via the API to other languages (JavaScript, Java, Go, Rust) and even export and natively load models with ONNX.\n",
|
||||
"\n",
|
||||
"What about other machine learning frameworks? Say we have an existing TF-IDF + Logistic Regression model that has been well tuned. Can this model be exported to ONNX and used in txtai for labeling and similarity queries? Or what about a simple PyTorch text classifier? Yes, both of these can be done!\n",
|
||||
"\n",
|
||||
"With the [onnxmltools](https://github.com/onnx/onnxmltools) library, traditional models from [scikit-learn](https://scikit-learn.org/stable/), [XGBoost](https://xgboost.readthedocs.io/en/latest/) and others can be exported to ONNX and loaded with txtai. Additionally, Hugging Face's trainer module can train generic PyTorch modules. This notebook will walk through all these examples.\n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dk31rbYjSTYm"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "XMQuuun2R06J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline,similarity] datasets"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "r6nmtieHdMfr"
|
||||
},
|
||||
"source": [
|
||||
"# Train a TF-IDF + Logistic Regression model\n",
|
||||
"\n",
|
||||
"For this example, we'll load the emotion dataset from Hugging Face datasets and build a TF-IDF + Logistic Regression model with scikit-learn.\n",
|
||||
"\n",
|
||||
"The emotion dataset has the following labels:\n",
|
||||
"\n",
|
||||
"- sadness (0)\n",
|
||||
"- joy (1)\n",
|
||||
"- love (2)\n",
|
||||
"- anger (3)\n",
|
||||
"- fear (4)\n",
|
||||
"- surprise (5)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "pg9-tUxEdRfk"
|
||||
},
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"from sklearn.feature_extraction.text import TfidfVectorizer\n",
|
||||
"from sklearn.linear_model import LogisticRegression\n",
|
||||
"from sklearn.pipeline import Pipeline\n",
|
||||
"\n",
|
||||
"ds = load_dataset(\"emotion\")\n",
|
||||
"\n",
|
||||
"# Train the model\n",
|
||||
"pipeline = Pipeline([\n",
|
||||
" ('tfidf', TfidfVectorizer()),\n",
|
||||
" ('lr', LogisticRegression(max_iter=250))\n",
|
||||
"])\n",
|
||||
"\n",
|
||||
"pipeline.fit(ds[\"train\"][\"text\"], ds[\"train\"][\"label\"])\n",
|
||||
"\n",
|
||||
"# Determine accuracy on validation set\n",
|
||||
"results = pipeline.predict(ds[\"validation\"][\"text\"])\n",
|
||||
"labels = ds[\"validation\"][\"label\"]\n",
|
||||
"\n",
|
||||
"results = [results[x] == label for x, label in enumerate(labels)]\n",
|
||||
"print(\"Accuracy =\", sum(results) / len(ds[\"validation\"]))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"Using custom data configuration default\n",
|
||||
"Reusing dataset emotion (/root/.cache/huggingface/datasets/emotion/default/0.0.0/348f63ca8e27b3713b6c04d723efe6d824a56fb3d1449794716c0f0296072705)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Accuracy = 0.8595\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "49jZD4jQgdBg"
|
||||
},
|
||||
"source": [
|
||||
"86% accuracy - not too bad! While we all get caught up in deep learning and advanced methods, good ole TF-IDF + Logistic Regression is still a solid performer and runs much faster. If that level of accuracy works, no reason to overcomplicate things."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "zZtHxNSwFNGC"
|
||||
},
|
||||
"source": [
|
||||
"# Export and load with txtai\n",
|
||||
"\n",
|
||||
"The next section exports this model to ONNX and shows how the model can be used for similarity queries. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "JBeScS5dFNeW",
|
||||
"outputId": "e1b5cbf4-87dd-4598-e7ee-14e36cf31a7c"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.pipeline import Labels, MLOnnx, Similarity\n",
|
||||
"\n",
|
||||
"def tokenize(inputs, **kwargs):\n",
|
||||
" if isinstance(inputs, str):\n",
|
||||
" inputs = [inputs]\n",
|
||||
"\n",
|
||||
" return {\"input_ids\": [[x] for x in inputs]}\n",
|
||||
"\n",
|
||||
"def query(model, tokenizer, multilabel=False):\n",
|
||||
" # Load models into similarity pipeline\n",
|
||||
" similarity = Similarity((model, tokenizer), dynamic=False)\n",
|
||||
"\n",
|
||||
" # Add labels to model\n",
|
||||
" similarity.pipeline.model.config.id2label = {0: \"sadness\", 1: \"joy\", 2: \"love\", 3: \"anger\", 4: \"fear\", 5: \"surprise\"}\n",
|
||||
" similarity.pipeline.model.config.label2id = dict((v, k) for k, v in similarity.pipeline.model.config.id2label.items())\n",
|
||||
"\n",
|
||||
" inputs = [\"that caught me off guard\", \"I didn t see that coming\", \"i feel bad\", \"What a wonderful goal!\"]\n",
|
||||
" scores = similarity(\"joy\", inputs, multilabel)\n",
|
||||
" for uid, score in scores[:5]:\n",
|
||||
" print(inputs[uid], score)\n",
|
||||
"\n",
|
||||
"# Export to ONNX\n",
|
||||
"onnx = MLOnnx()\n",
|
||||
"model = onnx(pipeline)\n",
|
||||
"\n",
|
||||
"# Create labels pipeline using scikit-learn ONNX model\n",
|
||||
"sklabels = Labels((model, tokenize), dynamic=False)\n",
|
||||
"\n",
|
||||
"# Add labels to model\n",
|
||||
"sklabels.pipeline.model.config.id2label = {0: \"sadness\", 1: \"joy\", 2: \"love\", 3: \"anger\", 4: \"fear\", 5: \"surprise\"}\n",
|
||||
"sklabels.pipeline.model.config.label2id = dict((v, k) for k, v in sklabels.pipeline.model.config.id2label.items())\n",
|
||||
"\n",
|
||||
"# Run test query using model\n",
|
||||
"query(model, tokenize, None)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"What a wonderful goal! 0.909473717212677\n",
|
||||
"I didn t see that coming 0.47113093733787537\n",
|
||||
"that caught me off guard 0.42067453265190125\n",
|
||||
"i feel bad 0.019547615200281143\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "d-y8gFJwCwKN"
|
||||
},
|
||||
"source": [
|
||||
"txtai can use a standard text classification model for similarity queries, where the label(s) are a list of fixed queries. The output above shows the best results for the query \"joy\"."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "cbqwX7GgKBkf"
|
||||
},
|
||||
"source": [
|
||||
"# Train a PyTorch model\n",
|
||||
"\n",
|
||||
"The next section defines a simple PyTorch text classifier. The transformers library has a trainer package that supports training PyTorch models, assuming some standard conventions/naming is used. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 239
|
||||
},
|
||||
"id": "k8PkTlBLKBTy",
|
||||
"outputId": "4f48bfb2-2f16-45e3-d3e6-f1a2a747fd09"
|
||||
},
|
||||
"source": [
|
||||
"# Set predictable seeds\n",
|
||||
"import os\n",
|
||||
"import random\n",
|
||||
"import torch\n",
|
||||
"\n",
|
||||
"import numpy as np\n",
|
||||
"\n",
|
||||
"from torch import nn\n",
|
||||
"from torch.nn import CrossEntropyLoss\n",
|
||||
"from transformers import AutoConfig, AutoTokenizer\n",
|
||||
"\n",
|
||||
"from txtai.models import Registry\n",
|
||||
"from txtai.pipeline import HFTrainer\n",
|
||||
"\n",
|
||||
"from transformers.modeling_outputs import SequenceClassifierOutput\n",
|
||||
"\n",
|
||||
"def seed(seed=42):\n",
|
||||
" random.seed(seed)\n",
|
||||
" os.environ['PYTHONHASHSEED'] = str(seed)\n",
|
||||
" np.random.seed(seed)\n",
|
||||
" torch.manual_seed(seed)\n",
|
||||
" torch.cuda.manual_seed(seed)\n",
|
||||
" torch.backends.cudnn.deterministic = True\n",
|
||||
"\n",
|
||||
"class Simple(nn.Module):\n",
|
||||
" def __init__(self, vocab, dimensions, labels):\n",
|
||||
" super().__init__()\n",
|
||||
"\n",
|
||||
" self.config = AutoConfig.from_pretrained(\"bert-base-uncased\")\n",
|
||||
" self.labels = labels\n",
|
||||
"\n",
|
||||
" self.embedding = nn.EmbeddingBag(vocab, dimensions)\n",
|
||||
" self.classifier = nn.Linear(dimensions, labels)\n",
|
||||
" self.init_weights()\n",
|
||||
"\n",
|
||||
" def init_weights(self):\n",
|
||||
" initrange = 0.5\n",
|
||||
" self.embedding.weight.data.uniform_(-initrange, initrange)\n",
|
||||
" self.classifier.weight.data.uniform_(-initrange, initrange)\n",
|
||||
" self.classifier.bias.data.zero_()\n",
|
||||
"\n",
|
||||
" def forward(self, input_ids=None, labels=None, **kwargs):\n",
|
||||
" embeddings = self.embedding(input_ids)\n",
|
||||
" logits = self.classifier(embeddings)\n",
|
||||
"\n",
|
||||
" loss = None\n",
|
||||
" if labels is not None:\n",
|
||||
" loss_fct = CrossEntropyLoss()\n",
|
||||
" loss = loss_fct(logits.view(-1, self.labels), labels.view(-1))\n",
|
||||
"\n",
|
||||
" return SequenceClassifierOutput(\n",
|
||||
" loss=loss,\n",
|
||||
" logits=logits,\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
"# Set seed for reproducibility\n",
|
||||
"seed()\n",
|
||||
"\n",
|
||||
"# Define model\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")\n",
|
||||
"model = Simple(tokenizer.vocab_size, 128, len(ds[\"train\"].unique(\"label\")))\n",
|
||||
"\n",
|
||||
"# Train model\n",
|
||||
"train = HFTrainer()\n",
|
||||
"model, tokenizer = train((model, tokenizer), ds[\"train\"], per_device_train_batch_size=8, learning_rate=1e-3, num_train_epochs=15, logging_steps=10000)\n",
|
||||
"\n",
|
||||
"# Register custom model to fully support pipelines\n",
|
||||
"Registry.register(model)\n",
|
||||
"\n",
|
||||
"# Create labels pipeline using PyTorch model\n",
|
||||
"thlabels = Labels((model, tokenizer), dynamic=False)\n",
|
||||
"\n",
|
||||
"# Determine accuracy on validation set\n",
|
||||
"results = [row[\"label\"] == thlabels(row[\"text\"])[0][0] for row in ds[\"validation\"]]\n",
|
||||
"print(\"Accuracy = \", sum(results) / len(ds[\"validation\"]))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"Loading cached processed dataset at /root/.cache/huggingface/datasets/emotion/default/0.0.0/348f63ca8e27b3713b6c04d723efe6d824a56fb3d1449794716c0f0296072705/cache-a983327c4471f5aa.arrow\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <div>\n",
|
||||
" \n",
|
||||
" <progress value='30000' max='30000' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
|
||||
" [30000/30000 02:28, Epoch 15/15]\n",
|
||||
" </div>\n",
|
||||
" <table border=\"1\" class=\"dataframe\">\n",
|
||||
" <thead>\n",
|
||||
" <tr style=\"text-align: left;\">\n",
|
||||
" <th>Step</th>\n",
|
||||
" <th>Training Loss</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" <tr>\n",
|
||||
" <td>10000</td>\n",
|
||||
" <td>1.017600</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <td>20000</td>\n",
|
||||
" <td>0.286200</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <td>30000</td>\n",
|
||||
" <td>0.152500</td>\n",
|
||||
" </tr>\n",
|
||||
" </tbody>\n",
|
||||
"</table><p>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Accuracy = 0.883\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "nHQoJnrj60Pz"
|
||||
},
|
||||
"source": [
|
||||
"88% accuracy this time. Pretty good for such a simple network and something that could definitely be improved upon. \n",
|
||||
"\n",
|
||||
"Once again let's run similarity queries using this model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "W5_NDInF5lFN",
|
||||
"outputId": "38a2c126-63e9-40dc-f309-a29826b5b937"
|
||||
},
|
||||
"source": [
|
||||
"query(model, tokenizer)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"What a wonderful goal! 1.0\n",
|
||||
"that caught me off guard 0.9998751878738403\n",
|
||||
"I didn t see that coming 0.7328283190727234\n",
|
||||
"i feel bad 5.2972134609891875e-19\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KmcsdIltDTwj"
|
||||
},
|
||||
"source": [
|
||||
"Same result order as with the scikit-learn model with scoring variations which is expected given this is a completely different model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-fNTi2jb68rv"
|
||||
},
|
||||
"source": [
|
||||
"# Pooled embeddings\n",
|
||||
"\n",
|
||||
"The PyTorch model above consists of an embeddings layer with a linear classifier on top of it. What if we take that embeddings layer and use it for similarity queries? Let's give it a try."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "J1yhfHKC7N7L",
|
||||
"outputId": "11567948-769a-44df-9057-9fe9837a73dd"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.embeddings import Embeddings\n",
|
||||
"\n",
|
||||
"class SimpleEmbeddings(nn.Module):\n",
|
||||
" def __init__(self, embeddings):\n",
|
||||
" super().__init__()\n",
|
||||
"\n",
|
||||
" self.embeddings = embeddings\n",
|
||||
"\n",
|
||||
" def forward(self, input_ids=None, **kwargs):\n",
|
||||
" return (self.embeddings(input_ids),)\n",
|
||||
"\n",
|
||||
"embeddings = Embeddings({\"method\": \"pooling\", \"path\": SimpleEmbeddings(model.embedding), \"tokenizer\": \"bert-base-uncased\"})\n",
|
||||
"print(embeddings.similarity(\"mad\", [\"Glad you found it\", \"Happy to see you\", \"I'm angry\"]))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"[(2, 0.8323876857757568), (1, -0.11010512709617615), (0, -0.16152513027191162)]\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "0kTUEIcmBNuV"
|
||||
},
|
||||
"source": [
|
||||
"Definitely looks like the embeddings have stored knowledge. Could these embeddings be good enough to build a semantic search index, especially for sentiment based data, given the training dataset? Possibly. It certainly would run faster than a standard transformer model (see below). "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "V7nAl3WtkBNK"
|
||||
},
|
||||
"source": [
|
||||
"# Train a transformer model and compare accuracy/speed\n",
|
||||
"\n",
|
||||
"Let's train a standard transformer sequence classifier and compare the accuracy/speed between the two. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 274
|
||||
},
|
||||
"id": "46fMiJrAIBu4",
|
||||
"outputId": "f0512cf8-3bc2-41ed-caff-e1541403f2a5"
|
||||
},
|
||||
"source": [
|
||||
"train = HFTrainer()\n",
|
||||
"model, tokenizer = train(\"microsoft/xtremedistil-l6-h384-uncased\", ds[\"train\"], logging_steps=2000)\n",
|
||||
"\n",
|
||||
"tflabels = Labels((model, tokenizer), dynamic=False)\n",
|
||||
"\n",
|
||||
"# Determine accuracy on validation set\n",
|
||||
"results = [row[\"label\"] == tflabels(row[\"text\"])[0][0] for row in ds[\"validation\"]]\n",
|
||||
"print(\"Accuracy = \", sum(results) / len(ds[\"validation\"]))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"Loading cached processed dataset at /root/.cache/huggingface/datasets/emotion/default/0.0.0/348f63ca8e27b3713b6c04d723efe6d824a56fb3d1449794716c0f0296072705/cache-98b7ef31bf6ca944.arrow\n",
|
||||
"Some weights of BertForSequenceClassification were not initialized from the model checkpoint at microsoft/xtremedistil-l6-h384-uncased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
|
||||
"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <div>\n",
|
||||
" \n",
|
||||
" <progress value='6000' max='6000' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
|
||||
" [6000/6000 07:13, Epoch 3/3]\n",
|
||||
" </div>\n",
|
||||
" <table border=\"1\" class=\"dataframe\">\n",
|
||||
" <thead>\n",
|
||||
" <tr style=\"text-align: left;\">\n",
|
||||
" <th>Step</th>\n",
|
||||
" <th>Training Loss</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" <tr>\n",
|
||||
" <td>2000</td>\n",
|
||||
" <td>0.635500</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <td>4000</td>\n",
|
||||
" <td>0.281700</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <td>6000</td>\n",
|
||||
" <td>0.192600</td>\n",
|
||||
" </tr>\n",
|
||||
" </tbody>\n",
|
||||
"</table><p>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Accuracy = 0.926\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ycvozGzPmlbS"
|
||||
},
|
||||
"source": [
|
||||
"As expected, the accuracy is better. The model above is a distilled model and even better accuracy can be obtained with a model like \"roberta-base\" with the tradeoff being increased training/inference time. \n",
|
||||
"\n",
|
||||
"Speaking of speed, let's compare the speed of these models."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "nWQMRQm0NwdN",
|
||||
"outputId": "4a49406c-b4eb-46b1-edab-de01c15fdccb"
|
||||
},
|
||||
"source": [
|
||||
"import time\n",
|
||||
"\n",
|
||||
"# Test inputs\n",
|
||||
"inputs = ds[\"test\"][\"text\"]\n",
|
||||
"print(\"Testing speed of %d items\" % len(inputs))\n",
|
||||
"\n",
|
||||
"start = time.time()\n",
|
||||
"r1 = sklabels(inputs, multilabel=None)\n",
|
||||
"print(\"TF-IDF + Logistic Regression time =\", time.time() - start)\n",
|
||||
"\n",
|
||||
"start = time.time()\n",
|
||||
"r2 = thlabels(inputs)\n",
|
||||
"print(\"PyTorch time =\", time.time() - start)\n",
|
||||
"\n",
|
||||
"start = time.time()\n",
|
||||
"r3 = tflabels(inputs)\n",
|
||||
"print(\"Transformers time =\", time.time() - start, \"\\n\")\n",
|
||||
"\n",
|
||||
"# Compare model results\n",
|
||||
"for x in range(5):\n",
|
||||
" print(\"index: %d\" % x)\n",
|
||||
" print(r1[x][0])\n",
|
||||
" print(r2[x][0])\n",
|
||||
" print(r3[x][0], \"\\n\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Testing speed of 2000 items\n",
|
||||
"TF-IDF + Logistic Regression time = 1.0483319759368896\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"PyTorch time = 2.0001697540283203\n",
|
||||
"Transformers time = 13.71584439277649 \n",
|
||||
"\n",
|
||||
"index: 0\n",
|
||||
"(0, 0.7258279323577881)\n",
|
||||
"(0, 1.0)\n",
|
||||
"(0, 0.998375654220581) \n",
|
||||
"\n",
|
||||
"index: 1\n",
|
||||
"(0, 0.854256272315979)\n",
|
||||
"(0, 1.0)\n",
|
||||
"(0, 0.9983494281768799) \n",
|
||||
"\n",
|
||||
"index: 2\n",
|
||||
"(0, 0.6306578516960144)\n",
|
||||
"(0, 0.9999700784683228)\n",
|
||||
"(0, 0.9982945322990417) \n",
|
||||
"\n",
|
||||
"index: 3\n",
|
||||
"(1, 0.554378092288971)\n",
|
||||
"(1, 0.9998960494995117)\n",
|
||||
"(1, 0.99846351146698) \n",
|
||||
"\n",
|
||||
"index: 4\n",
|
||||
"(0, 0.8961835503578186)\n",
|
||||
"(0, 1.0)\n",
|
||||
"(0, 0.9984095692634583) \n",
|
||||
"\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "1YMTyqIWDiOB"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook showed how frameworks outside of Transformers and ONNX can be used as models in txtai.\n",
|
||||
"\n",
|
||||
"As seen in the section above, TF-IDF + Logistic Regression is 16 times faster than a distilled Transformers model. A simple PyTorch network is 8 times faster. Depending on your accuracy requirements, it may make sense to use a simpler model to get better runtime performance."
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,951 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"name": "22 - Transform tabular data with composable workflows",
|
||||
"provenance": [],
|
||||
"collapsed_sections": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Pjmz-RORV8E"
|
||||
},
|
||||
"source": [
|
||||
"# Transform tabular data with composable workflows\n",
|
||||
"\n",
|
||||
"txtai has support for processing both unstructured and structured data. Structured or tabular data is grouped into rows and columns. This can be a spreadsheet, an API call that returns JSON or XML or even list of key-value pairs.\n",
|
||||
"\n",
|
||||
"This notebook will walk through examples on how to use workflows with the tabular pipeline to transform and index structured data.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dk31rbYjSTYm"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies. We will install the api, pipeline and workflow optional extras packages. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "XMQuuun2R06J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[api,pipeline,workflow] sacremoses"
|
||||
],
|
||||
"execution_count": 66,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "NSYrP0hjtR_E"
|
||||
},
|
||||
"source": [
|
||||
"# CSV Workflow\n",
|
||||
"\n",
|
||||
"The first example will transform and index a CSV file. The [COVID-19 Open Research Dataset](https://allenai.org/data/cord-19) (CORD-19) is a repository of medical articles covering COVID-19. This workflow reads the input CSV and builds a semantic search index.\n",
|
||||
"\n",
|
||||
"The first step is downloading the dataset locally."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "BoPJIKWoTibk"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"# Get CORD-19 metadata file\n",
|
||||
"!wget https://ai2-semanticscholar-cord-19.s3-us-west-2.amazonaws.com/2021-11-01/metadata.csv\n",
|
||||
"!head -1 metadata.csv > input.csv\n",
|
||||
"!tail -10000 metadata.csv >> input.csv"
|
||||
],
|
||||
"execution_count": 67,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Q1ivX4eBuU8T"
|
||||
},
|
||||
"source": [
|
||||
"The next section creates a simple workflow consisting of a tabular pipeline. The tabular pipeline builds a list of (id, text, tag) tuples that can be easily loaded into an Embeddings index. For this example, we'll use the `url` column as the id and the `title` column as the text column. The textcolumns parameter takes a list of columns to support indexing text content from multiple columns. \n",
|
||||
"\n",
|
||||
"The file input.csv is processed and the first 5 rows are shown."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "-pi2QU3TSlM_",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "a30ffeb2-1d85-467b-a0a3-134c8cbac19f"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.pipeline import Tabular\n",
|
||||
"from txtai.workflow import Task, Workflow\n",
|
||||
"\n",
|
||||
"# Create tabular instance mapping input.csv fields\n",
|
||||
"tabular = Tabular(\"url\", [\"title\"])\n",
|
||||
"\n",
|
||||
"# Create workflow\n",
|
||||
"workflow = Workflow([Task(tabular)])\n",
|
||||
"\n",
|
||||
"# Print 5 rows of input.csv via workflow\n",
|
||||
"list(workflow([\"input.csv\"]))[:5]"
|
||||
],
|
||||
"execution_count": 68,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[('https://doi.org/10.1016/j.cmpb.2021.106469; https://www.ncbi.nlm.nih.gov/pubmed/34715516/',\n",
|
||||
" 'Computer simulation of the dynamics of a spatial susceptible-infected-recovered epidemic model with time delays in transmission and treatment.',\n",
|
||||
" None),\n",
|
||||
" ('https://www.ncbi.nlm.nih.gov/pubmed/34232002/; https://doi.org/10.36849/jdd.5544',\n",
|
||||
" 'Understanding the Potential Role of Abrocitinib in the Time of SARS-CoV-2',\n",
|
||||
" None),\n",
|
||||
" ('https://doi.org/10.1186/1471-2458-8-42; https://www.ncbi.nlm.nih.gov/pubmed/18234083/',\n",
|
||||
" \"Can the concept of Health Promoting Schools help to improve students' health knowledge and practices to combat the challenge of communicable diseases: Case study in Hong Kong?\",\n",
|
||||
" None),\n",
|
||||
" ('https://www.ncbi.nlm.nih.gov/pubmed/32983582/; https://www.sciencedirect.com/science/article/pii/S2095809920302514?v=s5; https://api.elsevier.com/content/article/pii/S2095809920302514; https://doi.org/10.1016/j.eng.2020.07.018',\n",
|
||||
" 'Buying time for an effective epidemic response: The impact of a public holiday for outbreak control on COVID-19 epidemic spread',\n",
|
||||
" None),\n",
|
||||
" ('https://doi.org/10.1093/pcmedi/pbab016',\n",
|
||||
" 'The SARS-CoV-2 spike L452R-E484Q variant in the Indian B.1.617 strain showed significant reduction in the neutralization activity of immune sera',\n",
|
||||
" None)]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 68
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "UYYKnwNhu0hv"
|
||||
},
|
||||
"source": [
|
||||
"Next, we take the workflow output, build an Embeddings index and run a search query."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "G7M34puLWeZm",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "0e9e771b-d53d-4c4e-8ac8-257c33992f66"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.embeddings import Embeddings\n",
|
||||
"\n",
|
||||
"# Embeddings with sentence-transformers backend\n",
|
||||
"embeddings = Embeddings({\"method\": \"transformers\", \"path\": \"sentence-transformers/paraphrase-mpnet-base-v2\"})\n",
|
||||
"\n",
|
||||
"# Index subset of CORD-19 data\n",
|
||||
"data = list(workflow([\"input.csv\"]))\n",
|
||||
"embeddings.index(data)\n",
|
||||
"\n",
|
||||
"for uid, _ in embeddings.search(\"insulin\"):\n",
|
||||
" title = [text for url, text, _ in data if url == uid][0]\n",
|
||||
" print(title, uid)"
|
||||
],
|
||||
"execution_count": 69,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Importance of diabetes management during the COVID-19 pandemic. https://doi.org/10.1080/00325481.2021.1978704; https://www.ncbi.nlm.nih.gov/pubmed/34602003/\n",
|
||||
"Position Statement on How to Manage Patients with Diabetes and COVID-19 https://www.ncbi.nlm.nih.gov/pubmed/33442169/; https://doi.org/10.15605/jafes.035.01.03\n",
|
||||
"Successful blood glucose management of a severe COVID-19 patient with diabetes: A case report https://www.ncbi.nlm.nih.gov/pubmed/32590779/; https://doi.org/10.1097/md.0000000000020844\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "vPQv7GJSu9Uq"
|
||||
},
|
||||
"source": [
|
||||
"The example searched for the term `insulin`. The top results mention diabetes and blood glucose which are a closely associated terms for diabetes."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Workflow with stored content\n",
|
||||
"\n",
|
||||
"Next we'll re-run the same example adding in full content storage. Full content storage enables SQL queries."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "BQeLRC1md6GR"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import json\n",
|
||||
"\n",
|
||||
"# Create tabular instance mapping input.csv fields\n",
|
||||
"tabular = Tabular(\"url\", [\"title\"], True)\n",
|
||||
"\n",
|
||||
"# Create workflow\n",
|
||||
"workflow = Workflow([Task(tabular)])\n",
|
||||
"\n",
|
||||
"# Embeddings with sentence-transformers backend\n",
|
||||
"embeddings = Embeddings({\"method\": \"transformers\", \"path\": \"sentence-transformers/paraphrase-mpnet-base-v2\", \"content\": True})\n",
|
||||
"\n",
|
||||
"# Index subset of CORD-19 data\n",
|
||||
"data = list(workflow([\"input.csv\"]))\n",
|
||||
"embeddings.index(data)\n",
|
||||
"\n",
|
||||
"for result in embeddings.search(\"select title, abstract, authors, doi from txtai where similar('insulin')\"):\n",
|
||||
" print(json.dumps(result, default=str, indent=2))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "Ln61lB7QeDEq",
|
||||
"outputId": "ae364e17-36a5-480e-c614-c7e58d8b8462"
|
||||
},
|
||||
"execution_count": 70,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"{\n",
|
||||
" \"title\": \"Importance of diabetes management during the COVID-19 pandemic.\",\n",
|
||||
" \"abstract\": \"Uncontrolled diabetes and/or hyperglycemia is associated with severe COVID-19 disease and increased mortality. It is now known that poor glucose control before hospital admission can be associated with a high risk of in-hospital death. By achieving and maintaining glycemic control, primary care physicians (PCPs) play a critical role in limiting this potentially devastating outcome. Further, despite the hope that mass vaccination will help control the pandemic, genetic variants of the virus are causing surges in some countries. As such, PCPs will treat an increasing number of patients with diabetes who have symptoms of post-COVID-19 infection, or even have new-onset type 2 diabetes as a result of COVID-19 infection. However, much of the literature published focuses on the effects of COVID-19 in hospitalized patients, with few publications providing information and advice to those caring for people with diabetes in the primary care setting. This manuscript reviews the current knowledge of the risk and outcomes of individuals with diabetes who are infected with COVID-19 and provides information for PCPs on the importance of glucose control, appropriate treatment, and use of telemedicine and online prescription delivery systems to limit the potentially devastating effects of COVID-19 in people with hyperglycemia.\",\n",
|
||||
" \"authors\": \"Pettus, Jeremy; Skolnik, Neil\",\n",
|
||||
" \"doi\": \"10.1080/00325481.2021.1978704\"\n",
|
||||
"}\n",
|
||||
"{\n",
|
||||
" \"title\": \"Position Statement on How to Manage Patients with Diabetes and COVID-19\",\n",
|
||||
" \"abstract\": null,\n",
|
||||
" \"authors\": null,\n",
|
||||
" \"doi\": \"10.15605/jafes.035.01.03\"\n",
|
||||
"}\n",
|
||||
"{\n",
|
||||
" \"title\": \"Successful blood glucose management of a severe COVID-19 patient with diabetes: A case report\",\n",
|
||||
" \"abstract\": \"RATIONALE: Coronavirus disease 2019 (COVID-19) has emerged as a rapidly spreading communicable disease affecting individuals worldwide. Patients with diabetes are more vulnerable to the disease, and the mortality is higher than in those without diabetes. We reported a severe COVID-19 patient with diabetes and shared our experience with blood glucose management. PATIENT CONCERNS: A 64-year-old female diabetes patient was admitted to the intensive care unit due to productive coughing for 8 days without any obvious cause. The results of blood gas analysis indicated that the partial pressure of oxygen was 84 mm Hg with oxygen 8 L/min, and the oxygenation index was less than 200 mm Hg. In addition, postprandial blood glucose levels were abnormal (29.9 mmol/L). DIAGNOSES: The patient was diagnosed with COVID-19 (severe type) and type 2 diabetes. INTERVENTIONS: Comprehensive interventions including establishing a multidisciplinary team, closely monitoring her blood glucose level, an individualized diabetes diet, early activities, psychological care, etc, were performed to control blood glucose while actively treating COVID-19 infection. OUTCOMES: After the comprehensive measures, the patient's blood glucose level gradually became stable, and the patient was discharged after 20 days of hospitalization. LESSONS: This case indicated that the comprehensive measures performed by a multidisciplinary team achieved good treatment effects on a COVID-19 patient with diabetes. Targeted treatment and nursing methods should be performed based on patients\\u2019 actual situations in clinical practice.\",\n",
|
||||
" \"authors\": \"Hu, Rujun; Gao, Huiming; Huang, Di; Jiang, Deyu; Chen, Fang; Fu, Bao; Yuan, Xiaoli; Li, Jin; Jiang, Zhixia\",\n",
|
||||
" \"doi\": \"10.1097/md.0000000000020844\"\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Note how the same results are returned with additional content fields."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "yTvn0O1v_eHT"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "gzFmQDXfvniJ"
|
||||
},
|
||||
"source": [
|
||||
"# JSON Service Workflow\n",
|
||||
"\n",
|
||||
"The next example builds a workflow that runs a query against a remote URL, retrieves the results, then transforms and indexes the tabular data. This example gets the top results from the [Hacker News front page](https://news.ycombinator.com/). \n",
|
||||
"\n",
|
||||
"Below shows how to build the ServiceTask and prints the first JSON result. Details on how to configure the ServiceTask can be found in [txtai's documentation](https://neuml.github.io/txtai/workflows/)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "bA8SihkeZqbJ",
|
||||
"outputId": "05962fa0-7e79-49a8-b6ac-3d535ace3fad"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.workflow import ServiceTask\n",
|
||||
"\n",
|
||||
"service = ServiceTask(url=\"https://hn.algolia.com/api/v1/search\", method=\"get\", params={\"tags\": None}, batch=False, extract=\"hits\")\n",
|
||||
"workflow = Workflow([service])\n",
|
||||
"\n",
|
||||
"list(workflow([\"front_page\"]))[0][2]"
|
||||
],
|
||||
"execution_count": 71,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'_highlightResult': {'author': {'matchLevel': 'none',\n",
|
||||
" 'matchedWords': [],\n",
|
||||
" 'value': 'cheesestain'},\n",
|
||||
" 'title': {'matchLevel': 'none',\n",
|
||||
" 'matchedWords': [],\n",
|
||||
" 'value': 'Ante: A low-level functional language'},\n",
|
||||
" 'url': {'matchLevel': 'none',\n",
|
||||
" 'matchedWords': [],\n",
|
||||
" 'value': 'https://antelang.org/'}},\n",
|
||||
" '_tags': ['story', 'author_cheesestain', 'story_31775216', 'front_page'],\n",
|
||||
" 'author': 'cheesestain',\n",
|
||||
" 'comment_text': None,\n",
|
||||
" 'created_at': '2022-06-17T07:39:40.000Z',\n",
|
||||
" 'created_at_i': 1655451580,\n",
|
||||
" 'num_comments': 109,\n",
|
||||
" 'objectID': '31775216',\n",
|
||||
" 'parent_id': None,\n",
|
||||
" 'points': 207,\n",
|
||||
" 'story_id': None,\n",
|
||||
" 'story_text': None,\n",
|
||||
" 'story_title': None,\n",
|
||||
" 'story_url': None,\n",
|
||||
" 'title': 'Ante: A low-level functional language',\n",
|
||||
" 'url': 'https://antelang.org/'}"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 71
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Lv_ybw1VwK1N"
|
||||
},
|
||||
"source": [
|
||||
"Next we'll map the JSON data using the tabular pipeline. `url` will be used as the id column and `title` as the text to index."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "YAbwhsaveKo1",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "15c1b28a-3c16-4d38-c586-6b8eae244dbb"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.workflow import Task\n",
|
||||
"\n",
|
||||
"# Create tabular instance mapping input.csv fields\n",
|
||||
"tabular = Tabular(\"url\", [\"title\"])\n",
|
||||
"\n",
|
||||
"# Recreate service applying the tabular pipeline to each result\n",
|
||||
"service = ServiceTask(action=tabular, url=\"https://hn.algolia.com/api/v1/search\", method=\"get\", params={\"tags\": None}, batch=False, extract=\"hits\")\n",
|
||||
"workflow = Workflow([service])\n",
|
||||
"\n",
|
||||
"list(workflow([\"front_page\"]))[2]"
|
||||
],
|
||||
"execution_count": 72,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"('https://antelang.org/', 'Ante: A low-level functional language', None)"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 72
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "SbjMuN5lw63c"
|
||||
},
|
||||
"source": [
|
||||
"As we did previously, let's build an Embeddings index and run a search query."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "5lx9pa65e23E",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "99e05da7-b013-4d17-d82b-c7e13e79a563"
|
||||
},
|
||||
"source": [
|
||||
"# Embeddings with sentence-transformers backend\n",
|
||||
"embeddings = Embeddings({\"method\": \"transformers\", \"path\": \"sentence-transformers/paraphrase-mpnet-base-v2\"})\n",
|
||||
"\n",
|
||||
"# Index Hacker News front page\n",
|
||||
"data = list(workflow([\"front_page\"]))\n",
|
||||
"embeddings.index(data)\n",
|
||||
"\n",
|
||||
"for uid, _ in embeddings.search(\"programming\"):\n",
|
||||
" title = [text for url, text, _ in data if url == uid][0]\n",
|
||||
" print(title, uid)"
|
||||
],
|
||||
"execution_count": 73,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Bundling binary tools in Python wheels https://simonwillison.net/2022/May/23/bundling-binary-tools-in-python-wheels/\n",
|
||||
"Ante: A low-level functional language https://antelang.org/\n",
|
||||
"Adding a Rust compiler front end to GCC [video] https://www.youtube.com/watch?v=R8Pr21nlhig\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "yq26KqfgxJ6Q"
|
||||
},
|
||||
"source": [
|
||||
"# XML Service workflow\n",
|
||||
"\n",
|
||||
"txtai's ServiceTask can consume both JSON and XML. This example runs a query against the [arXiv API](https://arxiv.org/), transforms the results and indexes them for search.\n",
|
||||
"\n",
|
||||
"Below shows how to build the ServiceTask and prints the first XML result."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "K6CbS2QwltGi",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "2d9a0336-9680-4ca1-e2cd-68f6c5b4d1d3"
|
||||
},
|
||||
"source": [
|
||||
"service = ServiceTask(url=\"http://export.arxiv.org/api/query\", method=\"get\", params={\"search_query\": None, \"max_results\": 25}, batch=False, extract=[\"feed\", \"entry\"])\n",
|
||||
"workflow = Workflow([service])\n",
|
||||
"\n",
|
||||
"list(workflow([\"all:aliens\"]))[0][:1]"
|
||||
],
|
||||
"execution_count": 74,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'arxiv:comment': {'#text': 'To appear in Astrophysical Journal',\n",
|
||||
" '@xmlns:arxiv': 'http://arxiv.org/schemas/atom'},\n",
|
||||
" 'arxiv:doi': {'#text': '10.3847/1538-4357/ac2369',\n",
|
||||
" '@xmlns:arxiv': 'http://arxiv.org/schemas/atom'},\n",
|
||||
" 'arxiv:primary_category': {'@scheme': 'http://arxiv.org/schemas/atom',\n",
|
||||
" '@term': 'q-bio.OT',\n",
|
||||
" '@xmlns:arxiv': 'http://arxiv.org/schemas/atom'},\n",
|
||||
" 'author': [{'name': 'Robin Hanson'},\n",
|
||||
" {'name': 'Daniel Martin'},\n",
|
||||
" {'name': 'Calvin McCarter'},\n",
|
||||
" {'name': 'Jonathan Paulson'}],\n",
|
||||
" 'category': [{'@scheme': 'http://arxiv.org/schemas/atom',\n",
|
||||
" '@term': 'q-bio.OT'},\n",
|
||||
" {'@scheme': 'http://arxiv.org/schemas/atom', '@term': 'physics.pop-ph'}],\n",
|
||||
" 'id': 'http://arxiv.org/abs/2102.01522v3',\n",
|
||||
" 'link': [{'@href': 'http://dx.doi.org/10.3847/1538-4357/ac2369',\n",
|
||||
" '@rel': 'related',\n",
|
||||
" '@title': 'doi'},\n",
|
||||
" {'@href': 'http://arxiv.org/abs/2102.01522v3',\n",
|
||||
" '@rel': 'alternate',\n",
|
||||
" '@type': 'text/html'},\n",
|
||||
" {'@href': 'http://arxiv.org/pdf/2102.01522v3',\n",
|
||||
" '@rel': 'related',\n",
|
||||
" '@title': 'pdf',\n",
|
||||
" '@type': 'application/pdf'}],\n",
|
||||
" 'published': '2021-02-01T18:27:12Z',\n",
|
||||
" 'summary': \"If life on Earth had to achieve n 'hard steps' to reach humanity's level,\\nthen the chance of this event rose as time to the n-th power. Integrating this\\nover habitable star formation and planet lifetime distributions predicts >99%\\nof advanced life appears after today, unless n<3 and max planet duration\\n<50Gyr. That is, we seem early. We offer this explanation: a deadline is set by\\n'loud' aliens who are born according to a hard steps power law, expand at a\\ncommon rate, change their volumes' appearances, and prevent advanced life like\\nus from appearing in their volumes. 'Quiet' aliens, in contrast, are much\\nharder to see. We fit this three-parameter model of loud aliens to data: 1)\\nbirth power from the number of hard steps seen in Earth history, 2) birth\\nconstant by assuming a inform distribution over our rank among loud alien birth\\ndates, and 3) expansion speed from our not seeing alien volumes in our sky. We\\nestimate that loud alien civilizations now control 40-50% of universe volume,\\neach will later control ~10^5 - 3x10^7 galaxies, and we could meet them in\\n~200Myr - 2Gyr. If loud aliens arise from quiet ones, a depressingly low\\ntransition chance (~10^-4) is required to expect that even one other quiet\\nalien civilization has ever been active in our galaxy. Which seems bad news for\\nSETI. But perhaps alien volume appearances are subtle, and their expansion\\nspeed lower, in which case we predict many long circular arcs to find in our\\nsky.\",\n",
|
||||
" 'title': 'If Loud Aliens Explain Human Earliness, Quiet Aliens Are Also Rare',\n",
|
||||
" 'updated': '2021-09-06T14:18:23Z'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 74
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "aWFwkxjQyscc"
|
||||
},
|
||||
"source": [
|
||||
"Next we'll map the XML data using the tabular pipeline. `id` will be used as the id column and `title` as the text to index."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "DyIetJ7OmJjP",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "7e730b22-256d-4b1b-9064-2479e26d50c9"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.workflow import Task\n",
|
||||
"\n",
|
||||
"# Create tablular pipeline with new mapping\n",
|
||||
"tabular = Tabular(\"id\", [\"title\"])\n",
|
||||
"\n",
|
||||
"# Recreate service applying the tabular pipeline to each result\n",
|
||||
"service = ServiceTask(action=tabular, url=\"http://export.arxiv.org/api/query\", method=\"get\", params={\"search_query\": None, \"max_results\": 25}, batch=False, extract=[\"feed\", \"entry\"])\n",
|
||||
"workflow = Workflow([service])\n",
|
||||
"\n",
|
||||
"list(workflow([\"all:aliens\"]))[:1]"
|
||||
],
|
||||
"execution_count": 75,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[('http://arxiv.org/abs/2102.01522v3',\n",
|
||||
" 'If Loud Aliens Explain Human Earliness, Quiet Aliens Are Also Rare',\n",
|
||||
" None)]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 75
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "7pFnW7mCyycy"
|
||||
},
|
||||
"source": [
|
||||
"As we did previously, let's build an Embeddings index and run a search query."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "NX2oR5dhm_99",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "399c4316-bce6-443b-c383-aedd7ff9a1ec"
|
||||
},
|
||||
"source": [
|
||||
"# Embeddings with sentence-transformers backend\n",
|
||||
"embeddings = Embeddings({\"method\": \"transformers\", \"path\": \"sentence-transformers/paraphrase-mpnet-base-v2\"})\n",
|
||||
"\n",
|
||||
"# Index Hacker News front page\n",
|
||||
"data = list(workflow([\"all:aliens\"]))\n",
|
||||
"embeddings.index(data)\n",
|
||||
"\n",
|
||||
"for uid, _ in embeddings.search(\"alien radio signals\"):\n",
|
||||
" title = [text for url, text, _ in data if url == uid][0]\n",
|
||||
" print(title, uid)"
|
||||
],
|
||||
"execution_count": 76,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Calculating the probability of detecting radio signals from alien\n",
|
||||
" civilizations http://arxiv.org/abs/0707.0011v2\n",
|
||||
"Field Trial of Alien Wavelengths on GARR Optical Network http://arxiv.org/abs/1805.04278v1\n",
|
||||
"Aliens on Earth. Are reports of close encounters correct? http://arxiv.org/abs/1203.6805v2\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "8xknLo2ey0vZ"
|
||||
},
|
||||
"source": [
|
||||
"# Build a workflow with no code!\n",
|
||||
"\n",
|
||||
"The next example shows how one of the same workflows above can be constructed via API configuration. This is a no-code way to build a txtai indexing workflow!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "1eF5IJlzpNbw",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "4ff0bfdc-fb01-46a6-b61a-142a370dc273"
|
||||
},
|
||||
"source": [
|
||||
"%%writefile workflow.yml\n",
|
||||
"# Index settings\n",
|
||||
"writable: true\n",
|
||||
"embeddings:\n",
|
||||
" path: sentence-transformers/nli-mpnet-base-v2\n",
|
||||
"\n",
|
||||
"# Tabular pipeline\n",
|
||||
"tabular:\n",
|
||||
" idcolumn: id\n",
|
||||
" textcolumns: \n",
|
||||
" - title\n",
|
||||
"\n",
|
||||
"# Workflow definitions\n",
|
||||
"workflow:\n",
|
||||
" index:\n",
|
||||
" tasks:\n",
|
||||
" - task: service\n",
|
||||
" action: tabular\n",
|
||||
" url: http://export.arxiv.org/api/query?max_results=25\n",
|
||||
" method: get\n",
|
||||
" params:\n",
|
||||
" search_query: null\n",
|
||||
" batch: false\n",
|
||||
" extract: [feed, entry]\n",
|
||||
" - action: upsert"
|
||||
],
|
||||
"execution_count": 77,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing workflow.yml\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "dxi5w3IezR7Q"
|
||||
},
|
||||
"source": [
|
||||
"This workflow once again runs an arXiv query and indexes article titles. The workflow configures the same actions that were configured in Python previously. \n",
|
||||
"\n",
|
||||
"Let's start an API instance "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "B1DQyB5ErIzr",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "7e170445-bc93-4660-b944-8b10b8b02d98"
|
||||
},
|
||||
"source": [
|
||||
"!killall -9 uvicorn\n",
|
||||
"!CONFIG=workflow.yml nohup uvicorn \"txtai.api:app\" &> api.log &\n",
|
||||
"!sleep 30\n",
|
||||
"!cat api.log"
|
||||
],
|
||||
"execution_count": 78,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"INFO: Started server process [754]\n",
|
||||
"2022-06-17 15:05:58,554 [INFO] serve: Started server process [754]\n",
|
||||
"INFO: Waiting for application startup.\n",
|
||||
"2022-06-17 15:05:58,554 [INFO] startup: Waiting for application startup.\n",
|
||||
"INFO: Application startup complete.\n",
|
||||
"2022-06-17 15:06:07,707 [INFO] startup: Application startup complete.\n",
|
||||
"INFO: Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)\n",
|
||||
"2022-06-17 15:06:07,707 [INFO] _log_started_message: Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "45JaR7Nr0Zmg"
|
||||
},
|
||||
"source": [
|
||||
"Next we'll execute the workflow. txtai has API bindings for [JavaScript](https://github.com/neuml/txtai.js), [Java](https://github.com/neuml/txtai.java), [Rust](https://github.com/neuml/txtai.rs) and [Golang](https://github.com/neuml/txtai.go). But to keep things simple, we'll just run the commands via cURL. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "xt_qL6eA0SrS",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "3ca70a31-cf48-42b4-949f-f01aac14505c"
|
||||
},
|
||||
"source": [
|
||||
"# Execute workflow via API call\n",
|
||||
"!curl -X POST \"http://localhost:8000/workflow\" -H \"accept: application/json\" -H \"Content-Type: application/json\" -d \"{\\\"name\\\":\\\"index\\\",\\\"elements\\\":[\\\"all:aliens\\\"]}\""
|
||||
],
|
||||
"execution_count": 79,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"[[\"http://arxiv.org/abs/2102.01522v3\",\"If Loud Aliens Explain Human Earliness, Quiet Aliens Are Also Rare\",null],[\"http://arxiv.org/abs/cs/0306071v1\",\"AliEnFS - a Linux File System for the AliEn Grid Services\",null],[\"http://arxiv.org/abs/physics/0306103v1\",\"AliEn - EDG Interoperability in ALICE\",null],[\"http://arxiv.org/abs/2103.05559v1\",\"Oumuamua Is Not a Probe Sent to our Solar System by an Alien\\n Civilization\",null],[\"http://arxiv.org/abs/1403.3979v1\",\"Robust transitivity and density of periodic points of partially\\n hyperbolic diffeomorphisms\",null],[\"http://arxiv.org/abs/1712.09210v1\",\"Sampling alien species inside and outside protected areas: does it\\n matter?\",null],[\"http://arxiv.org/abs/cs/0306067v1\",\"The AliEn system, status and perspectives\",null],[\"http://arxiv.org/abs/0707.0011v2\",\"Calculating the probability of detecting radio signals from alien\\n civilizations\",null],[\"http://arxiv.org/abs/1805.04278v1\",\"Field Trial of Alien Wavelengths on GARR Optical Network\",null],[\"http://arxiv.org/abs/1808.00529v1\",\"Open Category Detection with PAC Guarantees\",null],[\"http://arxiv.org/abs/1206.3640v1\",\"The Study of Climate on Alien Worlds\",null],[\"http://arxiv.org/abs/1203.6805v2\",\"Aliens on Earth. Are reports of close encounters correct?\",null],[\"http://arxiv.org/abs/1604.05078v1\",\"The Imprecise Search for Habitability\",null],[\"http://arxiv.org/abs/1006.2613v1\",\"Resurgence, Stokes phenomenon and alien derivatives for level-one linear\\n differential systems\",null],[\"http://arxiv.org/abs/1307.0653v1\",\"General and alien solutions of a functional equation and of a functional\\n inequality\",null],[\"http://arxiv.org/abs/1705.03394v1\",\"That is not dead which can eternal lie: the aestivation hypothesis for\\n resolving Fermi's paradox\",null],[\"http://arxiv.org/abs/1701.02294v1\",\"Alien Calculus and non perturbative effects in Quantum Field Theory\",null],[\"http://arxiv.org/abs/1801.06180v1\",\"Are Alien Civilizations Technologically Advanced?\",null],[\"http://arxiv.org/abs/1902.05387v1\",\"Simultaneous x, y Pixel Estimation and Feature Extraction for Multiple\\n Small Objects in a Scene: A Description of the ALIEN Network\",null],[\"http://arxiv.org/abs/0711.4034v1\",\"The q-analogue of the wild fundamental group (II)\",null],[\"http://arxiv.org/abs/2111.07895v1\",\"Research Programs Arising from 'Oumuamua Considered as an Alien Craft\",null],[\"http://arxiv.org/abs/2112.15226v1\",\"Variations on the Resurgence of the Gamma Function\",null],[\"http://arxiv.org/abs/astro-ph/0501119v1\",\"Expanding advanced civilizations in the universe\",null],[\"http://arxiv.org/abs/cs/0306068v1\",\"AliEn Resource Brokers\",null],[\"http://arxiv.org/abs/hep-ph/9403231v2\",\"The Renormalization of Composite Operators in Yang-Mills Theories Using\\n General Covariant Gauge\",null]]"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "_bwn4KBt1Cos"
|
||||
},
|
||||
"source": [
|
||||
"The data is now indexed. Note that the index configuration has an `upsert` action. Each workflow call will insert new rows or update existing rows. This call could be scheduled with a system cron to execute periodically and build an index of arXiv article titles. \n",
|
||||
"\n",
|
||||
"Now that the index is ready, let's run a search."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "qbteIueJ1Fds",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "ad8ab84b-36a3-4346-d387-64321b933d25"
|
||||
},
|
||||
"source": [
|
||||
"# Run a search\n",
|
||||
"!curl -X GET \"http://localhost:8000/search?query=radio&limit=3\" -H \"accept: application/json\""
|
||||
],
|
||||
"execution_count": 80,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"[{\"id\":\"http://arxiv.org/abs/0707.0011v2\",\"score\":0.40350058674812317},{\"id\":\"http://arxiv.org/abs/1805.04278v1\",\"score\":0.3406212031841278},{\"id\":\"http://arxiv.org/abs/1902.05387v1\",\"score\":0.22262491285800934}]"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "C5_Tt6EA3Cxb"
|
||||
},
|
||||
"source": [
|
||||
"# Add a translation step to workflow\n",
|
||||
"\n",
|
||||
"Next we'll recreate the workflow, adding one additional step, translating the text into French before indexing. This workflow runs an arXiv query, translates the results and builds an semantic index of titles in French. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "j8rBVl17293q",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "215dc2cf-3a57-4deb-820d-e8ac03e3e041"
|
||||
},
|
||||
"source": [
|
||||
"%%writefile workflow.yml\n",
|
||||
"# Index settings\n",
|
||||
"writable: true\n",
|
||||
"embeddings:\n",
|
||||
" path: sentence-transformers/nli-mpnet-base-v2\n",
|
||||
"\n",
|
||||
"# Tabular pipeline\n",
|
||||
"tabular:\n",
|
||||
" idcolumn: id\n",
|
||||
" textcolumns: \n",
|
||||
" - title\n",
|
||||
"\n",
|
||||
"# Translation pipeline\n",
|
||||
"translation:\n",
|
||||
"\n",
|
||||
"# Workflow definitions\n",
|
||||
"workflow:\n",
|
||||
" index:\n",
|
||||
" tasks:\n",
|
||||
" - task: service\n",
|
||||
" action: tabular\n",
|
||||
" url: http://export.arxiv.org/api/query?max_results=25\n",
|
||||
" method: get\n",
|
||||
" params:\n",
|
||||
" search_query: null\n",
|
||||
" batch: false\n",
|
||||
" extract: [feed, entry]\n",
|
||||
" - action: translation\n",
|
||||
" args: [fr]\n",
|
||||
" - action: upsert"
|
||||
],
|
||||
"execution_count": 81,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Overwriting workflow.yml\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "UQWvvgb2CwgG",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "3b0c99d7-b5f6-4c81-eb6e-ac9055eaaae4"
|
||||
},
|
||||
"source": [
|
||||
"!killall -9 uvicorn\n",
|
||||
"!CONFIG=workflow.yml nohup uvicorn \"txtai.api:app\" &> api.log &\n",
|
||||
"!sleep 30\n",
|
||||
"!cat api.log"
|
||||
],
|
||||
"execution_count": 82,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"INFO: Started server process [775]\n",
|
||||
"2022-06-17 15:06:29,397 [INFO] serve: Started server process [775]\n",
|
||||
"INFO: Waiting for application startup.\n",
|
||||
"2022-06-17 15:06:29,397 [INFO] startup: Waiting for application startup.\n",
|
||||
"INFO: Application startup complete.\n",
|
||||
"2022-06-17 15:06:40,198 [INFO] startup: Application startup complete.\n",
|
||||
"INFO: Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)\n",
|
||||
"2022-06-17 15:06:40,199 [INFO] _log_started_message: Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "v1y4heYx679i"
|
||||
},
|
||||
"source": [
|
||||
"Same as before, we'll run the index workflow and a search"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "npW3rjCw6_nx",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "fba0bd30-2340-4197-d68c-4fc925ce80b8"
|
||||
},
|
||||
"source": [
|
||||
"# Execute workflow via API call\n",
|
||||
"!curl -s -X POST \"http://localhost:8000/workflow\" -H \"accept: application/json\" -H \"Content-Type: application/json\" -d \"{\\\"name\\\":\\\"index\\\",\\\"elements\\\":[\\\"all:aliens\\\"]}\" > /dev/null\n",
|
||||
"\n",
|
||||
"# Run a search\n",
|
||||
"!curl -X GET \"http://localhost:8000/search?query=radio&limit=3\" -H \"accept: application/json\""
|
||||
],
|
||||
"execution_count": 83,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"[{\"id\":\"http://arxiv.org/abs/0707.0011v2\",\"score\":0.532800555229187},{\"id\":\"http://arxiv.org/abs/0711.4034v1\",\"score\":0.24413327872753143},{\"id\":\"http://arxiv.org/abs/2102.01522v3\",\"score\":0.22881504893302917}]"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "OsQO5DG9zBQF"
|
||||
},
|
||||
"source": [
|
||||
"# Run YAML workflow in Python\n",
|
||||
"\n",
|
||||
"Workflow YAML files can also be directly executed in Python. In this case, all input data is passed locally in Python and not through network interfaces. The following section shows how to do this!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "iXmkLDTlzPT3",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "6b795963-2954-42e2-c42d-39c3f88c9ac3"
|
||||
},
|
||||
"source": [
|
||||
"import yaml\n",
|
||||
"\n",
|
||||
"from txtai.app import Application\n",
|
||||
"\n",
|
||||
"with open(\"workflow.yml\") as config:\n",
|
||||
" workflow = yaml.safe_load(config)\n",
|
||||
"\n",
|
||||
"app = Application(workflow)\n",
|
||||
"\n",
|
||||
"# Run the workflow\n",
|
||||
"data = list(app.workflow(\"index\", [\"all:aliens\"]))\n",
|
||||
"\n",
|
||||
"# Run a search\n",
|
||||
"for result in app.search(\"radio\", None):\n",
|
||||
" text = [row[1] for row in data if row[0] == result[\"id\"]][0]\n",
|
||||
" print(result[\"id\"], result[\"score\"], text)"
|
||||
],
|
||||
"execution_count": 84,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"http://arxiv.org/abs/0707.0011v2 0.532800555229187 Calcul de la probabilité de détection des signaux radio de l'étrangercivilisations\n",
|
||||
"http://arxiv.org/abs/0711.4034v1 0.24413327872753143 Le q-analogue du groupe fondamental sauvage (II)\n",
|
||||
"http://arxiv.org/abs/2102.01522v3 0.22881504893302917 Si les étrangers louds expliquent le début de l'humanité, les étrangers tranquilles sont aussi rares\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "EoQFEi_61P9O"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notework demonstrated how to transform, index and search tabular data from a variety of sources. txtai offers maximum flexibility in building composable workflows to maximize the number of ways data can be indexed for semantic search. "
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,503 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"name": "23 - Tensor workflows",
|
||||
"provenance": [],
|
||||
"collapsed_sections": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Pjmz-RORV8E"
|
||||
},
|
||||
"source": [
|
||||
"# Tensor workflows\n",
|
||||
"\n",
|
||||
"Many of the examples and use cases for txtai focus on transforming text. Makes sense as txt is even in the name! But that doesn't mean txtai only works with text.\n",
|
||||
"\n",
|
||||
"This notebook will cover examples of how to efficiently process tensors using txtai workflows."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Dk31rbYjSTYm"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies. We will install the api, pipeline and workflow optional extras packages, along with the datasets package. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "XMQuuun2R06J"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[api,pipeline,workflow] datasets"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "NSYrP0hjtR_E"
|
||||
},
|
||||
"source": [
|
||||
"# Transform large tensor arrays\n",
|
||||
"\n",
|
||||
"The first section attempts to apply a simple transform to a very large memory-mapped array (2,000,000 x 1024)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "BoPJIKWoTibk",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 220
|
||||
},
|
||||
"outputId": "143a6e4d-fe56-4353-e8ee-595ddfc12249"
|
||||
},
|
||||
"source": [
|
||||
"import numpy as np\n",
|
||||
"import torch\n",
|
||||
"\n",
|
||||
"# Generate large memory-mapped array\n",
|
||||
"rows, cols = 2000000, 1024\n",
|
||||
"data = np.memmap(\"data.npy\", dtype=np.float32, mode=\"w+\", shape=(rows, cols))\n",
|
||||
"del data\n",
|
||||
"\n",
|
||||
"# Open memory-mapped array\n",
|
||||
"data = np.memmap(\"data.npy\", dtype=np.float32, shape=(rows, cols))\n",
|
||||
"\n",
|
||||
"# Create tensor\n",
|
||||
"tensor = torch.from_numpy(data).to(\"cuda:0\")\n",
|
||||
"\n",
|
||||
"# Apply tanh transform to tensor\n",
|
||||
"torch.tanh(tensor).shape"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "error",
|
||||
"ename": "RuntimeError",
|
||||
"evalue": "ignored",
|
||||
"traceback": [
|
||||
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
|
||||
"\u001b[0;31mRuntimeError\u001b[0m Traceback (most recent call last)",
|
||||
"\u001b[0;32m<ipython-input-2-a1fc94fedb69>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m 14\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 15\u001b[0m \u001b[0;31m# Apply tanh transform to tensor\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 16\u001b[0;31m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtanh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
|
||||
"\u001b[0;31mRuntimeError\u001b[0m: CUDA out of memory. Tried to allocate 7.63 GiB (GPU 0; 11.17 GiB total capacity; 7.63 GiB already allocated; 3.04 GiB free; 7.63 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "O8mKzPP01d_m",
|
||||
"outputId": "929226a8-6948-4d17-ab70-025da2081abd"
|
||||
},
|
||||
"source": [
|
||||
"!ls -l --block-size=MB data.npy"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"-rw-r--r-- 1 root root 8192MB Dec 6 23:24 data.npy\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "vuObmAJ9FaJe"
|
||||
},
|
||||
"source": [
|
||||
"Not surprisingly this runs out of CUDA memory. The array needs `2,000,000 * 1024 * 4 = 8GB` which exceeds the amount of GPU memory available.\n",
|
||||
"\n",
|
||||
"One of the great things about NumPy and PyTorch arrays is that they can be sliced without having to copy data. Additionally, PyTorch has methods to work directly on NumPy arrays without copying data, in other words both NumPy arrays and PyTorch arrays can share the same memory. This opens the door to efficient processing of tensor data in place. \n",
|
||||
"\n",
|
||||
"Let's try applying a simple tanh transform in batches over the array."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ciD6unQYD-bJ",
|
||||
"outputId": "d3b6a0c5-aea5-451d-d3e3-60d04ef33a9e"
|
||||
},
|
||||
"source": [
|
||||
"def process(x):\n",
|
||||
" print(x.shape)\n",
|
||||
" return torch.tanh(torch.from_numpy(x).to(\"cuda:0\")).cpu().numpy()\n",
|
||||
"\n",
|
||||
"# Split into 250,000 rows per call\n",
|
||||
"batch = 250000\n",
|
||||
"count = 0\n",
|
||||
"for x in range(0, len(data), batch):\n",
|
||||
" for row in process(data[x : x + batch]):\n",
|
||||
" count += 1\n",
|
||||
"\n",
|
||||
"print(count)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"2000000\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "uZBzEMsRHpsi"
|
||||
},
|
||||
"source": [
|
||||
"Iterating over the data array and selecting slices to operate on allows the transform to complete successfully! Each `torch.from_numpy` call is building a view of a portion the existing large NumPy data array. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Oe7X17vbHJRV"
|
||||
},
|
||||
"source": [
|
||||
"# Enter workflows\n",
|
||||
"\n",
|
||||
"The next section takes the same array and shows how workflows can apply transformations to tensors. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ymqr92kW9hxd",
|
||||
"outputId": "e4a4c7d1-be54-46c2-bc7c-dd849adfeb7e"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.workflow import Task, Workflow\n",
|
||||
"\n",
|
||||
"# Create workflow with a single task calling process for each batch\n",
|
||||
"task = Task(process)\n",
|
||||
"workflow = Workflow([task], batch)\n",
|
||||
"\n",
|
||||
"# Run workflow\n",
|
||||
"count = 0\n",
|
||||
"for row in workflow(data):\n",
|
||||
" count += 1\n",
|
||||
"\n",
|
||||
"print(count)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"(250000, 1024)\n",
|
||||
"2000000\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "B9qC8qUbHjfk"
|
||||
},
|
||||
"source": [
|
||||
"Workflows process the data in the same fashion as the code in the previous section. On top of that, workflows can handle text, images, video, audio, document, tensors and more. Workflow graphs can also be connected together to handle complex use cases."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "wCRD9ERoJvsG"
|
||||
},
|
||||
"source": [
|
||||
"# Workflows with PyTorch models\n",
|
||||
"\n",
|
||||
"The next example applies a PyTorch model to the same data. The model applies a series of transforms and outputs a single float per row."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "N9UTfSTTIDaO",
|
||||
"outputId": "000c7e13-a249-43d0-f419-28ddd62e8ba1"
|
||||
},
|
||||
"source": [
|
||||
"from torch import nn\n",
|
||||
"\n",
|
||||
"class Model(nn.Module):\n",
|
||||
" def __init__(self):\n",
|
||||
" super().__init__()\n",
|
||||
"\n",
|
||||
" self.gelu = nn.ReLU()\n",
|
||||
" self.linear1 = nn.Linear(1024, 512)\n",
|
||||
" self.dropout = nn.Dropout(0.5)\n",
|
||||
" self.norm = nn.LayerNorm(512)\n",
|
||||
" self.linear2 = nn.Linear(512, 1)\n",
|
||||
"\n",
|
||||
" def forward(self, inputs):\n",
|
||||
" outputs = self.gelu(inputs)\n",
|
||||
" outputs = self.linear1(outputs)\n",
|
||||
" outputs = self.dropout(outputs)\n",
|
||||
" outputs = self.norm(outputs)\n",
|
||||
" outputs = self.linear2(outputs)\n",
|
||||
"\n",
|
||||
" return outputs\n",
|
||||
"\n",
|
||||
"model = Model().to(\"cuda:0\")\n",
|
||||
"\n",
|
||||
"def process(x):\n",
|
||||
" with torch.no_grad():\n",
|
||||
" outputs = model(torch.from_numpy(x).to(\"cuda:0\")).cpu().numpy()\n",
|
||||
" print(outputs.shape)\n",
|
||||
" return outputs\n",
|
||||
"\n",
|
||||
"# Create workflow with a single task calling model for each batch\n",
|
||||
"task = Task(process)\n",
|
||||
"workflow = Workflow([task], batch)\n",
|
||||
"\n",
|
||||
"# Run workflow\n",
|
||||
"count = 0\n",
|
||||
"for row in workflow(data):\n",
|
||||
" count += 1\n",
|
||||
"\n",
|
||||
"print(count)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"(250000, 1)\n",
|
||||
"(250000, 1)\n",
|
||||
"(250000, 1)\n",
|
||||
"(250000, 1)\n",
|
||||
"(250000, 1)\n",
|
||||
"(250000, 1)\n",
|
||||
"(250000, 1)\n",
|
||||
"(250000, 1)\n",
|
||||
"2000000\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Q1ivX4eBuU8T"
|
||||
},
|
||||
"source": [
|
||||
"Once again the data can be processed in batches using workflows, even with a more complex model. Let's try a more interesting example."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KoSB1mKzUnb0"
|
||||
},
|
||||
"source": [
|
||||
"# Workflows in parallel\n",
|
||||
"\n",
|
||||
"Workflows consist of a series of tasks. Each task can output one to many outputs per input element. Multi-output tasks have options available to [merge the data](https://neuml.github.io/txtai/workflow/task/#multi-action-task-merges) for downstream tasks.\n",
|
||||
"\n",
|
||||
"The following example builds a workflow with a task having three separate actions. Each action takes text as an input an applies a sentiment classifier. This is followed by a task that merges the three outputs for each row using a mean transform. Essentially, this workflow builds a weighted sentiment classifier using the outputs of three models. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "JlCdVgo_LXOl"
|
||||
},
|
||||
"source": [
|
||||
"import time\n",
|
||||
"\n",
|
||||
"from datasets import load_dataset\n",
|
||||
"from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
|
||||
"\n",
|
||||
"class Tokens:\n",
|
||||
" def __init__(self, texts):\n",
|
||||
" tokenizer = AutoTokenizer.from_pretrained(\"distilbert-base-uncased-finetuned-sst-2-english\")\n",
|
||||
" tokens = tokenizer(texts, padding=True, return_tensors=\"pt\").to(\"cuda:0\")\n",
|
||||
"\n",
|
||||
" self.inputs, self.attention = tokens[\"input_ids\"], tokens[\"attention_mask\"]\n",
|
||||
"\n",
|
||||
" def __len__(self):\n",
|
||||
" return len(self.inputs)\n",
|
||||
"\n",
|
||||
" def __getitem__(self, value):\n",
|
||||
" return (self.inputs[value], self.attention[value])\n",
|
||||
"\n",
|
||||
"class Classify:\n",
|
||||
" def __init__(self, model):\n",
|
||||
" self.model = model\n",
|
||||
"\n",
|
||||
" def __call__(self, tokens):\n",
|
||||
" with torch.no_grad():\n",
|
||||
" inputs, attention = tokens\n",
|
||||
" outputs = self.model(input_ids=inputs, attention_mask=attention)\n",
|
||||
" outputs = outputs[\"logits\"]\n",
|
||||
"\n",
|
||||
" return outputs\n",
|
||||
"\n",
|
||||
"# Load reviews from the rotten tomatoes dataset\n",
|
||||
"ds = load_dataset(\"rotten_tomatoes\")\n",
|
||||
"texts = ds[\"train\"][\"text\"]\n",
|
||||
"\n",
|
||||
"tokens = Tokens(texts)\n",
|
||||
"\n",
|
||||
"model1 = AutoModelForSequenceClassification.from_pretrained(\"M-FAC/bert-tiny-finetuned-sst2\")\n",
|
||||
"model1 = model1.to(\"cuda:0\")\n",
|
||||
"\n",
|
||||
"model2 = AutoModelForSequenceClassification.from_pretrained(\"howey/electra-base-sst2\")\n",
|
||||
"model2 = model2.to(\"cuda:0\")\n",
|
||||
"\n",
|
||||
"model3 = AutoModelForSequenceClassification.from_pretrained(\"philschmid/MiniLM-L6-H384-uncased-sst2\")\n",
|
||||
"model3 = model3.to(\"cuda:0\")\n",
|
||||
"\n",
|
||||
"task1 = Task([Classify(model1), Classify(model2), Classify(model3)])\n",
|
||||
"task2 = Task([lambda x: torch.sigmoid(x).mean(axis=1).cpu().numpy()])\n",
|
||||
"\n",
|
||||
"workflow = Workflow([task1, task2], 250)\n",
|
||||
"\n",
|
||||
"start = time.time()\n",
|
||||
"for x in workflow(tokens):\n",
|
||||
" pass\n",
|
||||
"\n",
|
||||
"print(f\"Took {time.time() - start} seconds\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"Using custom data configuration default\n",
|
||||
"Reusing dataset rotten_tomatoes_movie_review (/root/.cache/huggingface/datasets/rotten_tomatoes_movie_review/default/1.0.0/40d411e45a6ce3484deed7cc15b82a53dad9a72aafd9f86f8f227134bec5ca46)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Took 84.73194456100464 seconds\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "eBQzLrtAVUtB"
|
||||
},
|
||||
"source": [
|
||||
"Note that while the task actions are parallel, that doesn't necessarily mean the operations are concurrent. In the case above, the actions are are executed sequentially.\n",
|
||||
"\n",
|
||||
"Workflows have an additional option to run task actions concurrently. The two supported modes are \"thread\" and \"process\". I/O bound actions will do better with multithreading and CPU bound actions will do better with multiprocessing. More can be read in the [txtai documentation](https://neuml.github.io/txtai/workflow/task/#multi-action-task-concurrency). "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "AB0onoOlVT-e",
|
||||
"outputId": "a072d8a6-3b3b-4066-8881-8af9a8b96608"
|
||||
},
|
||||
"source": [
|
||||
"task1 = Task([Classify(model1), Classify(model2), Classify(model3)], concurrency=\"thread\")\n",
|
||||
"task2 = Task([lambda x: torch.sigmoid(x).mean(axis=1).cpu().numpy()])\n",
|
||||
"\n",
|
||||
"workflow = Workflow([task1, task2], 250)\n",
|
||||
"\n",
|
||||
"start = time.time()\n",
|
||||
"for x in workflow(tokens):\n",
|
||||
" pass\n",
|
||||
"\n",
|
||||
"print(f\"Took {time.time() - start} seconds\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Took 85.21102929115295 seconds\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "5s2KexhG_udx"
|
||||
},
|
||||
"source": [
|
||||
"In this case, concurrency doesn't improve performance. While the [GIL](https://wiki.python.org/moin/GlobalInterpreterLock) is a factor, a bigger factor is that the GPU is already fully loaded. This method would be more beneficial if the system had a second GPU or the primary GPU had idle cycles. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "EoQFEi_61P9O"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook introduced a number of different ways to work with large-scale tensor data and process it efficiently. This notebook purposely didn't cover embeddings and pipelines to demonstrate how workflows can stand on their own. In addition to workflows, this notebook covered efficient methods to work with large tensor arrays in PyTorch and NumPy."
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1,270 @@
|
||||
{
|
||||
"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": "26 - Entity extraction workflows",
|
||||
"provenance": [],
|
||||
"collapsed_sections": []
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "POWZoSJR6XzK"
|
||||
},
|
||||
"source": [
|
||||
"# Entity extraction workflows\n",
|
||||
"\n",
|
||||
"Entity extraction is the process of identifying names, locations, organizations and other entity-like tokens in unstructured text. Entity extraction can organize data into topics and/or feed downstream machine learning pipelines.\n",
|
||||
"\n",
|
||||
"This notebook will show how to use the entity extraction pipeline in txtai with workflows."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "qa_PPKVX6XzN"
|
||||
},
|
||||
"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"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Extract entities\n",
|
||||
"\n",
|
||||
"Let's get right to it! The following example creates an entity pipeline and extracts entities from text. \n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "0p3WCDniUths"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a",
|
||||
"_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
|
||||
"trusted": true,
|
||||
"id": "2j_CFGDR6Xzp",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "6276538e-5b06-4038-ae0f-0aa642e86814"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.pipeline import Entity\n",
|
||||
"\n",
|
||||
"data = [\"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",
|
||||
"entity = Entity()\n",
|
||||
"\n",
|
||||
"for x, e in enumerate(entity(data)):\n",
|
||||
" print(data[x])\n",
|
||||
" print(f\" {e}\", \"\\n\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"US tops 5 million confirmed virus cases\n",
|
||||
" [('US', 'LOC', 0.999273955821991)] \n",
|
||||
"\n",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\n",
|
||||
" [('Canada', 'LOC', 0.999609649181366), ('Manhattan', 'MISC', 0.651396632194519)] \n",
|
||||
"\n",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n",
|
||||
" [('Beijing', 'LOC', 0.9996659755706787), ('Taiwan', 'LOC', 0.9996755123138428)] \n",
|
||||
"\n",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack\n",
|
||||
" [('National Park Service', 'ORG', 0.9993489384651184)] \n",
|
||||
"\n",
|
||||
"Maine man wins $1M from $25 lottery ticket\n",
|
||||
" [('Maine', 'LOC', 0.9987521171569824)] \n",
|
||||
"\n",
|
||||
"Make huge profits without work, earn up to $100,000 a day\n",
|
||||
" [] \n",
|
||||
"\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"The section above is running an entity extraction pipeline for each row in data. The outputs are the token(s) identified as part of an entity, the type of entity and score or confidence in the prediction."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "bVd_qHh97IlJ"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Feed entities to a workflow\n",
|
||||
"\n",
|
||||
"The next section demonstrates how the entity extraction pipeline can be used as part of a workflow. This workflow uses the output entities and builds an embeddings index for each row. This effectively computes entity embeddings to compare the row similarity with a focus on mentioned entities."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "fY71Dyyt8Arv"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai.embeddings import Embeddings, Documents\n",
|
||||
"from txtai.workflow import Workflow, Task\n",
|
||||
"\n",
|
||||
"# Create workflow with an entity pipeline output into a documents collection\n",
|
||||
"documents = Documents()\n",
|
||||
"workflow = Workflow([Task(lambda x: entity(x, flatten=True, join=True)), Task(documents.add, unpack=False)])\n",
|
||||
"\n",
|
||||
"# Run workflow\n",
|
||||
"for _ in workflow([(x, row, None) for x, row in enumerate(data)]):\n",
|
||||
" pass\n",
|
||||
"\n",
|
||||
"embeddings = Embeddings({\"path\": \"sentence-transformers/nli-mpnet-base-v2\"})\n",
|
||||
"embeddings.index(documents)\n",
|
||||
"\n",
|
||||
"for query in [\"North America\", \"Asia Pacific\"]:\n",
|
||||
" index = embeddings.search(query, 1)[0][0]\n",
|
||||
" print(query, \"\\t\", data[index])"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "yB3e6HVM8dRJ",
|
||||
"outputId": "8424ce5d-9652-421d-d0bd-261b911ea7d6"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"North America \t Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\n",
|
||||
"Asia Pacific \t Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Run workflow YAML\n",
|
||||
"\n",
|
||||
"Below is the same example using workflow YAML."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "u_0WiV1NvbCK"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"workflow = \"\"\"\n",
|
||||
"writable: true\n",
|
||||
"embeddings:\n",
|
||||
" path: sentence-transformers/nli-mpnet-base-v2\n",
|
||||
"\n",
|
||||
"entity:\n",
|
||||
"\n",
|
||||
"workflow:\n",
|
||||
" index:\n",
|
||||
" tasks:\n",
|
||||
" - action: entity\n",
|
||||
" args: [null, \"simple\", true, true]\n",
|
||||
" - action: index\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"from txtai.app import Application\n",
|
||||
"\n",
|
||||
"# Create and run workflow\n",
|
||||
"app = Application(workflow)\n",
|
||||
"for _ in app.workflow(\"index\", [(x, row, None) for x, row in enumerate(data)]):\n",
|
||||
" pass\n",
|
||||
"\n",
|
||||
"# Run queries\n",
|
||||
"for query in [\"North America\", \"Asia Pacific\"]:\n",
|
||||
" index = app.search(query)[0][\"id\"]\n",
|
||||
" print(query, \"\\t\", data[index])"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "3qmRdU7LvhQ3",
|
||||
"outputId": "e01453fc-5bbf-41a6-e557-64fa453f80ba"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"North America \t Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\n",
|
||||
"Asia Pacific \t Beijing mobilises invasion craft along coast as Taiwan tensions escalate\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook introduced entity extraction pipelines with txtai. This pipeline supports a number of different configurations to help feed downstream systems and/or directly use the entities.\n",
|
||||
"\n",
|
||||
"As with other pipelines, the entity extraction pipeline can be used standalone in Python, as an API service or as part of a workflow!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "i3kkN5Vpx8ks"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,442 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "27 - Workflow scheduling",
|
||||
"provenance": [],
|
||||
"collapsed_sections": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "LjmhJ4ad9kBL"
|
||||
},
|
||||
"source": [
|
||||
"# Workflow scheduling\n",
|
||||
"\n",
|
||||
"Workflows are a simple yet powerful construct that takes a callable and returns elements. They are streaming and work on data in batches, allowing large volumes of data to be processed efficiently. When working with streaming data, workflows continually run until the data stream is exhausted. \n",
|
||||
"\n",
|
||||
"Workflows can also be scheduled to run. In this case, a static set of elements, dynamically expands. For example, an API service endpoint that returns items, or polling a directory with files coming in and out. \n",
|
||||
"\n",
|
||||
"This notebook will show how to use workflow scheduling in txtai."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "8tLWvo9v-Q0u"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "Fa5BCjMFqVKE"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install datasets git+https://github.com/neuml/txtai#egg=txtai[workflow]"
|
||||
],
|
||||
"execution_count": 6,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Create workflow action\n",
|
||||
"\n",
|
||||
"Workflows run a series of tasks to transform and process data. This section creates a callable object that can be used as a workflow action. The object iterates over a dataset, returning a batch of data."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "EJ_hHmQtRgQM"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "3hYRk9JnsM0J"
|
||||
},
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"class Stream:\n",
|
||||
" def __init__(self):\n",
|
||||
" self.dataset = load_dataset(\"ag_news\", split=\"train\")\n",
|
||||
" self.index, self.size = 0, 2500\n",
|
||||
"\n",
|
||||
" def __call__(self, fields):\n",
|
||||
" outputs = []\n",
|
||||
" for field in fields:\n",
|
||||
" output = []\n",
|
||||
" for row in self.dataset.select(range(self.index, self.index+self.size)):\n",
|
||||
" output.append((self.index, row[field], None))\n",
|
||||
" self.index += 1\n",
|
||||
"\n",
|
||||
" outputs.append(output)\n",
|
||||
"\n",
|
||||
" return outputs"
|
||||
],
|
||||
"execution_count": 7,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Build workflow\n",
|
||||
"\n",
|
||||
"Next we'll create the workflow. The workflow reads batches of data from a stream and loads it into an Embeddings index. We'll run this workflow four times on a scheduled interval to demonstrate a scheduled workflow."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "_B4YFu-1R2QC"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai.app import Application\n",
|
||||
"\n",
|
||||
"# Run up to every 5 seconds 4 times\n",
|
||||
"workflow = \"\"\"\n",
|
||||
"writable: true\n",
|
||||
"embeddings:\n",
|
||||
" path: sentence-transformers/nli-mpnet-base-v2\n",
|
||||
" content: true\n",
|
||||
"\n",
|
||||
"workflow:\n",
|
||||
" index:\n",
|
||||
" schedule:\n",
|
||||
" cron: '* * * * * 0/5'\n",
|
||||
" elements:\n",
|
||||
" - text\n",
|
||||
" iterations: 4\n",
|
||||
" tasks:\n",
|
||||
" - __main__.Stream\n",
|
||||
" - upsert\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"app = Application(workflow)\n",
|
||||
"app.wait()"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "1oZag3tKWkfe",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "54d6ce91-d782-421a-e092-7e78bb6ee1d0"
|
||||
},
|
||||
"execution_count": 8,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"2022-02-03 02:12:06,720 [WARNING] _create_builder_config: Using custom data configuration default\n",
|
||||
"2022-02-03 02:12:06,727 [WARNING] download_and_prepare: Reusing dataset ag_news (/root/.cache/huggingface/datasets/ag_news/default/0.0.0/bc2bcb40336ace1a0374767fc29bb0296cdaf8a6da7298436239c54d79180548)\n",
|
||||
"2022-02-03 02:12:06,751 [INFO] schedule: 'index' scheduler started with schedule * * * * * 0/5\n",
|
||||
"2022-02-03 02:12:06,757 [INFO] schedule: 'index' next run scheduled for 2022-02-03T02:12:10+00:00\n",
|
||||
"2022-02-03 02:12:34,937 [INFO] schedule: 'index' next run scheduled for 2022-02-03T02:12:35+00:00\n",
|
||||
"2022-02-03 02:12:59,967 [INFO] schedule: 'index' next run scheduled for 2022-02-03T02:13:00+00:00\n",
|
||||
"2022-02-03 02:13:23,349 [INFO] schedule: 'index' next run scheduled for 2022-02-03T02:13:25+00:00\n",
|
||||
"2022-02-03 02:13:49,621 [INFO] schedule: 'index' max iterations (4) reached\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Reviewing the log above, we see the `index` job ran four times. Now let's query the index and see what was loaded."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "SC3Yf9EgSHiK"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Run an embeddings search\n",
|
||||
"\n",
|
||||
"Let's run a search against the newly created index."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "5hBY2TsZSQY0"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import json\n",
|
||||
"\n",
|
||||
"# Show total number of records\n",
|
||||
"print(f\"Total records: {app.count()}\")\n",
|
||||
"\n",
|
||||
"# Run a search\n",
|
||||
"print(\"Search:\")\n",
|
||||
"print(json.dumps(app.search(\"life on mars\", limit=1), indent=2))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "3slP9p_KkbtD",
|
||||
"outputId": "9473aaea-eaa3-42e5-f9db-1f9f6e51f5c5"
|
||||
},
|
||||
"execution_count": 9,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Total records: 10000\n",
|
||||
"Search:\n",
|
||||
"[\n",
|
||||
" {\n",
|
||||
" \"id\": \"119\",\n",
|
||||
" \"text\": \"Life on Mars Likely, Scientist Claims (SPACE.com) SPACE.com - DENVER, COLORADO -- Those twin robots hard at work on Mars have transmitted teasing views that reinforce the prospect that microbial life may exist on the red planet.\",\n",
|
||||
" \"score\": 0.7236138582229614\n",
|
||||
" }\n",
|
||||
"]\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"The index has 10,000 records. We also see the top result for the query on `life on mars`."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "4i3H19ZmUSlK"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Run a scheduled embeddings search\n",
|
||||
"\n",
|
||||
"Now let's incrementally load the dataset with a scheduled workflow and run a scheduled search after each batch is loaded."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "-t-D7yu_WZl-"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai.app import Application\n",
|
||||
"\n",
|
||||
"# Run every 5 seconds up to 4 times\n",
|
||||
"workflow = \"\"\"\n",
|
||||
"writable: true\n",
|
||||
"embeddings:\n",
|
||||
" path: sentence-transformers/nli-mpnet-base-v2\n",
|
||||
" content: true\n",
|
||||
"\n",
|
||||
"workflow:\n",
|
||||
" index:\n",
|
||||
" schedule:\n",
|
||||
" cron: '* * * * * 0/5'\n",
|
||||
" elements:\n",
|
||||
" - text\n",
|
||||
" iterations: 4\n",
|
||||
" tasks:\n",
|
||||
" - __main__.Stream\n",
|
||||
" - upsert\n",
|
||||
" search:\n",
|
||||
" schedule:\n",
|
||||
" cron: '* * * * * 0/5'\n",
|
||||
" elements:\n",
|
||||
" - life on mars\n",
|
||||
" iterations: 4\n",
|
||||
" tasks:\n",
|
||||
" - action: search\n",
|
||||
" args: [3]\n",
|
||||
" task: console\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"app = Application(workflow)\n",
|
||||
"app.wait()"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "xpuwL9aCUJOd",
|
||||
"outputId": "d994aaaf-7315-4109-d024-3e09773cc538"
|
||||
},
|
||||
"execution_count": 10,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"2022-02-03 02:13:55,789 [WARNING] _create_builder_config: Using custom data configuration default\n",
|
||||
"2022-02-03 02:13:55,797 [WARNING] download_and_prepare: Reusing dataset ag_news (/root/.cache/huggingface/datasets/ag_news/default/0.0.0/bc2bcb40336ace1a0374767fc29bb0296cdaf8a6da7298436239c54d79180548)\n",
|
||||
"2022-02-03 02:13:55,808 [INFO] schedule: 'index' scheduler started with schedule * * * * * 0/5\n",
|
||||
"2022-02-03 02:13:55,808 [INFO] schedule: 'search' scheduler started with schedule * * * * * 0/5\n",
|
||||
"2022-02-03 02:13:55,810 [INFO] schedule: 'index' next run scheduled for 2022-02-03T02:14:00+00:00\n",
|
||||
"2022-02-03 02:13:55,814 [INFO] schedule: 'search' next run scheduled for 2022-02-03T02:14:00+00:00\n",
|
||||
"2022-02-03 02:14:00,001 [INFO] schedule: 'search' next run scheduled for 2022-02-03T02:14:05+00:00\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Inputs: [\n",
|
||||
" \"life on mars\"\n",
|
||||
"]\n",
|
||||
"Outputs: [\n",
|
||||
" null\n",
|
||||
"]\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"2022-02-03 02:14:24,500 [INFO] schedule: 'index' next run scheduled for 2022-02-03T02:14:25+00:00\n",
|
||||
"2022-02-03 02:14:24,522 [INFO] schedule: 'search' next run scheduled for 2022-02-03T02:14:25+00:00\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Inputs: [\n",
|
||||
" \"life on mars\"\n",
|
||||
"]\n",
|
||||
"Outputs: [\n",
|
||||
" {\n",
|
||||
" \"id\": \"119\",\n",
|
||||
" \"text\": \"Life on Mars Likely, Scientist Claims (SPACE.com) SPACE.com - DENVER, COLORADO -- Those twin robots hard at work on Mars have transmitted teasing views that reinforce the prospect that microbial life may exist on the red planet.\",\n",
|
||||
" \"score\": 0.7236138582229614\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" \"id\": \"271\",\n",
|
||||
" \"text\": \"Saturn's Moon Titan: Prebiotic Laboratory by Harry Bortman In this second and final part of the interview, Lunine explains how Huygens may help scientists understand the origin of life on Earth, even if it doesn't detect life on Titan. Astrobiology Magazine -- Titan is the only moon in our solar system with an atmosphere, and it is the organic chemistry that has been detected in that atmosphere that has sparked the imagination of planetary scientists like Lunine...\",\n",
|
||||
" \"score\": 0.4750666916370392\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" \"id\": \"1132\",\n",
|
||||
" \"text\": \"Is Mercury the Incredible Shrinking Planet? MESSENGER Spacecraft May Find Out (SPACE.com) SPACE.com - With a new spacecraft bound for Mercury, that tiny planet nbsp;near the heart of the solar system, researchers are hoping to solve a slew of riddles about the small world.\",\n",
|
||||
" \"score\": 0.47124743461608887\n",
|
||||
" }\n",
|
||||
"]\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"2022-02-03 02:14:25,496 [INFO] schedule: 'search' next run scheduled for 2022-02-03T02:14:30+00:00\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Inputs: [\n",
|
||||
" \"life on mars\"\n",
|
||||
"]\n",
|
||||
"Outputs: [\n",
|
||||
" {\n",
|
||||
" \"id\": \"119\",\n",
|
||||
" \"text\": \"Life on Mars Likely, Scientist Claims (SPACE.com) SPACE.com - DENVER, COLORADO -- Those twin robots hard at work on Mars have transmitted teasing views that reinforce the prospect that microbial life may exist on the red planet.\",\n",
|
||||
" \"score\": 0.7236138582229614\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" \"id\": \"271\",\n",
|
||||
" \"text\": \"Saturn's Moon Titan: Prebiotic Laboratory by Harry Bortman In this second and final part of the interview, Lunine explains how Huygens may help scientists understand the origin of life on Earth, even if it doesn't detect life on Titan. Astrobiology Magazine -- Titan is the only moon in our solar system with an atmosphere, and it is the organic chemistry that has been detected in that atmosphere that has sparked the imagination of planetary scientists like Lunine...\",\n",
|
||||
" \"score\": 0.4750666916370392\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" \"id\": \"1132\",\n",
|
||||
" \"text\": \"Is Mercury the Incredible Shrinking Planet? MESSENGER Spacecraft May Find Out (SPACE.com) SPACE.com - With a new spacecraft bound for Mercury, that tiny planet nbsp;near the heart of the solar system, researchers are hoping to solve a slew of riddles about the small world.\",\n",
|
||||
" \"score\": 0.47124743461608887\n",
|
||||
" }\n",
|
||||
"]\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"2022-02-03 02:14:50,112 [INFO] schedule: 'index' next run scheduled for 2022-02-03T02:14:55+00:00\n",
|
||||
"2022-02-03 02:14:50,138 [INFO] schedule: 'search' max iterations (4) reached\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Inputs: [\n",
|
||||
" \"life on mars\"\n",
|
||||
"]\n",
|
||||
"Outputs: [\n",
|
||||
" {\n",
|
||||
" \"id\": \"119\",\n",
|
||||
" \"text\": \"Life on Mars Likely, Scientist Claims (SPACE.com) SPACE.com - DENVER, COLORADO -- Those twin robots hard at work on Mars have transmitted teasing views that reinforce the prospect that microbial life may exist on the red planet.\",\n",
|
||||
" \"score\": 0.7236138582229614\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" \"id\": \"3300\",\n",
|
||||
" \"text\": \"Mars Hills, Crater Yield Evidence of Flowing Water LOS ANGELES (Reuters) - The hills of Mars yielded more tantalizing clues about how water shaped the Red Planet in tests by NASA #39;s robotic geologist, Spirit, while its twin, Opportunity, observed the deep crater it climbed into two months ...\",\n",
|
||||
" \"score\": 0.6666488647460938\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" \"id\": \"4201\",\n",
|
||||
" \"text\": \"Martian hill shows signs of ancient water LOS ANGELES - NASA #39;s Spirit rover has found more evidence of past water on the hills of Mars, while its twin, Opportunity, has observed a field of dunes inside a crater. \",\n",
|
||||
" \"score\": 0.6453495621681213\n",
|
||||
" }\n",
|
||||
"]\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stderr",
|
||||
"text": [
|
||||
"2022-02-03 02:15:18,333 [INFO] schedule: 'index' next run scheduled for 2022-02-03T02:15:20+00:00\n",
|
||||
"2022-02-03 02:15:44,592 [INFO] schedule: 'index' max iterations (4) reached\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"The workflow above runs up to every 5 seconds. Note that since the index job takes longer than 5 seconds, the time difference between jobs is longer.\n",
|
||||
"\n",
|
||||
"The index job loads the next batch of data and the search job runs a recurring search. \n",
|
||||
"\n",
|
||||
"See how the search results change over time as more relevant results are found."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "o9EX2NgxV23x"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook covered how to use workflow scheduling with txtai. While there are existing ways to schedule jobs (system cron, serverless, and so on), this is another easy and quick way to do it. "
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Fr99QHPtTMJt"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,373 @@
|
||||
{
|
||||
"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": {
|
||||
"provenance": []
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "POWZoSJR6XzK"
|
||||
},
|
||||
"source": [
|
||||
"# Embeddings SQL custom functions\n",
|
||||
"\n",
|
||||
"txtai 4.0 added support for SQL-based embeddings queries. This feature combines natural language queries for similarity with concrete filtering rules. txtai now has support for user-defined SQL functions, making this feature even more powerful."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "qa_PPKVX6XzN"
|
||||
},
|
||||
"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"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline]"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Create index\n",
|
||||
"Let's first recap how to create an index. We'll use the classic txtai example.\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "0p3WCDniUths"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a",
|
||||
"_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
|
||||
"trusted": true,
|
||||
"id": "2j_CFGDR6Xzp",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "f2488a78-6cae-4c25-985e-fb2dd674a534"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.embeddings import Embeddings\n",
|
||||
"\n",
|
||||
"data = [\"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",
|
||||
"# Create embeddings index with content enabled. The default behavior is to only store indexed vectors.\n",
|
||||
"embeddings = Embeddings({\"path\": \"sentence-transformers/nli-mpnet-base-v2\", \"content\": True})\n",
|
||||
"\n",
|
||||
"# Create an index for the list of text\n",
|
||||
"embeddings.index([(uid, text, None) for uid, text in enumerate(data)])\n",
|
||||
"\n",
|
||||
"# Run a search\n",
|
||||
"embeddings.search(\"feel good story\", 1)"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '4',\n",
|
||||
" 'score': 0.08329004049301147,\n",
|
||||
" 'text': 'Maine man wins $1M from $25 lottery ticket'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 14
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Custom SQL functions\n",
|
||||
"\n",
|
||||
"Next, we'll recreate the index adding user-defined SQL functions. These functions are simply Python callable objects or functions that take an input and return values. Pipelines, workflows, custom tasks and any other callable object is supported."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "QTee7YMNDD4R"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"def clength(text):\n",
|
||||
" return len(text) if text else 0\n",
|
||||
"\n",
|
||||
"# Create embeddings index with content enabled. The default behavior is to only store indexed vectors.\n",
|
||||
"embeddings = Embeddings({\"path\": \"sentence-transformers/nli-mpnet-base-v2\", \"content\": True, \"functions\": [clength]})\n",
|
||||
"\n",
|
||||
"# Create an index for the list of text\n",
|
||||
"embeddings.index([(uid, text, None) for uid, text in enumerate(data)])\n",
|
||||
"\n",
|
||||
"# Run a search using a custom SQL function\n",
|
||||
"embeddings.search(\"select clength(text) clength, length(text) length, text from txtai where similar('feel good story')\", 1)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "rbsEXtysDDNg",
|
||||
"outputId": "f966be17-086b-49b4-e1af-62b766f1c995"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'clength': 42,\n",
|
||||
" 'length': 42,\n",
|
||||
" 'text': 'Maine man wins $1M from $25 lottery ticket'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 15
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"The function itself is simple, it's just alternate length function. But this example is just warming us up to what is possible and what is more exciting. "
|
||||
],
|
||||
"metadata": {
|
||||
"id": "epIV58P1DyZa"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Pipelines in SQL\n",
|
||||
"\n",
|
||||
"As mentioned above, any callable can be registered as a custom SQL function. Let's add a translate SQL function."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "1Iw1WKR6FW3S"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai.pipeline import Translation\n",
|
||||
"\n",
|
||||
"# Translation pipeline\n",
|
||||
"translate = Translation()\n",
|
||||
"\n",
|
||||
"# Create embeddings index with content enabled. The default behavior is to only store indexed vectors.\n",
|
||||
"embeddings = Embeddings({\"path\": \"sentence-transformers/nli-mpnet-base-v2\", \"content\": True, \"functions\": [translate]})\n",
|
||||
"\n",
|
||||
"# Create an index for the list of text\n",
|
||||
"embeddings.index([(uid, text, None) for uid, text in enumerate(data)])\n",
|
||||
"\n",
|
||||
"query = \"\"\"\n",
|
||||
"select\n",
|
||||
" text,\n",
|
||||
" translation(text, 'de', null) 'text (DE)',\n",
|
||||
" translation(text, 'es', null) 'text (ES)',\n",
|
||||
" translation(text, 'fr', null) 'text (FR)'\n",
|
||||
"from txtai where similar('feel good story')\n",
|
||||
"limit 1\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"# Run a search using a custom SQL function\n",
|
||||
"embeddings.search(query)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "83e8yXpXFh4F",
|
||||
"outputId": "0b17e9be-8983-418d-9903-b1e72efc5918"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'text': 'Maine man wins $1M from $25 lottery ticket',\n",
|
||||
" 'text (DE)': 'Maine Mann gewinnt $1M von $25 Lotterie-Ticket',\n",
|
||||
" 'text (ES)': 'Maine hombre gana $1M de billete de lotería de $25',\n",
|
||||
" 'text (FR)': 'Maine homme gagne $1M à partir de $25 billet de loterie'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 16
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"And just like that we have translations through SQL! This is pretty 🔥🔥🔥\n",
|
||||
"\n",
|
||||
"We can do more to make this easier though. Let's define a helper function to not require as many parameters. The default logic will require all function parameters each call, including parameters with default values."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Ck_XTyBEQtaW"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"def translation(text, lang):\n",
|
||||
" return translate(text, lang)\n",
|
||||
"\n",
|
||||
"# Create embeddings index with content enabled. The default behavior is to only store indexed vectors.\n",
|
||||
"embeddings = Embeddings({\"path\": \"sentence-transformers/nli-mpnet-base-v2\", \"content\": True, \"functions\": [translation]})\n",
|
||||
"\n",
|
||||
"# Create an index for the list of text\n",
|
||||
"embeddings.index([(uid, text, None) for uid, text in enumerate(data)])\n",
|
||||
"\n",
|
||||
"query = \"\"\"\n",
|
||||
"select\n",
|
||||
" text,\n",
|
||||
" translation(text, 'de') 'text (DE)',\n",
|
||||
" translation(text, 'es') 'text (ES)',\n",
|
||||
" translation(text, 'fr') 'text (FR)'\n",
|
||||
"from txtai where similar('feel good story')\n",
|
||||
"limit 1\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"# Run a search using a custom SQL function\n",
|
||||
"embeddings.search(query)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "L2DDJrd0RAaN",
|
||||
"outputId": "0bb437ec-5c9b-4a0c-fe8a-07f641c94a49"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'text': 'Maine man wins $1M from $25 lottery ticket',\n",
|
||||
" 'text (DE)': 'Maine Mann gewinnt $1M von $25 Lotterie-Ticket',\n",
|
||||
" 'text (ES)': 'Maine hombre gana $1M de billete de lotería de $25',\n",
|
||||
" 'text (FR)': 'Maine homme gagne $1M à partir de $25 billet de loterie'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 17
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Custom SQL functions with applications\n",
|
||||
"\n",
|
||||
"Of course this is all available with YAML-configured applications."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "mTT8nopiRdVH"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"config = \"\"\"\n",
|
||||
"translation:\n",
|
||||
"\n",
|
||||
"writable: true\n",
|
||||
"embeddings:\n",
|
||||
" path: sentence-transformers/nli-mpnet-base-v2\n",
|
||||
" content: true\n",
|
||||
" functions:\n",
|
||||
" - {name: translation, argcount: 2, function: translation}\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"from txtai.app import Application\n",
|
||||
"\n",
|
||||
"# Build application and index data\n",
|
||||
"app = Application(config)\n",
|
||||
"app.add([{\"id\": x, \"text\": row} for x, row in enumerate(data)])\n",
|
||||
"app.index()\n",
|
||||
"\n",
|
||||
"# Run search with custom SQL\n",
|
||||
"app.search(query)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "FZ_7G6M4RUbz",
|
||||
"outputId": "4eca94f3-d2aa-4449-dc6f-f1091ad9dd67"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'text': 'Maine man wins $1M from $25 lottery ticket',\n",
|
||||
" 'text (DE)': 'Maine Mann gewinnt $1M von $25 Lotterie-Ticket',\n",
|
||||
" 'text (ES)': 'Maine hombre gana $1M de billete de lotería de $25',\n",
|
||||
" 'text (FR)': 'Maine homme gagne $1M à partir de $25 billet de loterie'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 18
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "aDIF3tYt6X0O"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook introduced running user-defined custom SQL functions through embeddings SQL. This powerful feature can be used with any callable function including pipelines, tasks and workflows in tandem with similarity and rules filters."
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1,508 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "POWZoSJR6XzK"
|
||||
},
|
||||
"source": [
|
||||
"# Query translation\n",
|
||||
"\n",
|
||||
"txtai supports two main types of queries: natural language statements and SQL statements. Natural language queries handles a search engine like query. SQL statements enable more complex filtering, sorting and column selection. Query translation bridges the gap between the two and enables filtering for natural language queries.\n",
|
||||
"\n",
|
||||
"For example, the query:\n",
|
||||
"\n",
|
||||
"```\n",
|
||||
"Tell me a feel good story since yesterday\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"becomes\n",
|
||||
"\n",
|
||||
"```sql\n",
|
||||
"select * from txtai where similar(\"Tell me a feel good story\") and\n",
|
||||
"entry >= date('now', '-1 day')\n",
|
||||
"```"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "qa_PPKVX6XzN"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
|
||||
"_kg_hide-output": true,
|
||||
"_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
|
||||
"id": "24q-1n5i6XzQ",
|
||||
"trusted": true
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "0p3WCDniUths"
|
||||
},
|
||||
"source": [
|
||||
"# Create index\n",
|
||||
"Let's first recap how to create an index. We'll use the classic txtai example.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
|
||||
"_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "2j_CFGDR6Xzp",
|
||||
"outputId": "e65238c5-d67f-4fe4-9cbf-c88b0ff3a8bb",
|
||||
"trusted": true
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '4',\n",
|
||||
" 'text': 'Maine man wins $1M from $25 lottery ticket',\n",
|
||||
" 'score': 0.08329025655984879}]"
|
||||
]
|
||||
},
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"data = [\"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",
|
||||
"# Create embeddings index with content enabled. The default behavior is to only store indexed vectors.\n",
|
||||
"embeddings = Embeddings(path=\"sentence-transformers/nli-mpnet-base-v2\", content=True)\n",
|
||||
"\n",
|
||||
"# Create an index for the list of text\n",
|
||||
"embeddings.index(data)\n",
|
||||
"\n",
|
||||
"# Run a search\n",
|
||||
"embeddings.search(\"feel good story\", 1)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "QTee7YMNDD4R"
|
||||
},
|
||||
"source": [
|
||||
"# Query translation models\n",
|
||||
"\n",
|
||||
"Next we'll explore how query translation models work with examples. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "04iOgP_ojSfK",
|
||||
"outputId": "41e73130-75e6-4ae9-b690-685972a13565"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Input: feel good story\n",
|
||||
"SQL: select id, text, score from txtai where similar('feel good story')\n",
|
||||
"\n",
|
||||
"Input: feel good story since yesterday\n",
|
||||
"SQL: select id, text, score from txtai where similar('feel good story') and entry >= date('now', '-1 day')\n",
|
||||
"\n",
|
||||
"Input: feel good story with 'lottery' in text\n",
|
||||
"SQL: select id, text, score from txtai where similar('feel good story') and text like '%lottery%'\n",
|
||||
"\n",
|
||||
"Input: how many feel good story\n",
|
||||
"SQL: select count(*) from txtai where similar('feel good story')\n",
|
||||
"\n",
|
||||
"Input: feel good story translated to fr\n",
|
||||
"SQL: select id, translate(text, 'fr') text, score from txtai where similar('feel good story')\n",
|
||||
"\n",
|
||||
"Input: feel good story summarized\n",
|
||||
"SQL: select id, summary(text) text, score from txtai where similar('feel good story')\n",
|
||||
"\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import LLM\n",
|
||||
"\n",
|
||||
"llm = LLM(\"NeuML/t5-small-txtsql\", template=\"translate English to SQL: {text}\")\n",
|
||||
"\n",
|
||||
"queries = [\n",
|
||||
" \"feel good story\",\n",
|
||||
" \"feel good story since yesterday\",\n",
|
||||
" \"feel good story with 'lottery' in text\",\n",
|
||||
" \"how many feel good story\",\n",
|
||||
" \"feel good story translated to fr\",\n",
|
||||
" \"feel good story summarized\"\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"for query in queries:\n",
|
||||
" print(f\"Input: {query}\")\n",
|
||||
" print(f\"SQL: {llm(query)}\")\n",
|
||||
" print()\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "rAnEMaiWlOXm"
|
||||
},
|
||||
"source": [
|
||||
"Looking at the query translations above gives an idea on how this model works.\n",
|
||||
"\n",
|
||||
"[t5-small-txtsql](https://huggingface.co/NeuML/t5-small-txtsql) is the default model. Custom domain query syntax languages can be created using this same methodology, including for other languages. Natural language can be translated to functions, query clauses, column selection and more!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "P9hOcgNfjSyL"
|
||||
},
|
||||
"source": [
|
||||
"# Natural language filtering\n",
|
||||
"\n",
|
||||
"Now it's time for this in action! Let's first initialize the embeddings index with the appropriate settings."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "L2DDJrd0RAaN",
|
||||
"outputId": "c60bfc29-187b-4926-8656-04063b2ca85c"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'id': '4',\n",
|
||||
" 'score': 0.08329025655984879,\n",
|
||||
" 'text': 'Maine Mann gewinnt $1M von $25 Lotterie-Ticket'}"
|
||||
]
|
||||
},
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai.pipeline import Translation\n",
|
||||
"\n",
|
||||
"def translate(text, lang):\n",
|
||||
" return translation(text, lang)\n",
|
||||
"\n",
|
||||
"translation = Translation()\n",
|
||||
"\n",
|
||||
"# Create embeddings index with content enabled. The default behavior is to only store indexed vectors.\n",
|
||||
"embeddings = Embeddings(\n",
|
||||
" path=\"sentence-transformers/nli-mpnet-base-v2\",\n",
|
||||
" content=True,\n",
|
||||
" query={\"path\": \"NeuML/t5-small-txtsql\"},\n",
|
||||
" functions=[translate]\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Create an index for the list of text\n",
|
||||
"embeddings.index(data)\n",
|
||||
"\n",
|
||||
"query = \"select id, score, translate(text, 'de') 'text' from txtai where similar('feel good story')\"\n",
|
||||
"\n",
|
||||
"# Run a search using a custom SQL function\n",
|
||||
"embeddings.search(query)[0]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "MNd7QmFmnh-f"
|
||||
},
|
||||
"source": [
|
||||
"Note how the query model was provided as a embeddings index configuration parameter. Custom SQL functions were also added in. Let's now see if the same SQL statement can be run with a natural language query."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "mnGSBNMonur2",
|
||||
"outputId": "dc14d898-b46f-42e6-88e0-434863cc15d6"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'id': '4',\n",
|
||||
" 'text': 'Maine Mann gewinnt $1M von $25 Lotterie-Ticket',\n",
|
||||
" 'score': 0.08329025655984879}"
|
||||
]
|
||||
},
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings.search(\"feel good story translated to de\")[0]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "OnmVUK2DoZkh"
|
||||
},
|
||||
"source": [
|
||||
"Same result. Let's try a few more."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "Js2b_M61oitg",
|
||||
"outputId": "fab080cb-d082-4cf3-8314-12d23acc3c02"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'id': '4',\n",
|
||||
" 'text': 'Maine man wins $1M from $25 lottery ticket',\n",
|
||||
" 'score': 0.08329025655984879}"
|
||||
]
|
||||
},
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings.search(\"feel good story since yesterday\")[0]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "pcxPgriwomWv",
|
||||
"outputId": "7e71dd09-7c4f-4147-ba9e-47a45c1cf90f"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '4',\n",
|
||||
" 'text': 'Maine man wins $1M from $25 lottery ticket',\n",
|
||||
" 'score': 0.08329025655984879}]"
|
||||
]
|
||||
},
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings.search(\"feel good story with 'lottery' in text\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "1ob4Q_CLpGBm"
|
||||
},
|
||||
"source": [
|
||||
"For good measure, a couple queries with filters that return no results."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "LMaf7JJzonAC",
|
||||
"outputId": "71080bd1-2d9e-41ba-9501-2cdedaf26d6d"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[]"
|
||||
]
|
||||
},
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings.search(\"feel good story with missing in text\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ADFNLhvto_0J",
|
||||
"outputId": "12e08ea8-7203-4fb0-cb84-1133b5db9dc0"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[]"
|
||||
]
|
||||
},
|
||||
"execution_count": 8,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings.search(\"feel good story with field equal 14\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "mTT8nopiRdVH"
|
||||
},
|
||||
"source": [
|
||||
"# Query translation with applications\n",
|
||||
"\n",
|
||||
"Of course this is all available with YAML-configured applications."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "FZ_7G6M4RUbz",
|
||||
"outputId": "b00ce1c2-8df2-4289-8d03-991174a0e74f"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'id': '4',\n",
|
||||
" 'text': 'Maine Mann gewinnt $1M von $25 Lotterie-Ticket',\n",
|
||||
" 'score': 0.08329025655984879}"
|
||||
]
|
||||
},
|
||||
"execution_count": 9,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"config = \"\"\"\n",
|
||||
"translation:\n",
|
||||
"\n",
|
||||
"writable: true\n",
|
||||
"embeddings:\n",
|
||||
" path: sentence-transformers/nli-mpnet-base-v2\n",
|
||||
" content: true\n",
|
||||
" query:\n",
|
||||
" path: NeuML/t5-small-txtsql\n",
|
||||
" functions:\n",
|
||||
" - {name: translate, argcount: 2, function: translation}\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"from txtai.app import Application\n",
|
||||
"\n",
|
||||
"# Build application and index data\n",
|
||||
"app = Application(config)\n",
|
||||
"app.add([{\"id\": x, \"text\": row} for x, row in enumerate(data)])\n",
|
||||
"app.index()\n",
|
||||
"\n",
|
||||
"# Run search query\n",
|
||||
"app.search(\"feel good story translated to de\")[0]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "aDIF3tYt6X0O"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook introduced natural language filtering with query translation models. This powerful feature adds filtering and pipelines to natural language statements. Custom domain-specific query languages can be created to enable rich queries natively in txtai."
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.10.19"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
@@ -0,0 +1,455 @@
|
||||
{
|
||||
"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": "34 - Build a QA database",
|
||||
"provenance": [],
|
||||
"collapsed_sections": []
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "POWZoSJR6XzK"
|
||||
},
|
||||
"source": [
|
||||
"# Build a QA database\n",
|
||||
"\n",
|
||||
"Conversational AI is a growing field that could potentially automate much of the customer service industry. Full automation is still a ways away (most of us have been on a call with an automated agent and just want to get to a person) but it certainly can be a solid first line before human intervention.\n",
|
||||
"\n",
|
||||
"This notebook presents a process to answer user questions using a txtai embeddings instance. It's not conversational AI but instead looks to find the closest existing question to a user question. This is useful in cases where there are a list of frequently asked questions. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "qa_PPKVX6XzN"
|
||||
},
|
||||
"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"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai datasets"
|
||||
],
|
||||
"execution_count": 75,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Load the dataset\n",
|
||||
"\n",
|
||||
"We'll use a Hugging Face dataset of web questions for this example. The dataset has a list of questions and answers. The code below loads the dataset and prints a couple examples to get an idea of how the data is formatted."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "QAEVXfOiVzEB"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"ds = load_dataset(\"web_questions\", split=\"train\")\n",
|
||||
"\n",
|
||||
"for row in ds.select(range(5)):\n",
|
||||
" print(row[\"question\"], row[\"answers\"])"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "koM4vYHXL82P",
|
||||
"outputId": "91926a1b-8b9c-46be-f450-f76a1e3aab82"
|
||||
},
|
||||
"execution_count": 76,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"what is the name of justin bieber brother? ['Jazmyn Bieber', 'Jaxon Bieber']\n",
|
||||
"what character did natalie portman play in star wars? ['Padmé Amidala']\n",
|
||||
"what state does selena gomez? ['New York City']\n",
|
||||
"what country is the grand bahama island in? ['Bahamas']\n",
|
||||
"what kind of money to take to bahamas? ['Bahamian dollar']\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Create index\n",
|
||||
"\n",
|
||||
"Next, we'll create a txtai index. The question will be the indexed text. We'll also store full content so we can access the answer at query time.\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "0p3WCDniUths"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a",
|
||||
"_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
|
||||
"trusted": true,
|
||||
"id": "2j_CFGDR6Xzp"
|
||||
},
|
||||
"source": [
|
||||
"from txtai.embeddings import Embeddings\n",
|
||||
"\n",
|
||||
"# Create embeddings index with content enabled. The default behavior is to only store indexed vectors.\n",
|
||||
"embeddings = Embeddings({\"path\": \"sentence-transformers/nli-mpnet-base-v2\", \"content\": True})\n",
|
||||
"\n",
|
||||
"# Map question to text and store content\n",
|
||||
"embeddings.index([(uid, {\"url\": row[\"url\"], \"text\": row[\"question\"], \"answer\": \", \".join(row[\"answers\"])}, None) for uid, row in enumerate(ds)])"
|
||||
],
|
||||
"execution_count": 77,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Asking questions\n",
|
||||
"\n",
|
||||
"Now that the index is built, let's ask some questions! We'll use txtai SQL to select the fields we want to return.\n",
|
||||
"\n",
|
||||
"See the list of questions asked and best matching question-answer combo."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "dqx-PmpzWMez"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"def question(text):\n",
|
||||
" return embeddings.search(f\"select text, answer, score from txtai where similar('{text}') limit 1\")\n",
|
||||
"\n",
|
||||
"question(\"What is the timezone of NYC?\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "Gatp4mDQPA_v",
|
||||
"outputId": "93efaec7-c41e-4865-f042-371c55170ffe"
|
||||
},
|
||||
"execution_count": 78,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'answer': 'North American Eastern Time Zone',\n",
|
||||
" 'score': 0.8904051184654236,\n",
|
||||
" 'text': 'what time zone is new york under?'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 78
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"question(\"Things to do in New York\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "WUvdAvvlSPDT",
|
||||
"outputId": "00ed87fe-b659-406e-d03a-97b7b5ea9773"
|
||||
},
|
||||
"execution_count": 79,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'answer': \"Chelsea Art Museum, Brooklyn Bridge, Empire State Building, The Broadway Theatre, American Museum of Natural History, Central Park, St. Patrick's Cathedral, Japan Society of New York, FusionArts Museum, American Folk Art Museum\",\n",
|
||||
" 'score': 0.8308358192443848,\n",
|
||||
" 'text': 'what are some places to visit in new york?'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 79
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"question(\"Microsoft founder\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "vvsSj2qKRrrY",
|
||||
"outputId": "1c4a94d7-3bf6-4cae-8b40-13fbfef35205"
|
||||
},
|
||||
"execution_count": 80,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'answer': 'Bill Gates',\n",
|
||||
" 'score': 0.6617322564125061,\n",
|
||||
" 'text': 'who created microsoft windows?'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 80
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"question(\"Apple founder university\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "UUD56XSRR1jc",
|
||||
"outputId": "fb94e89d-1580-471b-9e0e-541e5208d937"
|
||||
},
|
||||
"execution_count": 81,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'answer': 'Reed College',\n",
|
||||
" 'score': 0.5137897729873657,\n",
|
||||
" 'text': 'what college did steve jobs attend?'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 81
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"question(\"What country uses the Yen?\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "D6Ur0wZWSBfd",
|
||||
"outputId": "caf6e0aa-af62-4f33-fed1-6e1a111fed1a"
|
||||
},
|
||||
"execution_count": 82,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'answer': 'Japanese yen',\n",
|
||||
" 'score': 0.6663530468940735,\n",
|
||||
" 'text': 'what money do japanese use?'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 82
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"question(\"Show me a list of Pixar movies\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "2POCUWrqSKGP",
|
||||
"outputId": "f8069bf4-a135-47df-a571-198f168c03fb"
|
||||
},
|
||||
"execution_count": 83,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'answer': \"A Bug's Life, Toy Story 2, Ratatouille, Cars, Up, Toy Story, Monsters, Inc., The Incredibles, Finding Nemo, WALL-E\",\n",
|
||||
" 'score': 0.653051495552063,\n",
|
||||
" 'text': 'what does pixar produce?'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 83
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"question(\"What is the timezone of Florida?\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "xll4a1ChTaVg",
|
||||
"outputId": "00cb967f-753e-4d15-cb87-308c08dfde59"
|
||||
},
|
||||
"execution_count": 84,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'answer': 'North American Eastern Time Zone',\n",
|
||||
" 'score': 0.9672279357910156,\n",
|
||||
" 'text': 'where is the time zone in florida?'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 84
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"question(\"Tell me an animal found offshore in Florida\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "4EyAOhWXUcKe",
|
||||
"outputId": "ef3f83b1-9314-4458-843e-16afb3364cd9"
|
||||
},
|
||||
"execution_count": 85,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'answer': 'Largemouth bass',\n",
|
||||
" 'score': 0.6526554822921753,\n",
|
||||
" 'text': 'what kind of fish do you catch in florida?'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 85
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Not too bad! This database only has over 6,000 question-answer pairs. To improve quality a score filter could be put on the query to only return highly confident answers. But this gives an idea of what is possible."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "KFxsjtsnWgpe"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Run as an application\n",
|
||||
"\n",
|
||||
"This can also be run as an application. See below."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "2x9awoKNZfZg"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai.app import Application\n",
|
||||
"\n",
|
||||
"# Save index\n",
|
||||
"embeddings.save(\"questions.tar.gz\")\n",
|
||||
"\n",
|
||||
"# Build application and index data\n",
|
||||
"app = Application(\"path: questions.tar.gz\")\n",
|
||||
"\n",
|
||||
"# Run search query\n",
|
||||
"app.search(\"select text, answer, score from txtai where similar('Tell me an animal found offshore in Florida') limit 1\")[0]"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "0lH9cf1bZokt",
|
||||
"outputId": "77afcb03-c834-46c9-98f2-b1a3c34c0e3b"
|
||||
},
|
||||
"execution_count": 86,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'answer': 'Largemouth bass',\n",
|
||||
" 'score': 0.6526554822921753,\n",
|
||||
" 'text': 'what kind of fish do you catch in florida?'}"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 86
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "aDIF3tYt6X0O"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook introduced a simple question matching service. This could be the foundation of an automated customer service agent and/or an online FAQ.\n",
|
||||
"\n",
|
||||
"For a full example, see [codequestion](https://github.com/neuml/codequestion), which is an application that matches user questions to Stack Overflow question-answer pairs."
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,728 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
},
|
||||
"gpuClass": "standard",
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Run txtai in native code\n",
|
||||
"\n",
|
||||
"txtai currently has two main methods of execution: Python or via a HTTP API. There are API bindings for [JavaScript](https://github.com/neuml/txtai.js), [Java](https://github.com/neuml/txtai.java), [Rust](https://github.com/neuml/txtai.rs) and [Go](https://github.com/neuml/txtai.go).\n",
|
||||
"\n",
|
||||
"This notebook presents a way to run txtai as part of a native executable with the [Python C API](https://docs.python.org/3/c-api/index.html). We'll run an example in C and even call txtai from assembly code!\n",
|
||||
"\n",
|
||||
"Before diving into this notebook, it's important to emphasize that connecting to txtai via the HTTP API has a number of major advantages. This includes decoupling from Python, the ability to offload txtai to a different machine and scaling with cloud compute. With that being said, this notebook demonstrates an additional way to integrate txtai along with providing an informative and perhaps academic programming exercise."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "-xU9P9iSR-Cy"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "shlUi2kKS7KT"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "xEvX9vCpn4E0"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline] sacremoses\n",
|
||||
"\n",
|
||||
"# Remove tensorflow as it's not used and prints noisy log messages\n",
|
||||
"!pip uninstall -y tensorflow\n",
|
||||
"\n",
|
||||
"# Install python3.7-dev and nasm\n",
|
||||
"!apt-get install python3.7-dev nasm"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Workflow configuration\n",
|
||||
"\n",
|
||||
"This configuration builds a workflow to translate input text to French. More information on workflows can be found in [txtai's documentation](https://neuml.github.io/txtai/workflow).\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "AtEdP7Utw3mk"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"%%writefile config.yml\n",
|
||||
"summary:\n",
|
||||
" path: sshleifer/distilbart-cnn-12-6\n",
|
||||
"\n",
|
||||
"textractor:\n",
|
||||
" join: true\n",
|
||||
" lines: false\n",
|
||||
" minlength: 100\n",
|
||||
" paragraphs: true\n",
|
||||
" sentences: false\n",
|
||||
" tika: false\n",
|
||||
"\n",
|
||||
"translation:\n",
|
||||
"\n",
|
||||
"workflow:\n",
|
||||
" summary:\n",
|
||||
" tasks:\n",
|
||||
" - action: textractor\n",
|
||||
" task: url\n",
|
||||
" - action: summary\n",
|
||||
"\n",
|
||||
" translate:\n",
|
||||
" tasks:\n",
|
||||
" - action: translation\n",
|
||||
" args: \n",
|
||||
" - fr"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "DPWrubv5oOn7",
|
||||
"outputId": "bdfba131-c36e-421d-8c17-a726be4615d2"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing config.yml\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Python C API\n",
|
||||
"\n",
|
||||
"Next we'll build an interface to txtai workflows with the Python C API. This logic will load Python, create a txtai application instance and add methods to run workflows.\n",
|
||||
"\n",
|
||||
"Some assumptions are made:\n",
|
||||
"\n",
|
||||
"- txtai is installed and available\n",
|
||||
"- A workflow is available in a file named `config.yml`\n",
|
||||
"- The workflow only returns the first element\n",
|
||||
"\n",
|
||||
"These assumptions are for brevity. This example could be expanded on and built into a more robust, full-fledged library.\n",
|
||||
"\n",
|
||||
"While this example is in C, Rust has a well-maintained and popular library for interfacing with Python, [PyO3](https://github.com/PyO3/pyo3). Interfacing with the Python C API is also possible in Java, JavaScript and Go but not as straighforward."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "KIg_QDVRo-6-"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"%%writefile workflow.c\n",
|
||||
"#include <Python.h>\n",
|
||||
"\n",
|
||||
"// Global instances\n",
|
||||
"PyObject *module = NULL, *app = NULL;\n",
|
||||
"\n",
|
||||
"/**\n",
|
||||
" * Create txtai module.\n",
|
||||
" */\n",
|
||||
"PyObject* txtai() {\n",
|
||||
" PyObject* module = NULL;\n",
|
||||
" module = PyImport_ImportModule(\"txtai.app\");\n",
|
||||
" return module;\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"/**\n",
|
||||
" * Create txtai application instance.\n",
|
||||
" */\n",
|
||||
"PyObject* application() {\n",
|
||||
" PyObject* app = NULL;\n",
|
||||
" app = PyObject_CallMethod(module, \"Application\", \"z\", \"config.yml\");\n",
|
||||
" return app;\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"/**\n",
|
||||
" * Run txtai workflow.\n",
|
||||
" */\n",
|
||||
"PyObject* run(char** args) {\n",
|
||||
" PyObject* result = NULL;\n",
|
||||
" result = PyObject_CallMethod(app, \"workflow\", \"z[z]\", args[0], args[1]);\n",
|
||||
" return result;\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"/**\n",
|
||||
" * Cleanup Python objects.\n",
|
||||
" */\n",
|
||||
"void cleanup() {\n",
|
||||
" // Ensure Python instance exists\n",
|
||||
" if (Py_IsInitialized()) {\n",
|
||||
" PyErr_Print();\n",
|
||||
"\n",
|
||||
" Py_CLEAR(app);\n",
|
||||
" Py_CLEAR(module);\n",
|
||||
"\n",
|
||||
" Py_FinalizeEx();\n",
|
||||
" }\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"/**\n",
|
||||
" * Initialize a txtai application and run a workflow.\n",
|
||||
" */\n",
|
||||
"const char* workflow(char** args) {\n",
|
||||
" PyObject* result = NULL;\n",
|
||||
"\n",
|
||||
" // Create application instance if it doesn't already exist\n",
|
||||
" if (!Py_IsInitialized()) {\n",
|
||||
" // Start Python Interpreter\n",
|
||||
" Py_Initialize();\n",
|
||||
"\n",
|
||||
" // Create txtai module\n",
|
||||
" module = txtai();\n",
|
||||
"\n",
|
||||
" // Handle errors\n",
|
||||
" if (!module) {\n",
|
||||
" cleanup();\n",
|
||||
" return NULL;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" // Create txtai application\n",
|
||||
" app = application();\n",
|
||||
"\n",
|
||||
" // Handle errors\n",
|
||||
" if (!app) {\n",
|
||||
" cleanup();\n",
|
||||
" return NULL;\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" // Run workflow\n",
|
||||
" result = run(args);\n",
|
||||
"\n",
|
||||
" // Handle errors\n",
|
||||
" if (!result) {\n",
|
||||
" cleanup();\n",
|
||||
" return NULL;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" // Get first result\n",
|
||||
" const char *text = PyUnicode_AsUTF8(PyIter_Next(result));\n",
|
||||
"\n",
|
||||
" // Cleanup result\n",
|
||||
" Py_CLEAR(result);\n",
|
||||
"\n",
|
||||
" return text;\n",
|
||||
"}"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "rv64SWDkpmIx",
|
||||
"outputId": "1c19ae3c-7f06-43ac-8638-37ee7b5e4eab"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing workflow.c\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Run txtai workflow in C\n",
|
||||
"\n",
|
||||
"Let's now write a C program to run a workflow using command line arguments as input. "
|
||||
],
|
||||
"metadata": {
|
||||
"id": "K19FZK5StVbf"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"%%writefile main.c\n",
|
||||
"#include <stdio.h>\n",
|
||||
"\n",
|
||||
"extern char* workflow(char** argv);\n",
|
||||
"extern void cleanup();\n",
|
||||
"\n",
|
||||
"/**\n",
|
||||
" * Run a txtai workflow and print results.\n",
|
||||
" */\n",
|
||||
"int main(int argc, char** argv) {\n",
|
||||
" if (argc < 3) {\n",
|
||||
" printf(\"Usage: workflow <name> <element>\\n\");\n",
|
||||
" return 1;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" // Run workflow using command line arguments\n",
|
||||
" char* text = workflow(argv + 1);\n",
|
||||
" if (text) {\n",
|
||||
" printf(\"%s\\n\", text);\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" // Cleanup\n",
|
||||
" cleanup();\n",
|
||||
"\n",
|
||||
" return 0;\n",
|
||||
"}"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "apBnMta8pepD",
|
||||
"outputId": "e6573cec-6b92-424c-e594-aa46c7cd18b6"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing main.c\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Compile and run\n",
|
||||
"\n",
|
||||
"Time to compile this all into an executable and run!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "5kVxaL0FuOgu"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!cc -c main.c -I/usr/include/python3.7m\n",
|
||||
"!cc -c workflow.c -I/usr/include/python3.7m\n",
|
||||
"!cc -o workflow workflow.o main.o -lpython3.7m"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "KjQYctCqp1Vi"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!./workflow translate \"I'm running machine translation using a transformers model in C!\""
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "t5IWOF8ap-a7",
|
||||
"outputId": "071f3d26-955e-4e10-cce4-5ed9774960d5"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"J'exécute la traduction automatique à l'aide d'un modèle de transformateurs en C!\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"And there it is, a translation workflow from English to French in a native executable, all backed by Transformers models. Any workflow YAML can be loaded and run in C using this method, which is pretty powerful.\n",
|
||||
"\n",
|
||||
"Embedding txtai in native executable adds libpython as a dependency (libraries from 3rd party modules such as PyTorch and NumPy also load dynamically). See output of ldd below.\n",
|
||||
"This opens up an avenue to embed txtai in native code provided it is acceptable to add libpython as a project dependency. \n",
|
||||
"\n",
|
||||
"As mentioned above, connecting to a txtai HTTP API instance is a less tightly coupled way to accomplish the same thing."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "4eNk20mCxp3r"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!ldd workflow | grep python"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "SLq5ix5mGPAb",
|
||||
"outputId": "a46720e6-c721-4df5-92a7-a1bd897511d0"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"\tlibpython3.7m.so.1.0 => /usr/lib/x86_64-linux-gnu/libpython3.7m.so.1.0 (0x00007efcba85e000)\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Machine learning in Assembly?\n",
|
||||
"\n",
|
||||
"Now for a more academic exercise perhaps bringing you back to a computer organization/logic class from college. Let's see if we can run the same program in assembly!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "n6E-yHIjy0UY"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"%%writefile main.asm\n",
|
||||
"global main\n",
|
||||
"\n",
|
||||
"; External C library functions\n",
|
||||
"extern puts\n",
|
||||
"\n",
|
||||
"; External txtai functions\n",
|
||||
"extern workflow, cleanup\n",
|
||||
"\n",
|
||||
"; Default to REL mode\n",
|
||||
"default REL\n",
|
||||
"\n",
|
||||
"section .data\n",
|
||||
" message: db \"Usage: workflow <name> <element>\", 0\n",
|
||||
"\n",
|
||||
"section .text\n",
|
||||
"\n",
|
||||
"; Print a usage message\n",
|
||||
"usage:\n",
|
||||
" mov rdi, message\n",
|
||||
" call puts\n",
|
||||
" jmp done\n",
|
||||
"\n",
|
||||
"; Main function\n",
|
||||
"main:\n",
|
||||
" ; Enter\n",
|
||||
" sub rsp, 8\n",
|
||||
"\n",
|
||||
" ; Read argc - require workflow name and element (plus program name)\n",
|
||||
" cmp rdi, 3\n",
|
||||
" jl usage\n",
|
||||
"\n",
|
||||
" ; Run txtai workflow with argv params (skip program name) and print result\n",
|
||||
" lea rdi, [rsi + 8]\n",
|
||||
" call workflow\n",
|
||||
" mov rdi, rax\n",
|
||||
" call puts\n",
|
||||
"\n",
|
||||
"done:\n",
|
||||
" ; Close txtai application instance\n",
|
||||
" call cleanup\n",
|
||||
"\n",
|
||||
" ; Exit\n",
|
||||
" add rsp, 8\n",
|
||||
" ret"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "osrjXJonqcTm",
|
||||
"outputId": "b2d662d0-2a71-4b53-a956-c9277c4a7d2b"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing main.asm\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Build workflow executable\n",
|
||||
"!nasm -felf64 main.asm\n",
|
||||
"!cc -c workflow.c -I/usr/include/python3.7m\n",
|
||||
"!cc -o workflow -no-pie workflow.o main.o -lpython3.7m"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "sQMEx0LFq80Z"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!./workflow translate \"I'm running machine translation using a transformers model with assembler!\""
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "YG6faU7HrCrH",
|
||||
"outputId": "a7a77cbf-5806-4241-cd67-03e79438e6f2"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"J'exécute la traduction automatique à l'aide d'un modèle de transformateurs avec assembleur!\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Just as before, the input text is translated to French using a machine translation model. But this time the code executing the logic was in assembly!\n",
|
||||
"\n",
|
||||
"Probably not terribly useful but using the lowest level of code possible proves that any higher-level native code can do the same. "
|
||||
],
|
||||
"metadata": {
|
||||
"id": "SNXAqFm6bNSw"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Multiple workflow calls\n",
|
||||
"\n",
|
||||
"Everything up to this point has been a single workflow call. Much of the run time is spent on loading models as part of the txtai workflow. The next example will run a series of workflow calls and compare how long it takes vs a single workflow command line call. Once again in assembly."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "dB_5UWoDH0nY"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"%%writefile main.asm\n",
|
||||
"global main\n",
|
||||
"\n",
|
||||
"; External C library functions\n",
|
||||
"extern printf\n",
|
||||
"\n",
|
||||
"; External txtai functions\n",
|
||||
"extern workflow, cleanup\n",
|
||||
"\n",
|
||||
"; Default to REL mode\n",
|
||||
"default REL\n",
|
||||
"\n",
|
||||
"section .data\n",
|
||||
" format: db \"action: %s\", 10, \"input: %s\", 10, \"output: %s\", 10, 10, 0\n",
|
||||
" summary: db \"summary\", 0\n",
|
||||
" translate: db \"translate\", 0\n",
|
||||
" text1: db \"txtai executes machine-learning workflows to transform data and build AI-powered semantic search applications.\", 0\n",
|
||||
" text2: db \"Traditional search systems use keywords to find data\", 0\n",
|
||||
" url1: db \"https://github.com/neuml/txtai\", 0\n",
|
||||
" url2: db \"https://github.com/neuml/paperai\", 0\n",
|
||||
"\n",
|
||||
"section .text\n",
|
||||
"\n",
|
||||
"; Run txtai workflow and print results\n",
|
||||
"%macro txtai 2\n",
|
||||
" ; Workflow name and element\n",
|
||||
" push %2\n",
|
||||
" push %1\n",
|
||||
"\n",
|
||||
" ; Run workflow\n",
|
||||
" lea rdi, [rsp]\n",
|
||||
" call workflow\n",
|
||||
"\n",
|
||||
" ; Print action-input-output\n",
|
||||
" mov rdi, format\n",
|
||||
" mov rsi, [rsp]\n",
|
||||
" mov rdx, [rsp + 8]\n",
|
||||
" mov rcx, rax\n",
|
||||
" call printf\n",
|
||||
"\n",
|
||||
" ; Restore stack\n",
|
||||
" add rsp, 16\n",
|
||||
"%endmacro\n",
|
||||
"\n",
|
||||
"; Main function\n",
|
||||
"main:\n",
|
||||
" ; Enter\n",
|
||||
" sub rsp, 8\n",
|
||||
"\n",
|
||||
" ; Run workflows\n",
|
||||
" txtai translate, text1\t\n",
|
||||
" txtai translate, text2\n",
|
||||
" txtai summary, url1\n",
|
||||
" txtai summary, url2\n",
|
||||
"\n",
|
||||
"done:\n",
|
||||
" ; Close txtai application instance\n",
|
||||
" call cleanup\n",
|
||||
"\n",
|
||||
" ; Exit\n",
|
||||
" add rsp, 8\n",
|
||||
" ret"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "wzTYCSaqCSZu",
|
||||
"outputId": "1f48015f-7654-4d86-9f2a-2e462c22f4e6"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Overwriting main.asm\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!time ./workflow translate \"I'm running machine translation using a transformers model with assembler!\""
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "BSmdBuPhIIbw",
|
||||
"outputId": "b45c1ea5-e0fb-4331-f033-5370cdc2ead0"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"J'exécute la traduction automatique à l'aide d'un modèle de transformateurs avec assembleur!\n",
|
||||
"\n",
|
||||
"real\t0m19.208s\n",
|
||||
"user\t0m11.256s\n",
|
||||
"sys\t0m3.224s\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Build workflow executable\n",
|
||||
"!nasm -felf64 main.asm\n",
|
||||
"!cc -c workflow.c -I/usr/include/python3.7m\n",
|
||||
"!cc -no-pie -o workflow workflow.o main.o -lpython3.7m"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "HoRsPeorC5rC"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!time ./workflow"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "7Da6TZlDC8Mj",
|
||||
"outputId": "e2b46994-62d1-42d1-ab77-0d690a7d7d67"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"action: translate\n",
|
||||
"input: txtai executes machine-learning workflows to transform data and build AI-powered semantic search applications.\n",
|
||||
"output: txtai exécute des workflows d'apprentissage automatique pour transformer les données et construire des applications de recherche sémantique alimentées par l'IA.\n",
|
||||
"\n",
|
||||
"action: translate\n",
|
||||
"input: Traditional search systems use keywords to find data\n",
|
||||
"output: Les systèmes de recherche traditionnels utilisent des mots-clés pour trouver des données\n",
|
||||
"\n",
|
||||
"action: summary\n",
|
||||
"input: https://github.com/neuml/txtai\n",
|
||||
"output: txtai executes machine-learning workflows to transform data and build AI-powered semantic search applications. Semantic search applications have an understanding of natural language and identify results that have the same meaning, not necessarily the same keywords. API bindings for JavaScript, Java, Rust and Go. Cloud-native architecture scales out with container orchestration systems (e. g. Kubernetes)\n",
|
||||
"\n",
|
||||
"action: summary\n",
|
||||
"input: https://github.com/neuml/paperai\n",
|
||||
"output: paperai is an AI-powered literature discovery and review engine for medical/scientific papers. Paperai was used to analyze the COVID-19 Open Research Dataset (CORD-19) paperai and NeuML have been recognized in the following articles: Cord-19 Kaggle Challenge Awards Machine-Learning Experts Delve Into 47,000 Papers on Coronavirus Family.\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"real\t0m22.478s\n",
|
||||
"user\t0m13.776s\n",
|
||||
"sys\t0m3.218s\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"As we can see, running 4 workflow actions is about the same runtime as a single action when accounting for model load times."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "zNalTEg5cmhX"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook walked through an example on how to run txtai with native code. While the HTTP API is a better route to go, this is another way to work with txtai!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "4L8smyyXc8q8"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,566 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
},
|
||||
"gpuClass": "standard",
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Embeddings index components\n",
|
||||
"\n",
|
||||
"The main components of txtai are `embeddings`, `pipeline`, `workflow` and an `api`. The following shows the top level view of the txtai src tree.\n",
|
||||
"\n",
|
||||
"```\n",
|
||||
"Abbreviated listing of src/txtai\n",
|
||||
" ann\n",
|
||||
" api\n",
|
||||
" database\n",
|
||||
" embeddings\n",
|
||||
" pipeline\n",
|
||||
" scoring\n",
|
||||
" vectors\n",
|
||||
" workflow\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"One might ask, why are `ann`, `database`, `scoring` and `vectors` top level packages and not under the `embeddings` package? The `embeddings` package provides the glue between these components, making everything easy to use. The reason is that each of these packages are modular and can be used on their own! \n",
|
||||
"\n",
|
||||
"This notebook will go through a series of examples demonstrating how these components can be used standalone as well as combined together to build custom search indexes.\n",
|
||||
"\n",
|
||||
"_Note: This is intended as a deep dive into txtai `embeddings` components. There are much simpler high-level APIs for standard use cases._"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "-xU9P9iSR-Cy"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "shlUi2kKS7KT"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "xEvX9vCpn4E0"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai datasets"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Load dataset\n",
|
||||
"\n",
|
||||
"This example will use the `ag_news` dataset, which is a collection of news article headlines."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "408IyXzKFSiG"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"dataset = load_dataset(\"ag_news\", split=\"train\")"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "IQ_ns6YvFRm1"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Approximate nearest neighbor (ANN) and Vectors\n",
|
||||
"\n",
|
||||
"In this section, we'll use the `ann` and `vectors` package to build a similarity index over the `ag_news` dataset.\n",
|
||||
"\n",
|
||||
"The first step is vectorizing the text. We'll use a `sentence-transformers` model. "
|
||||
],
|
||||
"metadata": {
|
||||
"id": "AtEdP7Utw3mk"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import numpy as np\n",
|
||||
"\n",
|
||||
"from txtai.vectors import VectorsFactory\n",
|
||||
"\n",
|
||||
"model = VectorsFactory.create({\"path\": \"sentence-transformers/all-MiniLM-L6-v2\"}, None)\n",
|
||||
"\n",
|
||||
"embeddings = []\n",
|
||||
"\n",
|
||||
"# List of all text elements\n",
|
||||
"texts = dataset[\"text\"]\n",
|
||||
"\n",
|
||||
"# Create embeddings buffer, vector model has 384 features\n",
|
||||
"embeddings = np.zeros(dtype=np.float32, shape=(len(texts), 384))\n",
|
||||
"\n",
|
||||
"# Vectorize text in batches\n",
|
||||
"batch, index, batchsize = [], 0, 128\n",
|
||||
"for text in texts:\n",
|
||||
" batch.append(text)\n",
|
||||
"\n",
|
||||
" if len(batch) == batchsize:\n",
|
||||
" vectors = model.encode(batch)\n",
|
||||
" embeddings[index : index + vectors.shape[0]] = vectors\n",
|
||||
" index += vectors.shape[0]\n",
|
||||
" batch = []\n",
|
||||
"\n",
|
||||
"# Last batch\n",
|
||||
"if batch:\n",
|
||||
" vectors = model.encode(batch)\n",
|
||||
" embeddings[index : index + vectors.shape[0]] = vectors\n",
|
||||
"\n",
|
||||
"# Normalize embeddings\n",
|
||||
"embeddings /= np.linalg.norm(embeddings, axis=1)[:, np.newaxis]\n",
|
||||
"\n",
|
||||
"# Print shape\n",
|
||||
"embeddings.shape"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "DPWrubv5oOn7",
|
||||
"outputId": "972c1837-2404-42f4-9a5e-51f3c3f149ee"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"(120000, 384)"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 3
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Next we'll build a vector index using these embeddings!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "SDaDLMyXLGe1"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai.ann import ANNFactory\n",
|
||||
"\n",
|
||||
"# Create Faiss index using normalized embeddings\n",
|
||||
"ann = ANNFactory.create({\"backend\": \"faiss\"})\n",
|
||||
"ann.index(embeddings)\n",
|
||||
"\n",
|
||||
"# Show total\n",
|
||||
"ann.count()"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ILSfWHxVHex0",
|
||||
"outputId": "b9da6a79-778f-4338-a6b5-d693772fcdae"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"120000"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 4
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Now let's run a search."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "B_XnpIpXNKSP"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"query = model.encode([\"best planets to explore for life\"])\n",
|
||||
"query /= np.linalg.norm(query)\n",
|
||||
"\n",
|
||||
"for uid, score in ann.search(query, 3)[0]:\n",
|
||||
" print(uid, texts[uid], score)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "c2FVQlxSLKgP",
|
||||
"outputId": "825ca5de-d765-4c67-fdde-c7fe06557095"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"17752 Rocky Road: Planet hunting gets closer to Earth Astronomers have discovered the three lightest planets known outside the solar system, moving researchers closer to the goal of finding extrasolar planets that resemble Earth. 0.599043607711792\n",
|
||||
"16158 Earth #39;s #39;big brothers #39; floating around stars Washington - A new class of planets has been found orbiting stars besides our sun, in a possible giant leap forward in the search for Earth-like planets that might harbour life. 0.5688529014587402\n",
|
||||
"45029 Coming Soon: \"Good\" Jupiters Most of the extrasolar planets discovered to date are gas giants like Jupiter, but their orbits are either much closer to their parent stars or are highly eccentric. Planet hunters are on the verge of confirming the discovery of Jupiter-size planets with Jupiter-like orbits. Solar systems that contain these \"good\" Jupiters may harbor habitable Earth-like planets as well. 0.5606889724731445\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"And there it is, a full vector search system without using the `embeddings` package.\n",
|
||||
"\n",
|
||||
"Just as a reminder, the following much simpler code does the same thing with an Embeddings instance."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "00dnum6fNNM0"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai.embeddings import Embeddings\n",
|
||||
"\n",
|
||||
"embeddings = Embeddings({\"path\": \"sentence-transformers/all-MiniLM-L6-v2\"})\n",
|
||||
"embeddings.index((x, text, None) for x, text in enumerate(texts))\n",
|
||||
"\n",
|
||||
"for uid, score in embeddings.search(\"best planets to explore for life\"):\n",
|
||||
" print(uid, texts[uid], score)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "OYAqPoTmNaNN",
|
||||
"outputId": "30b6305a-11da-4439-e0f5-646aef8d96f6"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"17752 Rocky Road: Planet hunting gets closer to Earth Astronomers have discovered the three lightest planets known outside the solar system, moving researchers closer to the goal of finding extrasolar planets that resemble Earth. 0.599043607711792\n",
|
||||
"16158 Earth #39;s #39;big brothers #39; floating around stars Washington - A new class of planets has been found orbiting stars besides our sun, in a possible giant leap forward in the search for Earth-like planets that might harbour life. 0.568852961063385\n",
|
||||
"45029 Coming Soon: \"Good\" Jupiters Most of the extrasolar planets discovered to date are gas giants like Jupiter, but their orbits are either much closer to their parent stars or are highly eccentric. Planet hunters are on the verge of confirming the discovery of Jupiter-size planets with Jupiter-like orbits. Solar systems that contain these \"good\" Jupiters may harbor habitable Earth-like planets as well. 0.560688853263855\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Database\n",
|
||||
"\n",
|
||||
"When the `content` parameter is enabled, an Embeddings instance stores both vector content and raw content in a database. But the `database` package can be used standalone too."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "KfGc7iNRO0Tw"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai.database import DatabaseFactory\n",
|
||||
"\n",
|
||||
"# Load content into database\n",
|
||||
"database = DatabaseFactory.create({\"content\": True})\n",
|
||||
"database.insert((x, row, None) for x, row in enumerate(dataset))\n",
|
||||
"\n",
|
||||
"# Show total\n",
|
||||
"database.search(\"select count(*) from txtai\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "9IwcIgocPSDr",
|
||||
"outputId": "eeceee2f-cf35-414f-b9e4-ad975c118e42"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'count(*)': 120000}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 7
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"The full txtai [SQL query syntax](https://neuml.github.io/txtai/embeddings/query/#sql) is available, including working with dynamically created fields."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "l18dapNgSqnA"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"database.search(\"select count(*), label from txtai group by label\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "7yPjcT3peDnZ",
|
||||
"outputId": "241c98df-9c24-47a1-8ca5-e5bace19abfb"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'count(*)': 30000, 'label': 0},\n",
|
||||
" {'count(*)': 30000, 'label': 1},\n",
|
||||
" {'count(*)': 30000, 'label': 2},\n",
|
||||
" {'count(*)': 30000, 'label': 3}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 8
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Let's run a query to find text containing the word planets."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "S8G09Ib0kRB9"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"for row in database.search(\"select id, text from txtai where text like '%planets%' limit 3\"):\n",
|
||||
" print(row[\"id\"], row[\"text\"])"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "Aq1ZvOhvQHRO",
|
||||
"outputId": "3c2f9c94-57b1-489d-c854-60b909187ff0"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"100 Comets, Asteroids and Planets around a Nearby Star (SPACE.com) SPACE.com - A nearby star thought to harbor comets and asteroids now appears to be home to planets, too. The presumed worlds are smaller than Jupiter and could be as tiny as Pluto, new observations suggest.\n",
|
||||
"102 Redesigning Rockets: NASA Space Propulsion Finds a New Home (SPACE.com) SPACE.com - While the exploration of the Moon and other planets in our solar system is nbsp;exciting, the first task for astronauts and robots alike is to actually nbsp;get to those destinations.\n",
|
||||
"272 Sharpest Image Ever Obtained of a Circumstellar Disk Reveals Signs of Young Planets MAUNA KEA, Hawaii -- The sharpest image ever taken of a dust disk around another star has revealed structures in the disk which are signs of unseen planets. Dr...\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Since this is just a SQL database, text search is quite limited. The query above just retrieved results with the word planets in it."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "xhND31tnkOY2"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Scoring\n",
|
||||
"\n",
|
||||
"Since the original txtai release, there has been a `scoring` package. The main use case for this package is building a weighted sentence embeddings vector when using word vector models. But this package can also be used standalone to build BM25, TF-IDF and/or SIF text indexes."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "-vNVSA2FQnKj"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai.scoring import ScoringFactory\n",
|
||||
"\n",
|
||||
"# Build index\n",
|
||||
"scoring = ScoringFactory.create({\"method\": \"bm25\", \"terms\": True, \"content\": True})\n",
|
||||
"scoring.index((x, text, None) for x, text in enumerate(texts))\n",
|
||||
"\n",
|
||||
"# Show total\n",
|
||||
"scoring.count()"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "LJ2FskiiQ_l_",
|
||||
"outputId": "d023191a-b9bc-47fa-be99-375b81f02e8e"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"120000"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 10
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"for row in scoring.search(\"planets explore life earth\", 3):\n",
|
||||
" print(row[\"id\"], row[\"text\"], row[\"score\"])"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "slLtmzfbRuf6",
|
||||
"outputId": "be7ba16e-fd6d-4c30-fb43-a4617bac21ca"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"16327 3 Planets Are Found Close in Size to Earth, Making Scientists Think 'Life' A trio of newly discovered worlds are much smaller than any other planets previously discovered outside of the solar system. 17.768332448130707\n",
|
||||
"16158 Earth #39;s #39;big brothers #39; floating around stars Washington - A new class of planets has been found orbiting stars besides our sun, in a possible giant leap forward in the search for Earth-like planets that might harbour life. 17.65941968170793\n",
|
||||
"16620 New Planets could advance search for Life Astronomers in Europe and the United States have found two new planets about 20 times the size of Earth beyond the solar system. The discovery might be a giant leap forward in 17.65941968170793\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"The search above ran a BM25 search across the dataset. The search will return more keyword/literal results. With proper query construction, the results can be decent.\n",
|
||||
"\n",
|
||||
"Comparing the vector search results earlier and these results are a good lesson in the differences between keyword and vector search."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "BOW5JlxYS3Rm"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Database and Scoring\n",
|
||||
"\n",
|
||||
"Earlier we showed how the `ann` and `vectors` components can be combined to build a vector search engine. Can we combine the `database` and `scoring` components to add keyword search to a database? Yes!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "gqEBeBEoTrup"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"def search(query, limit=3):\n",
|
||||
" # Get similar clauses, if any\n",
|
||||
" similar = database.parse(query).get(\"similar\")\n",
|
||||
" return database.search(query, [scoring.search(args[0], limit * 10) for args in similar] if similar else None, limit)\n",
|
||||
"\n",
|
||||
"# Rebuild scoring - only need terms index\n",
|
||||
"scoring = ScoringFactory.create({\"method\": \"bm25\", \"terms\": True})\n",
|
||||
"scoring.index((x, text, None) for x, text in enumerate(texts))\n",
|
||||
"\n",
|
||||
"for row in search(\"select id, text, score from txtai where similar('planets explore life earth') and label = 0\"):\n",
|
||||
" print(row[\"id\"], row[\"text\"], row[\"score\"])"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "MEpZAr0TUMCK",
|
||||
"outputId": "b320604b-f2a0-4546-c91d-2d2b0c449382"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"15363 NASA to Announce New Class of Planets Astronomers have discovered four new planets in a week's time, an exciting end-of-summer flurry that signals a sharper era in the hunt for new worlds. While none of these new bodies would be mistaken as Earth's twin, some appear to be noticeably smaller and more solid - more like Earth and Mars - than the gargantuan, gaseous giants identified before... 12.582923259697132\n",
|
||||
"15900 Astronomers Spot Smallest Planets Yet American astronomers say they have discovered the two smallest planets yet orbiting nearby stars, trumping a small planet discovery by European scientists five days ago and capping the latest round in a frenzied hunt for other worlds like Earth. All three of these smaller planets belong to a new class of \"exoplanets\" - those that orbit stars other than our sun, the scientists said in a briefing Tuesday... 12.563928231067155\n",
|
||||
"15879 Astronomers see two new planets US astronomers find the smallest worlds detected circling other stars and say it is a breakthrough in the search for life in space. 12.078383982352994\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"And there it is, scoring-based similarity search with the same syntax as standard txtai vector queries, including additional filters!\n",
|
||||
"\n",
|
||||
"txtai is built on vector search, machine learning and finding results based on semantic meaning. It's been well-discussed from a functionality standpoint how vector search has many advantages over keyword search. The one advantage keyword search has is speed. "
|
||||
],
|
||||
"metadata": {
|
||||
"id": "aPaZsoxnYW8I"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook walked through each of the packages used by an Embeddings index. The Embeddings index makes this all transparent and easy to use. But each of the components do stand on their own and can be individually integrated into a project!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "4L8smyyXc8q8"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,454 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
},
|
||||
"gpuClass": "standard"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Classic Topic Modeling with BM25\n",
|
||||
"\n",
|
||||
"txtai 5.0 introduced topic modeling via [semantic graphs](https://neuml.hashnode.dev/introducing-the-semantic-graph). Semantic graphs can be easily integrated into an embeddings instance to add topic modeling to a txtai index.\n",
|
||||
"\n",
|
||||
"In addition to transformers-backed models, txtai also has support for traditional indexing methods. Given the modular design of txtai, traditional scoring methods like BM25 can be combined with graphs to build topic models. \n",
|
||||
"\n",
|
||||
"This notebook is all classic Python code on the CPU. No GPUs or machine learning models required!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "-xU9P9iSR-Cy"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "shlUi2kKS7KT"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "xEvX9vCpn4E0"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[graph] datasets"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Load dataset\n",
|
||||
"\n",
|
||||
"This example will use the `ag_news` dataset, which is a collection of news article headlines."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "408IyXzKFSiG"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"dataset = load_dataset(\"ag_news\", split=\"train\")"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "IQ_ns6YvFRm1"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Build BM25 Index\n",
|
||||
"\n",
|
||||
"Since the original txtai release, there has been a `scoring` package. This package supports building standalone BM25, TF-IDF and/or SIF text indexes."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "-vNVSA2FQnKj"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai.scoring import ScoringFactory\n",
|
||||
"\n",
|
||||
"# List of all text elements\n",
|
||||
"texts = dataset[\"text\"]\n",
|
||||
"\n",
|
||||
"# Build index\n",
|
||||
"scoring = ScoringFactory.create({\"method\": \"bm25\", \"terms\": True})\n",
|
||||
"scoring.index((x, text, None) for x, text in enumerate(texts))\n",
|
||||
"\n",
|
||||
"# Show total\n",
|
||||
"scoring.count()"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "LJ2FskiiQ_l_",
|
||||
"outputId": "fc4685d9-5ec2-4fbe-a347-857a74b9b509"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"120000"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 3
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Let's test the index."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "BOW5JlxYS3Rm"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"for id, score in scoring.search(\"planets explore life earth\", 3):\n",
|
||||
" print(id, texts[id], score)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "slLtmzfbRuf6",
|
||||
"outputId": "27aca9cb-8704-475c-c38d-01da65328686"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"16327 3 Planets Are Found Close in Size to Earth, Making Scientists Think 'Life' A trio of newly discovered worlds are much smaller than any other planets previously discovered outside of the solar system. 20.72295380862701\n",
|
||||
"16158 Earth #39;s #39;big brothers #39; floating around stars Washington - A new class of planets has been found orbiting stars besides our sun, in a possible giant leap forward in the search for Earth-like planets that might harbour life. 19.917461045326878\n",
|
||||
"16620 New Planets could advance search for Life Astronomers in Europe and the United States have found two new planets about 20 times the size of Earth beyond the solar system. The discovery might be a giant leap forward in 19.917461045326878\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Results look as expected. BM25 returns keyword-based results vs contextual matches."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "XnAhNb8Td6Wd"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Build topic model\n",
|
||||
"\n",
|
||||
"Now that we have a scoring index, we'll use it to build a graph.\n",
|
||||
"\n",
|
||||
"Graphs have built-in methods to insert nodes and build a relationship index between the nodes. The `index` method takes a search parameter that can be any function that returns (id, score) pairs. This logic is built into embeddings instances. \n",
|
||||
"\n",
|
||||
"Graphs constructed via a BM25 index will have more literal relationships. In other words, it will be keyword-driven. Semantic graphs backed by embeddings will have contextual relationships.\n",
|
||||
"\n",
|
||||
"The next section builds a graph to support topic modeling. We'll use a multiprocessing pool to maximize CPU usage."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "gqEBeBEoTrup"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import os\n",
|
||||
"\n",
|
||||
"from multiprocessing import Pool\n",
|
||||
"\n",
|
||||
"from txtai.graph import GraphFactory\n",
|
||||
"\n",
|
||||
"# Multiprocessing helper methods\n",
|
||||
"SCORING = None\n",
|
||||
"\n",
|
||||
"def create(search):\n",
|
||||
" global SCORING\n",
|
||||
"\n",
|
||||
" # Create a global scoring object\n",
|
||||
" SCORING = search\n",
|
||||
"\n",
|
||||
"def run(params):\n",
|
||||
" query, limit = params\n",
|
||||
" return SCORING.search(query, limit)\n",
|
||||
"\n",
|
||||
"def batchsearch(queries, limit):\n",
|
||||
" return pool.imap(run, [(query, limit) for query in queries])\n",
|
||||
"\n",
|
||||
"# Build the graph\n",
|
||||
"pool = None\n",
|
||||
"with Pool(os.cpu_count(), initializer=create, initargs=(scoring,)) as pool:\n",
|
||||
" graph = GraphFactory.create({\"topics\": {}})\n",
|
||||
" graph.insert((x, text, None) for x, text in enumerate(texts))\n",
|
||||
" graph.index(batchsearch, None)"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "MEpZAr0TUMCK"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Let's list the top 10 topics. Keep in mind this dataset is from 2004."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "aPaZsoxnYW8I"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"list(graph.topics)[:10]"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "HyKiVFBerQfw",
|
||||
"outputId": "fdba468b-bac1-4f63-c915-5fcad78880f7"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"['kerry_bush_john_president',\n",
|
||||
" 'nhl_players_league_lockout',\n",
|
||||
" 'arafat_yasser_palestinian_leader',\n",
|
||||
" 'sharon_ariel_prime_minister',\n",
|
||||
" 'blair_tony_minister_prime',\n",
|
||||
" 'xp_windows_microsoft_sp2',\n",
|
||||
" 'athens_gold_medal_olympic',\n",
|
||||
" 'space_prize_million_spaceshipone',\n",
|
||||
" 'nikkei_tokyo_reuters_average',\n",
|
||||
" 'hostage_british_bigley_iraq']"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 10
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Topics map a list of ids for each matching text element ordered by topic relevance. Let's print the most relevant text element for a topic."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "BKPJuq6br-ru"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"uid = graph.topics[\"xp_windows_microsoft_sp2\"][0]\n",
|
||||
"graph.attribute(uid, \"text\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 54
|
||||
},
|
||||
"id": "sYd9nKqYrwhe",
|
||||
"outputId": "0db63a74-2c01-433a-d754-6f0988894ffe"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'Microsoft continues Windows XP SP2 distribution Continuing the roll-out of Windows XP Service Pack 2 (SP2), Microsoft Corp. on Wednesday began pushing the security-focused update to PCs running Windows XP Professional Edition '"
|
||||
],
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 44
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Graph analysis\n",
|
||||
"\n",
|
||||
"Given this is a standard txtai graph, analysis methods such as centrality and pagerank are available."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "muN33RNB0Kob"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"centrality = list(graph.centrality().keys())\n",
|
||||
"print(\"Top connection count:\", [len(graph.edges(uid)) for uid in centrality[:5]], \"\\n\")\n",
|
||||
"\n",
|
||||
"# Print most central node/topic\n",
|
||||
"print(\"Most central node:\", graph.attribute(centrality[0], \"text\"))\n",
|
||||
"\n",
|
||||
"topic = graph.attribute(centrality[0], \"topic\")\n",
|
||||
"for uid in graph.topics[topic][:3]:\n",
|
||||
" print(\"->\", graph.attribute(uid, \"text\"))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "hpwfJFcRv19j",
|
||||
"outputId": "0bb4d53d-27c2-42f5-cfaf-89c785cb73da"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Top connection count: [30, 30, 28, 28, 28] \n",
|
||||
"\n",
|
||||
"Most central node: Manning Gets Chance to Start Giants Coach Tom Coughlin announced that rookie quarterback Eli Manning will start ahead of two-time M.V.P. Kurt Warner in Thursday's preseason game against Carolina.\n",
|
||||
"-> Manning Replaces Warner As Giants QB (AP) AP - Eli Manning has replaced Kurt Warner as the New York Giants' starting quarterback.\n",
|
||||
"-> Eli Manning replaces Warner at quarterback Eli Manning, the top pick in this year #39;s NFL draft, has been named the starting quarterback of the New York Giants. Coach Tom Coughlin made the announcement at a Monday news conference.\n",
|
||||
"-> Giants to Start Manning Against Carolina (AP) AP - Eli Manning is going to get a chance to open the season as the New York Giants' starting quarterback.\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Notice the correlation between the number of connections and centrality.\n",
|
||||
"\n",
|
||||
"Given that BM25 is keyword-driven, we expect that the most central node would be text that is duplicative in nature. And that is the case here."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "lqOdJR6a0ftB"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Walk the graph\n",
|
||||
"\n",
|
||||
"Just like semantic graphs, relationship paths can be explored."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "8u_tTNrq2EoN"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from IPython.display import HTML\n",
|
||||
"\n",
|
||||
"def showpath(source, target):\n",
|
||||
" path = graph.showpath(source, target)\n",
|
||||
" path = [graph.attribute(p, \"text\") for p in path]\n",
|
||||
"\n",
|
||||
" sections = []\n",
|
||||
" for x, p in enumerate(path):\n",
|
||||
" # Print start node\n",
|
||||
" sections.append(f\"{x + 1}. {p}\")\n",
|
||||
"\n",
|
||||
" return HTML(\"<br/><br/>\".join(sections))"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "oodIUmwrsZRd"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"showpath(83978, 8107)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 208
|
||||
},
|
||||
"id": "gZzStey9tEkl",
|
||||
"outputId": "8e6398d5-0fd8-465c-ce95-45b98eb2c195"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
],
|
||||
"text/html": [
|
||||
"1. NFL Game Summary - NY Jets at Buffalo Orchard Park, NY (Sports Network) - Willis McGahee ran for 132 yards and a touchdown to lead the Buffalo Bills to a 22-17 victory over the New York Jets at Ralph Wilson Stadium.<br/><br/>2. NCAA Game Summary - Marshall at Georgia Athens, GA (Sports Network) - Michael Cooper ran for the only touchdown of the game, as third-ranked Georgia rode its defense to a 13-3 victory over Marshall at Sanford Stadium.<br/><br/>3. NCAA Game Summary - Northwestern at Wisconsin Madison, WI (Sports Network) - Anthony Davis ran for 122 yards and two touchdowns to lead No. 6 Wisconsin over Northwestern, 24-12, to celebrate Homecoming weekend at Camp Randall Stadium.<br/><br/>4. NCAA Top 25 Game Summary - Northwestern at Minnesota The last time Minnesota won four games to start three consecutive seasons was 1934-36...Chris Malleo replaced Basanez for two series in the third quarter for his first career appearance.<br/><br/>5. UConn ousts Marist Sophomore Steve Sealy netted his third winning goal in the last four games, giving Connecticut a 2-1 overtime victory over Marist yesterday in an NCAA Division 1 first-round men's soccer playoff game at Morrone Stadium in Storrs, Conn.<br/><br/>6. United States upsets Germany to move to soccer semifinals Deep into overtime, and maybe the last time for the Fab Five of US women #39;s soccer, the breaks were going against them. A last-gasp goal that stole victory in regulation, a wide-open shot that bounced off the goal post."
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 43
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Notice how the data pivots from the start node to the end node. If you've read the [Introducing the Semantic Graph](https://neuml.hashnode.dev/introducing-the-semantic-graph) article, you'll notice how this traversal is more literal in nature. In other words, the relationships are keyword-driven vs contextual."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "5qv1CbY35uUy"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook demonstrated how graphs can index traditional indexes such as BM25. This method can also be applied to an external index provided a search function is available to build connections.\n",
|
||||
"\n",
|
||||
"Semantic graphs backed by embeddings instances have a number of advantages and are recommended in most cases. But this is a classic way to do it - no machine learning models required!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "4L8smyyXc8q8"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,499 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
},
|
||||
"gpuClass": "standard",
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Train a language model from scratch\n",
|
||||
"\n",
|
||||
"txtai has a robust training pipeline that can fine-tune large language models (LLMs) for downstream tasks such as labeling text. txtai also has the ability to train language models from scratch.\n",
|
||||
"\n",
|
||||
"The vast majority of time, fine-tuning a LLM yields the best results. But when making significant changes to the structure of a model, training from scratch is often required.\n",
|
||||
"\n",
|
||||
"Examples of significant changes are:\n",
|
||||
"\n",
|
||||
"- Changing the vocabulary size\n",
|
||||
"- Changing the number of hidden dimensions\n",
|
||||
"- Changing the number of attention heads or layers\n",
|
||||
"- Create a custom model architecture\n",
|
||||
"\n",
|
||||
"This notebook will show how to build a new tokenizer and train a small language model (known as a micromodel) from scratch.\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "-xU9P9iSR-Cy"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "shlUi2kKS7KT"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "xEvX9vCpn4E0"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline-train] datasets sentence-transformers onnxruntime onnx"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Load dataset\n",
|
||||
"\n",
|
||||
"This example will use the `ag_news` dataset, which is a collection of news article headlines."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "408IyXzKFSiG"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"dataset = load_dataset(\"ag_news\", split=\"train\")"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "IQ_ns6YvFRm1"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Train the tokenizer\n",
|
||||
"\n",
|
||||
"The first step is to train the tokenizer. We could use an existing tokenizer but in this case, we want a smaller vocabulary.\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "-vNVSA2FQnKj"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from transformers import AutoTokenizer\n",
|
||||
"\n",
|
||||
"def stream(batch=10000):\n",
|
||||
" for x in range(0, len(dataset), batch):\n",
|
||||
" yield dataset[x: x + batch][\"text\"]\n",
|
||||
"\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")\n",
|
||||
"tokenizer = tokenizer.train_new_from_iterator(stream(), vocab_size=500, length=len(dataset))\n",
|
||||
"tokenizer.model_max_length = 512\n",
|
||||
"\n",
|
||||
"tokenizer.save_pretrained(\"bert\")"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "LJ2FskiiQ_l_"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Let's test the tokenizer."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "BOW5JlxYS3Rm"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"print(tokenizer.tokenize(\"Red Sox defeat Yankees 5-3\"))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "slLtmzfbRuf6",
|
||||
"outputId": "2391b58b-7428-49a2-9225-0268a1f2ad64"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"['re', '##d', 'so', '##x', 'de', '##f', '##e', '##at', 'y', '##ank', '##e', '##es', '5', '-', '3']\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"With a limited vocabulary size of 500, most words require multiple tokens. This limited vocabulary lowers the number of token representations the model needs to learn."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "IozRdXdEqegD"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Train the language model\n",
|
||||
"\n",
|
||||
"Now it's time to train the model. We'll train a micromodel, which is an extremely small language model with a limited vocabulary. Micromodels, when paired with a limited vocabulary have the potential to work in limited compute environments like edge devices and microcontrollers."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "gqEBeBEoTrup"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from transformers import AutoTokenizer, BertConfig, BertForMaskedLM\n",
|
||||
"\n",
|
||||
"from txtai.pipeline import HFTrainer\n",
|
||||
"\n",
|
||||
"config = BertConfig(\n",
|
||||
" vocab_size = 500,\n",
|
||||
" hidden_size = 50,\n",
|
||||
" num_hidden_layers = 2,\n",
|
||||
" num_attention_heads = 2,\n",
|
||||
" intermediate_size = 100,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"model = BertForMaskedLM(config)\n",
|
||||
"model.save_pretrained(\"bert\")\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(\"bert\")\n",
|
||||
"\n",
|
||||
"train = HFTrainer()\n",
|
||||
"\n",
|
||||
"# Train model\n",
|
||||
"train((model, tokenizer), dataset, task=\"language-modeling\", output_dir=\"bert\",\n",
|
||||
" fp16=True, per_device_train_batch_size=128, num_train_epochs=10,\n",
|
||||
" dataloader_num_workers=2)"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "MEpZAr0TUMCK"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Sentence embeddings\n",
|
||||
"\n",
|
||||
"Next let's take the language model and fine-tune it to build sentence embeddings. "
|
||||
],
|
||||
"metadata": {
|
||||
"id": "53bvB9c6MbPS"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!wget https://raw.githubusercontent.com/UKPLab/sentence-transformers/master/examples/training/nli/training_nli_v2.py\n",
|
||||
"!python training_nli_v2.py bert\n",
|
||||
"!mv output/* bert-nli"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "f11f5tjfS85m"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Embeddings search\n",
|
||||
"\n",
|
||||
"Now we'll build a txtai embeddings index using the fine-tuned model. We'll index the `ag_news` dataset. "
|
||||
],
|
||||
"metadata": {
|
||||
"id": "FTOm5ofaMmcv"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai.embeddings import Embeddings\n",
|
||||
"\n",
|
||||
"# Get list of all text\n",
|
||||
"texts = dataset[\"text\"]\n",
|
||||
"\n",
|
||||
"embeddings = Embeddings({\"path\": \"bert-nli\", \"content\": True})\n",
|
||||
"embeddings.index((x, text, None) for x, text in enumerate(texts))"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "_kKe5kRnVRhM"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Let's run a search and see how much the model has learned."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Rh9yA6ZJM47H"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"embeddings.search(\"Boston Red Sox Cardinals World Series\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "XCRhpfLmV1-q",
|
||||
"outputId": "662f1b1d-fcf5-4383-fd8a-369081a77501"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '76733',\n",
|
||||
" 'text': 'Red Sox sweep Cardinals to win World Series The Boston Red Sox ended their 86-year championship drought with a 3-0 win over the St. Louis Cardinals in Game Four of the World Series.',\n",
|
||||
" 'score': 0.8008379936218262},\n",
|
||||
" {'id': '71169',\n",
|
||||
" 'text': 'Red Sox lead 2-0 over Cardinals of World Series The host Boston Red Sox scored a 6-2 victory over the St. Louis Cardinals, helped by Curt Schilling #39;s pitching through pain and seeping blood, in World Series Game 2 on Sunday night.',\n",
|
||||
" 'score': 0.7896029353141785},\n",
|
||||
" {'id': '70100',\n",
|
||||
" 'text': 'Sports: Red Sox 9 Cardinals 7 after 7 innings BOSTON Boston has scored twice in the seventh inning to take an 9-to-7 lead over the St. Louis Cardinals in the World Series opener at Fenway Park.',\n",
|
||||
" 'score': 0.7735188603401184}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 49
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Not too bad. It's far from perfect but we can tell that it has some knowledge! This model was trained for 5 minutes, there is certainly room for improvement in training longer and/or with a larger dataset.\n",
|
||||
"\n",
|
||||
"The standard `bert-base-uncased` model has 110M parameters and is around 440MB. Let's see how many parameters this model has."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "M5Pk1spcM72L"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Show number of parameters\n",
|
||||
"parameters = sum(p.numel() for p in embeddings.model.model.parameters())\n",
|
||||
"print(f\"Number of parameters:\\t\\t{parameters:,}\")\n",
|
||||
"print(f\"% of bert-base-uncased\\t\\t{(parameters / 110000000) * 100:.2f}%\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "RIEnDwuxeakq",
|
||||
"outputId": "131930f1-e6cd-4b23-b9dc-a62e670eb8e4"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Number of parameters:\t\t94,450\n",
|
||||
"% of bert-base-uncased\t\t0.09%\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!ls -lh bert-nli/pytorch_model.bin"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ATg4aInQeRN-",
|
||||
"outputId": "a6181fbc-ee6c-426e-882e-1d6cbefa22a0"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"-rw-r--r-- 1 root root 386K Jan 11 20:52 bert-nli/pytorch_model.bin\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"This model is 386KB and has only 0.1% of the parameters. With proper vocabulary selection, a small language model has potential."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "7GfsF8ziNFJa"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Quantization\n",
|
||||
"\n",
|
||||
"If 386KB isn't small enough, we can quantize the model to get it down even further. "
|
||||
],
|
||||
"metadata": {
|
||||
"id": "CcbJNidNwuXt"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai.pipeline import HFOnnx\n",
|
||||
"\n",
|
||||
"onnx = HFOnnx()\n",
|
||||
"onnx(\"bert-nli\", task=\"pooling\", output=\"bert-nli.onnx\", quantize=True)"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "IYZnex9kRcb0"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"embeddings = Embeddings({\"path\": \"bert-nli.onnx\", \"tokenizer\": \"bert-nli\", \"content\": True})\n",
|
||||
"embeddings.index((x, text, None) for x, text in enumerate(texts))\n",
|
||||
"embeddings.search(\"Boston Red Sox Cardinals World Series\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "QL_1UosIVkZ7",
|
||||
"outputId": "6b2bedb8-5cc6-44b6-855a-617a6a07478c"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '76733',\n",
|
||||
" 'text': 'Red Sox sweep Cardinals to win World Series The Boston Red Sox ended their 86-year championship drought with a 3-0 win over the St. Louis Cardinals in Game Four of the World Series.',\n",
|
||||
" 'score': 0.8008379936218262},\n",
|
||||
" {'id': '71169',\n",
|
||||
" 'text': 'Red Sox lead 2-0 over Cardinals of World Series The host Boston Red Sox scored a 6-2 victory over the St. Louis Cardinals, helped by Curt Schilling #39;s pitching through pain and seeping blood, in World Series Game 2 on Sunday night.',\n",
|
||||
" 'score': 0.7896029353141785},\n",
|
||||
" {'id': '70100',\n",
|
||||
" 'text': 'Sports: Red Sox 9 Cardinals 7 after 7 innings BOSTON Boston has scored twice in the seventh inning to take an 9-to-7 lead over the St. Louis Cardinals in the World Series opener at Fenway Park.',\n",
|
||||
" 'score': 0.7735188603401184}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 50
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!ls -lh bert-nli.onnx"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "oCa3RDeInCkN",
|
||||
"outputId": "89cf379a-09cf-4fbc-8568-4d66085450f3"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"-rw-r--r-- 1 root root 187K Jan 11 20:53 bert-nli.onnx\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"We're down to 187KB with a quantized model!\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "r95T1HZXVhnZ"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Train on BERT dataset\n",
|
||||
"\n",
|
||||
"The [BERT paper](https://arxiv.org/abs/1810.04805) has all the information regarding training parameters and datasets used. Hugging Face Datasets hosts the `bookcorpus` and `wikipedia` datasets.\n",
|
||||
"\n",
|
||||
"Training on this size of a dataset is out of scope for this notebook but example code is shown below on how to build the BERT dataset.\n",
|
||||
"\n",
|
||||
"```python\n",
|
||||
"bookcorpus = load_dataset(\"bookcorpus\", split=\"train\")\n",
|
||||
"wiki = load_dataset(\"wikipedia\", \"20220301.en\", split=\"train\")\n",
|
||||
"wiki = wiki.remove_columns([col for col in wiki.column_names if col != \"text\"])\n",
|
||||
"dataset = concatenate_datasets([bookcorpus, wiki])\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"Then the same steps to train the tokenizer and model can be run. The dataset is 25GB compressed, so it will take some space and time to process!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "aPaZsoxnYW8I"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook covered how to build micromodels from scratch with txtai. Micromodels can be fully rebuilt in hours using the most up-to-date knowledge available. If properly constructed, prepared and trained, micromodels have the potential to be a viable choice for limited resource environments. They can also help when realtime response is more important than having the highest accuracy scores.\n",
|
||||
"\n",
|
||||
"It's our hope that further research and exploration into micromodels leads to productive and useful models."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "4L8smyyXc8q8"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,401 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "vwELCooy4ljr"
|
||||
},
|
||||
"source": [
|
||||
"# Prompt-driven search with LLMs\n",
|
||||
"\n",
|
||||
"This notebook revisits the RAG pipeline, which has been covered in a number of previous notebooks. This pipeline is a combination of a similarity instance (embeddings or similarity pipeline) to build a question context and a model that answers questions.\n",
|
||||
"\n",
|
||||
"The RAG pipeline recently underwent a number of major upgrades to support the following.\n",
|
||||
"\n",
|
||||
"- Ability to run embeddings searches. Given that content is supported, text can be retrieved from the embeddings instance.\n",
|
||||
"- In addition to extractive qa, support text generation models, sequence to sequence models and custom pipelines\n",
|
||||
"\n",
|
||||
"These changes enable embeddings-guided and prompt-driven search with Large Language Models (LLMs) 🔥🔥🔥"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ew7orE2O441o"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "LPQTb25tASIG"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai datasets"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "_YnqorRKAbLu"
|
||||
},
|
||||
"source": [
|
||||
"# Create Embeddings and RAG instances\n",
|
||||
"\n",
|
||||
"An Embeddings instance defines methods to represent text as vectors and build vector indexes for search.\n",
|
||||
"\n",
|
||||
"The RAG pipeline is a combination of a similarity instance (embeddings or similarity pipeline) to build a question context and a model that answers questions. The model can be a prompt-driven large language model (LLM), an extractive question-answering model or a custom pipeline.\n",
|
||||
"\n",
|
||||
"Let's run a basic example.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"id": "OUc9gqTyAYnm"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"\n",
|
||||
"from txtai import Embeddings, RAG\n",
|
||||
"\n",
|
||||
"# Create embeddings model with content support\n",
|
||||
"embeddings = Embeddings(path=\"sentence-transformers/all-MiniLM-L6-v2\", content=True)\n",
|
||||
"\n",
|
||||
"# Create the RAG pipeline\n",
|
||||
"rag = RAG(embeddings, \"Qwen/Qwen3-4B-Instruct-2507\", template=\"\"\"\n",
|
||||
" Answer the following question using the provided context.\n",
|
||||
"\n",
|
||||
" Question:\n",
|
||||
" {question}\n",
|
||||
"\n",
|
||||
" Context:\n",
|
||||
" {context}\n",
|
||||
"\"\"\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "4X5z3UjnAGe7",
|
||||
"outputId": "cacb7e9d-471a-437d-c68d-a8b51f876413"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"---- Red Sox - Blue Jays ----\n",
|
||||
"{'answer': 'The Blue Jays won the game.'}\n",
|
||||
"{'answer': 'The score was 2-1 in favor of the Blue Jays.'}\n",
|
||||
"\n",
|
||||
"---- Phillies - Braves ----\n",
|
||||
"{'answer': 'The Phillies won the game.'}\n",
|
||||
"{'answer': 'The score was 5-0 in favor of the Phillies.'}\n",
|
||||
"\n",
|
||||
"---- Dodgers - Giants ----\n",
|
||||
"{'answer': 'The Giants won the game.'}\n",
|
||||
"{'answer': 'The score was Giants 5, Dodgers 4.'}\n",
|
||||
"\n",
|
||||
"---- Flyers - Lightning ----\n",
|
||||
"{'answer': 'The Flyers won the game.'}\n",
|
||||
"{'answer': 'The score was Flyers 4, Lightning 1.'}\n",
|
||||
"\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"data = [\"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 beat Red Sox final score 2-1\",\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",
|
||||
" \"Lightning goaltender pulled, lose to Flyers 4-1\",\n",
|
||||
" \"Flyers 4 Lightning 1 final\",\n",
|
||||
" \"Flyers win 4-1\"]\n",
|
||||
"\n",
|
||||
"questions = [\"What team won the game?\", \"What was score?\"]\n",
|
||||
"\n",
|
||||
"for query in [\"Red Sox - Blue Jays\", \"Phillies - Braves\", \"Dodgers - Giants\", \"Flyers - Lightning\"]:\n",
|
||||
" print(\"----\", query, \"----\")\n",
|
||||
" for answer in rag([f\"{query} {x}\" for x in questions], data):\n",
|
||||
" print(answer)\n",
|
||||
" print()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "7AnPvSeM3N1Z"
|
||||
},
|
||||
"source": [
|
||||
"This code runs a series of questions. First it runs an embeddings filtering query to find the most relevant text. For example, `Red Sox - Blue Jays` finds text related to those teams. Then `What team won the game?` and `What was the score?` are asked.\n",
|
||||
"\n",
|
||||
"This logic is the same logic found in Notebook 5 - Extractive QA with txtai but uses prompt-based QA vs extractive QA. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Aj8GoDk331cS"
|
||||
},
|
||||
"source": [
|
||||
"# Embeddings-guided and Prompt-driven Search\n",
|
||||
"\n",
|
||||
"Now for the fun stuff. Let's build an embeddings index for the `ag_news` dataset (a set of news stories from the mid 2000s). Then we'll use prompts to ask questions with embeddings results as the context."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "yL716oEZ43t-",
|
||||
"outputId": "23f4b0e7-a60a-4e89-fb57-06966e6612f8"
|
||||
},
|
||||
"outputs": [
|
||||
],
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"dataset = load_dataset(\"ag_news\", split=\"train\")\n",
|
||||
"\n",
|
||||
"# Create an embeddings index over the dataset\n",
|
||||
"embeddings = Embeddings(path=\"sentence-transformers/all-MiniLM-L6-v2\", content=True)\n",
|
||||
"embeddings.index(dataset[\"text\"])\n",
|
||||
"\n",
|
||||
"# Create RAG instance\n",
|
||||
"rag = RAG(embeddings, \"Qwen/Qwen3-4B-Instruct-2507\", template=\"\"\"\n",
|
||||
" Answer the following question using the provided context.\n",
|
||||
"\n",
|
||||
" Question:\n",
|
||||
" {question}\n",
|
||||
"\n",
|
||||
" Context:\n",
|
||||
" {context}\n",
|
||||
"\"\"\", output=\"flatten\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Ifl8JwLDBL7k"
|
||||
},
|
||||
"source": [
|
||||
"Now let's run a prompt-driven search!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 22,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "5O1WBJ8153Mo",
|
||||
"outputId": "ddf09da2-7d4c-4fd3-b0da-df5631bacd13"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Who won the 2004 presidential election? George W. Bush won the 2004 presidential election.\n",
|
||||
"Who did the candidate beat? George W. Bush beat John F. Kerry in the 2004 presidential election.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"question = \"Who won the 2004 presidential election?\"\n",
|
||||
"answer = rag(question)\n",
|
||||
"print(question, answer)\n",
|
||||
"\n",
|
||||
"nquestion = \"Who did the candidate beat?\"\n",
|
||||
"print(nquestion, rag(f\"{question} {answer}. {nquestion}\"))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AhViFXH_BZSo"
|
||||
},
|
||||
"source": [
|
||||
"And there are the answers. Let's unpack how this works.\n",
|
||||
"\n",
|
||||
"The first thing the RAG pipeline does is run an embeddings search to find the most relevant text within the index. A context string is then built using those search results.\n",
|
||||
"\n",
|
||||
"After that, a prompt is generated, run and the answer printed."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "JtDcVPdOB0Rv"
|
||||
},
|
||||
"source": [
|
||||
"# Additional examples\n",
|
||||
"\n",
|
||||
"Before moving on, a couple more example questions."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 23,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "0NNLBwC-83MM",
|
||||
"outputId": "4b9fabd7-baf9-4f7d-bb8b-c9c60c5101e4"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Who won the World Series in 2004? The Boston Red Sox won the World Series in 2004.\n",
|
||||
"What team did the Red Sox beat in the World Series? The Boston Red Sox beat the St. Louis Cardinals in the World Series.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"question = \"Who won the World Series in 2004?\"\n",
|
||||
"answer = rag(question)\n",
|
||||
"print(question, answer)\n",
|
||||
"\n",
|
||||
"nquestion = \"What team did the Red Sox beat in the World Series?\"\n",
|
||||
"print(nquestion, rag(f\"{question} {answer}. {nquestion}\"))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 29,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 36
|
||||
},
|
||||
"id": "1P0zqkTW9cZW",
|
||||
"outputId": "4f2232db-761b-464f-b47d-6695c84ffb80"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'An interesting fact is that herrings communicate by farting—a quirky and unusual discovery that was honored with an Ig Nobel Prize for its oddball research.'"
|
||||
]
|
||||
},
|
||||
"execution_count": 29,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"rag(\"Tell me something interesting\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ygFFcwWPGI9p"
|
||||
},
|
||||
"source": [
|
||||
"Whhaaaattt??? Is this a model hallucination?\n",
|
||||
"\n",
|
||||
"Let's run an embeddings query and see if that text is in the results."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 28,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "qZPhLqSxGMbK",
|
||||
"outputId": "e3b2909a-1ee8-480f-afa3-201a8e27cb08"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"herrings communicate by farting\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"answer = \"herrings communicate by farting\"\n",
|
||||
"for x in embeddings.search(\"Tell me something interesting\"):\n",
|
||||
" if answer in x[\"text\"]:\n",
|
||||
" start = x[\"text\"].find(answer)\n",
|
||||
" print(x[\"text\"][start:start + len(answer)])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "IpgxMc1DZcds"
|
||||
},
|
||||
"source": [
|
||||
"Sure enough it is 😃"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KqfvCXp2B3li"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook covered how to run embeddings-guided and prompt-driven search with LLMs. This functionality is a major step forward towards `Generative Semantic Search` for txtai. More to come, stay tuned!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"gpuClass": "standard",
|
||||
"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.10.19"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
@@ -0,0 +1,496 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"accelerator": "GPU",
|
||||
"gpuClass": "standard"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "vwELCooy4ljr"
|
||||
},
|
||||
"source": [
|
||||
"# Embeddings in the Cloud\n",
|
||||
"\n",
|
||||
"Embeddings is the engine that delivers semantic search. Data is transformed into embeddings vectors where similar concepts will produce similar vectors. Indexes both large and small are built with these vectors. The indexes are used to find results that have the same meaning, not necessarily the same keywords. \n",
|
||||
"\n",
|
||||
"In addition to local storage, embeddings can be synced with cloud storage. Given that txtai is a fully encapsulated index format, cloud sync is simply a matter of moving a group of files to and from cloud storage. This can be object storage such as AWS S3/Azure Blob/Google Cloud or the [Hugging Face Hub](https://hf.co/models). More details on available options can be found in the [documentation](https://neuml.github.io/txtai/embeddings/configuration/cloud/). There is also an [article](https://medium.com/neuml/serverless-vector-search-with-txtai-96f6163ab972) available that covers how to build and store indexes in cloud object storage.\n",
|
||||
"\n",
|
||||
"This notebook will cover an example of loading embeddings indexes from the Hugging Face Hub."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ew7orE2O441o"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "LPQTb25tASIG"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "_YnqorRKAbLu"
|
||||
},
|
||||
"source": [
|
||||
"# Integration with the Hugging Face Hub\n",
|
||||
"\n",
|
||||
"The Hugging Face Hub has a vast array of models, datasets and example applications available to jumpstart your project. This now includes txtai indexes 🔥🔥🔥\n",
|
||||
"\n",
|
||||
"Let's load the embeddings used in the standard [Introducing txtai](https://huggingface.co/NeuML/txtai-intro) example.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "OUc9gqTyAYnm"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"from txtai.embeddings import Embeddings\n",
|
||||
"\n",
|
||||
"# Load the index from the Hub\n",
|
||||
"embeddings = Embeddings()\n",
|
||||
"embeddings.load(provider=\"huggingface-hub\", container=\"neuml/txtai-intro\")"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Notice the two fields, `provider` and `container`. The `provider` field tells txtai to look for an index in the Hugging Face Hub. The `container` field sets the target repository."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "oW5juxjBbrQ-"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "4X5z3UjnAGe7",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "7a052744-c727-4bc9-d892-35b23775011d"
|
||||
},
|
||||
"source": [
|
||||
"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\", \"public health story\", \"war\", \"wildlife\", \"asia\", \"lucky\", \"dishonest junk\"):\n",
|
||||
" # Get to the top result\n",
|
||||
" result = embeddings.search(query, 1)[0]\n",
|
||||
"\n",
|
||||
" # Print text\n",
|
||||
" print(\"%-20s %s\" % (query, result[\"text\"]))"
|
||||
],
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"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",
|
||||
"public health story 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",
|
||||
"lucky Maine man wins $1M from $25 lottery ticket\n",
|
||||
"dishonest junk Make huge profits without work, earn up to $100,000 a day\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"If you've seen txtai before, this is the classic example. The big difference though is the index was loaded from the Hugging Face Hub instead of being built dynamically."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "7AnPvSeM3N1Z"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wikipedia search with txtai\n",
|
||||
"\n",
|
||||
"Let's try something more interesting using the [Wikipedia index available on the Hugging Face Hub](https://huggingface.co/NeuML/txtai-wikipedia)"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Aj8GoDk331cS"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"from txtai.embeddings import Embeddings\n",
|
||||
"\n",
|
||||
"# Load the index from the Hub\n",
|
||||
"embeddings = Embeddings()\n",
|
||||
"embeddings.load(provider=\"huggingface-hub\", container=\"neuml/txtai-wikipedia\")"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "yL716oEZ43t-"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from IPython.display import HTML, display\n",
|
||||
"\n",
|
||||
"def wrap():\n",
|
||||
" display(HTML(\"\"\"<style>pre { white-space: pre-wrap; }</style>\"\"\"))\n",
|
||||
"\n",
|
||||
"get_ipython().events.register('pre_run_cell', wrap)"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "1V8nrm9IHPQc"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Now run a series of searches to show the kind of data available in this index."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Ifl8JwLDBL7k"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import json\n",
|
||||
"\n",
|
||||
"for x in embeddings.search(\"Roman Empire\", 1):\n",
|
||||
" print(json.dumps(x, indent=2))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 196
|
||||
},
|
||||
"id": "5O1WBJ8153Mo",
|
||||
"outputId": "f20e0688-54e7-43ce-e2f9-6f11f94026bc"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
],
|
||||
"text/html": [
|
||||
"<style>pre { white-space: pre-wrap; }</style>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"{\n",
|
||||
" \"id\": \"Roman Empire\",\n",
|
||||
" \"text\": \"The Roman Empire ( ; ) was the post-Republican period of ancient Rome. As a polity, it included large territorial holdings around the Mediterranean Sea in Europe, North Africa, and Western Asia, and was ruled by emperors. From the accession of Caesar Augustus as the first Roman emperor to the military anarchy of the 3rd century, it was a Principate with Italia as the metropole of its provinces and the city of Rome as its sole capital. The Empire was later ruled by multiple emperors who shared control over the Western Roman Empire and the Eastern Roman Empire. The city of Rome remained the nominal capital of both parts until AD 476 when the imperial insignia were sent to Constantinople following the capture of the Western capital of Ravenna by the Germanic barbarians. The adoption of Christianity as the state church of the Roman Empire in AD 380 and the fall of the Western Roman Empire to Germanic kings conventionally marks the end of classical antiquity and the beginning of the Middle Ages. Because of these events, along with the gradual Hellenization of the Eastern Roman Empire, historians distinguish the medieval Roman Empire that remained in the Eastern provinces as the Byzantine Empire.\",\n",
|
||||
" \"score\": 0.8913329243659973\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"for x in embeddings.search(\"How does a car engine work\", 1):\n",
|
||||
" print(json.dumps(x, indent=2))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 160
|
||||
},
|
||||
"id": "V4_tI2urWyZ3",
|
||||
"outputId": "3101db94-2ad8-4ea9-f646-411f57b3027b"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
],
|
||||
"text/html": [
|
||||
"<style>pre { white-space: pre-wrap; }</style>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"{\n",
|
||||
" \"id\": \"Internal combustion engine\",\n",
|
||||
" \"text\": \"An internal combustion engine (ICE or IC engine) is a heat engine in which the combustion of a fuel occurs with an oxidizer (usually air) in a combustion chamber that is an integral part of the working fluid flow circuit. In an internal combustion engine, the expansion of the high-temperature and high-pressure gases produced by combustion applies direct force to some component of the engine. The force is typically applied to pistons (piston engine), turbine blades (gas turbine), a rotor (Wankel engine), or a nozzle (jet engine). This force moves the component over a distance, transforming chemical energy into kinetic energy which is used to propel, move or power whatever the engine is attached to. This replaced the external combustion engine for applications where the weight or size of an engine was more important.\",\n",
|
||||
" \"score\": 0.8664469122886658\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"for x in embeddings.search(\"Who won the World Series in 2022?\", 1):\n",
|
||||
" print(json.dumps(x, indent=2))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 142
|
||||
},
|
||||
"id": "N1owMgfqW6ZO",
|
||||
"outputId": "04699cc3-6127-4a76-83c0-b105ce1924b2"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
],
|
||||
"text/html": [
|
||||
"<style>pre { white-space: pre-wrap; }</style>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"{\n",
|
||||
" \"id\": \"2022 World Series\",\n",
|
||||
" \"text\": \"The 2022 World Series was the championship series of Major League Baseball's (MLB) 2022 season. The 118th edition of the World Series, it was a best-of-seven playoff between the American League (AL) champion Houston Astros and the National League (NL) champion Philadelphia Phillies. The Astros defeated the Phillies in six games to earn their second championship. The series was broadcast in the United States on Fox television and ESPN Radio. \",\n",
|
||||
" \"score\": 0.8889098167419434\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"for x in embeddings.search(\"What was New York called under the Dutch?\", 1):\n",
|
||||
" print(json.dumps(x, indent=2))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 178
|
||||
},
|
||||
"id": "wqIAc67RXdpR",
|
||||
"outputId": "37c273fa-c3e5-4d91-dbfe-3d09d31be127"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
],
|
||||
"text/html": [
|
||||
"<style>pre { white-space: pre-wrap; }</style>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"{\n",
|
||||
" \"id\": \"Dutch Americans in New York City\",\n",
|
||||
" \"text\": \"Dutch people have had a continuous presence in New York City for nearly 400 years, being the earliest European settlers. New York City traces its origins to a trading post founded on the southern tip of Manhattan Island by Dutch colonists in 1624. The settlement was named New Amsterdam in 1626 and was chartered as a city in 1653. Because of the history of Dutch colonization, Dutch culture, politics, law, architecture, and language played a formative role in shaping the culture of the city. The Dutch were the majority in New York City until the early 1700s and the Dutch language was commonly spoken until the mid to late-1700s. Many places and institutions in New York City still bear a colonial Dutch toponymy, including Brooklyn (Breukelen), Harlem (Haarlem), Wall Street (Waal Straat), The Bowery (bouwerij (\\u201cfarm\\u201d), and Coney Island (conyne).\",\n",
|
||||
" \"score\": 0.8840358853340149\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"It's now probably clear how these results can be combined with another component (such as a LLM prompt) to build a conversational QA-based system!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "RwrCxFJJczW-"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Filter by popularity\n",
|
||||
"\n",
|
||||
"Let's try one last query. This is a generic query where there are a lot of matching results with similarity search alone."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "kg6P_NAhY2FL"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"for x in embeddings.search(\"Boston\", 1):\n",
|
||||
" print(json.dumps(x, indent=2))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 124
|
||||
},
|
||||
"id": "h7BsoJC3CFOq",
|
||||
"outputId": "1e8f2ff2-6847-460f-b5b8-9e546fd7bf62"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
],
|
||||
"text/html": [
|
||||
"<style>pre { white-space: pre-wrap; }</style>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"{\n",
|
||||
" \"id\": \"Boston (song)\",\n",
|
||||
" \"text\": \"\\\"Boston\\\" is a song by American rock band Augustana, from their debut album All the Stars and Boulevards (2005). It was originally produced in 2003 by Jon King for their demo, Midwest Skies and Sleepless Mondays, and was later re-recorded with producer Brendan O'Brien for All the Stars and Boulevards.\",\n",
|
||||
" \"score\": 0.8729256987571716\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"While the result is about Boston, it's not the most popular result. This is where the `percentile` field comes in to help. The results can be filtered based on the number of page views.\n",
|
||||
"\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "EeC8PjQLY8G5"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"for x in embeddings.search(\"SELECT id, text, score, percentile FROM txtai WHERE similar('Boston') AND percentile >= 0.99\", 1):\n",
|
||||
" print(json.dumps(x, indent=2))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 196
|
||||
},
|
||||
"id": "NWvENgWNU23V",
|
||||
"outputId": "7008c6e5-8fbc-4275-fa6a-8d32d7ea3fbd"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "display_data",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
],
|
||||
"text/html": [
|
||||
"<style>pre { white-space: pre-wrap; }</style>"
|
||||
]
|
||||
},
|
||||
"metadata": {}
|
||||
},
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"{\n",
|
||||
" \"id\": \"Boston\",\n",
|
||||
" \"text\": \"Boston, officially the City of Boston, is the state capital and most populous city of the Commonwealth of Massachusetts, as well as the cultural and financial center of the New England region of the United States. It is the 24th-most populous city in the country. The city boundaries encompass an area of about and a population of 675,647 as of 2020. It is the seat of Suffolk County (although the county government was disbanded on July 1, 1999). The city is the economic and cultural anchor of a substantially larger metropolitan area known as Greater Boston, a metropolitan statistical area (MSA) home to a census-estimated 4.8\\u00a0million people in 2016 and ranking as the tenth-largest MSA in the country. A broader combined statistical area (CSA), generally corresponding to the commuting area and including Providence, Rhode Island, is home to approximately 8.2\\u00a0million people, making it the sixth most populous in the United States.\",\n",
|
||||
" \"score\": 0.8668985366821289,\n",
|
||||
" \"percentile\": 0.9999025135905505\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"This query adds an additional filter to only match results for the Top 1% of visited Wikipedia pages. "
|
||||
],
|
||||
"metadata": {
|
||||
"id": "dn2wGwTuZJ4A"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook covered how to load embeddings indexes from cloud storage. The Hugging Face Hub is a great resource for sharing models, datasets, example applications and now txtai embeddings indexes. This is especially useful when indexing time is long or requires significant GPU resources.\n",
|
||||
"\n",
|
||||
"Looking forward to seeing what embeddings indexes the community shares in the coming months!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "KqfvCXp2B3li"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,439 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "vwELCooy4ljr"
|
||||
},
|
||||
"source": [
|
||||
"# Prompt templates and task chains\n",
|
||||
"\n",
|
||||
"txtai has long had support for workflows. Workflows connect the input and outputs of machine learning models together to create powerful transformation and processing functions.\n",
|
||||
"\n",
|
||||
"There has been a recent surge in interest in \"model prompting\", which is the process of building a natural language description of a task and passing it to a large language model (LLM). txtai has recently improved support for task templating, which builds string outputs from a set of parameters.\n",
|
||||
"\n",
|
||||
"This notebook demonstrates how txtai workflows can be used to apply prompt templates and chain those tasks together."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ew7orE2O441o"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "LPQTb25tASIG"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[api]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "_YnqorRKAbLu"
|
||||
},
|
||||
"source": [
|
||||
"# Prompt workflow\n",
|
||||
"\n",
|
||||
"First, we'll look at building a workflow with a series of model prompts. This workflow creates a conditional translation using a statement and target language. Another task reads that output text and detects the language.\n",
|
||||
"\n",
|
||||
"This workflow uses a LLM pipeline. The LLM pipeline loads a local model for inference, in this case [Qwen3-4B-Instruct-2507](https://huggingface.co/Qwen/Qwen3-4B-Instruct-2507). The [LLM pipeline](https://neuml.github.io/txtai/pipeline/llm/llm) supports local transformers models, llama.cpp models and LLM APIs such as Ollama, vLLM, OpenAI, Claude etc. \n",
|
||||
"\n",
|
||||
"It's important to note that a pipeline is simply a callable function. It can easily be replaced with a call to an external API."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "OUc9gqTyAYnm",
|
||||
"outputId": "83300311-736c-47c8-bc16-ec0303274054"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"['French', 'German']\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import LLM, Workflow\n",
|
||||
"from txtai.workflow import TemplateTask\n",
|
||||
"\n",
|
||||
"# Create LLM\n",
|
||||
"llm = LLM(\"Qwen/Qwen3-4B-Instruct-2507\")\n",
|
||||
"\n",
|
||||
"# Define workflow or chaining of tasks together.\n",
|
||||
"workflow = Workflow([\n",
|
||||
" TemplateTask(\n",
|
||||
" template=\"Translate text '{statement}' to {language} if the text is English, otherwise keep the original text\",\n",
|
||||
" action=llm\n",
|
||||
" ),\n",
|
||||
" TemplateTask(\n",
|
||||
" template=\"What language is the following text. Only print the answer? {text}\",\n",
|
||||
" action=llm\n",
|
||||
" )\n",
|
||||
"])\n",
|
||||
"\n",
|
||||
"inputs = [\n",
|
||||
" {\"statement\": \"Hello, how are you\", \"language\": \"French\"},\n",
|
||||
" {\"statement\": \"Hallo, wie geht's dir\", \"language\": \"French\"}\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"print(list(workflow(inputs)))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "_zz4Do8BV-Lk"
|
||||
},
|
||||
"source": [
|
||||
"Let's recap what happened here. The first workflow task conditionally translates text to a language if it's English.\n",
|
||||
"\n",
|
||||
"The first statement is `Hello, how are you` with a target language of French. So the statement is translated to French.\n",
|
||||
"\n",
|
||||
"The second statement is German, so it's not converted to French.\n",
|
||||
"\n",
|
||||
"The next step asks the model what the language is and it correctly prints `French` and `German`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "iXDAKP4CX0W9"
|
||||
},
|
||||
"source": [
|
||||
"# Prompt Workflow as YAML\n",
|
||||
"\n",
|
||||
"The same workflow above can be created with YAML configuration."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "GwV5A9xRYtYs",
|
||||
"outputId": "ffe6ee65-95a7-46c6-e6b9-5324eab26ca8"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Writing workflow.yml\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"%%writefile workflow.yml\n",
|
||||
"\n",
|
||||
"llm:\n",
|
||||
" path: Qwen/Qwen3-4B-Instruct-2507\n",
|
||||
"\n",
|
||||
"workflow:\n",
|
||||
" chain:\n",
|
||||
" tasks:\n",
|
||||
" - task: template\n",
|
||||
" template: Translate text '{statement}' to {language} if the text is English, otherwise keep the original text\n",
|
||||
" action: llm\n",
|
||||
" - task: template\n",
|
||||
" template: What language is the following text. Only print the answer? {text}\n",
|
||||
" action: llm"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "dr7Lv5S5X98e",
|
||||
"outputId": "d6ac0427-671d-4525-aa21-664430109af3"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"['French', 'German']\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import Application\n",
|
||||
"\n",
|
||||
"app = Application(\"workflow.yml\")\n",
|
||||
"print(list(app.workflow(\"chain\", inputs)))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "EGqiV45fYVse"
|
||||
},
|
||||
"source": [
|
||||
"As expected, the same result! This is a matter of preference on how you want to create a workflow. One advantage of YAML workflows is that an API can easily be created from the workflow file."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "9PqMU0bNYinf"
|
||||
},
|
||||
"source": [
|
||||
"# Prompt Workflow via an API call\n",
|
||||
"\n",
|
||||
"Let's say you want the workflow to be available via an API call. Well good news, txtai has a built in API mechanism using FastAPI. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "vDxQj1ZIYsz3"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Start an API service\n",
|
||||
"!CONFIG=workflow.yml nohup uvicorn \"txtai.api:app\" &> api.log &\n",
|
||||
"!sleep 60"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "R1o08SVtZW7h",
|
||||
"outputId": "99875acd-18a8-4c2c-ead3-cb6975a4b2d2"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"['French', 'German']"
|
||||
]
|
||||
},
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import requests\n",
|
||||
"\n",
|
||||
"# Run API request\n",
|
||||
"requests.post(\"http://localhost:8000/workflow\", json={\"name\": \"chain\", \"elements\": inputs}).json()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "B88mCrGFl5W-"
|
||||
},
|
||||
"source": [
|
||||
"Just like the previous steps, except through an API call. Let's run via cURL for good measure."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "hRUyh0cQl_P2",
|
||||
"outputId": "9db8481d-0b6e-4a31-bdf6-5443df5f768a"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[\"French\",\"German\"]"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"%%bash\n",
|
||||
"\n",
|
||||
"curl -s -X POST \"http://localhost:8000/workflow\" \\\n",
|
||||
" -H \"Content-Type: application/json\" \\\n",
|
||||
" --data @- << EOF\n",
|
||||
"{\n",
|
||||
" \"name\": \"chain\",\n",
|
||||
" \"elements\": [\n",
|
||||
" {\"statement\": \"Hello, how are you\", \"language\": \"French\"},\n",
|
||||
" {\"statement\": \"Hallo, wie geht's dir\", \"language\": \"French\"}\n",
|
||||
" ]\n",
|
||||
"}\n",
|
||||
"EOF"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "W0zL93WPoaCo"
|
||||
},
|
||||
"source": [
|
||||
"One last time, the same output is shown.\n",
|
||||
"\n",
|
||||
"If your primary development environment isn't Python, txtai does have API bindings for [JavaScript](https://github.com/neuml/txtai.js), [Rust](https://github.com/neuml/txtai.rs), [Go](https://github.com/neuml/txtai.go) and [Java](https://github.com/neuml/txtai.java).\n",
|
||||
"\n",
|
||||
"More information on the API is available [here](https://neuml.github.io/txtai/api/)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "q9WiFG6fpzw5"
|
||||
},
|
||||
"source": [
|
||||
"# Chat with your data\n",
|
||||
"\n",
|
||||
"\"Chat with your data\" is a popular entry point into the AI space. Let's run an example."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "rM3Y551LqF-J",
|
||||
"outputId": "85623785-c15f-4996-9460-0644f69cf5bf"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Writing search.yml\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"%%writefile search.yml\n",
|
||||
"\n",
|
||||
"writable: false\n",
|
||||
"cloud:\n",
|
||||
" provider: huggingface-hub\n",
|
||||
" container: neuml/txtai-intro\n",
|
||||
"\n",
|
||||
"rag:\n",
|
||||
" path: Qwen/Qwen3-4B-Instruct-2507\n",
|
||||
" output: reference\n",
|
||||
" template: |\n",
|
||||
" Answer the following question using only the context below.\n",
|
||||
"\n",
|
||||
" Question: {question}\n",
|
||||
" Context: {context}\n",
|
||||
"\n",
|
||||
"workflow:\n",
|
||||
" search:\n",
|
||||
" tasks:\n",
|
||||
" - action: rag\n",
|
||||
" - task: template\n",
|
||||
" template: \"{answer}\\n\\nReference: {reference}\"\n",
|
||||
" rules:\n",
|
||||
" answer: I don't have data on that"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "1Elb8JANqpwX",
|
||||
"outputId": "b1f1ffa1-6c47-4d90-b6f1-8098d4dc45f8"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[\"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg.\\n\\nReference: 1\"]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"app = Application(\"search.yml\")\n",
|
||||
"print(list(app.workflow(\"search\", [\"Find something about North America\"])))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4r49V4c9s5nf"
|
||||
},
|
||||
"source": [
|
||||
"The first thing the code above does is run an embeddings search to build a conversational context. That context is then used to build a prompt and inference is run against the LLM. \n",
|
||||
"\n",
|
||||
"The next task formats the outputs with a reference to the best matching record. In this case, it's only an id of 1. But this can be much more useful if the id is a URL or there is logic to format the id back to a unique reference string."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KqfvCXp2B3li"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook covered how to build prompt templates and task chains through a series of results. txtai has long had a robust and efficient workflow framework for connecting models together. This can be small and simple models and/or prompting with large models. Go ahead and give it a try!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"gpuClass": "standard",
|
||||
"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.10.19"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,773 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": [],
|
||||
"gpuType": "T4"
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# 💡 What's new in txtai 6.0\n",
|
||||
"\n",
|
||||
"txtai 6.0 brings a number of major feature enhancements. Highlights include:\n",
|
||||
"\n",
|
||||
"- Embeddings\n",
|
||||
" - Sparse/keyword indexes\n",
|
||||
" - Hybrid search\n",
|
||||
" - Subindexes\n",
|
||||
" - Streamlined methods\n",
|
||||
"\n",
|
||||
"- Large Language Models (LLMs)\n",
|
||||
" - Automatically instantiate the best available underlying model\n",
|
||||
" - Pass through parameters enabling immediate support as features are released upstream\n",
|
||||
"\n",
|
||||
"These are just the big, high level changes. There are also many improvements and bug fixes.\n",
|
||||
"\n",
|
||||
"This notebook will cover all the changes with examples.\n",
|
||||
"\n",
|
||||
"**Standard upgrade disclaimer below**\n",
|
||||
"\n",
|
||||
"6.0 is one of the largest, if not largest releases to date! While almost everything is backwards compatible, it's prudent to backup production indexes before upgrading and test before deploying."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "e3wdiK5fGUoZ"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "p8BbfjrhH-V2"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "-OXsTQgaGQPM"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[graph] datasets"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Sparse indexes\n",
|
||||
"\n",
|
||||
"While dense vector indexes are by far the best option for semantic search systems, sparse keyword indexes can still add value. There may be cases where finding an exact match is important or we just want a fast index to quickly do an initial scan of the dataset.\n",
|
||||
"\n",
|
||||
"Unfortunately, there aren't a ton of great options for a local Python-based keyword index library. Most of the options available don't scale and are highly inefficient, designed only for simple situations. With 6.0, txtai has added a performant sparse index component with speed and accuracy on par with Apache Lucene. A future article will discuss the engineering behind this.\n",
|
||||
"\n",
|
||||
"Let's take a look. We'll use a [prompt dataset on the Hugging Face Hub](https://huggingface.co/datasets/fka/awesome-chatgpt-prompts) for all examples."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "n4EXrtcYIIYE"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"import txtai\n",
|
||||
"\n",
|
||||
"# Load dataset\n",
|
||||
"ds = load_dataset(\"fka/awesome-chatgpt-prompts\", split=\"train\")\n",
|
||||
"\n",
|
||||
"def stream():\n",
|
||||
" for row in ds:\n",
|
||||
" yield f\"{row['act']} {row['prompt']}\"\n",
|
||||
"\n",
|
||||
"# Build sparse keyword index\n",
|
||||
"embeddings = txtai.Embeddings(keyword=True, content=True)\n",
|
||||
"embeddings.index(stream())\n",
|
||||
"\n",
|
||||
"embeddings.search(\"Linux terminal\", 1)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "hzPD8_cQJNtN",
|
||||
"outputId": "22059d6e-9022-4fa9-8f54-bd2609508f1e"
|
||||
},
|
||||
"execution_count": 34,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '0',\n",
|
||||
" 'text': 'Linux Terminal I want you to act as a linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. do not write explanations. do not type commands unless I instruct you to do so. when i need to tell you something in english, i will do so by putting text inside curly brackets {like this}. my first command is pwd',\n",
|
||||
" 'score': 0.5932681465337526}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 34
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"And there it is, a keyword index!\n",
|
||||
"\n",
|
||||
"Couple things to unpack here. First, for those familar with txtai, notice that only a text field was yielded in the `stream` method. With 6.0, when ids aren't provided, they are automatically generated.\n",
|
||||
"\n",
|
||||
"Next notice the score. Those familar with keyword scores (TF-IDF, BM25) will notice that the score seems low. That is because with a keyword index, the default score is normalized between 0 and 1.\n",
|
||||
"\n",
|
||||
"More on these items later."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "caDteWUVK5jI"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Hybrid Search\n",
|
||||
"\n",
|
||||
"The addition of sparse indexes enables hybrid search. Hybrid search combines the results from sparse and dense vector indexes for the best of both worlds."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "m7lGqfvDMPQA"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Build hybrid index\n",
|
||||
"embeddings = txtai.Embeddings(hybrid=True, content=True)\n",
|
||||
"embeddings.index(stream())\n",
|
||||
"\n",
|
||||
"embeddings.search(\"Linux terminal\", 1)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "cRXnGeqJMWql",
|
||||
"outputId": "e4c38ad5-59c0-4fd1-ec1d-d64372b19902"
|
||||
},
|
||||
"execution_count": 19,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '0',\n",
|
||||
" 'text': 'Linux Terminal I want you to act as a linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. do not write explanations. do not type commands unless I instruct you to do so. when i need to tell you something in english, i will do so by putting text inside curly brackets {like this}. my first command is pwd',\n",
|
||||
" 'score': 0.6078515601252442}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 19
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Simple change with big impacts. This new index now has both a sparse and dense (using default `sentence-transformers/all-MiniLM-L6-v2` model) index. These scores are combined into a single score as seen above.\n",
|
||||
"\n",
|
||||
"The scoring weights (also known as alpha) control the weighting between the sparse and dense index."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "6bHDxRm5NXI7"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"embeddings.search(\"Linux terminal\", 1, weights=1)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "RWFnXVTRNPHC",
|
||||
"outputId": "b644cbb0-e4d5-41d6-ea72-6a139a160f0c"
|
||||
},
|
||||
"execution_count": 20,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '0',\n",
|
||||
" 'text': 'Linux Terminal I want you to act as a linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. do not write explanations. do not type commands unless I instruct you to do so. when i need to tell you something in english, i will do so by putting text inside curly brackets {like this}. my first command is pwd',\n",
|
||||
" 'score': 0.6224349737167358}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 20
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"embeddings.search(\"Linux terminal\", 1, weights=0)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "EYcv21vENRtp",
|
||||
"outputId": "25b7516f-e23e-4f67-bcdc-c1ff3e157009"
|
||||
},
|
||||
"execution_count": 21,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '0',\n",
|
||||
" 'text': 'Linux Terminal I want you to act as a linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. do not write explanations. do not type commands unless I instruct you to do so. when i need to tell you something in english, i will do so by putting text inside curly brackets {like this}. my first command is pwd',\n",
|
||||
" 'score': 0.5932681465337526}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 21
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"A weight of 1 only uses the dense index and 0 only uses the sparse index. Notice the score with `weight = 0` is the same as the sparse index query earlier."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "cVwf7KeqNydT"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Subindexes\n",
|
||||
"\n",
|
||||
"While sparse and hybrid indexes are great new features, the prize of this release is the addition of subindexes. Subindexes will add a host of new ways to build txtai embeddings instances. Let's give a brief intro here."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "4gUIB_vIODSg"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Build index with subindexes\n",
|
||||
"embeddings = txtai.Embeddings(\n",
|
||||
" content=True,\n",
|
||||
" defaults=False,\n",
|
||||
" indexes={\n",
|
||||
" \"sparse\": {\n",
|
||||
" \"keyword\": True\n",
|
||||
" },\n",
|
||||
" \"dense\":{\n",
|
||||
"\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
")\n",
|
||||
"embeddings.index(stream())\n",
|
||||
"\n",
|
||||
"# Run search\n",
|
||||
"embeddings.search(\"select id, text, score from txtai where similar('Linux terminal', 'sparse') and similar('Linux terminal', 'dense')\", 1)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "9nkna32MObBi",
|
||||
"outputId": "3f57c63e-8b4b-4d75-f4db-3455ed667ec2"
|
||||
},
|
||||
"execution_count": 22,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '0',\n",
|
||||
" 'text': 'Linux Terminal I want you to act as a linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. do not write explanations. do not type commands unless I instruct you to do so. when i need to tell you something in english, i will do so by putting text inside curly brackets {like this}. my first command is pwd',\n",
|
||||
" 'score': 0.6078515601252442}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 22
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"embeddings.search(\"select id, text, score from txtai where similar('Linux terminal', 'dense')\", 1)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "cPysPMHRR2IS",
|
||||
"outputId": "77b19baa-0643-4958-dd59-363ce05017a8"
|
||||
},
|
||||
"execution_count": 23,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '0',\n",
|
||||
" 'text': 'Linux Terminal I want you to act as a linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. do not write explanations. do not type commands unless I instruct you to do so. when i need to tell you something in english, i will do so by putting text inside curly brackets {like this}. my first command is pwd',\n",
|
||||
" 'score': 0.6224349737167358}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 23
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"embeddings.search(\"select id, text, score from txtai where similar('Linux terminal', 'sparse')\", 1)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "KYEFkDRDSJ4J",
|
||||
"outputId": "bcb2460d-9cdd-4808-f861-71405044849d"
|
||||
},
|
||||
"execution_count": 24,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '0',\n",
|
||||
" 'text': 'Linux Terminal I want you to act as a linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. do not write explanations. do not type commands unless I instruct you to do so. when i need to tell you something in english, i will do so by putting text inside curly brackets {like this}. my first command is pwd',\n",
|
||||
" 'score': 0.5932681465337526}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 24
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Notice how the scores are the same as above. The three searches above run a hybrid search, dense and sparse search. This time though it's using subindexes. The top-level Embeddings only has an associated database.\n",
|
||||
"\n",
|
||||
"Each of the sections in the `indexes` is a full embeddings index supporting all available options. For example, let's add a graph subindex."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "hJGAjnpbSTRu"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Build index with graph subindex\n",
|
||||
"embeddings = txtai.Embeddings(\n",
|
||||
" content=True,\n",
|
||||
" defaults=False,\n",
|
||||
" functions=[\n",
|
||||
" {\"name\": \"graph\", \"function\": \"indexes.act.graph.attribute\"}\n",
|
||||
" ],\n",
|
||||
" expressions=[\n",
|
||||
" {\"name\": \"topic\", \"expression\": \"graph(indexid, 'topic')\"},\n",
|
||||
" ],\n",
|
||||
" indexes={\n",
|
||||
" \"act\": {\n",
|
||||
" \"keyword\": True,\n",
|
||||
" \"columns\": {\n",
|
||||
" \"text\": \"act\"\n",
|
||||
" },\n",
|
||||
" \"graph\": {\n",
|
||||
" \"topics\": {}\n",
|
||||
" }\n",
|
||||
" },\n",
|
||||
" \"prompt\":{\n",
|
||||
" \"columns\": {\n",
|
||||
" \"text\": \"prompt\"\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
")\n",
|
||||
"embeddings.index(ds)\n",
|
||||
"\n",
|
||||
"# Run search\n",
|
||||
"embeddings.search(\"select id, act, prompt, score, topic from txtai where similar('Linux terminal')\", 1)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "D35uFFnESqP3",
|
||||
"outputId": "d5ee3775-d488-4e46-a44b-3ae4a12a6c1f"
|
||||
},
|
||||
"execution_count": 25,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '0',\n",
|
||||
" 'act': 'Linux Terminal',\n",
|
||||
" 'prompt': 'I want you to act as a linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. do not write explanations. do not type commands unless I instruct you to do so. when i need to tell you something in english, i will do so by putting text inside curly brackets {like this}. my first command is pwd',\n",
|
||||
" 'score': 0.6382951796072414,\n",
|
||||
" 'topic': 'terminal_linux_sql'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 25
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Notice the new `topic` field added to this query. That comes from the graph index, which runs topic modeling. Also notice that two indexes for two different columns are added.\n",
|
||||
"\n",
|
||||
"Note that graph indexes are different in that they depend on a sparse or dense index being available. That is how the graph is automatically constructed. For good measure, let's add the graph to a dense index."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "t_OIM6tPUnC_"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Build index with graph subindex\n",
|
||||
"embeddings = txtai.Embeddings(\n",
|
||||
" content=True,\n",
|
||||
" defaults=False,\n",
|
||||
" functions=[\n",
|
||||
" {\"name\": \"graph\", \"function\": \"indexes.act.graph.attribute\"}\n",
|
||||
" ],\n",
|
||||
" expressions=[\n",
|
||||
" {\"name\": \"topic\", \"expression\": \"graph(indexid, 'topic')\"},\n",
|
||||
" ],\n",
|
||||
" indexes={\n",
|
||||
" \"act\": {\n",
|
||||
" \"path\": \"intfloat/e5-small-v2\",\n",
|
||||
" \"columns\": {\n",
|
||||
" \"text\": \"act\"\n",
|
||||
" },\n",
|
||||
" \"graph\": {\n",
|
||||
" \"topics\": {}\n",
|
||||
" }\n",
|
||||
" },\n",
|
||||
" \"prompt\":{\n",
|
||||
" \"path\": \"sentence-transformers/all-MiniLM-L6-v2\",\n",
|
||||
" \"columns\": {\n",
|
||||
" \"text\": \"prompt\"\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
")\n",
|
||||
"embeddings.index(ds)\n",
|
||||
"\n",
|
||||
"# Run search\n",
|
||||
"embeddings.search(\"select id, act, prompt, score, topic from txtai where similar('Linux terminal')\", 1)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "Me8zJse7WbdM",
|
||||
"outputId": "14433e6b-9eed-4cd2-e590-4c9c03be390a"
|
||||
},
|
||||
"execution_count": 26,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '0',\n",
|
||||
" 'act': 'Linux Terminal',\n",
|
||||
" 'prompt': 'I want you to act as a linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. do not write explanations. do not type commands unless I instruct you to do so. when i need to tell you something in english, i will do so by putting text inside curly brackets {like this}. my first command is pwd',\n",
|
||||
" 'score': 1.0,\n",
|
||||
" 'topic': 'linux_terminal'}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 26
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Almost the same as above except the topic is different. This is due to the grouping of the vector index. Notice how the `act` column and `prompt` column are both vector indexes but specify different vector models. This opens up another possibility of weighting not only sparse vs vector but different vector models."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "yKHVU_r4W9HD"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"embeddings.search(\"select id, act, prompt, score from txtai where similar('Linux terminal', 'act') and similar('Linux terminal', 'prompt')\", 1)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "2Yqky3ntYKUJ",
|
||||
"outputId": "02672b7b-9e28-4d6d-dd74-acaf475ddf72"
|
||||
},
|
||||
"execution_count": 27,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '0',\n",
|
||||
" 'act': 'Linux Terminal',\n",
|
||||
" 'prompt': 'I want you to act as a linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. do not write explanations. do not type commands unless I instruct you to do so. when i need to tell you something in english, i will do so by putting text inside curly brackets {like this}. my first command is pwd',\n",
|
||||
" 'score': 0.7881423830986023}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 27
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"As always, everything discussed so far is also supported with txtai application instances."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "obZhFSIKZT4u"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Build index with graph subindex\n",
|
||||
"app = txtai.Application(\"\"\"\n",
|
||||
"writable: True\n",
|
||||
"embeddings:\n",
|
||||
" content: True\n",
|
||||
" defaults: False\n",
|
||||
" functions:\n",
|
||||
" - name: graph\n",
|
||||
" function: indexes.act.graph.attribute\n",
|
||||
" expressions:\n",
|
||||
" - name: topic\n",
|
||||
" expression: graph(indexid, 'topic')\n",
|
||||
" indexes:\n",
|
||||
" act:\n",
|
||||
" path: intfloat/e5-small-v2\n",
|
||||
" columns:\n",
|
||||
" text: act\n",
|
||||
" graph:\n",
|
||||
" topics:\n",
|
||||
" prompt:\n",
|
||||
" path: sentence-transformers/all-MiniLM-L6-v2\n",
|
||||
" columns:\n",
|
||||
" text: prompt\n",
|
||||
"\"\"\")\n",
|
||||
"\n",
|
||||
"app.add(ds)\n",
|
||||
"app.index()\n",
|
||||
"\n",
|
||||
"app.search(\"select id, act, prompt, topic, score from txtai where similar('Linux terminal', 'act') and similar('Linux terminal', 'prompt')\", 1)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "3MDvZoaGZZhI",
|
||||
"outputId": "981072e4-70bf-4f2e-bc92-ee92083ec21a"
|
||||
},
|
||||
"execution_count": 28,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '0',\n",
|
||||
" 'act': 'Linux Terminal',\n",
|
||||
" 'prompt': 'I want you to act as a linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. do not write explanations. do not type commands unless I instruct you to do so. when i need to tell you something in english, i will do so by putting text inside curly brackets {like this}. my first command is pwd',\n",
|
||||
" 'topic': 'linux_terminal',\n",
|
||||
" 'score': 0.7881423830986023}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 28
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Streamlined methods\n",
|
||||
"\n",
|
||||
"Much of this has been covered already but a number of changes were added to make it easier to search and index data. The existing interfaces are all still supported, this is all about ease of use.\n",
|
||||
"\n",
|
||||
"See the code explanations below."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "WNSJ6tR1Y_KI"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Top-level import includes Application and Embeddings\n",
|
||||
"import txtai\n",
|
||||
"\n",
|
||||
"app = txtai.Application(\"\"\"writable: False\"\"\")\n",
|
||||
"embeddings = txtai.Embeddings()"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "MoW5QkxNZKBL"
|
||||
},
|
||||
"execution_count": 29,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Ids are automatically generated when omitted\n",
|
||||
"embeddings.index([\"test\"])\n",
|
||||
"print(embeddings.search(\"test\"))\n",
|
||||
"\n",
|
||||
"# UUID ids are also supported - use any of the methods in https://docs.python.org/3/library/uuid.html\n",
|
||||
"embeddings = txtai.Embeddings(autoid=\"uuid5\")\n",
|
||||
"embeddings.index([\"test\"])\n",
|
||||
"embeddings.search(\"test\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "6CckssTWbD3y",
|
||||
"outputId": "e826ea96-e6aa-4902-d75d-4ed2df031a65"
|
||||
},
|
||||
"execution_count": 30,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"[(0, 0.9999998807907104)]\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[('4be0643f-1d98-573b-97cd-ca98a65347dd', 0.9999998807907104)]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 30
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Large Language Models (LLMs)\n",
|
||||
"\n",
|
||||
"While the bulk of the changes in this release came with the embeddings package, LLMs also have important changes that make it easier to use."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "v_Dt9kltcRg8"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import torch\n",
|
||||
"\n",
|
||||
"from txtai import LLM\n",
|
||||
"\n",
|
||||
"# Create model and set dtype to use 16-bit floats\n",
|
||||
"llm = LLM(\"tiiuae/falcon-rw-1b\", torch_dtype=torch.bfloat16)"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "HL1caw21cfxO"
|
||||
},
|
||||
"execution_count": 31,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"print(llm(\"Write a short list of things to do in Paris\", maxlength=55))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "G7OgcAUXg2RA",
|
||||
"outputId": "83e86070-9212-4d21-b5b1-2712389cac09"
|
||||
},
|
||||
"execution_count": 32,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
".\n",
|
||||
"- Visit the Eiffel Tower.\n",
|
||||
"- Visit the Louvre.\n",
|
||||
"- Visit the Arc de Triomphe.\n",
|
||||
"- Visit the Notre Dame Cathedral.\n",
|
||||
"- Visit the Sacre Coeur Basilica\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"The new `LLM` pipeline automatically detects the type of model and loads it using the best available method.\n",
|
||||
"\n",
|
||||
"The pipeline framework now passes through keyword arguments to the underlying methods, which adds support for new Hugging Face features automatically as they are released."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "j06rZByEhwnO"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook gave a quick overview of txtai 6.0. Updated documentation and more examples will be forthcoming. There is much to cover and much to build on!\n",
|
||||
"\n",
|
||||
"See the following links for more information.\n",
|
||||
"\n",
|
||||
"- [6.0 Release on GitHub](https://github.com/neuml/txtai/releases/tag/v6.0.0)\n",
|
||||
"- [Documentation site](https://neuml.github.io/txtai)"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "tvnMO1Eai6Gy"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,802 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# External database integration\n",
|
||||
"\n",
|
||||
"txtai provides many default settings to help a developer quickly get started. For example, metadata is stored in SQLite, dense vectors in Faiss, sparse vectors in a terms index and graph data with NetworkX.\n",
|
||||
"\n",
|
||||
"Each of these components is customizable and can be swapped with alternate implementations. This has been covered in [several previous notebooks](https://neuml.github.io/txtai/examples/#architecture).\n",
|
||||
"\n",
|
||||
"This notebook will introduce how to store metadata in client-server RDBMS systems. In addition to SQLite and DuckDB, any [SQLAlchemy-supported database](https://docs.sqlalchemy.org/en/20/dialects/) with [JSON support](https://docs.sqlalchemy.org/en/20/core/type_basics.html#sqlalchemy.types.JSON) can now be used."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "781Mrl0FtFR8"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "VG8sklTPwlJa"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "XrqaDHiCnFnr"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[database] elasticsearch==7.10.1 datasets"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Install Postgres\n",
|
||||
"\n",
|
||||
"Next, we'll install Postgres and start a Postgres instance."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "6pc469NKwyMu"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "1JZhUIMmkmil"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"\n",
|
||||
"# Install and start Postgres\n",
|
||||
"!apt-get update && apt-get install postgresql\n",
|
||||
"!service postgresql start\n",
|
||||
"!sudo -u postgres psql -U postgres -c \"ALTER USER postgres PASSWORD 'postgres';\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Load a dataset\n",
|
||||
"\n",
|
||||
"Now we're ready to load a dataset. We'll use the `ag_news` dataset. This dataset consists of 120,000 news headlines."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "lS16J4QBxDJP"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "TMVkdPSgNlD0"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"# Load dataset\n",
|
||||
"ds = load_dataset(\"ag_news\", split=\"train\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Build an Embeddings instance with Postgres\n",
|
||||
"\n",
|
||||
"Let's load this dataset into an embeddings database. We'll configure this instance to store metadata in Postgres. Note that the content parameter below is a [SQLAlchemy connection string](https://docs.sqlalchemy.org/en/20/core/engines.html).\n",
|
||||
"\n",
|
||||
"This embeddings database will use the default vector settings and build that index locally."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "iERZH9w1xQGC"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "OmnobNLGNhW4"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import txtai\n",
|
||||
"\n",
|
||||
"# Create embeddings\n",
|
||||
"embeddings = txtai.Embeddings(\n",
|
||||
" content=\"postgresql+psycopg2://postgres:postgres@localhost/postgres\",\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Index dataset\n",
|
||||
"embeddings.index(ds[\"text\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Let's run a search query and see what comes back."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "GqOMCOT0xnnC"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "qD4tBdMKI4MX",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "109fe906-ca74-49c9-ed4e-922eaefd398e"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '63561',\n",
|
||||
" 'text': 'Red Sox Beat Yankees 6-4 in 12 Innings BOSTON - Down to their last three outs of the season, the Boston Red Sox rallied - against Mariano Rivera, the New York Yankees and decades of disappointment. Bill Mueller singled home the tying run off Rivera in the ninth inning and David Ortiz homered against Paul Quantrill in the 12th, leading Boston to a 6-4 victory Sunday over the Yankees that avoided a four-game sweep in the AL championship series...',\n",
|
||||
" 'score': 0.8104304671287537},\n",
|
||||
" {'id': '63221',\n",
|
||||
" 'text': 'Red Sox Beat Yankees 6-4 in 12 Innings BOSTON - Down to their last three outs of the season, the Boston Red Sox rallied - against Mariano Rivera, the New York Yankees and decades of disappointment. Bill Mueller singled home the tying run off Rivera in the ninth inning and David Ortiz homered against Paul Quantrill in the 12th, leading Boston to a 6-4 victory over the Yankees on Sunday night that avoided a four-game sweep in the AL championship series...',\n",
|
||||
" 'score': 0.8097385168075562},\n",
|
||||
" {'id': '66861',\n",
|
||||
" 'text': 'Record-Breaking Red Sox Clinch World Series Berth NEW YORK (Reuters) - The Boston Red Sox crushed the New York Yankees 10-3 Wednesday to complete an historic comeback victory over their arch-rivals by four games to three in the American League Championship Series.',\n",
|
||||
" 'score': 0.8003846406936646}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 24
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings.search(\"red sox defeat yankees\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"As expected, we get the standard `id, text, score` fields with the top matches for the query. The difference though is that all the database metadata normally stored in a local SQLite file is now stored in a Postgres server.\n",
|
||||
"\n",
|
||||
"This opens up several possibilities such as row-level security. If a row isn't returned by the database, it won't be shown here. Alternatively, this search could optionally return only the ids and scores, which lets the user know a record exists they don't have access to.\n",
|
||||
"\n",
|
||||
"As with other supported databases, underlying database functions can be called from txtai SQL."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "SeV3i6vixs-e"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"embeddings.search(\"SELECT id, text, md5(text), score FROM txtai WHERE similar('red sox defeat yankees')\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "3PQAoPfvSyP1",
|
||||
"outputId": "1eed6cf0-6402-4760-b673-5798fc329d0e"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '63561',\n",
|
||||
" 'text': 'Red Sox Beat Yankees 6-4 in 12 Innings BOSTON - Down to their last three outs of the season, the Boston Red Sox rallied - against Mariano Rivera, the New York Yankees and decades of disappointment. Bill Mueller singled home the tying run off Rivera in the ninth inning and David Ortiz homered against Paul Quantrill in the 12th, leading Boston to a 6-4 victory Sunday over the Yankees that avoided a four-game sweep in the AL championship series...',\n",
|
||||
" 'md5': '1e55a78fdf0cb3be3ef61df650f0a50f',\n",
|
||||
" 'score': 0.8104304671287537},\n",
|
||||
" {'id': '63221',\n",
|
||||
" 'text': 'Red Sox Beat Yankees 6-4 in 12 Innings BOSTON - Down to their last three outs of the season, the Boston Red Sox rallied - against Mariano Rivera, the New York Yankees and decades of disappointment. Bill Mueller singled home the tying run off Rivera in the ninth inning and David Ortiz homered against Paul Quantrill in the 12th, leading Boston to a 6-4 victory over the Yankees on Sunday night that avoided a four-game sweep in the AL championship series...',\n",
|
||||
" 'md5': 'a0417e1fc503a5a2945c8755b6fb18d5',\n",
|
||||
" 'score': 0.8097385168075562},\n",
|
||||
" {'id': '66861',\n",
|
||||
" 'text': 'Record-Breaking Red Sox Clinch World Series Berth NEW YORK (Reuters) - The Boston Red Sox crushed the New York Yankees 10-3 Wednesday to complete an historic comeback victory over their arch-rivals by four games to three in the American League Championship Series.',\n",
|
||||
" 'md5': '398a8508692aed109bd8c56f067a8083',\n",
|
||||
" 'score': 0.8003846406936646}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 25
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Note the addition of the Postgres `md5` function to the query.\n",
|
||||
"\n",
|
||||
"Let's save and show the files in the embeddings database."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "lqkfoknYziiU"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"embeddings.save(\"vectors\")\n",
|
||||
"!ls -l vectors"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "gKTsjD3uS0YG",
|
||||
"outputId": "aa2ec6af-509b-468e-a7e9-11b550707073"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"total 183032\n",
|
||||
"-rw-r--r-- 1 root root 355 Sep 7 16:38 config\n",
|
||||
"-rw-r--r-- 1 root root 187420123 Sep 7 16:38 embeddings\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Only the configuration and the local vectors index are stored in this case."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "ylGVrmL1ztG0"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# External indexing\n",
|
||||
"\n",
|
||||
"As mentioned previously, all of the main components of txtai can be replaced with custom components. For example, there are external integrations for storing dense vectors in [Weaviate](https://github.com/hsm207/weaviate-txtai) and [Qdrant](https://github.com/qdrant/qdrant-txtai) to name a few.\n",
|
||||
"\n",
|
||||
"Next, we'll build an example that stores metadata in Postgres and builds a sparse index with Elasticsearch."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "RUqaj-tVz19K"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"## Scoring component for Elasticsearch\n",
|
||||
"\n",
|
||||
"First, we need to define a custom scoring component for Elasticsearch. While could have used an existing integration, it's important to show that creating a new component isn't a large LOE (~70 lines of code). See below."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "OZLCQceF3US9"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "fbPd8uXZrekP"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from elasticsearch import Elasticsearch\n",
|
||||
"from elasticsearch.helpers import bulk\n",
|
||||
"\n",
|
||||
"from txtai.scoring import Scoring\n",
|
||||
"\n",
|
||||
"class Elastic(Scoring):\n",
|
||||
" def __init__(self, config=None):\n",
|
||||
" # Scoring configuration\n",
|
||||
" self.config = config if config else {}\n",
|
||||
"\n",
|
||||
" # Server parameters\n",
|
||||
" self.url = self.config.get(\"url\", \"http://localhost:9200\")\n",
|
||||
" self.indexname = self.config.get(\"indexname\", \"testindex\")\n",
|
||||
"\n",
|
||||
" # Elasticsearch connection\n",
|
||||
" self.connection = Elasticsearch(self.url)\n",
|
||||
"\n",
|
||||
" self.terms = True\n",
|
||||
" self.normalize = self.config.get(\"normalize\")\n",
|
||||
"\n",
|
||||
" def insert(self, documents, index=None):\n",
|
||||
" rows = []\n",
|
||||
" for uid, document, tags in documents:\n",
|
||||
" rows.append((index, document))\n",
|
||||
"\n",
|
||||
" # Increment index\n",
|
||||
" index = index + 1\n",
|
||||
"\n",
|
||||
" bulk(self.connection, ({\"_index\": self.indexname, \"_id\": uid, \"text\": text} for uid, text in rows))\n",
|
||||
"\n",
|
||||
" def index(self, documents=None):\n",
|
||||
" self.connection.indices.refresh(index=self.indexname)\n",
|
||||
"\n",
|
||||
" def search(self, query, limit=3):\n",
|
||||
" return self.batchsearch([query], limit)\n",
|
||||
"\n",
|
||||
" def batchsearch(self, queries, limit=3):\n",
|
||||
" # Generate bulk queries\n",
|
||||
" request = []\n",
|
||||
" for query in queries:\n",
|
||||
" req_head = {\"index\": self.indexname, \"search_type\": \"dfs_query_then_fetch\"}\n",
|
||||
" req_body = {\n",
|
||||
" \"_source\": False,\n",
|
||||
" \"query\": {\"multi_match\": {\"query\": query, \"type\": \"best_fields\", \"fields\": [\"text\"], \"tie_breaker\": 0.5}},\n",
|
||||
" \"size\": limit,\n",
|
||||
" }\n",
|
||||
" request.extend([req_head, req_body])\n",
|
||||
"\n",
|
||||
" # Run ES query\n",
|
||||
" response = self.connection.msearch(body=request, request_timeout=600)\n",
|
||||
"\n",
|
||||
" # Read responses\n",
|
||||
" results = []\n",
|
||||
" for resp in response[\"responses\"]:\n",
|
||||
" result = resp[\"hits\"][\"hits\"]\n",
|
||||
" results.append([(r[\"_id\"], r[\"_score\"]) for r in result])\n",
|
||||
"\n",
|
||||
" return results\n",
|
||||
"\n",
|
||||
" def count(self):\n",
|
||||
" response = self.connection.cat.count(self.indexname, params={\"format\": \"json\"})\n",
|
||||
" return int(response[0][\"count\"])\n",
|
||||
"\n",
|
||||
" def load(self, path):\n",
|
||||
" # No local storage\n",
|
||||
" pass\n",
|
||||
"\n",
|
||||
" def save(self, path):\n",
|
||||
" # No local storage\n",
|
||||
" pass"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"## Elasticsearch server\n",
|
||||
"\n",
|
||||
"As with Postgres, we'll install and start an Elasticsearch instance."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "3Zu1en9A2fP7"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "bZFJGE9RpRTt"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"import os\n",
|
||||
"\n",
|
||||
"# Download and extract elasticsearch\n",
|
||||
"os.system(\"wget https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.10.1-linux-x86_64.tar.gz\")\n",
|
||||
"os.system(\"tar -xzf elasticsearch-7.10.1-linux-x86_64.tar.gz\")\n",
|
||||
"os.system(\"chown -R daemon:daemon elasticsearch-7.10.1\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "oeCGDn5NpR9p"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from subprocess import Popen, PIPE, STDOUT\n",
|
||||
"\n",
|
||||
"# Start and wait for serverw\n",
|
||||
"server = Popen(['elasticsearch-7.10.1/bin/elasticsearch'], stdout=PIPE, stderr=STDOUT, preexec_fn=lambda: os.setuid(1))\n",
|
||||
"!sleep 30"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Let's build the index. The only difference from the previous example is setting the custom `scoring` component."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "l9WI7NJy4TXo"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "ZCQhwUVRnnEY"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import txtai\n",
|
||||
"\n",
|
||||
"# Create embeddings\n",
|
||||
"embeddings = txtai.Embeddings(\n",
|
||||
" keyword=True,\n",
|
||||
" content=\"postgresql+psycopg2://postgres:postgres@localhost/postgres\",\n",
|
||||
" scoring= \"__main__.Elastic\"\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Index dataset\n",
|
||||
"embeddings.index(ds[\"text\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Below is the same search as shown before."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "JrfZTEyE4Q-i"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"embeddings.search(\"red sox defeat yankees\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "VFgIV_S8Stby",
|
||||
"outputId": "75f3b9ab-2805-4a3c-ca93-b843e21af439"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '66954',\n",
|
||||
" 'text': 'Boston Red Sox make history Believe it, New England -- the Boston Red Sox are in the World Series. And they got there with the most unbelievable comeback of all, with four sweet swings after decades of defeat, shaming the dreaded New York Yankees.',\n",
|
||||
" 'score': 21.451942},\n",
|
||||
" {'id': '69577',\n",
|
||||
" 'text': 'Passing thoughts on Yankees-Red Sox series The Red Sox beat the Yankees at Yankee Stadium in a season-deciding game. The Red Sox beat the Yankees at Yankee Stadium in a season-deciding game and it wasn #39;t even close.',\n",
|
||||
" 'score': 20.923117},\n",
|
||||
" {'id': '67253',\n",
|
||||
" 'text': 'Sox Victorious At Last!! BOSTON -- After suffering decades of defeat and disappointment, the 2004 Boston Red Sox made history Wednesday night, beating the Yankees in the house that Ruth built and claiming the American League championship trophy.',\n",
|
||||
" 'score': 20.865997}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 31
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"And once again we get the top matches. This time though the index is in Elasticsearch. Why are results and scores different? This is because this is a keyword index and it's using Elasticsearch's raw BM25 scores.\n",
|
||||
"\n",
|
||||
"One enhancement to this component would be adding score normalization as seen in the standard scoring components.\n",
|
||||
"\n",
|
||||
"For good measure, let's also show that the `md5` function can be called here too."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "PE9SAEDn4gDE"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"embeddings.search(\"SELECT id, text, md5(text), score FROM txtai WHERE similar('red sox defeat yankees')\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "4oYAIFPqWbSq",
|
||||
"outputId": "98178870-393e-42ca-c1d6-9c91239ddabe"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '66954',\n",
|
||||
" 'text': 'Boston Red Sox make history Believe it, New England -- the Boston Red Sox are in the World Series. And they got there with the most unbelievable comeback of all, with four sweet swings after decades of defeat, shaming the dreaded New York Yankees.',\n",
|
||||
" 'md5': '29084f8640d4d72e402e991bc9fdbfa0',\n",
|
||||
" 'score': 21.451942},\n",
|
||||
" {'id': '69577',\n",
|
||||
" 'text': 'Passing thoughts on Yankees-Red Sox series The Red Sox beat the Yankees at Yankee Stadium in a season-deciding game. The Red Sox beat the Yankees at Yankee Stadium in a season-deciding game and it wasn #39;t even close.',\n",
|
||||
" 'md5': '056983d301975084b49a5987185f2ddf',\n",
|
||||
" 'score': 20.923117},\n",
|
||||
" {'id': '67253',\n",
|
||||
" 'text': 'Sox Victorious At Last!! BOSTON -- After suffering decades of defeat and disappointment, the 2004 Boston Red Sox made history Wednesday night, beating the Yankees in the house that Ruth built and claiming the American League championship trophy.',\n",
|
||||
" 'md5': '7838fcf610f0b569829c9bafdf9012f2',\n",
|
||||
" 'score': 20.865997}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 32
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Same results with the additional `md5` column, as expected."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "zW1NkPpA5V0b"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Explore the data stores\n",
|
||||
"\n",
|
||||
"The last thing we'll do is see where and how this data is stored in Postgres and Elasticsearch.\n",
|
||||
"\n",
|
||||
"Let's connect to the local Postgres instance and sample content from the `sections` table."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "54R2s6CM5a-6"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "pWRJ4vNkn2QL",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "fab25634-474f-441d-ef82-cf903b8c8762"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"env: DATABASE_URL=postgresql://postgres:postgres@localhost:5432/postgres\n",
|
||||
"The sql extension is already loaded. To reload it, use:\n",
|
||||
" %reload_ext sql\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"%env DATABASE_URL=postgresql://postgres:postgres@localhost:5432/postgres\n",
|
||||
"%load_ext sql"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "axy-m3JMn9H8",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 175
|
||||
},
|
||||
"outputId": "6d7a523f-722e-4139-a4f2-db30f8a8f877"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
" * postgresql://postgres:***@localhost:5432/postgres\n",
|
||||
"3 rows affected.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[('66954', 'Boston Red Sox make history Believe it, New England -- the Boston Red Sox are in the World Series. And they got there with the most unbelievable comeback of all, with four sweet swings after decades of defeat, shaming the dreaded New York Yankees.'),\n",
|
||||
" ('62732', \"BoSox, Astros Play for Crucial Game 4 Wins (AP) AP - The Boston Red Sox entered this AL championship series hoping to finally overcome their bitter r ... (50 characters truncated) ... n-game defeat last October. Instead, they've been reduced to trying to prevent the Yankees from completing a humiliating sweep in their own ballpark.\"),\n",
|
||||
" ('62752', \"BoSox, Astros Play for Crucial Game 4 Wins The Boston Red Sox entered this AL championship series hoping to finally overcome their bitter rivals from ... (42 characters truncated) ... game defeat last October. Instead, they've been reduced to trying to prevent the Yankees from completing a humiliating sweep in their own ballpark...\")]"
|
||||
],
|
||||
"text/html": [
|
||||
"<table>\n",
|
||||
" <thead>\n",
|
||||
" <tr>\n",
|
||||
" <th>id</th>\n",
|
||||
" <th>text</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" <tr>\n",
|
||||
" <td>66954</td>\n",
|
||||
" <td>Boston Red Sox make history Believe it, New England -- the Boston Red Sox are in the World Series. And they got there with the most unbelievable comeback of all, with four sweet swings after decades of defeat, shaming the dreaded New York Yankees.</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <td>62732</td>\n",
|
||||
" <td>BoSox, Astros Play for Crucial Game 4 Wins (AP) AP - The Boston Red Sox entered this AL championship series hoping to finally overcome their bitter rivals from New York following a heartbreaking seven-game defeat last October. Instead, they've been reduced to trying to prevent the Yankees from completing a humiliating sweep in their own ballpark.</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <td>62752</td>\n",
|
||||
" <td>BoSox, Astros Play for Crucial Game 4 Wins The Boston Red Sox entered this AL championship series hoping to finally overcome their bitter rivals from New York following a heartbreaking seven-game defeat last October. Instead, they've been reduced to trying to prevent the Yankees from completing a humiliating sweep in their own ballpark...</td>\n",
|
||||
" </tr>\n",
|
||||
" </tbody>\n",
|
||||
"</table>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 34
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"%%sql\n",
|
||||
"select id, text from sections where text like '%Red Sox%' and text like '%Yankees%' and text like '%defeat%' limit 3;"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"As expected, we can see content stored directly in Postgres!\n",
|
||||
"\n",
|
||||
"Now let's check Elasticsearch."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "6JfMoZPB5uZo"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "X6A7GP7ZBm_v",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "151df22c-6d43-4599-eba6-a579fe437677"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"{\n",
|
||||
" \"took\": 13,\n",
|
||||
" \"timed_out\": false,\n",
|
||||
" \"_shards\": {\n",
|
||||
" \"total\": 1,\n",
|
||||
" \"successful\": 1,\n",
|
||||
" \"skipped\": 0,\n",
|
||||
" \"failed\": 0\n",
|
||||
" },\n",
|
||||
" \"hits\": {\n",
|
||||
" \"total\": {\n",
|
||||
" \"value\": 3297,\n",
|
||||
" \"relation\": \"eq\"\n",
|
||||
" },\n",
|
||||
" \"max_score\": 21.451942,\n",
|
||||
" \"hits\": [\n",
|
||||
" {\n",
|
||||
" \"_index\": \"testindex\",\n",
|
||||
" \"_type\": \"_doc\",\n",
|
||||
" \"_id\": \"66954\",\n",
|
||||
" \"_score\": 21.451942,\n",
|
||||
" \"_source\": {\n",
|
||||
" \"text\": \"Boston Red Sox make history Believe it, New England -- the Boston Red Sox are in the World Series. And they got there with the most unbelievable comeback of all, with four sweet swings after decades of defeat, shaming the dreaded New York Yankees.\"\n",
|
||||
" }\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" \"_index\": \"testindex\",\n",
|
||||
" \"_type\": \"_doc\",\n",
|
||||
" \"_id\": \"69577\",\n",
|
||||
" \"_score\": 20.923117,\n",
|
||||
" \"_source\": {\n",
|
||||
" \"text\": \"Passing thoughts on Yankees-Red Sox series The Red Sox beat the Yankees at Yankee Stadium in a season-deciding game. The Red Sox beat the Yankees at Yankee Stadium in a season-deciding game and it wasn #39;t even close.\"\n",
|
||||
" }\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" \"_index\": \"testindex\",\n",
|
||||
" \"_type\": \"_doc\",\n",
|
||||
" \"_id\": \"67253\",\n",
|
||||
" \"_score\": 20.865997,\n",
|
||||
" \"_source\": {\n",
|
||||
" \"text\": \"Sox Victorious At Last!! BOSTON -- After suffering decades of defeat and disappointment, the 2004 Boston Red Sox made history Wednesday night, beating the Yankees in the house that Ruth built and claiming the American League championship trophy.\"\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" ]\n",
|
||||
" }\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import json\n",
|
||||
"import requests\n",
|
||||
"\n",
|
||||
"response = requests.get(\"http://localhost:9200/_search?q=red+sox+defeat+yankees&size=3\")\n",
|
||||
"print(json.dumps(response.json(), indent=2))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Same query results as what was run through the embeddings database.\n",
|
||||
"\n",
|
||||
"Let's save the embeddings database and review what's stored."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "6sS0ym5Y59w9"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "BsFNRvRzoj9T",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "38bef15f-2ce3-4c8a-9bd0-3d7d0c84d794"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"total 4\n",
|
||||
"-rw-r--r-- 1 root root 155 Sep 7 16:39 config\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings.save(\"elastic\")\n",
|
||||
"!ls -l elastic"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"And all we have is the configuration. No `database`, `embeddings` or `scoring` files. That data is in Postgres and Elasticsearch!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "OAtoEmV36FWp"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook showed how external databases and other external integrations can be used with embeddings databases. This architecture ensures that as new ways to index and store data become available, txtai can easily adapt.\n",
|
||||
"\n",
|
||||
"This notebook also showed how creating a custom component is a low level of effort and can easily be done for a component without an existing integration.\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "SQxnieiT6RTe"
|
||||
}
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,454 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": [],
|
||||
"gpuType": "T4"
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Custom API Endpoints\n",
|
||||
"\n",
|
||||
"The [txtai API](https://neuml.github.io/txtai/api/) is a web-based service backed by [FastAPI](https://fastapi.tiangolo.com/). Semantic search, LLM orchestration and Language Model Workflows can all run through the API.\n",
|
||||
"\n",
|
||||
"While the API is extremely flexible and complex logic can be executed through YAML-driven workflows, some may prefer to create an endpoint in Python.\n",
|
||||
"\n",
|
||||
"This notebook introduces API extensions and shows how they can be used to define custom Python endpoints that interact with txtai applications."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "VGeVB8M41jqW"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "ZQrHIw351lwE"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"id": "R0AqRP7v1hdr"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[api] datasets"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Define the extension\n",
|
||||
"\n",
|
||||
"First, we'll create an application that defines a persistent embeddings database and LLM. Then we'll combine those two into a RAG endpoint through the API."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "xmPN8RDF1pXd"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"%%writefile app.yml\n",
|
||||
"\n",
|
||||
"# Embeddings index\n",
|
||||
"writable: true\n",
|
||||
"embeddings:\n",
|
||||
" hybrid: true\n",
|
||||
" content: true\n",
|
||||
"\n",
|
||||
"# LLM pipeline\n",
|
||||
"llm:\n",
|
||||
" path: Qwen/Qwen3-4B-Instruct-2507"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "XZ7vPBIs1rGZ",
|
||||
"outputId": "b5cf95f1-1a99-4839-ae9b-9141922bd248"
|
||||
},
|
||||
"execution_count": 2,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing app.yml\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"The code below creates an API endpoint at `/rag`. This is a `GET` endpoint that takes a `text` parameter as input."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "syd1PZ621sok"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"%%writefile rag.py\n",
|
||||
"from fastapi import APIRouter\n",
|
||||
"from txtai.api import application, Extension\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"class RAG(Extension):\n",
|
||||
" \"\"\"\n",
|
||||
" API extension\n",
|
||||
" \"\"\"\n",
|
||||
"\n",
|
||||
" def __call__(self, app):\n",
|
||||
" app.include_router(RAGRouter().router)\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"class RAGRouter:\n",
|
||||
" \"\"\"\n",
|
||||
" API router\n",
|
||||
" \"\"\"\n",
|
||||
"\n",
|
||||
" router = APIRouter()\n",
|
||||
"\n",
|
||||
" @staticmethod\n",
|
||||
" @router.get(\"/rag\")\n",
|
||||
" def rag(text: str):\n",
|
||||
" \"\"\"\n",
|
||||
" Runs a retrieval augmented generation (RAG) pipeline.\n",
|
||||
"\n",
|
||||
" Args:\n",
|
||||
" text: input text\n",
|
||||
"\n",
|
||||
" Returns:\n",
|
||||
" response\n",
|
||||
" \"\"\"\n",
|
||||
"\n",
|
||||
" # Run embeddings search\n",
|
||||
" results = application.get().search(text, 3)\n",
|
||||
" context = \" \".join([x[\"text\"] for x in results])\n",
|
||||
"\n",
|
||||
" prompt = f\"\"\"\n",
|
||||
" Answer the following question using only the context below.\n",
|
||||
"\n",
|
||||
" Question: {text}\n",
|
||||
" Context: {context}\n",
|
||||
" \"\"\"\n",
|
||||
"\n",
|
||||
" return {\n",
|
||||
" \"response\": application.get().pipeline(\"llm\", (prompt,))\n",
|
||||
" }"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "zXERt7Vw1ujq",
|
||||
"outputId": "2c680298-895b-419c-967d-70030265f5a6"
|
||||
},
|
||||
"execution_count": 3,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing rag.py\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Start the API instance\n",
|
||||
"\n",
|
||||
"Let's start the API with the RAG extension."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "p7vl6_9i1w39"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!CONFIG=app.yml EXTENSIONS=rag.RAG nohup uvicorn \"txtai.api:app\" &> api.log &\n",
|
||||
"!sleep 60"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "FRif4lhW1y8m"
|
||||
},
|
||||
"execution_count": 4,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Create the embeddings database\n",
|
||||
"\n",
|
||||
"Next, we'll create the embeddings database using the `ag_news` dataset. This is a set of news stories from the mid 2000s."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "FTdkEDa0106G"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"import requests\n",
|
||||
"\n",
|
||||
"ds = load_dataset(\"ag_news\", split=\"train\")\n",
|
||||
"\n",
|
||||
"# API endpoint\n",
|
||||
"url = \"http://localhost:8000\"\n",
|
||||
"headers = {\"Content-Type\": \"application/json\"}\n",
|
||||
"\n",
|
||||
"# Add data\n",
|
||||
"batch = []\n",
|
||||
"for text in ds[\"text\"]:\n",
|
||||
" batch.append({\"text\": text})\n",
|
||||
" if len(batch) == 4096:\n",
|
||||
" requests.post(f\"{url}/add\", headers=headers, json=batch, timeout=120)\n",
|
||||
" batch = []\n",
|
||||
"\n",
|
||||
"if batch:\n",
|
||||
" requests.post(f\"{url}/add\", headers=headers, json=batch, timeout=120)\n",
|
||||
"\n",
|
||||
"# Build index\n",
|
||||
"index = requests.get(f\"{url}/index\")"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Ns6BKNQQ13FA"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Run queries\n",
|
||||
"\n",
|
||||
"Now that we have a knowledge source indexed, let's run a set of queries. The code below defines a method that calls the `/rag` endpoint and retrieves the response. Keep in mind this dataset is from 2004.\n",
|
||||
"\n",
|
||||
"While the Python Requests library is used in this notebook, this is a simple web endpoint that can be called from any programming language."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "_wGvCWsP17it"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"def rag(text):\n",
|
||||
" return requests.get(f\"{url}/rag?text={text}\").json()[\"response\"]\n",
|
||||
"\n",
|
||||
"rag(\"Who is the current President?\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 39
|
||||
},
|
||||
"id": "_WbFu64L15Ch",
|
||||
"outputId": "3d631fe8-d1d3-4437-bf64-9248599caff9"
|
||||
},
|
||||
"execution_count": 14,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'George W. Bush'"
|
||||
],
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 14
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"rag(\"Who lost the presidential election?\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 39
|
||||
},
|
||||
"id": "YtJ7LJ_819vw",
|
||||
"outputId": "e102b060-edb3-483c-98f9-50892e5e6c70"
|
||||
},
|
||||
"execution_count": 15,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'John Kerry'"
|
||||
],
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 15
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"rag(\"Who won the World Series?\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 39
|
||||
},
|
||||
"id": "BlYDMTj41_QL",
|
||||
"outputId": "4f58fb40-2e75-4248-8065-5efc969fdd0e"
|
||||
},
|
||||
"execution_count": 16,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'Boston'"
|
||||
],
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 16
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"rag(\"Who did the Red Sox beat to win the world series?\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 39
|
||||
},
|
||||
"id": "XMHLmQ532ApE",
|
||||
"outputId": "4bf5c7fa-dd42-43e3-b473-2df9d2c64d29"
|
||||
},
|
||||
"execution_count": 17,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'Cardinals'"
|
||||
],
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 17
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"rag(\"What major hurricane hit the USA?\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 39
|
||||
},
|
||||
"id": "pTMqQSx82B_h",
|
||||
"outputId": "9ee72bc9-664b-407d-ac65-95f1a09a2cb2"
|
||||
},
|
||||
"execution_count": 18,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'Charley'"
|
||||
],
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 18
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"rag(\"What mobile phone manufacturer has the largest current marketshare?\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 39
|
||||
},
|
||||
"id": "BV99h7272DVj",
|
||||
"outputId": "20602f12-09fe-4a44-a3f4-1797885e9d22"
|
||||
},
|
||||
"execution_count": 19,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'Nokia'"
|
||||
],
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 19
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook showed how a txtai application can be extended with custom endpoints in Python. While applications have a robust workflow framework, it may be preferable to write complex logic in Python and this method enables that."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "oPwgCgBc2Er2"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,592 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "VGeVB8M41jqW"
|
||||
},
|
||||
"source": [
|
||||
"# Build RAG pipelines with txtai\n",
|
||||
"\n",
|
||||
"Large Language Models (LLMs) have completely dominated the tech space in recent years. The results have been amazing and the public imagination is almost endless.\n",
|
||||
"\n",
|
||||
"While LLMs have been impressive, they are not problem free. The biggest challenge is with hallucinations. Hallucinations is the term for when a LLM generates output that is factually incorrect. The alarming part of this is that on a cursory glance, it actually sounds like good content. The default behavior of LLMs is to produce plausible answers even when no plausible answer exists. LLMs are not great at saying I don't know.\n",
|
||||
"\n",
|
||||
"Retrieval augmented generation (RAG) helps reduce the risk of hallucinations by limiting the context in which a LLM can generate answers. This is typically done with a vector search query that hydrates a prompt with a relevant context. RAG is one of the most practical and production-ready use cases for *Generative AI*. It's so popular now, that some are creating their entire companies around it.\n",
|
||||
"\n",
|
||||
"[txtai](https://github.com/neuml/txtai) has long had question-answering pipelines, which employ the same process of retrieving a relevant context. LLMs are now the preferred approach for analyzing that context and RAG pipelines are one of the main features of txtai. One of the other main features of txtai is that it's a vector database! You can build your prompts and limit your context all with one library. Hence the phrase *all-in-one AI framework*.\n",
|
||||
"\n",
|
||||
"This notebook shows how to build RAG pipelines with txtai."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ZQrHIw351lwE"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies. Since this notebook is using optional pipelines, we need to install the pipeline extras package."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "R0AqRP7v1hdr"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline]\n",
|
||||
"\n",
|
||||
"# Get test data\n",
|
||||
"!wget -N https://github.com/neuml/txtai/releases/download/v6.2.0/tests.tar.gz\n",
|
||||
"!tar -xvzf tests.tar.gz\n",
|
||||
"\n",
|
||||
"# Install NLTK\n",
|
||||
"import nltk\n",
|
||||
"nltk.download(['punkt', 'punkt_tab'])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "xmPN8RDF1pXd"
|
||||
},
|
||||
"source": [
|
||||
"# Start with the basics\n",
|
||||
"\n",
|
||||
"Let's jump right in and start with a simple LLM pipeline. The [LLM pipeline](https://neuml.github.io/txtai/pipeline/text/llm/) supports local LLM models via [Hugging Face Transformers](https://github.com/huggingface/transformers) and [llama.cpp](https://github.com/abetlen/llama-cpp-python).\n",
|
||||
"\n",
|
||||
"The LLM pipeline also supports [API services (i.e. OpenAI, Claude, Bedrock etc) via LiteLLM](https://github.com/BerriAI/litellm). The LLM pipeline automatically detects the underlying LLM framework from the `path` parameter.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"id": "XZ7vPBIs1rGZ"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import LLM\n",
|
||||
"\n",
|
||||
"# Create LLM\n",
|
||||
"llm = LLM(\"Qwen/Qwen3-4B-Instruct-2507\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "9rmTWMxAH3Vx"
|
||||
},
|
||||
"source": [
|
||||
"Next, we'll load a document to query. The [Textractor pipeline](https://neuml.github.io/txtai/pipeline/data/textractor/) has support for extracting text from common document formats (docx, pdf, xlsx, web)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "nifStGtOHuyc",
|
||||
"outputId": "5a4010e0-75f9-4095-a24c-cd4c859847d0"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"# txtai – the all-in-one embeddings database\n",
|
||||
"txtai is an all-in-one embeddings database for semantic search, LLM orchestration and language model workflows.\n",
|
||||
"\n",
|
||||
"Summary of txtai features:\n",
|
||||
"· *Vector search* with SQL, object storage, topic modeling\n",
|
||||
"· Create *embeddings* for text, documents, audio, images and video\n",
|
||||
"· *Pipelines* powered by language models that run LLM prompts\n",
|
||||
"· *Workflows* to join pipelines together and aggregate business logic\n",
|
||||
"· Build with *Python* or *YAML* . API bindings available for JavaScript, Java, Rust and Go.\n",
|
||||
"· *Run local or scale out with container orchestration* \n",
|
||||
"\n",
|
||||
"\n",
|
||||
"## Examples\n",
|
||||
"List of example notebooks.\n",
|
||||
"|Notebook|Description|\n",
|
||||
"|---|---|\n",
|
||||
"|Introducing txtai |Overview of the functionality provided by txtai|\n",
|
||||
"|Similarity search with images|Embed images and text into the same space for search|\n",
|
||||
"|Build a QA database|Question matching with semantic search|\n",
|
||||
"|Semantic Graphs|Explore topics, data connectivity and run network analysis|\n",
|
||||
"\n",
|
||||
"## Install\n",
|
||||
"The easiest way to install is via pip and PyPI\n",
|
||||
"pip install txtai\n",
|
||||
"Python 3.10+ is supported. Using a Python virtual environment is **recommended** .\n",
|
||||
"See the detailed install instructions for more information covering optional dependencies, environment specific prerequisites, installing from source, conda support and how to run with containers.\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"## Model guide\n",
|
||||
"The following shows a list of suggested models.\n",
|
||||
"|Component|Model(s)|\n",
|
||||
"|---|---|\n",
|
||||
"|Embeddings|all-MiniLM-L6-v2|\n",
|
||||
"||E5-base-v2|\n",
|
||||
"|Image Captions|BLIP|\n",
|
||||
"|Labels - Zero Shot|BART-Large-MNLI|\n",
|
||||
"|Labels - Fixed|Fine-tune with training pipeline|\n",
|
||||
"|Large Language Model (LLM)|Flan T5 XL|\n",
|
||||
"||Mistral 7B OpenOrca|\n",
|
||||
"|Summarization|DistilBART|\n",
|
||||
"|Text-to-Speech|ESPnet JETS|\n",
|
||||
"|Transcription|Whisper|\n",
|
||||
"|Translation|OPUS Model Series|\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import Textractor\n",
|
||||
"\n",
|
||||
"# Create Textractor\n",
|
||||
"textractor = Textractor()\n",
|
||||
"text = textractor(\"txtai/document.docx\")\n",
|
||||
"print(text)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "2jkamwgdIgEp"
|
||||
},
|
||||
"source": [
|
||||
"Now we'll define a simple LLM pipeline. It takes a question and context (which in this case is the whole file), creates a prompt and runs it with the LLM."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 18,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 59
|
||||
},
|
||||
"id": "9HU6C0OIIAKn",
|
||||
"outputId": "f9d556c4-cd7a-4774-ef62-1f3fff90aa47"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'txtai is an all-in-one embeddings database that supports semantic search, LLM orchestration, and language model workflows with features like vector search, embeddings for text, audio, images, and video, pipelines powered by language models, and scalable workflows available via Python or YAML with API bindings for multiple languages.'"
|
||||
]
|
||||
},
|
||||
"execution_count": 18,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"def execute(question, text):\n",
|
||||
" return llm([\n",
|
||||
" {\"role\": \"system\", \"content\": \"You are a friendly assistant. You answer questions from users.\"},\n",
|
||||
" {\"role\": \"user\", \"content\": f\"\"\"\n",
|
||||
" Answer the following question using only the context below. Only include information specifically discussed.\n",
|
||||
"\n",
|
||||
" question: {question}\n",
|
||||
" context: {text} \n",
|
||||
" \"\"\"}\n",
|
||||
" ], maxlength=4096)\n",
|
||||
"\n",
|
||||
"execute(\"Tell me about txtai in one sentence\", text)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 19,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 39
|
||||
},
|
||||
"id": "xxF7ajCPJP5_",
|
||||
"outputId": "0e6f6dbb-c784-4841-fe3f-82754ef478eb"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'txtai recommends using Whisper for transcription.'"
|
||||
]
|
||||
},
|
||||
"execution_count": 19,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"execute(\"What model does txtai recommend for transcription?\", text)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 20,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 97
|
||||
},
|
||||
"id": "AKmmTqsnJa5X",
|
||||
"outputId": "834bf3ee-b7ed-4e38-e2ef-d5950f99ed9a"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'The best thing to read if you don\\'t know anything about txtai would be the \"Introducing txtai\" notebook, as it provides an overview of the functionality offered by txtai.'"
|
||||
]
|
||||
},
|
||||
"execution_count": 20,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"execute(\"I don't know anything about txtai, what would be the best thing to read?\", text)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "WaVeEHrIMpFr"
|
||||
},
|
||||
"source": [
|
||||
"If this is the first time you've seen *Generative AI*, then these statements are 🤯. Even if you've been in the space a while, it's still amazing how much a language model can understand and the high level of quality in it's answers.\n",
|
||||
"\n",
|
||||
"While this use case is fun, lets try to scale it to a larger set of documents."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "viVVft59NbKv"
|
||||
},
|
||||
"source": [
|
||||
"# Build a RAG pipeline with vector search\n",
|
||||
"\n",
|
||||
"Let's say we have a large number of documents, hundreds/thousands etc. We can't just put all those documents into a single prompt, we'll run out of GPU memory fast!\n",
|
||||
"\n",
|
||||
"This is where retrieval augmented generation enters the picture. We can use a query step that finds the best candidates to add to the prompt.\n",
|
||||
"\n",
|
||||
"Typically, this candidate query uses vector search but it can be anything that runs a search and returns results. In fact, many complex production systems have customized retrieval pipelines that feed a context into LLM prompts.\n",
|
||||
"\n",
|
||||
"The first step in building our RAG pipeline is creating the knowledge store. In this case, it's a vector database of file content. The files will be split into paragraphs with each paragraph stored as a separate row."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ipmsmtN1NahT",
|
||||
"outputId": "64733e1f-fb7b-4a2d-bf02-8930478a8ee8"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Indexing txtai/article.pdf\n",
|
||||
"Indexing txtai/document.docx\n",
|
||||
"Indexing txtai/document.pdf\n",
|
||||
"Indexing txtai/spreadsheet.xlsx\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"\n",
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"def stream(path):\n",
|
||||
" for f in sorted(os.listdir(path)):\n",
|
||||
" fpath = os.path.join(path, f)\n",
|
||||
"\n",
|
||||
" # Only accept documents\n",
|
||||
" if f.endswith((\"docx\", \"xlsx\", \"pdf\")):\n",
|
||||
" print(f\"Indexing {fpath}\")\n",
|
||||
" for paragraph in textractor(fpath):\n",
|
||||
" yield paragraph\n",
|
||||
"\n",
|
||||
"# Document text extraction, split into paragraphs\n",
|
||||
"textractor = Textractor(paragraphs=True)\n",
|
||||
"\n",
|
||||
"# Vector Database\n",
|
||||
"embeddings = Embeddings(content=True)\n",
|
||||
"embeddings.index(stream(\"txtai\"))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ASlmAaR3nBPN"
|
||||
},
|
||||
"source": [
|
||||
"The next step is defining the RAG pipeline. This pipeline takes the input question, runs a vector search and builds a context using the search results. The context is then inserted into a prompt template and run with the LLM."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 23,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 39
|
||||
},
|
||||
"id": "_-9SW6r4P5ha",
|
||||
"outputId": "6a7bcd69-bcd0-4f6e-81c1-e32c323b3ffb"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'txtai recommends using BLIP for image captioning.'"
|
||||
]
|
||||
},
|
||||
"execution_count": 23,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"def context(question):\n",
|
||||
" context = \"\\n\".join(x[\"text\"] for x in embeddings.search(question))\n",
|
||||
" return context\n",
|
||||
"\n",
|
||||
"def rag(question):\n",
|
||||
" return execute(question, context(question))\n",
|
||||
"\n",
|
||||
"rag(\"What model does txtai recommend for image captioning?\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 24,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "NbQhSunPQtB0",
|
||||
"outputId": "de0caf04-4cdb-48e8-aadf-37283be9909a"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"The BLIP model was added for image captioning on 2022-03-17.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"result = rag(\"When was the BLIP model added for image captioning?\")\n",
|
||||
"print(result)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "D6HW-3GtnTFl"
|
||||
},
|
||||
"source": [
|
||||
"As we can see, the result is similar to what we had before without vector search. The difference is that we only used a relevant portion of the documents to generate the answer.\n",
|
||||
"\n",
|
||||
"As we discussed before, this is important when dealing with large volumes of data. Not all of the data can be added to a LLM prompt. Additionally, having only the most relevant context helps the LLM generate higher quality answers."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "FZ-yPC-xiUqa"
|
||||
},
|
||||
"source": [
|
||||
"# Citations for LLMs\n",
|
||||
"\n",
|
||||
"A healthy level of skepticism should be applied to answers generated by AI. We're far from the day where we can blindly trust answers from an AI model.\n",
|
||||
"\n",
|
||||
"txtai has a couple approaches for generating citations. The basic approach is to take the answer and search the vector database for the closest match."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 25,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "obttSg_dSFT5",
|
||||
"outputId": "c7ae7675-6959-4bcb-ad06-065ea8609c31"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"E5-base-v2\n",
|
||||
"Image Captions BLIP\n",
|
||||
"Labels - Zero Shot BART-Large-MNLI\n",
|
||||
"# Model Guide\n",
|
||||
"|Component |Model(s)|Date Added|\n",
|
||||
"|---|---|---|\n",
|
||||
"|Embeddings |all-MiniLM-L6-v2|2022-04-15|\n",
|
||||
"|Image Captions |BLIP|2022-03-17|\n",
|
||||
"|Labels - Zero Shot |BART-Large-MNLI|2022-01-01|\n",
|
||||
"|Large Language Model (LLM) |Mistral 7B OpenOrca|2023-10-01|\n",
|
||||
"|Summarization |DistilBART|2021-02-22|\n",
|
||||
"|Text-to-Speech |ESPnet JETS|2022-08-01|\n",
|
||||
"|Transcription |Whisper|2022-08-01|\n",
|
||||
"|Translation |OPUS Model Series|2021-04-06|\n",
|
||||
"&\"Times New Roman,Regular\"&12&A\n",
|
||||
"## Model guide\n",
|
||||
"The following shows a list of suggested models.\n",
|
||||
"|Component|Model(s)|\n",
|
||||
"|---|---|\n",
|
||||
"|Embeddings|all-MiniLM-L6-v2|\n",
|
||||
"||E5-base-v2|\n",
|
||||
"|Image Captions|BLIP|\n",
|
||||
"|Labels - Zero Shot|BART-Large-MNLI|\n",
|
||||
"|Labels - Fixed|Fine-tune with training pipeline|\n",
|
||||
"|Large Language Model (LLM)|Flan T5 XL|\n",
|
||||
"||Mistral 7B OpenOrca|\n",
|
||||
"|Summarization|DistilBART|\n",
|
||||
"|Text-to-Speech|ESPnet JETS|\n",
|
||||
"|Transcription|Whisper|\n",
|
||||
"|Translation|OPUS Model Series|\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"for x in embeddings.search(result):\n",
|
||||
" print(x[\"text\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "MDcB7GCWY6TO"
|
||||
},
|
||||
"source": [
|
||||
"While the basic approach above works in this case, txtai has a more robust pipeline to handle citations and references.\n",
|
||||
"\n",
|
||||
"The RAG pipeline is defined below. A RAG pipeline works in the same way as a LLM + Vector Search pipeline, except it has special logic for generating citations. This pipeline takes the answers and compares it to the context passed to the LLM to determine the most likely reference."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 26,
|
||||
"metadata": {
|
||||
"id": "Lm6gg85_Y7ot"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import RAG\n",
|
||||
"\n",
|
||||
"# Create the RAG pipeline\n",
|
||||
"rag = RAG(embeddings, \"Qwen/Qwen3-4B-Instruct-2507\", template=\"\"\"\n",
|
||||
" Answer the following question using the provided context.\n",
|
||||
"\n",
|
||||
" Question:\n",
|
||||
" {question}\n",
|
||||
"\n",
|
||||
" Context:\n",
|
||||
" {context}\n",
|
||||
"\"\"\", output=\"reference\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 27,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "4pOfE5paZatH",
|
||||
"outputId": "2bed2de5-22ff-4f7b-dba5-41b8e4cc6c75"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"ANSWER: Python 3.10 and later versions (Python 3.10+) are supported.\n",
|
||||
"CITATION: [{'id': '24', 'text': 'Python 3.10+ is supported. Using a Python virtual environment is recommended.'}]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"result = rag(\"What version of Python is supported?\", maxlength=4096)\n",
|
||||
"print(\"ANSWER:\", result[\"answer\"])\n",
|
||||
"print(\"CITATION:\", embeddings.search(\"select id, text from txtai where id = :id\", limit=1, parameters={\"id\": result[\"reference\"]}))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "vHdE2Q59jNnF"
|
||||
},
|
||||
"source": [
|
||||
"And as we can see, not only is the answer to the statement shown, the RAG pipeline also provides a citation. This step is crucial in any line of work where answers must be verified (which is most lines of work)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "oPwgCgBc2Er2"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook introduced retrieval augmented generation (RAG), explained why we need it and showed the options available for running RAG pipelines with txtai.\n",
|
||||
"\n",
|
||||
"The advantages of building RAG pipelines with txtai are:\n",
|
||||
"\n",
|
||||
"- **All-in-one AI framework** - one library can handle LLM inference and vector search retrieval\n",
|
||||
"- **Generating citations** - generating answers is useful but referencing where those answers came from is crucial in gaining the trust of users\n",
|
||||
"- **Simple yet powerful** - building pipelines can be done in a small amount of Python. Options are available to build pipelines in YAML and/or run through the API"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "T4",
|
||||
"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.10.19"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
@@ -0,0 +1,334 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": [],
|
||||
"gpuType": "T4"
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Integrate LLM Frameworks\n",
|
||||
"\n",
|
||||
"The release of [BERT](https://arxiv.org/abs/1810.04805) in 2018 kicked off the language model revolution. The Transformers architecture succeeded RNNs and LSTMs to become the architecture of choice. Unbelievable progress was made in a number of areas: summarization, translation, text classification, entity classification and more. 2023 tooks things to another level with the rise of large language models (LLMs). Models with billions of parameters showed an amazing ability to generate coherent dialogue.\n",
|
||||
"\n",
|
||||
"Looking ahead towards the next wave of innovation, we're due for another shift in model architecture. For example, the [Mamba paper](https://arxiv.org/abs/2312.00752) previews a possible future after Transformers.\n",
|
||||
"\n",
|
||||
"With that in mind, [txtai](https://github.com/neuml/txtai) now has the capability to easily integrate additional LLM frameworks. While local models through Hugging Face Transformers continues to be the default choice, these additional LLM frameworks broaden the number of options available.\n",
|
||||
"\n",
|
||||
"This notebook will demonstrate how txtai can integrate with [llama.cpp](https://github.com/ggerganov/llama.cpp), [LiteLLM](https://github.com/BerriAI/litellm) and custom generation methods. For custom generation, we'll show how to run inference with a `Mamba` model."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "VGeVB8M41jqW"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies. Since this notebook is using optional libraries, we need to install the `pipeline-llm` extras package."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "ZQrHIw351lwE"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"id": "R0AqRP7v1hdr"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline-llm]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# llama.cpp\n",
|
||||
"\n",
|
||||
"First, we'll demonstrate how to load a model with llama.cpp. This framework is an extremely popular method with those who run local LLMs. It provides a number of innovations in running LLMs on CPUs, especially on Mac's.\n",
|
||||
"\n",
|
||||
"The following example shows a retrieval augmented generation (RAG) pipeline with llama.cpp. txtai automatically loads llama.cpp models when working with GGUF files."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "32xg8L1JHd3d"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai import Embeddings, RAG, LLM\n",
|
||||
"\n",
|
||||
"data = [\n",
|
||||
" \"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",
|
||||
"\n",
|
||||
"# Create embeddings\n",
|
||||
"embeddings = Embeddings(content=True, autoid=\"uuid5\")\n",
|
||||
"\n",
|
||||
"# Create an index for the list of text\n",
|
||||
"embeddings.index(data)\n",
|
||||
"\n",
|
||||
"# Create LLM with llama.cpp - GGUF file is automatically downloaded\n",
|
||||
"llm = LLM(\"unsloth/Qwen3-4B-Instruct-2507-GGUF/Qwen3-4B-Instruct-2507-Q4_K_M.gguf\", verbose=True)\n",
|
||||
"\n",
|
||||
"template = \"\"\"<|im_start|>system\n",
|
||||
"You are a friendly assistant. You answer questions from users.<|im_end|>\n",
|
||||
"<|im_start|>user\n",
|
||||
"Find the best matching text in the context for the question. The response should be the text from the context only.\n",
|
||||
"\n",
|
||||
"Question:\n",
|
||||
"{question}\n",
|
||||
"\n",
|
||||
"Context:\n",
|
||||
"{context}\n",
|
||||
"\n",
|
||||
"Text:\n",
|
||||
"<|im_end|>\n",
|
||||
"<|im_start|>assistant\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"# Create and run RAG instance\n",
|
||||
"rag = RAG(embeddings, llm, output=\"reference\", separator=\"\\n\", template=template)\n",
|
||||
"result = rag(\"Tell me about someone lucky\")\n",
|
||||
"\n",
|
||||
"print(\"ANSWER:\", result[\"answer\"])\n",
|
||||
"print(\"REFERENCE:\", embeddings.search(\"select id, text from txtai where id = :id\", parameters={\"id\": result[\"reference\"]}))\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "XZ7vPBIs1rGZ",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "18554fc7-1383-46dd-fd7d-87d816747fad"
|
||||
},
|
||||
"execution_count": 7,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"ANSWER: Maine man wins $1M from $25 lottery ticket\n",
|
||||
"REFERENCE: [{'id': '37e5fae7-74c2-5f1c-bf69-2962dd7470d1', 'text': 'Maine man wins $1M from $25 lottery ticket'}]\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"The code above builds an embeddings database, runs a vector search and passes those results to a LLM prompt. As expected, it prints the best answer and reference. The difference is that LLM inference is run through `llama.cpp` vs `transformers`."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "_qDllo4AIky2"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# LiteLLM\n",
|
||||
"\n",
|
||||
"LiteLLM is an abstraction framework designed to run with API-based LLMs. At the time of writing this article, LiteLLM supports over 100+ LLMs. See the [full list of providers](https://docs.litellm.ai/docs/providers) for all the options.\n",
|
||||
"\n",
|
||||
"The following example shows a LLM call with the Hugging Face Inference API. This method automatically detects that this is a LiteLLM model string."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "83BSRNPAI6Q9"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Hugging Face Inference API\n",
|
||||
"llm = LLM(\"huggingface/roneneldan/TinyStories-1M\")\n",
|
||||
"print(llm(\"The cat and the dog.\", maxlength=5))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ye0DshJF4MLj",
|
||||
"outputId": "6bc74aa4-201b-4b68-9545-a1110778e0b2"
|
||||
},
|
||||
"execution_count": 14,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"They are friends.\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Given that all these APIs require a paid account, we'll leave it to you to try other API models using your own authentication."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "bCpx6BtOJgw2"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Custom Generation\n",
|
||||
"\n",
|
||||
"Last but certainly not least, we'll demonstrate how to add a custom generation framework. For this example, we'll use the recently released [mamba-chat](https://huggingface.co/havenhq/mamba-chat) model to build a RAG pipeline. You can read more about the model in this [GitHub Repository](https://github.com/havenhq/mamba-chat)\n",
|
||||
"\n",
|
||||
"The following sections install support for Mamba models, define a Mamba Generation instance and run a Mamba-based RAG pipeline."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "43x3YPZVJ2SA"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install mamba-ssm\n",
|
||||
"\n",
|
||||
"# Link CUDA libraries into environment\n",
|
||||
"!export LC_ALL=\"en_US.UTF-8\"\n",
|
||||
"!export LD_LIBRARY_PATH=\"/usr/lib64-nvidia\"\n",
|
||||
"!export LIBRARY_PATH=\"/usr/local/cuda/lib64/stubs\"\n",
|
||||
"!ldconfig /usr/lib64-nvidia"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "j6GQo1_J_HpU"
|
||||
},
|
||||
"execution_count": 4,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import torch\n",
|
||||
"\n",
|
||||
"from transformers import AutoTokenizer\n",
|
||||
"from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel\n",
|
||||
"\n",
|
||||
"from txtai.pipeline import Generation\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"class MambaGeneration(Generation):\n",
|
||||
" def __init__(self, path, template=None, **kwargs):\n",
|
||||
" super().__init__(path, template, **kwargs)\n",
|
||||
"\n",
|
||||
" self.tokenizer = AutoTokenizer.from_pretrained(path)\n",
|
||||
" self.tokenizer.eos_token = \"<|endoftext|>\"\n",
|
||||
" self.tokenizer.pad_token = self.tokenizer.eos_token\n",
|
||||
"\n",
|
||||
" self.model = MambaLMHeadModel.from_pretrained(path, device=\"cuda\", dtype=torch.float16)\n",
|
||||
"\n",
|
||||
" def execute(self, texts, maxlength, **kwargs):\n",
|
||||
" results = []\n",
|
||||
" for text in texts:\n",
|
||||
" # Tokenize prompt\n",
|
||||
" tokens = self.tokenizer(text, return_tensors=\"pt\").to(\"cuda\")[\"input_ids\"]\n",
|
||||
"\n",
|
||||
" # Run inference\n",
|
||||
" output = self.model.generate(input_ids=tokens, max_length=maxlength, eos_token_id=self.tokenizer.eos_token_id, **kwargs)\n",
|
||||
"\n",
|
||||
" # Decode results\n",
|
||||
" output = self.tokenizer.batch_decode(output)\n",
|
||||
" output = output[0].split(\"<|assistant|>\\n\")[-1].replace(\"<|endoftext|>\", \"\").strip()\n",
|
||||
" results.append(output)\n",
|
||||
"\n",
|
||||
" return results"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "1uYgXT_q_U9U"
|
||||
},
|
||||
"execution_count": 5,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"llm = LLM(\"havenhq/mamba-chat\", method=\"__main__.MambaGeneration\")\n",
|
||||
"\n",
|
||||
"template = \"\"\"<|system|>You are a friendly assistant. You answer questions from users.</s>\n",
|
||||
"<|user|>\n",
|
||||
"Find the best matching text in the context for the question. The response should be the text from the context only.\n",
|
||||
"\n",
|
||||
"Question:\n",
|
||||
"{question}\n",
|
||||
"\n",
|
||||
"Context:\n",
|
||||
"{context}\n",
|
||||
"</s>\n",
|
||||
"<|assistant|>\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"# Create and run RAG instance\n",
|
||||
"rag = RAG(embeddings, llm, output=\"reference\", separator=\"\\n\", template=template)\n",
|
||||
"result = rag(\"Tell me something about about wildlife\")\n",
|
||||
"\n",
|
||||
"print(\"ANSWER:\", result[\"answer\"])\n",
|
||||
"print(\"REFERENCE:\", embeddings.search(\"select id, text from txtai where id = :id\", parameters={\"id\": result[\"reference\"]}))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "4vZUYLJB_jvb",
|
||||
"outputId": "cb8f81f5-6744-4450-9132-dadb16b9096c"
|
||||
},
|
||||
"execution_count": 6,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"ANSWER: The National Park Service warns against sacrificing slower friends in a bear attack.\n",
|
||||
"REFERENCE: [{'id': '7224f159-658b-5891-b06c-9a96cfa6a54d', 'text': 'The National Park Service warns against sacrificing slower friends in a bear attack'}]\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"As expected, the best answer and reference is shown.\n",
|
||||
"\n",
|
||||
"There is much to learn and validate about Mamba but it's important to note this model is only 2.8B parameters. The Mamba architecture is one to watch moving forward!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "i_7fgP4qLveq"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook demonstrated how to run LLMs through txtai using alternate LLM frameworks. It's an exciting time in AI/NLP/Machine Learning. What new innovations will 2024 bring? Time will tell but txtai is ready to integrate them in!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "oPwgCgBc2Er2"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,510 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# API Authorization and Authentication\n",
|
||||
"\n",
|
||||
"txtai can run in Python, with YAML configuration and through an API service.\n",
|
||||
"\n",
|
||||
"The default API service implementation runs without any security. This may be OK for a local prototype or if it's run on a small internal network. But in most cases, additional security measures should be taken.\n",
|
||||
"\n",
|
||||
"This notebook will demonstrate how to add authorization, authentication and middleware dependencies to a txtai API service."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "VGeVB8M41jqW"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies. Since this notebook uses the API, we need to install the api extras package."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "ZQrHIw351lwE"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "R0AqRP7v1hdr"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[api]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Create an API Service\n",
|
||||
"\n",
|
||||
"For this example, we'll load an existing txtai index from the Hugging Face Hub."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "32xg8L1JHd3d"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"%%writefile config.yml\n",
|
||||
"cloud:\n",
|
||||
" provider: huggingface-hub\n",
|
||||
" container: neuml/txtai-intro\n",
|
||||
"\n",
|
||||
"embeddings:"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "XZ7vPBIs1rGZ",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "468b09cb-ba01-426d-9640-142acf1e3dd9"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing config.yml\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Next, we'll generate a test token to use for this notebook."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Eu8rtjB_QZOg"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import uuid\n",
|
||||
"str(uuid.uuid5(uuid.NAMESPACE_DNS, \"TokenTest\"))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 38
|
||||
},
|
||||
"id": "Np18NtytQG8k",
|
||||
"outputId": "5bf60d69-311d-4a02-f3dd-d210aedc8b4f"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'edd590d3-bfab-5425-8a85-79b01e3127ee'"
|
||||
],
|
||||
"application/vnd.google.colaboratory.intrinsic+json": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 3
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"txtai has a default API token authorization method built-in. We'll set a token and start the service.\n",
|
||||
"\n",
|
||||
"**It's important to note that this service is running via HTTP as this is only for demonstration purposes. HTTPS must be added either with a proxy service like NGINX or by passing a SSL cert to Uvicorn. See [this link](https://neuml.github.io/txtai/api/security/) for more.**"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "ehZSh97_Nr7h"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!CONFIG=config.yml TOKEN=`echo -n 'edd590d3-bfab-5425-8a85-79b01e3127ee' | sha256sum | head -c 64` uvicorn \"txtai.api:app\" &> api.log &\n",
|
||||
"!sleep 60"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "PjKT3vOuNkfX"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Connect to Service\n",
|
||||
"\n",
|
||||
"First, we'll try a request with no token to see what happens."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "NfOn1vgwRF2C"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!curl -X GET -I 'http://localhost:8000/search?query=feel+good+story&limit=1'"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "9Ljn13HWOb26",
|
||||
"outputId": "4e437ebb-cf92-4e7d-c804-25bb42f6cb30"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"HTTP/1.1 401 Unauthorized\r\n",
|
||||
"\u001b[1mdate\u001b[0m: Thu, 04 Jan 2024 15:08:38 GMT\r\n",
|
||||
"\u001b[1mserver\u001b[0m: uvicorn\r\n",
|
||||
"\u001b[1mcontent-length\u001b[0m: 40\r\n",
|
||||
"\u001b[1mcontent-type\u001b[0m: application/json\r\n",
|
||||
"\r\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"As expected, we received a HTTP 401 saying the request is not authorized.\n",
|
||||
"\n",
|
||||
"Now let's try an invalid token."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Rvi4PRSTRTjt"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!curl -X GET -I 'http://localhost:8000/search?query=feel+good+story&limit=1' -H 'Authorization: Bearer junk'"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "vJvQHtcJdRXb",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "099577d1-2ed0-4051-827d-2ab72d5929fa"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"HTTP/1.1 401 Unauthorized\r\n",
|
||||
"\u001b[1mdate\u001b[0m: Thu, 04 Jan 2024 15:08:38 GMT\r\n",
|
||||
"\u001b[1mserver\u001b[0m: uvicorn\r\n",
|
||||
"\u001b[1mcontent-length\u001b[0m: 40\r\n",
|
||||
"\u001b[1mcontent-type\u001b[0m: application/json\r\n",
|
||||
"\r\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Once again, the request is rejected.\n",
|
||||
"\n",
|
||||
"Let's try again, this time passing a valid API token."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "7oPf-Crb6ETH"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!curl -X GET 'http://localhost:8000/search?query=feel+good+story&limit=1' -H 'Authorization: Bearer edd590d3-bfab-5425-8a85-79b01e3127ee'"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "HNxzVN1CSE1C",
|
||||
"outputId": "77962561-52fa-4697-9c0c-bc2d630ea8b4"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"[{\"id\":\"4\",\"text\":\"Maine man wins $1M from $25 lottery ticket\",\"score\":0.08329025655984879}]"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"This time we get search results!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "H025YmZ4TSSS"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Dependencies\n",
|
||||
"\n",
|
||||
"Next, let's add a custom dependency to test out authentication. A dependency could integrate with external identity providers to validate user credentials such as OAuth, Active Directory, LDAP or another identity management service.\n",
|
||||
"\n",
|
||||
"For this simple example, we'll validate user credentials using basic HTTP authentication. The code below checks if a specific username and password are provided. It's based on [this FastAPI example](https://fastapi.tiangolo.com/advanced/security/http-basic-auth/)."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "3JjhlfhsUDNU"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"%%writefile authentication.py\n",
|
||||
"\n",
|
||||
"import secrets\n",
|
||||
"\n",
|
||||
"from fastapi import Depends, HTTPException, status\n",
|
||||
"from fastapi.security import HTTPBasic, HTTPBasicCredentials\n",
|
||||
"\n",
|
||||
"security = HTTPBasic()\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"class Authentication:\n",
|
||||
" def __call__(self, credentials: HTTPBasicCredentials = Depends(security)):\n",
|
||||
" user = credentials.username.encode(\"utf8\")\n",
|
||||
" validuser = secrets.compare_digest(user, b\"txtai\")\n",
|
||||
"\n",
|
||||
" password = credentials.password.encode(\"utf8\")\n",
|
||||
" validpassword = secrets.compare_digest(password, b\"theembeddingsdb\")\n",
|
||||
"\n",
|
||||
" if not (validuser and validpassword):\n",
|
||||
" raise HTTPException(\n",
|
||||
" status_code=status.HTTP_401_UNAUTHORIZED,\n",
|
||||
" detail=\"Incorrect user or password\",\n",
|
||||
" headers={\"WWW-Authenticate\": \"Basic\"},\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" return credentials.username"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "1W9Vw0PMUa83",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "8bece467-588d-4d80-98a0-2451d7ebf85e"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing authentication.py\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Now let's restart the application and add this dependency in.\n",
|
||||
"\n",
|
||||
"**Once again, same note as above, this demonstration uses HTTP. Real use-cases must use HTTPS.**"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "VNYHD20xhVAS"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!killall -9 uvicorn\n",
|
||||
"!CONFIG=config.yml DEPENDENCIES=authentication.Authentication uvicorn \"txtai.api:app\" &> api.log &\n",
|
||||
"!sleep 30"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "F60-Aakyg1x_"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!curl -X GET -I 'http://localhost:8000/search?query=feel+good+story&limit=1'"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "jda9RugAhcNT",
|
||||
"outputId": "89111cdc-bee4-4cec-83ce-de3a4bf8e21a"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"HTTP/1.1 401 Unauthorized\r\n",
|
||||
"\u001b[1mdate\u001b[0m: Thu, 04 Jan 2024 15:09:10 GMT\r\n",
|
||||
"\u001b[1mserver\u001b[0m: uvicorn\r\n",
|
||||
"\u001b[1mwww-authenticate\u001b[0m: Basic\r\n",
|
||||
"\u001b[1mcontent-length\u001b[0m: 30\r\n",
|
||||
"\u001b[1mcontent-type\u001b[0m: application/json\r\n",
|
||||
"\r\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"The request is rejected as expected. Next let's try an invalid username/password."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "XZCsCDOG6VBS"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!curl -X GET -I 'http://localhost:8000/search?query=feel+good+story&limit=1' -H \"Authorization: Basic junk\""
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "VFj1UPuC6c4P",
|
||||
"outputId": "4fa7d6df-c06a-4cfe-88ff-4c341783704d"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"HTTP/1.1 401 Unauthorized\r\n",
|
||||
"\u001b[1mdate\u001b[0m: Thu, 04 Jan 2024 15:09:10 GMT\r\n",
|
||||
"\u001b[1mserver\u001b[0m: uvicorn\r\n",
|
||||
"\u001b[1mwww-authenticate\u001b[0m: Basic\r\n",
|
||||
"\u001b[1mcontent-length\u001b[0m: 47\r\n",
|
||||
"\u001b[1mcontent-type\u001b[0m: application/json\r\n",
|
||||
"\r\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Once again, the request is rejected.\n",
|
||||
"\n",
|
||||
"Now we'll add the expected username/password to the request. [HTTP basic authentication](https://en.wikipedia.org/wiki/Basic_access_authentication) simply concats the username-password separated by a colon and then base64 encodes it."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "jwZYyTDKh5pB"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!echo -n txtai:theembeddingsdb | base64"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "uL943ry3h2hR",
|
||||
"outputId": "aa2bc9a6-cd0a-42aa-933c-3cb5aa579588"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"dHh0YWk6dGhlZW1iZWRkaW5nc2Ri\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!curl -X GET 'http://localhost:8000/search?query=feel+good+story&limit=1' -H \"Authorization: Basic dHh0YWk6dGhlZW1iZWRkaW5nc2Ri\""
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "X9vfyRsphxiA",
|
||||
"outputId": "01cfe7b9-80ee-4745-eae9-7db46b2ec554"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"[{\"id\":\"4\",\"text\":\"Maine man wins $1M from $25 lottery ticket\",\"score\":0.08329025655984879}]"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Now that we have the correct username/password, a response is returned!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "M-NSiDGuiigq"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook introduced how to add authorization, authentication and middleware dependencies to a txtai API service. As noted multiple times, ensure that HTTPS is enabled when using this in production environments.\n",
|
||||
"\n",
|
||||
"For more advanced authentication methods, check out the [FastAPI security documentation](https://fastapi.tiangolo.com/tutorial/security/).\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "ZivWat7-r3ID"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,275 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": [],
|
||||
"gpuType": "T4"
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# External vectorization\n",
|
||||
"\n",
|
||||
"Vectorization is the process of transforming data into numbers using machine learning models. Input data is run through a model and fixed dimension vectors are returned. These vectors can then be loaded into a vector database for similarity search.\n",
|
||||
"\n",
|
||||
"txtai is an open-source first system. Given it's own open-source roots, like-minded projects such as [sentence-transformers](https://github.com/UKPLab/sentence-transformers) are prioritized during development. But that doesn't mean txtai can't work with Embeddings API services.\n",
|
||||
"\n",
|
||||
"This notebook will show to use txtai with external vectorization."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "SwgRD_NGutB9"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "68Iw-nPbhYIJ"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"id": "R0AqRP7v1hdr"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Create an Embeddings dataset\n",
|
||||
"\n",
|
||||
"The first thing we'll do is pre-compute an embeddings dataset. In addition to Embeddings APIs, this can also be used during internal testing to tune index and database settings."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "_F84mqRHwpCP"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"# Load dataset\n",
|
||||
"wikipedia = Embeddings()\n",
|
||||
"wikipedia.load(provider=\"huggingface-hub\", container=\"neuml/txtai-wikipedia\")\n",
|
||||
"\n",
|
||||
"# Query for Top 10,000 most popular articles\n",
|
||||
"query = \"\"\"\n",
|
||||
"SELECT id, text FROM txtai\n",
|
||||
"order by percentile desc\n",
|
||||
"LIMIT 10000\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"data = wikipedia.search(query)\n",
|
||||
"\n",
|
||||
"# Encode vectors using same vector model as Wikipedia\n",
|
||||
"vectors = wikipedia.batchtransform(x[\"text\"] for x in data)\n",
|
||||
"\n",
|
||||
"# Build dataset of id, text, embeddings\n",
|
||||
"dataset = []\n",
|
||||
"for i, row in enumerate(data):\n",
|
||||
" dataset.append({\"id\": row[\"id\"], \"article\": row[\"text\"], \"embeddings\": vectors[i]})\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "U52gN-vUxjky"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Build an Embeddings index with external vectors\n",
|
||||
"\n",
|
||||
"Next, we'll create an Embedding index with an external transform function set.\n",
|
||||
"\n",
|
||||
"The external transform function can be any function or callable object. This function takes an array of data and returns an array of embeddings."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Ggo-N9iOQJ4d"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"def transform(inputs):\n",
|
||||
" return wikipedia.batchtransform(inputs)\n",
|
||||
"\n",
|
||||
"def stream():\n",
|
||||
" for row in dataset:\n",
|
||||
" # Index vector\n",
|
||||
" yield row[\"id\"], row[\"embeddings\"]\n",
|
||||
"\n",
|
||||
" # Index metadata\n",
|
||||
" yield {\"id\": row[\"id\"], \"article\": row[\"article\"]}\n",
|
||||
"\n",
|
||||
"embeddings = Embeddings(transform=\"__main__.transform\", content=True)\n",
|
||||
"embeddings.index(stream())"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "iujXcHzMCd0B"
|
||||
},
|
||||
"execution_count": 64,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"🚀 Notice how fast creating the index was compared to indexing. This is because there is no vectorization! Now let's run a query."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "B0Se6nG7JxOG"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"embeddings.search(\"select id, article, score from txtai where similar(:x)\", parameters={\"x\": \"operating system\"})"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "Bi0WnvosEL8E",
|
||||
"outputId": "d190e8c1-cb08-4d40-989c-aed4a67edc06"
|
||||
},
|
||||
"execution_count": 63,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': 'Operating system',\n",
|
||||
" 'article': 'An operating system (OS) is system software that manages computer hardware and software resources, and provides common services for computer programs.',\n",
|
||||
" 'score': 0.8955847024917603},\n",
|
||||
" {'id': 'MacOS',\n",
|
||||
" 'article': \"macOS (;), originally Mac\\xa0OS\\xa0X, previously shortened as OS\\xa0X, is an operating system developed and marketed by Apple Inc. since 2001. It is the primary operating system for Apple's Mac computers. Within the market of desktop and laptop computers, it is the second most widely used desktop OS, after Microsoft Windows and ahead of all Linux distributions, including ChromeOS.\",\n",
|
||||
" 'score': 0.8666583299636841},\n",
|
||||
" {'id': 'Linux',\n",
|
||||
" 'article': 'Linux is a family of open-source Unix-like operating systems based on the Linux kernel, an operating system kernel first released on September 17, 1991, by Linus Torvalds. Linux is typically packaged as a Linux distribution (distro), which includes the kernel and supporting system software and libraries, many of which are provided by the GNU Project. Many Linux distributions use the word \"Linux\" in their name, but the Free Software Foundation uses and recommends the name \"GNU/Linux\" to emphasize the use and importance of GNU software in many distributions, causing some controversy.',\n",
|
||||
" 'score': 0.839817225933075}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 63
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"All as expected! This method can also be used with existing datasets on the Hugging Face Hub."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "jUxUBP0vKUZo"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Integrate with Embeddings API services\n",
|
||||
"\n",
|
||||
"Next, we'll integrate with an Embeddings API service to build vectors.\n",
|
||||
"\n",
|
||||
"The code below interfaces with the Hugging Face Inference API. This can easily be switched to OpenAI, Cohere or even your own local API."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "lUU8WlIbKnOi"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import numpy as np\n",
|
||||
"import requests\n",
|
||||
"\n",
|
||||
"BASE = \"https://api-inference.huggingface.co/pipeline/feature-extraction\"\n",
|
||||
"\n",
|
||||
"def transform(inputs):\n",
|
||||
" # Your API provider of choice\n",
|
||||
" response = requests.post(f\"{BASE}/sentence-transformers/nli-mpnet-base-v2\", json={\"inputs\": inputs})\n",
|
||||
" return np.array(response.json(), dtype=np.float32)\n",
|
||||
"\n",
|
||||
"data = [\n",
|
||||
" \"US tops 5 million confirmed virus cases\",\n",
|
||||
" \"Canada's last fully intact ice shelf has suddenly collapsed, \" +\n",
|
||||
" \"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 \" +\n",
|
||||
" \"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",
|
||||
"\n",
|
||||
"embeddings = Embeddings({\"transform\": transform, \"backend\": \"numpy\", \"content\": True})\n",
|
||||
"embeddings.index(data)\n",
|
||||
"embeddings.search(\"feel good story\", 1)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "9-V6OIEtLEIs",
|
||||
"outputId": "e2492e33-27f7-48f3-e409-5c63e50c7f35"
|
||||
},
|
||||
"execution_count": 74,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '4',\n",
|
||||
" 'text': 'Maine man wins $1M from $25 lottery ticket',\n",
|
||||
" 'score': 0.08329013735055923}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 74
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"This is the classic txtai tutorial example. Except this time, vectorization is run with an external API service!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "uW8pDvD5Ms2r"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook showed how txtai can integrate with external vectorization. This can be a dataset with pre-computed embeddings and/or an Embeddings API service.\n",
|
||||
"\n",
|
||||
"Each of txtai's components can be fully customized and vectorization is no exception. Flexibility and customization for the win!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "8wGf3O_2YGbd"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1,351 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "e3wdiK5fGUoZ"
|
||||
},
|
||||
"source": [
|
||||
"# Advanced RAG with guided generation\n",
|
||||
"\n",
|
||||
"A standard RAG process typically runs a single vector search query and returns the closest matches. Those matches are then passed into a LLM prompt and used to limit the context and help ensure more factually correct answers are generated. This works well with most simple cases. More complex use cases, require a more advanced approach.\n",
|
||||
"\n",
|
||||
"This notebook will demonstrate how constrained or guided generation can be applied to better control LLM output."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "p8BbfjrhH-V2"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "-OXsTQgaGQPM"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai outlines"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Define the RAG process\n",
|
||||
"\n",
|
||||
"The first step we'll take is to define the RAG process. The following code creates a LLM instance, defines method that takes a question and context then prompts an LLM."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Wj7SlK_ZdwN9"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "zC_yDuyKk8ZY"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import LLM\n",
|
||||
"\n",
|
||||
"llm = LLM(\"Qwen/Qwen3-4B-Instruct-2507\")\n",
|
||||
"\n",
|
||||
"def rag(question, text):\n",
|
||||
" prompt = f\"\"\"<|im_start|>system\n",
|
||||
" You are a friendly assistant. You answer questions from users.<|im_end|>\n",
|
||||
" <|im_start|>user\n",
|
||||
" Answer the following question using only the context below. Only include information specifically discussed.\n",
|
||||
"\n",
|
||||
" question: {question}\n",
|
||||
" context: {text} <|im_end|>\n",
|
||||
" <|im_start|>assistant\n",
|
||||
" \"\"\"\n",
|
||||
"\n",
|
||||
" return llm(prompt, maxlength=4096)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Let's run a simple RAG call to get the idea of the default behavior."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "M2NQHT8nfZWJ"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Manually generated context. Replace with an Embedding search or other request. See prior examples on txtai's documentation site for more.\n",
|
||||
"context = \"\"\"\n",
|
||||
"England's terrain chiefly consists of low hills and plains, especially in the centre and south.\n",
|
||||
"The Battle of Hastings was fought on 14 October 1066 between the Norman army of William, the Duke of Normandy, and an English army under the Anglo-Saxon King Harold Godwinson\n",
|
||||
"Bounded by the Atlantic Ocean on the east, Brazil has a coastline of 7,491 kilometers (4,655 mi).\n",
|
||||
"Spain pioneered the exploration of the New World and the first circumnavigation of the globe.\n",
|
||||
"Christopher Columbus lands in the Caribbean in 1492.\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"print(rag(\"List the countries discussed\", context))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "yVXLDj5wfkqF",
|
||||
"outputId": "984f5e15-1949-413a-d4f0-4e55c991ede2"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"1. England\n",
|
||||
"2. Brazil\n",
|
||||
"3. Spain\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Guided Generation\n",
|
||||
"\n",
|
||||
"The next step is defining how to guide generation. For this step, we'll use the [Outlines](https://github.com/outlines-dev/outlines) library. Outlines is a library for controlling how tokens are generated. It applies logic to enforce schemas, regular expressions and/or specific output formats such as JSON.\n",
|
||||
"\n",
|
||||
"For our first example, we'll guide generation with a model that has answers and citations. With this multi-answer and multi-citation model, we can generate multiple answers along with associated references on how those answers were derived."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "co2ZA6hyeD3e"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from typing import List\n",
|
||||
"\n",
|
||||
"from outlines.integrations.transformers import JSONPrefixAllowedTokens\n",
|
||||
"from pydantic import BaseModel\n",
|
||||
"\n",
|
||||
"class Response(BaseModel):\n",
|
||||
" answers: List[str]\n",
|
||||
" citations: List[str]\n",
|
||||
"\n",
|
||||
"# Define method that guides LLM generation\n",
|
||||
"prefix_allowed_tokens_fn=JSONPrefixAllowedTokens(\n",
|
||||
" schema=Response,\n",
|
||||
" tokenizer_or_pipe=llm.generator.llm.pipeline.tokenizer,\n",
|
||||
" whitespace_pattern=r\" ?\"\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"def rag(question, text):\n",
|
||||
" prompt = f\"\"\"<|im_start|>system\n",
|
||||
" You are a friendly assistant. You answer questions from users.<|im_end|>\n",
|
||||
" <|im_start|>user\n",
|
||||
" Answer the following question using only the context below. Only include information specifically discussed.\n",
|
||||
"\n",
|
||||
" question: {question}\n",
|
||||
" context: {text} <|im_end|>\n",
|
||||
" <|im_start|>assistant\n",
|
||||
" \"\"\"\n",
|
||||
"\n",
|
||||
" return llm(prompt, maxlength=4096, prefix_allowed_tokens_fn=prefix_allowed_tokens_fn)\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "5jB_9SMhEZut"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Couple things to unpack here.\n",
|
||||
"\n",
|
||||
"First, note the method `prefix_allowed_tokens_fn`. This method applies a [Pydantic](https://github.com/pydantic/pydantic) model to constrain/guide how the LLM generates tokens. Next, see how that constrain can be applied to txtai's LLM pipeline.\n",
|
||||
"\n",
|
||||
"Let's try it out."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "ifCN1dpjgvTi"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import json\n",
|
||||
"\n",
|
||||
"json.loads(rag(\"List the countries discussed\", context))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "Uf7_Co4VL0uu",
|
||||
"outputId": "ed752198-4074-4220-83f7-358ab9237bce"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'answers': ['England', 'Brazil', 'Spain'],\n",
|
||||
" 'citations': [\"England's terrain chiefly consists of low hills and plains, especially in the centre and south.\",\n",
|
||||
" 'The Battle of Hastings was fought on 14 October 1066 between the Norman army of William, the Duke of Normandy, and an English army under the Anglo-Saxon King Harold Godwinson.',\n",
|
||||
" 'Bounded by the Atlantic Ocean on the east, Brazil has a coastline of 7,491 kilometers (4,655 mi).',\n",
|
||||
" 'Spain pioneered the exploration of the New World and the first circumnavigation of the globe.',\n",
|
||||
" 'Christopher Columbus lands in the Caribbean in 1492.']}"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 157
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"This is pretty 🔥\n",
|
||||
"\n",
|
||||
"See how not only are the answers generated as they were previously but the answers are now list of answers. And there is a list of citations supporting how the answers were generated! This is also valid JSON."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "rGbvrtZmhi4Y"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Extracting information models\n",
|
||||
"\n",
|
||||
"In our last example, we'll define a more complex model to help with extracting structured information."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "W0rsUfwniOgn"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"class Response(BaseModel):\n",
|
||||
" countries: List[str]\n",
|
||||
" geography: List[str]\n",
|
||||
" years: List[str]\n",
|
||||
" people: List[str]\n",
|
||||
"\n",
|
||||
"prefix_allowed_tokens_fn=JSONPrefixAllowedTokens(\n",
|
||||
" schema=Response,\n",
|
||||
" tokenizer_or_pipe=llm.generator.llm.pipeline.tokenizer,\n",
|
||||
" whitespace_pattern=r\" ?\"\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"def rag(question, text):\n",
|
||||
" prompt = f\"\"\"<|im_start|>system\n",
|
||||
" You are a friendly assistant. You answer questions from users.<|im_end|>\n",
|
||||
" <|im_start|>user\n",
|
||||
" Answer the following question using only the context below. Only include information specifically discussed.\n",
|
||||
"\n",
|
||||
" question: {question}\n",
|
||||
" context: {text} <|im_end|>\n",
|
||||
" <|im_start|>assistant\n",
|
||||
" \"\"\"\n",
|
||||
"\n",
|
||||
" return llm(prompt, maxlength=4096, prefix_allowed_tokens_fn=prefix_allowed_tokens_fn)\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "XHieIoxRgPSU"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"json.loads(rag(\"List the entities discussed\", context))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "7_RpO3v4ib_-",
|
||||
"outputId": "c3d394db-3f2b-4f60-d2b6-754f8247f929"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'countries': ['England', 'Brazil', 'Spain'],\n",
|
||||
" 'geography': ['low hills and plains', 'Atlantic Ocean', 'New World'],\n",
|
||||
" 'years': ['1066', '1492'],\n",
|
||||
" 'people': ['William, the Duke of Normandy',\n",
|
||||
" 'Anglo-Saxon King Harold Godwinson',\n",
|
||||
" 'Christopher Columbus']}"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 159
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"This is extremely guided generation. This is constraining the output of the LLM into a very specific model of information. It's quite impressive and simple to get started with!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "fDCFeFjSlveR"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"Guided generation adds a number of different options to the RAG toolkit. It's a flexible approach that can enable more advanced functionality and/or fine-tuned control over how content is created.\n",
|
||||
"\n",
|
||||
"It does add overhead which may or may not be acceptable depending on the use case. Expect new methods with improved efficiency and accuracy coming in the future. The space continues to advance forward fast!"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "2LQTekhsl75H"
|
||||
}
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "T4",
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"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.8.18"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,432 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "za-eT0AZAVEi"
|
||||
},
|
||||
"source": [
|
||||
"# RAG with llama.cpp and external API services\n",
|
||||
"\n",
|
||||
"txtai has been and always will be a local-first framework. It was originally designed to run models on local hardware using Hugging Face Transformers. As the AI space has evolved over the last year, so has txtai. Additional LLM inference frameworks have been available for a while using llama.cpp and external API services (via LiteLLM). Recent changes have added the ability to use these frameworks for vectorization and made it easier to use for LLM inference.\n",
|
||||
"\n",
|
||||
"This notebook will demonstrate how to run retrieval-augmented-generation (RAG) processes (vectorization and LLM inference) with llama.cpp and external API services."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "2e6iFXweAVEk"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "-UwwKDHCAVEk"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"\n",
|
||||
"# Install txtai and dependencies\n",
|
||||
"!pip install llama-cpp-python[server] --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu121\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline-llm]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "YGu8d_niAVEk"
|
||||
},
|
||||
"source": [
|
||||
"# Embeddings with llama.cpp vectorization\n",
|
||||
"\n",
|
||||
"The first example will build an Embeddings database backed by [llama.cpp](https://github.com/ggerganov/llama.cpp) vectorization.\n",
|
||||
"\n",
|
||||
"The llama.cpp project states: _The main goal of llama.cpp is to enable LLM inference with minimal setup and state-of-the-art performance on a wide variety of hardware - locally and in the cloud_.\n",
|
||||
"\n",
|
||||
"Let's give it a try."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "gZgObbS0AVEl"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"# Create Embeddings with llama.cpp GGUF model\n",
|
||||
"embeddings = Embeddings(\n",
|
||||
" path=\"second-state/All-MiniLM-L6-v2-Embedding-GGUF/all-MiniLM-L6-v2-Q4_K_M.gguf\",\n",
|
||||
" content=True\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Load dataset\n",
|
||||
"wikipedia = Embeddings()\n",
|
||||
"wikipedia.load(provider=\"huggingface-hub\", container=\"neuml/txtai-wikipedia\")\n",
|
||||
"\n",
|
||||
"query = \"\"\"\n",
|
||||
"SELECT id, text FROM txtai\n",
|
||||
"order by percentile desc\n",
|
||||
"LIMIT 10000\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"# Index dataset\n",
|
||||
"embeddings.index(wikipedia.search(query))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "qD2jypWnAVEl"
|
||||
},
|
||||
"source": [
|
||||
"Now that the Embeddings database is ready, let's run a search query."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "QhGrDP7AAVEl",
|
||||
"outputId": "42223d8d-a33e-47f6-9de3-e24cec4b31b0"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': 'Thomas Edison',\n",
|
||||
" 'text': 'Thomas Alva Edison (February 11, 1847October 18, 1931) was an American inventor and businessman. He developed many devices in fields such as electric power generation, mass communication, sound recording, and motion pictures. These inventions, which include the phonograph, the motion picture camera, and early versions of the electric light bulb, have had a widespread impact on the modern industrialized world. He was one of the first inventors to apply the principles of organized science and teamwork to the process of invention, working with many researchers and employees. He established the first industrial research laboratory.',\n",
|
||||
" 'score': 0.6758285164833069},\n",
|
||||
" {'id': 'Nikola Tesla',\n",
|
||||
" 'text': 'Nikola Tesla (; , ; 1856\\xa0– 7 January 1943) was a Serbian-American inventor, electrical engineer, mechanical engineer, and futurist. He is best-known for his contributions to the design of the modern alternating current (AC) electricity supply system.',\n",
|
||||
" 'score': 0.6077840328216553},\n",
|
||||
" {'id': 'Alexander Graham Bell',\n",
|
||||
" 'text': 'Alexander Graham Bell (, born Alexander Bell; March 3, 1847 – August 2, 1922) was a Scottish-born Canadian-American inventor, scientist and engineer who is credited with patenting the first practical telephone. He also co-founded the American Telephone and Telegraph Company (AT&T) in 1885.',\n",
|
||||
" 'score': 0.4573010802268982}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 3
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings.search(\"Inventors of electric-powered devices\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "QyaqgY8QAVEm"
|
||||
},
|
||||
"source": [
|
||||
"As we can see, this Embeddings database works just like any other Embeddings database. The difference is that it's using a llama.cpp model for vectorization instead of PyTorch."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "cewP6JCEAVEm"
|
||||
},
|
||||
"source": [
|
||||
"# RAG with llama.cpp\n",
|
||||
"\n",
|
||||
"LLM inference with llama.cpp is not a new txtai feature. A recent change added support for conversational messages in additional to standard prompts. This abstracts away having to understand prompting formats.\n",
|
||||
"\n",
|
||||
"Let's run a retrieval-augmented-generation (RAG) process fully backed by llama.cpp models.\n",
|
||||
"\n",
|
||||
"_It's important to note that conversational messages work with all LLM backends supported by txtai (transformers, llama.cpp, litellm)._"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "XTNX1mtMAVEm",
|
||||
"outputId": "62d92e44-02eb-4695-a507-0fc9729f0b57"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Based on the given context, here's a list of invented electric-powered devices:\n",
|
||||
"\n",
|
||||
"1. Electric light bulb by Thomas Edison\n",
|
||||
"2. Phonograph by Thomas Edison\n",
|
||||
"3. Motion picture camera by Thomas Edison\n",
|
||||
"4. Alternating current (AC) electricity supply system by Nikola Tesla\n",
|
||||
"5. Telephone by Alexander Graham Bell\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import LLM\n",
|
||||
"\n",
|
||||
"# LLM instance\n",
|
||||
"llm = LLM(path=\"unsloth/Qwen3-4B-Instruct-2507-GGUF/Qwen3-4B-Instruct-2507-Q4_K_M.gguf\")\n",
|
||||
"\n",
|
||||
"# Question and context\n",
|
||||
"question = \"Write a list of invented electric-powered devices\"\n",
|
||||
"context = \"\\n\".join(x[\"text\"] for x in embeddings.search(question))\n",
|
||||
"\n",
|
||||
"# Pass messages to LLM\n",
|
||||
"response = llm([\n",
|
||||
" {\"role\": \"system\", \"content\": \"You are a friendly assistant. You answer questions from users.\"},\n",
|
||||
" {\"role\": \"user\", \"content\": f\"\"\"\n",
|
||||
"Answer the following question using only the context below. Only include information specifically discussed.\n",
|
||||
"\n",
|
||||
"question: {question}\n",
|
||||
"context: {context}\n",
|
||||
"\"\"\"}\n",
|
||||
"])\n",
|
||||
"print(response)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "RQSApBO_AVEm"
|
||||
},
|
||||
"source": [
|
||||
"And just like that, RAG with llama.cpp🦙!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Qrf44ulzAVEn"
|
||||
},
|
||||
"source": [
|
||||
"# Embeddings with external vectorization\n",
|
||||
"\n",
|
||||
"Next, we'll show how an Embeddings database can integrate with external API services via [LiteLLM](https://github.com/BerriAI/litellm) .\n",
|
||||
"\n",
|
||||
"In the LiteLLM project's own words: _LiteLLM handles loadbalancing, fallbacks and spend tracking across 100+ LLMs. All in the OpenAI format._\n",
|
||||
"\n",
|
||||
"Let's first startup a local API service to use for this demo."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "iRoh9rPEAVEn"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"\n",
|
||||
"# Download models\n",
|
||||
"!wget https://huggingface.co/second-state/All-MiniLM-L6-v2-Embedding-GGUF/resolve/main/all-MiniLM-L6-v2-Q4_K_M.gguf\n",
|
||||
"!wget https://huggingface.co/unsloth/Qwen3-4B-Instruct-2507-GGUF/blob/main/Qwen3-4B-Instruct-2507-Q4_K_M.gguf\n",
|
||||
"\n",
|
||||
"# Start local API services\n",
|
||||
"!nohup python -m llama_cpp.server --n_gpu_layers -1 --model all-MiniLM-L6-v2-Q4_K_M.gguf --host 127.0.0.1 --port 8000 &> vector.log &\n",
|
||||
"!nohup python -m llama_cpp.server --n_gpu_layers -1 --model Qwen3-4B-Instruct-2507-Q4_K_M.gguf --chat_format chatml --host 127.0.0.1 --port 8001 &> llm.log &\n",
|
||||
"!sleep 30"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "fPnpAMnvAVEo"
|
||||
},
|
||||
"source": [
|
||||
"Now let's connect and use this local service to generate vectors for a new Embeddings database. Note that the local service responds in OpenAI's response format, hence the `path` setting below."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "EpSLnVfwAVEo"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"# Create Embeddings instance with external vectorization\n",
|
||||
"embeddings = Embeddings(\n",
|
||||
" path=\"openai/gpt-4-turbo\",\n",
|
||||
" content=True,\n",
|
||||
" vectors={\n",
|
||||
" \"api_base\": \"http://localhost:8000/v1\",\n",
|
||||
" \"api_key\": \"sk-1234\"\n",
|
||||
" }\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Load dataset\n",
|
||||
"wikipedia = Embeddings()\n",
|
||||
"wikipedia.load(provider=\"huggingface-hub\", container=\"neuml/txtai-wikipedia\")\n",
|
||||
"\n",
|
||||
"query = \"\"\"\n",
|
||||
"SELECT id, text FROM txtai\n",
|
||||
"order by percentile desc\n",
|
||||
"LIMIT 10000\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"# Index dataset\n",
|
||||
"embeddings.index(wikipedia.search(query))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "teoKbfzZAVEo",
|
||||
"outputId": "55594186-c989-419e-974a-37a150ba4734",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': 'Thomas Edison',\n",
|
||||
" 'text': 'Thomas Alva Edison (February 11, 1847October 18, 1931) was an American inventor and businessman. He developed many devices in fields such as electric power generation, mass communication, sound recording, and motion pictures. These inventions, which include the phonograph, the motion picture camera, and early versions of the electric light bulb, have had a widespread impact on the modern industrialized world. He was one of the first inventors to apply the principles of organized science and teamwork to the process of invention, working with many researchers and employees. He established the first industrial research laboratory.',\n",
|
||||
" 'score': 0.6758285164833069},\n",
|
||||
" {'id': 'Nikola Tesla',\n",
|
||||
" 'text': 'Nikola Tesla (; , ; 1856\\xa0– 7 January 1943) was a Serbian-American inventor, electrical engineer, mechanical engineer, and futurist. He is best-known for his contributions to the design of the modern alternating current (AC) electricity supply system.',\n",
|
||||
" 'score': 0.6077840328216553},\n",
|
||||
" {'id': 'Alexander Graham Bell',\n",
|
||||
" 'text': 'Alexander Graham Bell (, born Alexander Bell; March 3, 1847 – August 2, 1922) was a Scottish-born Canadian-American inventor, scientist and engineer who is credited with patenting the first practical telephone. He also co-founded the American Telephone and Telegraph Company (AT&T) in 1885.',\n",
|
||||
" 'score': 0.4573010802268982}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 7
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings.search(\"Inventors of electric-powered devices\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "0X1F1KcNAVEo"
|
||||
},
|
||||
"source": [
|
||||
"Like the previous example with llama.cpp, this Embeddings database behaves exactly the same. The main difference is that content is sent to an external service for vectorization."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "4Dr4HjhOAVEo"
|
||||
},
|
||||
"source": [
|
||||
"# RAG with External API services\n",
|
||||
"\n",
|
||||
"For our last task, we'll run a retrieval-augmented-generation (RAG) process fully backed by an external API service."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "42sPo6HCAVEo",
|
||||
"outputId": "7e42e28a-0bdd-4f68-a1ee-403527a241fd",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
}
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Based on the given context, a list of invented electric-powered devices includes:\n",
|
||||
"\n",
|
||||
"1. Phonograph by Thomas Edison\n",
|
||||
"2. Motion Picture Camera by Thomas Edison\n",
|
||||
"3. Early versions of the Electric Light Bulb by Thomas Edison\n",
|
||||
"4. AC (Alternating Current) Electricity Supply System by Nikola Tesla\n",
|
||||
"5. Telephone by Alexander Graham Bell\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import LLM\n",
|
||||
"\n",
|
||||
"# LLM instance\n",
|
||||
"llm = LLM(path=\"openai/gpt-4-turbo\", api_base=\"http://localhost:8001/v1\", api_key=\"sk-1234\")\n",
|
||||
"\n",
|
||||
"# Question and context\n",
|
||||
"question = \"Write a list of invented electric-powered devices\"\n",
|
||||
"context = \"\\n\".join(x[\"text\"] for x in embeddings.search(question))\n",
|
||||
"\n",
|
||||
"# Pass messages to LLM\n",
|
||||
"response = llm([\n",
|
||||
" {\"role\": \"system\", \"content\": \"You are a friendly assistant. You answer questions from users.\"},\n",
|
||||
" {\"role\": \"user\", \"content\": f\"\"\"\n",
|
||||
"Answer the following question using only the context below. Only include information specifically discussed.\n",
|
||||
"\n",
|
||||
"question: {question}\n",
|
||||
"context: {context}\n",
|
||||
"\"\"\"}\n",
|
||||
"])\n",
|
||||
"print(response)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "kFcqP5A7AVEp"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"txtai supports a number of different vector and LLM backends. The default method uses PyTorch models via the Hugging Face Transformers library. This notebook demonstrated how llama.cpp and external API services can also be used.\n",
|
||||
"\n",
|
||||
"These additional vector and LLM backends enable maximum flexibility and scalability. For example, vectorization can be fully offloaded to an external API service or another local service. llama.cpp has great support for macOS devices, alternate accelerators such AMD ROCm / Intel GPUs and has been known to run on Raspberry Pi devices.\n",
|
||||
"\n",
|
||||
"It's exciting to see the confluence of all these new advances coming together. Stay tuned for more!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "T4",
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"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.8.19"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
@@ -0,0 +1,527 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "VGeVB8M41jqW"
|
||||
},
|
||||
"source": [
|
||||
"# How RAG with txtai works\n",
|
||||
"\n",
|
||||
"Large Language Models (LLMs) have captured the public's attention with their impressive capabilities. The Generative AI era has reached a fever pitch with some predicting the coming rise of superintelligence.\n",
|
||||
"\n",
|
||||
"LLMs are far from perfect though and we're still a ways away from true AI. The biggest challenge is with hallucinations. Hallucinations is the term for when a LLM generates output that is factually incorrect. The alarming part of this is that on a cursory glance, it actually sounds like factual content. The default behavior of LLMs is to produce plausible answers even when no plausible answer exists. LLMs are not great at saying I don't know.\n",
|
||||
"\n",
|
||||
"Retrieval Augmented Generation (RAG) helps reduce the risk of hallucinations by limiting the context in which a LLM can generate answers. This is typically done with a search query that hydrates a prompt with a relevant context. RAG has been one of the most practical use cases of the Generative AI era.\n",
|
||||
"\n",
|
||||
"txtai has a multiple ways to run RAG pipelines as follows.\n",
|
||||
"\n",
|
||||
"- Embeddings instance and LLM. Run the embeddings search and plug the search results into a LLM prompt.\n",
|
||||
"- RAG (aka Extractor) pipeline which automatically adds a search context to LLM prompts.\n",
|
||||
"- RAG FastAPI service with YAML\n",
|
||||
"\n",
|
||||
"This notebook will cover all these methods and shows how RAG with txtai works."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ZQrHIw351lwE"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"id": "R0AqRP7v1hdr"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[api,pipeline]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "xmPN8RDF1pXd"
|
||||
},
|
||||
"source": [
|
||||
"# Components of a RAG pipeline\n",
|
||||
"\n",
|
||||
"Before using txtai's RAG pipeline, we'll show how each of the underlying components work together. In this example, we'll load the [txtai Wikipedia embeddings database](https://huggingface.co/NeuML/txtai-wikipedia) and a LLM. From there, we'll run a RAG process."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "XZ7vPBIs1rGZ"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import Embeddings, LLM\n",
|
||||
"\n",
|
||||
"# Load Wikipedia Embeddings database\n",
|
||||
"embeddings = Embeddings()\n",
|
||||
"embeddings.load(provider=\"huggingface-hub\", container=\"neuml/txtai-wikipedia\")\n",
|
||||
"\n",
|
||||
"# Create LLM\n",
|
||||
"llm = LLM(\"Qwen/Qwen3-4B-Instruct-2507\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Next, we'll create a prompt template to use for the RAG pipeline. The prompt has a placeholder for the question and context."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "ufUxz8MXLA8n"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Prompt template\n",
|
||||
"prompt = \"\"\"<|im_start|>system\n",
|
||||
"You are a friendly assistant. You answer questions from users.<|im_end|>\n",
|
||||
"<|im_start|>user\n",
|
||||
"Answer the following question using only the context below. Only include information\n",
|
||||
"specifically discussed.\n",
|
||||
"\n",
|
||||
"question: {question}\n",
|
||||
"context: {context} <|im_end|>\n",
|
||||
"<|im_start|>assistant\n",
|
||||
"\"\"\""
|
||||
],
|
||||
"metadata": {
|
||||
"id": "4f8DbjzMKwaE"
|
||||
},
|
||||
"execution_count": 3,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"After that, we'll generate the context using an embeddings (aka vector) query. This query finds the top 3 most similar matches to the question **\"How do you make beer 🍺?\"**"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "x8tp7IoBL27y"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"question = \"How do you make beer?\"\n",
|
||||
"\n",
|
||||
"# Generate context\n",
|
||||
"context = \"\\n\".join([x[\"text\"] for x in embeddings.search(question)])\n",
|
||||
"print(context)"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "Uyi_WgpULLFs",
|
||||
"outputId": "a4c7ae0c-287c-409a-ab55-c8c353125405"
|
||||
},
|
||||
"execution_count": 4,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Brewing is the production of beer by steeping a starch source (commonly cereal grains, the most popular of which is barley) in water and fermenting the resulting sweet liquid with yeast. It may be done in a brewery by a commercial brewer, at home by a homebrewer, or communally. Brewing has taken place since around the 6th millennium BC, and archaeological evidence suggests that emerging civilizations, including ancient Egypt, China, and Mesopotamia, brewed beer. Since the nineteenth century the brewing industry has been part of most western economies.\n",
|
||||
"Beer is produced through steeping a sugar source (commonly Malted cereal grains) in water and then fermenting with yeast. Brewing has taken place since around the 6th millennium BC, and archeological evidence suggests that this technique was used in ancient Egypt. Descriptions of various beer recipes can be found in Sumerian writings, some of the oldest known writing of any sort. Brewing is done in a brewery by a brewer, and the brewing industry is part of most western economies. In 19th century Britain, technological discoveries and improvements such as Burtonisation and the Burton Union system significantly changed beer brewing.\n",
|
||||
"Craft beer is a beer that has been made by craft breweries, which typically produce smaller amounts of beer, than larger \"macro\" breweries, and are often independently owned. Such breweries are generally perceived and marketed as emphasising enthusiasm, new flavours, and varied brewing techniques.\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Now we'll take the question and context and put that into the prompt."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "x_EPbqvFL_Uw"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"print(llm(prompt.format(question=question, context=context)))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "QCh1IaJ1L1Kt",
|
||||
"outputId": "93fbeab4-e477-4834-d01c-366fbf86361f"
|
||||
},
|
||||
"execution_count": 5,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"To make beer, you need to steep a starch source, such as malted cereal grains (commonly barley), in water. This process creates a sweet liquid called wort. Then, yeast is added to the wort, which ferments the liquid and produces alcohol and carbon dioxide. The beer is then aged, filtered, and packaged for consumption. This process has been used since around the 6th millennium BC and has been a part of most western economies since the 19th century.\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Looking at the generated answer, we can see it's based on the context above. The LLM generates a paragraph of text using the context as input. While this same answer could be directly asked of the LLM, this helps ensure the answer is based on known factual data."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "CSNbxxIMN_26"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"_Before continuing, it's important to note that txtai has multiple ways to run LLM inference. In the past, prior to \"Chat Templates\", it was expected that the submitted text had all the required chat tokens embedded. The same prompt above can also be written with chat messages. This is especially important when working with LLM APIs (i.e. OpenAI, Claude, Bedrock etc)._\n",
|
||||
"\n",
|
||||
"```python\n",
|
||||
"llm([\n",
|
||||
" {\"role\": \"system\": \"You are a friendly assistant. You answer questions from users.\"}\n",
|
||||
" {\"role\": \"user\", \"content\": f\"\"\"\n",
|
||||
" Answer the following question using only the context below. Only include information specifically discussed.\n",
|
||||
"\n",
|
||||
" question: {question}\n",
|
||||
" context: {text} \n",
|
||||
" \"\"\"}\n",
|
||||
"])\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"_See the [LLM pipeline documentation](https://neuml.github.io/txtai/pipeline/text/llm/) for more information._"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# The RAG Pipeline\n",
|
||||
"\n",
|
||||
"txtai has a RAG pipeline that makes this even easier. The logic to generate the context and join it context with the prompt is built in. Let's try that."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "tgLA-61iOP0r"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai import RAG\n",
|
||||
"\n",
|
||||
"# Create RAG pipeline using existing components. LLM parameter can also be a model path.\n",
|
||||
"rag = RAG(embeddings, llm, template=prompt)"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "eWMQUTxFOf_5"
|
||||
},
|
||||
"execution_count": 6,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Let's ask a question similar to the last one. This time we'll ask **\"How do you make wine🍷?\"**"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "8ePe71D0O0j_"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"print(rag(\"How do you make wine?\", maxlength=2048)[\"answer\"])"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "3ZaUe2YcOz_B",
|
||||
"outputId": "ada69201-9f94-498b-b33a-971b4d2e4fdb"
|
||||
},
|
||||
"execution_count": 7,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"To make wine, follow these steps:\n",
|
||||
"\n",
|
||||
"1. Select the fruit: Choose high-quality grapes or other fruit for wine production.\n",
|
||||
"\n",
|
||||
"2. Fermentation: Introduce yeast to the fruit, which will consume the sugar present in the juice and convert it into ethanol and carbon dioxide.\n",
|
||||
"\n",
|
||||
"3. Monitor temperature and oxygen levels: Control the temperature and speed of fermentation, as well as the levels of oxygen present in the must at the start of fermentation.\n",
|
||||
"\n",
|
||||
"4. Primary fermentation: This stage lasts from 5 to 14 days, during which the yeast consumes the sugar and produces alcohol and carbon dioxide.\n",
|
||||
"\n",
|
||||
"5. Secondary fermentation (optional): If desired, allow the wine to undergo a secondary fermentation, which can last another 5 to 10 days.\n",
|
||||
"\n",
|
||||
"6. Fermentation location: Choose the appropriate fermentation vessel, such as stainless steel tanks, open wooden vats, wine barrels, or wine bottles for sparkling wines.\n",
|
||||
"\n",
|
||||
"7. Bottle and age the wine: Transfer the finished wine into bottles and allow it to age, if desired, to develop flavors and complexity.\n",
|
||||
"\n",
|
||||
"Remember that wine can be made from various fruits, but grapes are most commonly used, and the term \"wine\" generally refers to grape wine when used without a qualifier.\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"_As with the LLM pipeline, the RAG pipeline also supports chat messages. See the [RAG pipeline documentation](https://neuml.github.io/txtai/pipeline/text/rag/) for more._"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# RAG API Endpoint\n",
|
||||
"\n",
|
||||
"Did you know that txtai has a built-in framework for automatically generating FastAPI services? This can be done with a YAML configuration file."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "L5udDj5dMrjY"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"%%writefile config.yml\n",
|
||||
"\n",
|
||||
"# Load Wikipedia Embeddings index\n",
|
||||
"cloud:\n",
|
||||
" provider: huggingface-hub\n",
|
||||
" container: neuml/txtai-wikipedia\n",
|
||||
"\n",
|
||||
"# RAG pipeline configuration\n",
|
||||
"rag:\n",
|
||||
" path: Qwen/Qwen3-4B-Instruct-2507\n",
|
||||
" output: flatten\n",
|
||||
" template: |\n",
|
||||
" <|im_start|>system\n",
|
||||
" You are a friendly assistant. You answer questions from users.<|im_end|>\n",
|
||||
" <|im_start|>user\n",
|
||||
" Answer the following question using only the context below. Only include information\n",
|
||||
" specifically discussed.\n",
|
||||
"\n",
|
||||
" question: {question}\n",
|
||||
" context: {context} <|im_end|>\n",
|
||||
" <|im_start|>assistant"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "NxEQqdzmNA0T",
|
||||
"outputId": "268d6027-d3d7-46f1-acc4-2016d3059f5e"
|
||||
},
|
||||
"execution_count": 8,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Writing config.yml\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Note how the same prompt template and models are set. This time instead of doing that with Python, it's done with a YAML configuration file 🔥\n",
|
||||
"\n",
|
||||
"Now let's start the API service using this configuration."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "rAmN_3aWNNyK"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!CONFIG=config.yml nohup uvicorn \"txtai.api:app\" &> api.log &\n",
|
||||
"!sleep 90"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "gzHPH5GKNXOQ"
|
||||
},
|
||||
"execution_count": 7,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Now let's run a RAG query using the API service. Keeping with the theme, we'll ask **\"How do you make whisky 🥃?\"**"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "bqIqEsAjPqsa"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!curl \"http://localhost:8000/rag?query=how+do+you+make+whisky&maxlength=2048\""
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "f14clGi1PeQW",
|
||||
"outputId": "0a056edc-e5c2-464d-fef0-c766369a2590"
|
||||
},
|
||||
"execution_count": 22,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"To make whisky, follow these steps:\n\n",
|
||||
"1. Choose the grains: Select the grains you want to use for your whisky, such as barley, corn, rye, or wheat.\n\n",
|
||||
"2. Malt the grains (optional): If using barley, malt the grains by soaking them in water and allowing them to germinate. This process releases enzymes that help break down starches into fermentable sugars.\n\n",
|
||||
"3. Mill the grains: Grind the grains to create a coarse flour, which will be mixed with water to create a mash.\n\n",
|
||||
"4. Create the mash: Combine the milled grains with hot water in a large vessel, and let it sit for several hours to allow fermentation to occur. The mash should have a temperature of around 65°C (149°F) to encourage the growth of yeast.\n\n",
|
||||
"5. Add yeast: Once the mash has cooled to around 30°C (86°F), add yeast to the mixture. The yeast will ferment the sugars in the mash, producing alcohol.\n\n",
|
||||
"6. Fermentation: Allow the mixture to ferment for several days, during which the yeast will consume the sugars and produce alcohol and carbon dioxide.\n\n",
|
||||
"7. Distillation: Transfer the fermented liquid, called \"wash\" to a copper still. Heat the wash in the still, and the alcohol will vaporize and rise through the still's neck. The vapors are then condensed back into a liquid form, creating a high-proof spirit.\n\n",
|
||||
"8. Maturation: Transfer the distilled spirit to wooden casks, typically made of charred white oak. The spirit will mature in the casks for a specified period, usually ranging from 3 to 25 years. During this time, the wood imparts flavors and color to the whisky.\n\n",
|
||||
"9. Bottling: Once the whisky has reached the desired maturity, it is bottled and ready for consumption."
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"And as before, we get an answer bound by the search context provided to the LLM. This time it comes from an API service vs a direct Python method."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "L8xTq1cTVTn7"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# RAG API Service with Docker\n",
|
||||
"\n",
|
||||
"txtai builds Docker images with each release. There are also Docker files available to help configure API services.\n",
|
||||
"\n",
|
||||
"The Dockerfile below builds an API service using the same config.yml."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "_UszA8iLQzGr"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"```dockerfile\n",
|
||||
"# Set base image\n",
|
||||
"ARG BASE_IMAGE=neuml/txtai-gpu\n",
|
||||
"FROM $BASE_IMAGE\n",
|
||||
"\n",
|
||||
"# Copy configuration\n",
|
||||
"COPY config.yml .\n",
|
||||
"\n",
|
||||
"# Install latest version of txtai from GitHub\n",
|
||||
"RUN \\\n",
|
||||
" apt-get update && \\\n",
|
||||
" apt-get -y --no-install-recommends install git && \\\n",
|
||||
" rm -rf /var/lib/apt/lists && \\\n",
|
||||
" python -m pip install git+https://github.com/neuml/txtai\n",
|
||||
"\n",
|
||||
"# Run local API instance to cache models in container\n",
|
||||
"RUN python -c \"from txtai.api import API; API('config.yml')\"\n",
|
||||
"\n",
|
||||
"# Start server and listen on all interfaces\n",
|
||||
"ENV CONFIG \"config.yml\"\n",
|
||||
"ENTRYPOINT [\"uvicorn\", \"--host\", \"0.0.0.0\", \"txtai.api:app\"]\n",
|
||||
"```"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "z5BQu2UPQ-wV"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"The following commands build and start a Docker API service."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "rhADZQUyVmDL"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"```bash\n",
|
||||
"docker build -t txtai-wikipedia --build-arg BASE_IMAGE=neuml/txtai-gpu .\n",
|
||||
"docker run -d --gpus=all -it -p 8000:8000 txtai-wikipedia\n",
|
||||
"```"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "fMyuCMcDRPzW"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"This creates the same API service just this time it's through Docker. RAG queries can be run the same way."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "TF1joce_R4nx"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"```bash\n",
|
||||
"curl \"http://localhost:8000/rag?query=how+do+you+make+whisky&maxlength=2048\"\n",
|
||||
"```"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "oN6l0G-gSDBb"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "oPwgCgBc2Er2"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook covered the various ways to run retrieval augmented generation (RAG) with txtai. We hope you find txtai is one of the easiest and most flexible ways to get up and running fast!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "T4",
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
@@ -0,0 +1,569 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
},
|
||||
"gpuClass": "standard",
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Embeddings index format for open data access\n",
|
||||
"\n",
|
||||
"The main programming language with txtai is Python. A key tenet is that the underlying data in an embeddings index is accessible without txtai.\n",
|
||||
"\n",
|
||||
"This notebook will demonstrate this through a series of examples.\n"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "-xU9P9iSR-Cy"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "shlUi2kKS7KT"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"id": "xEvX9vCpn4E0"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[graph] datasets sqlite-vec"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Load dataset\n",
|
||||
"\n",
|
||||
"This example will use the `chatgpt-prompts` dataset."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "408IyXzKFSiG"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"dataset = load_dataset(\"fka/awesome-chatgpt-prompts\", split=\"train\")"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "IQ_ns6YvFRm1"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Build an Embeddings index\n",
|
||||
"\n",
|
||||
"Let's first build an embeddings index using txtai."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "AtEdP7Utw3mk"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"embeddings = Embeddings()\n",
|
||||
"embeddings.index((x[\"act\"], x[\"prompt\"]) for x in dataset)\n",
|
||||
"embeddings.save(\"txtai-index\")"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "DPWrubv5oOn7"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Let's take a look at the index that was created"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "0zrjOl6JtPI2"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!ls -l txtai-index\n",
|
||||
"!echo\n",
|
||||
"!file txtai-index/*"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "i7PN9itKtIyx",
|
||||
"outputId": "ac904afb-83a8-4c5f-cf62-f5340225b8ac"
|
||||
},
|
||||
"execution_count": 4,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"total 268\n",
|
||||
"-rw-r--r-- 1 root root 342 Sep 6 15:21 config.json\n",
|
||||
"-rw-r--r-- 1 root root 262570 Sep 6 15:21 embeddings\n",
|
||||
"-rw-r--r-- 1 root root 2988 Sep 6 15:21 ids\n",
|
||||
"\n",
|
||||
"txtai-index/config.json: JSON data\n",
|
||||
"txtai-index/embeddings: data\n",
|
||||
"txtai-index/ids: data\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"The txtai embeddings index format is documented [here](https://neuml.github.io/txtai/embeddings/format/). Looking at the files above, we have configuration, embeddings data and ids storage. Ids storage is only used when content is disabled.\n",
|
||||
"\n",
|
||||
"Let's inspect each file."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "hx8H0dpXtX5b"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import json\n",
|
||||
"\n",
|
||||
"with open(\"txtai-index/config.json\") as f:\n",
|
||||
" print(json.dumps(json.load(f), sort_keys=True, indent=2))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "h0yrQ8rqtnzh",
|
||||
"outputId": "3379ab18-1805-4ef1-b946-1a61d18522db"
|
||||
},
|
||||
"execution_count": 5,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"{\n",
|
||||
" \"backend\": \"faiss\",\n",
|
||||
" \"build\": {\n",
|
||||
" \"create\": \"2024-09-06T15:21:11Z\",\n",
|
||||
" \"python\": \"3.10.12\",\n",
|
||||
" \"settings\": {\n",
|
||||
" \"components\": \"IDMap,Flat\"\n",
|
||||
" },\n",
|
||||
" \"system\": \"Linux (x86_64)\",\n",
|
||||
" \"txtai\": \"7.5.0\"\n",
|
||||
" },\n",
|
||||
" \"dimensions\": 384,\n",
|
||||
" \"offset\": 170,\n",
|
||||
" \"path\": \"sentence-transformers/all-MiniLM-L6-v2\",\n",
|
||||
" \"update\": \"2024-09-06T15:21:11Z\"\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import faiss\n",
|
||||
"\n",
|
||||
"index = faiss.read_index(\"txtai-index/embeddings\")\n",
|
||||
"print(f\"Total records {index.ntotal}\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "-aqiqfqeuM5p",
|
||||
"outputId": "4b7856f0-ec8d-4960-a419-e88a14d988a8"
|
||||
},
|
||||
"execution_count": 6,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Total records 170\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import msgpack\n",
|
||||
"\n",
|
||||
"with open(\"txtai-index/ids\", \"rb\") as f:\n",
|
||||
" print(msgpack.unpack(f)[5:10])"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "I0ffDyJsunrW",
|
||||
"outputId": "d8c5072d-d181-46e1-fc1a-cc713a78e69f"
|
||||
},
|
||||
"execution_count": 18,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"['JavaScript Console', 'Excel Sheet', 'English Pronunciation Helper', 'Spoken English Teacher and Improver', 'Travel Guide']\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Each file can be read without txtai. [JSON](https://www.json.org/json-en.html), [MessagePack](https://msgpack.org/index.html) and [Faiss](https://github.com/facebookresearch/faiss) all have libraries in multiple programming languages."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "0e3dpAFuvP-_"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Embeddings index with SQLite\n",
|
||||
"\n",
|
||||
"In the next example, we'll use SQLite to store content and vectors courtesy of the [sqlite-vec](https://github.com/asg017/sqlite-vec) library."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "UUwk13mzwUTS"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"embeddings = Embeddings(content=True, backend=\"sqlite\")\n",
|
||||
"embeddings.index((x[\"act\"], x[\"prompt\"]) for x in dataset)\n",
|
||||
"embeddings.save(\"txtai-sqlite\")"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Z80FWhuNwj14"
|
||||
},
|
||||
"execution_count": 8,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Let's once again explore the generated index files."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Kxcm42rixAH0"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!ls -l txtai-sqlite\n",
|
||||
"!echo\n",
|
||||
"!file txtai-sqlite/*"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "PEHT6LqHw_lw",
|
||||
"outputId": "81d707dd-9760-40cd-e0b2-be43c26ba2d1"
|
||||
},
|
||||
"execution_count": 9,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"total 1696\n",
|
||||
"-rw-r--r-- 1 root root 384 Sep 6 15:21 config.json\n",
|
||||
"-rw-r--r-- 1 root root 126976 Sep 6 15:21 documents\n",
|
||||
"-rw-r--r-- 1 root root 1605632 Sep 6 15:21 embeddings\n",
|
||||
"\n",
|
||||
"txtai-sqlite/config.json: JSON data\n",
|
||||
"txtai-sqlite/documents: SQLite 3.x database, last written using SQLite version 3037002, file counter 1, database pages 31, cookie 0x1, schema 4, UTF-8, version-valid-for 1\n",
|
||||
"txtai-sqlite/embeddings: SQLite 3.x database, last written using SQLite version 3037002, file counter 1, database pages 392, cookie 0x1, schema 4, UTF-8, version-valid-for 1\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"This time note how there is a documents file with content stored in SQLite and a separate SQLite file for embeddings. Let's test it out."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "UXiKSG0JxLPo"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"embeddings.search(\"teacher\")"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "VYMTXtUpxSLd",
|
||||
"outputId": "6259f3c2-6a76-434d-d40f-f11cb2e63c80"
|
||||
},
|
||||
"execution_count": 10,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "execute_result",
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': 'Math Teacher',\n",
|
||||
" 'text': 'I want you to act as a math teacher. I will provide some mathematical equations or concepts, and it will be your job to explain them in easy-to-understand terms. This could include providing step-by-step instructions for solving a problem, demonstrating various techniques with visuals or suggesting online resources for further study. My first request is \"I need help understanding how probability works.\"',\n",
|
||||
" 'score': 0.3421396017074585},\n",
|
||||
" {'id': 'Educational Content Creator',\n",
|
||||
" 'text': 'I want you to act as an educational content creator. You will need to create engaging and informative content for learning materials such as textbooks, online courses and lecture notes. My first suggestion request is \"I need help developing a lesson plan on renewable energy sources for high school students.\"',\n",
|
||||
" 'score': 0.3267676830291748},\n",
|
||||
" {'id': 'Philosophy Teacher',\n",
|
||||
" 'text': 'I want you to act as a philosophy teacher. I will provide some topics related to the study of philosophy, and it will be your job to explain these concepts in an easy-to-understand manner. This could include providing examples, posing questions or breaking down complex ideas into smaller pieces that are easier to comprehend. My first request is \"I need help understanding how different philosophical theories can be applied in everyday life.\"',\n",
|
||||
" 'score': 0.30780404806137085}]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"execution_count": 10
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"The top N results as expected. Let's again inspect the files."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "R0zsQx37xjWz"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import json\n",
|
||||
"\n",
|
||||
"with open(\"txtai-sqlite/config.json\") as f:\n",
|
||||
" print(json.dumps(json.load(f), sort_keys=True, indent=2))"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "jDvH6dHrxqwf",
|
||||
"outputId": "dcce1647-0488-4f3e-dba0-667a932bad27"
|
||||
},
|
||||
"execution_count": 11,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"{\n",
|
||||
" \"backend\": \"sqlite\",\n",
|
||||
" \"build\": {\n",
|
||||
" \"create\": \"2024-09-06T15:21:13Z\",\n",
|
||||
" \"python\": \"3.10.12\",\n",
|
||||
" \"settings\": {\n",
|
||||
" \"sqlite\": \"3.37.2\",\n",
|
||||
" \"sqlite-vec\": \"v0.1.1\"\n",
|
||||
" },\n",
|
||||
" \"system\": \"Linux (x86_64)\",\n",
|
||||
" \"txtai\": \"7.5.0\"\n",
|
||||
" },\n",
|
||||
" \"content\": true,\n",
|
||||
" \"dimensions\": 384,\n",
|
||||
" \"offset\": 170,\n",
|
||||
" \"path\": \"sentence-transformers/all-MiniLM-L6-v2\",\n",
|
||||
" \"update\": \"2024-09-06T15:21:13Z\"\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import sqlite3, sqlite_vec\n",
|
||||
"\n",
|
||||
"db = sqlite3.connect(\"txtai-sqlite/documents\")\n",
|
||||
"print(db.execute(\"SELECT COUNT(*) FROM sections\").fetchone()[0])\n",
|
||||
"\n",
|
||||
"db = sqlite3.connect(\"txtai-sqlite/embeddings\")\n",
|
||||
"db.enable_load_extension(True)\n",
|
||||
"sqlite_vec.load(db)\n",
|
||||
"print(db.execute(\"SELECT COUNT(*) FROM vectors\").fetchone()[0])"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "AsalQ-uUxxO0",
|
||||
"outputId": "2a0e303e-e10e-424c-a952-e473d86cd2db"
|
||||
},
|
||||
"execution_count": 12,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"170\n",
|
||||
"170\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"As in the previous example, each file can be read without txtai. [JSON](https://www.json.org/json-en.html), [SQLite](https://www.sqlite.org/) and [sqlite-vec](https://github.com/asg017/sqlite-vec) all have libraries in multiple programming languages."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Fo7XEUNny6s5"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Graph storage\n",
|
||||
"\n",
|
||||
"Starting with txtai 7.4, graphs are stored using MessagePack. The indexed file has a list of nodes and edges that can easily be imported."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Ipu08WhL0j6p"
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"embeddings = Embeddings(content=True, backend=\"sqlite\", graph={\"approximate\": False})\n",
|
||||
"embeddings.index((x[\"act\"], x[\"prompt\"]) for x in dataset)\n",
|
||||
"embeddings.save(\"txtai-graph\")"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "cmFiae6j0wHz"
|
||||
},
|
||||
"execution_count": 13,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"!ls -l txtai-graph\n",
|
||||
"!echo\n",
|
||||
"!file txtai-graph/*"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "hnv82EAB059Q",
|
||||
"outputId": "454b263b-8733-4997-eb6e-7d585cfa32c6"
|
||||
},
|
||||
"execution_count": 14,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"total 1816\n",
|
||||
"-rw-r--r-- 1 root root 454 Sep 6 15:21 config.json\n",
|
||||
"-rw-r--r-- 1 root root 126976 Sep 6 15:21 documents\n",
|
||||
"-rw-r--r-- 1 root root 1605632 Sep 6 15:21 embeddings\n",
|
||||
"-rw-r--r-- 1 root root 119970 Sep 6 15:21 graph\n",
|
||||
"\n",
|
||||
"txtai-graph/config.json: JSON data\n",
|
||||
"txtai-graph/documents: SQLite 3.x database, last written using SQLite version 3037002, file counter 1, database pages 31, cookie 0x1, schema 4, UTF-8, version-valid-for 1\n",
|
||||
"txtai-graph/embeddings: SQLite 3.x database, last written using SQLite version 3037002, file counter 1, database pages 392, cookie 0x1, schema 4, UTF-8, version-valid-for 1\n",
|
||||
"txtai-graph/graph: data\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"import msgpack\n",
|
||||
"\n",
|
||||
"with open(\"txtai-graph/graph\", \"rb\") as f:\n",
|
||||
" data = msgpack.unpack(f)\n",
|
||||
" print(data.keys())\n",
|
||||
"\n",
|
||||
" for key in data:\n",
|
||||
" if data[key]:\n",
|
||||
" print(key, data[key][100])"
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "Bx3Tzt3l08R3",
|
||||
"outputId": "662d157c-6912-4626-db64-d959c2719331"
|
||||
},
|
||||
"execution_count": 15,
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"dict_keys(['nodes', 'edges', 'categories', 'topics'])\n",
|
||||
"nodes [100, {'id': 'Ascii Artist', 'text': 'I want you to act as an ascii artist. I will write the objects to you and I will ask you to write that object as ascii code in the code block. Write only ascii code. Do not explain about the object you wrote. I will say the objects in double quotes. My first object is \"cat\"'}]\n",
|
||||
"edges [5, 100, {'weight': 0.39010339975357056}]\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook gave an overview of the txtai embeddings index file format and how it supports open data access.\n",
|
||||
"\n",
|
||||
"While txtai can be used as an all-in-one embeddings database, it can also be used for only one part of the stack such as data ingestion. For example, it can be used to populate a Postgres or SQLite database for downstream use. The options are there."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "y7N-YZlR5S-0"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,260 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Speech to Speech RAG\n",
|
||||
"\n",
|
||||
"There are many articles, notebooks and examples covering how to perform vector search and/or retrieval augmented generation (RAG) with txtai. A lesser known component of txtai is it's built-in workflow component.\n",
|
||||
"\n",
|
||||
"Workflows are a simple yet powerful construct that takes a callable and returns elements. Workflows enable efficient processing of pipeline data. Workflows are streaming by nature and work on data in batches. This allows large volumes of data to be processed efficiently.\n",
|
||||
"\n",
|
||||
"This notebook will demonstrate how to to build a Speech to Speech (S2S) workflow with txtai.\n",
|
||||
"\n",
|
||||
"_Note: This process is intended to run on local machines due to it's use of input and output audio devices._"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline-audio]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Define the S2S RAG Workflow\n",
|
||||
"\n",
|
||||
"The next section defines the Speech to Speech (S2S) RAG workflow. The objective of this workflow is to respond to a user request in near real-time.\n",
|
||||
"\n",
|
||||
"txtai supports workflow definitions in Python and with YAML. We'll cover both methods.\n",
|
||||
"\n",
|
||||
"The S2S workflow below starts with a microphone pipeline, which streams and processes input audio. The microphone pipeline has voice activity detection (VAD) built-in. When speech is detected, the pipeline returns the captured audio data. Next, the speech is transcribed to text and then passed to a RAG pipeline prompt. Finally, the RAG result is run through a text to speech (TTS) pipeline and streamed to an output audio device."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import logging\n",
|
||||
"\n",
|
||||
"from txtai import Embeddings, RAG\n",
|
||||
"from txtai.pipeline import AudioStream, Microphone, TextToSpeech, Transcription\n",
|
||||
"from txtai.workflow import Workflow, StreamTask, Task\n",
|
||||
"\n",
|
||||
"# Enable DEBUG logging\n",
|
||||
"logging.basicConfig()\n",
|
||||
"logging.getLogger().setLevel(logging.DEBUG)\n",
|
||||
"\n",
|
||||
"# Microphone\n",
|
||||
"microphone = Microphone()\n",
|
||||
"\n",
|
||||
"# Transcription\n",
|
||||
"transcribe = Transcription(\"distil-whisper/distil-large-v3\")\n",
|
||||
"\n",
|
||||
"# Embeddings database\n",
|
||||
"embeddings = Embeddings()\n",
|
||||
"embeddings.load(provider=\"huggingface-hub\", container=\"neuml/txtai-wikipedia\")\n",
|
||||
"\n",
|
||||
"# Define prompt template\n",
|
||||
"template = \"\"\"\n",
|
||||
"Answer the following question using only the context below. Only include information\n",
|
||||
"specifically discussed. Answer the question without explaining how you found the answer.\n",
|
||||
"\n",
|
||||
"question: {question}\n",
|
||||
"context: {context}\"\"\"\n",
|
||||
"\n",
|
||||
"# Create RAG pipeline\n",
|
||||
"rag = RAG(\n",
|
||||
" embeddings,\n",
|
||||
" \"Qwen/Qwen3-4B-Instruct-2507\",\n",
|
||||
" system=\"You are a friendly assistant. You answer questions from users.\",\n",
|
||||
" template=template,\n",
|
||||
" context=10\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Text to speech\n",
|
||||
"tts = TextToSpeech(\"neuml/vctk-vits-onnx\")\n",
|
||||
"\n",
|
||||
"# Audio stream\n",
|
||||
"audiostream = AudioStream()\n",
|
||||
"\n",
|
||||
"# Define speech to speech workflow\n",
|
||||
"workflow = Workflow(tasks=[\n",
|
||||
" Task(action=microphone),\n",
|
||||
" Task(action=transcribe, unpack=False),\n",
|
||||
" StreamTask(action=lambda x: rag(x, maxlength=4096, stream=True), batch=True),\n",
|
||||
" StreamTask(action=lambda x: tts(x, stream=True, speaker=15), batch=True),\n",
|
||||
" StreamTask(action=audiostream, batch=True)\n",
|
||||
"])\n",
|
||||
"\n",
|
||||
"while True:\n",
|
||||
" print(\"Waiting for input...\")\n",
|
||||
" list(workflow([None]))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Given that the input and outputs are audio, you'll have to use your imagination if you're reading this as an article.\n",
|
||||
"\n",
|
||||
"[Check out this video](https://www.youtube.com/watch?v=tH8QWwkVMKA) to see the workflow in action! The following examples are run:\n",
|
||||
"\n",
|
||||
"- Tell me about the Roman Empire\n",
|
||||
"- Explain how faster than light travel could work\n",
|
||||
"- Write a short poem about the Vikings\n",
|
||||
"- Tell me about the Roman Empire in French"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# S2S Workflow in YAML\n",
|
||||
"\n",
|
||||
"A crucial feature of txtai workflows is that they can be defined with YAML. This enables building workflows in a low-code and/or no-code setting. These YAML workflows can then be \"dockerized\" and run.\n",
|
||||
"\n",
|
||||
"Let's define the same workflow below."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%writefile s2s.yml\n",
|
||||
"# Microphone\n",
|
||||
"microphone:\n",
|
||||
"\n",
|
||||
"# Transcription\n",
|
||||
"transcription:\n",
|
||||
" path: distil-whisper/distil-large-v3\n",
|
||||
"\n",
|
||||
"# Embeddings database\n",
|
||||
"cloud:\n",
|
||||
" provider: huggingface-hub\n",
|
||||
" container: neuml/txtai-wikipedia\n",
|
||||
"\n",
|
||||
"embeddings:\n",
|
||||
"\n",
|
||||
"# RAG\n",
|
||||
"rag:\n",
|
||||
" path: \"Qwen/Qwen3-4B-Instruct-2507\"\n",
|
||||
" system: You are a friendly assistant. You answer questions from users.\n",
|
||||
" template: |\n",
|
||||
" Answer the following question using only the context below. Only include information\n",
|
||||
" specifically discussed. Answer the question without explaining how you found the answer.\n",
|
||||
"\n",
|
||||
" question: {question}\n",
|
||||
" context: {context}\n",
|
||||
" context: 10\n",
|
||||
"\n",
|
||||
"# TTS\n",
|
||||
"texttospeech:\n",
|
||||
" path: neuml/vctk-vits-onnx\n",
|
||||
"\n",
|
||||
"# AudioStream\n",
|
||||
"audiostream:\n",
|
||||
"\n",
|
||||
"# Speech to Speech Chat workflow\n",
|
||||
"workflow:\n",
|
||||
" s2s:\n",
|
||||
" tasks:\n",
|
||||
" - microphone\n",
|
||||
" - action: transcription\n",
|
||||
" unpack: False\n",
|
||||
" - task: stream\n",
|
||||
" action: rag\n",
|
||||
" args:\n",
|
||||
" maxlength: 4096\n",
|
||||
" stream: True\n",
|
||||
" batch: True\n",
|
||||
" - task: stream\n",
|
||||
" action: texttospeech\n",
|
||||
" args:\n",
|
||||
" stream: True\n",
|
||||
" speaker: 15\n",
|
||||
" batch: True\n",
|
||||
" - task: stream\n",
|
||||
" action: audiostream\n",
|
||||
" batch: True"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import Application\n",
|
||||
"\n",
|
||||
"app = Application(\"s2s.yml\")\n",
|
||||
"while True:\n",
|
||||
" print(\"Waiting for input...\")\n",
|
||||
" list(app.workflow(\"s2s\", [None]))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Once again, the same idea, just a different way to do it. In the video demo, the following query was asked.\n",
|
||||
"\n",
|
||||
"- As a Patriots fan, who would you guess is my favorite quarterback of all time is?\n",
|
||||
"- I'm tall and run fast, what do you think the best soccer position for me is?\n",
|
||||
"- I run slow, what do you think the best soccer position for me is?\n",
|
||||
"\n",
|
||||
"With YAML workflows, it's possible to fully define the process outside of code such as with a web interface. Perhaps someday we'll see this with [txtai.cloud](https://txtai.cloud) 😀"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook demonstrated how to build a Speech to Speech (S2S) workflow with txtai. While the workflow uses an off-the-shelf embeddings database, a custom embeddings database can easily be swapped in. From there, we have S2S with our own data!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"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.8.20"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,556 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "e3wdiK5fGUoZ"
|
||||
},
|
||||
"source": [
|
||||
"# 💡 What's new in txtai 8.0\n",
|
||||
"\n",
|
||||
"The 8.0 release brings a major new feature: Agents 🚀\n",
|
||||
"\n",
|
||||
"Agents automatically create workflows to answer multi-faceted user requests. Agents iteratively prompt and/or interface with tools to\n",
|
||||
"step through a process and ultimately come to an answer for a request.\n",
|
||||
"\n",
|
||||
"This release also adds support for Model2Vec vectorization.\n",
|
||||
"\n",
|
||||
"**Standard upgrade disclaimer below**\n",
|
||||
"\n",
|
||||
"While everything is backwards compatible, it's prudent to backup production indexes before upgrading and test before deploying."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "p8BbfjrhH-V2"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "-OXsTQgaGQPM"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[agent] model2vec"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "n4EXrtcYIIYE"
|
||||
},
|
||||
"source": [
|
||||
"# Agents with txtai\n",
|
||||
"\n",
|
||||
"The biggest change and reason this is a major release is the addition of agents. The following defines a basic agent. This agent has access to two embeddings databases (Wikipedia and ArXiv) and the web. Given the user's input request, the agent decides the best tool to solve the task."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "hzPD8_cQJNtN"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from datetime import datetime\n",
|
||||
"\n",
|
||||
"from txtai import Agent\n",
|
||||
"\n",
|
||||
"wikipedia = {\n",
|
||||
" \"name\": \"wikipedia\",\n",
|
||||
" \"description\": \"Searches a Wikipedia database\",\n",
|
||||
" \"provider\": \"huggingface-hub\",\n",
|
||||
" \"container\": \"neuml/txtai-wikipedia\"\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"arxiv = {\n",
|
||||
" \"name\": \"arxiv\",\n",
|
||||
" \"description\": \"Searches a database of scientific papers\",\n",
|
||||
" \"provider\": \"huggingface-hub\",\n",
|
||||
" \"container\": \"neuml/txtai-arxiv\"\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"def today() -> str:\n",
|
||||
" \"\"\"\n",
|
||||
" Gets the current date and time\n",
|
||||
"\n",
|
||||
" Returns:\n",
|
||||
" current date and time\n",
|
||||
" \"\"\"\n",
|
||||
"\n",
|
||||
" return datetime.today().isoformat()\n",
|
||||
"\n",
|
||||
"agent = Agent(\n",
|
||||
" tools=[today, wikipedia, arxiv, \"websearch\"],\n",
|
||||
" llm=\"Qwen/Qwen3-4B-Instruct-2507\",\n",
|
||||
" max_iterations=10,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"\u001b[32;20;1m======== New task ========\u001b[0m\n",
|
||||
"\u001b[37;1mWhich city has the highest population, Boston or New York?\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: I will use the tool 'web_search' to find the population of both cities.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'web_search' with arguments: {'query': 'Population of Boston and New York'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: The results from the web search indicate that the population of Boston is 7,948 and the population of New York is 653,833. However, it seems that the question is asking about the city of Boston in New York, not the city of Boston in Massachusetts. I will use the tool 'wikipedia' to find the population of the city of New York.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'wikipedia' with arguments: {'query': 'Population of New York City'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: The results from the Wikipedia search indicate that the population of New York City is approximately 8,804,190. However, this is the population of the entire city, not just the city of New York. I will use the tool 'wikipedia' to find the population of the city of New York, not the state of New York.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'wikipedia' with arguments: {'query': 'Population of New York City (borough)'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: The results from the Wikipedia search indicate that the population of New York City is approximately 8,804,190, and the population of Brooklyn is approximately 2,592,149. Since Brooklyn is one of the five boroughs of New York City, I will assume that the population of the city of New York is approximately 8,804,190. However, this is still not the correct answer to the original question. I will use the tool 'wikipedia' to find the population of the city of Boston in Massachusetts, which is the other city mentioned in the original question.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'wikipedia' with arguments: {'query': 'Population of Boston, Massachusetts'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: The results from the Wikipedia search indicate that the population of Boston, Massachusetts is approximately 675,647. This is the correct answer to the original question, which was to determine which city has the highest population, Boston or New York. Since the population of Boston is significantly lower than the population of New York City, which is approximately 8,804,190, the answer to the original question is New York City.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'final_answer' with arguments: {'answer': 'New York City'}\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'New York City'"
|
||||
]
|
||||
},
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"agent(\"Which city has the highest population, Boston or New York?\", maxlength=16000)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"How about that 🔥? \n",
|
||||
"\n",
|
||||
"Look at how the agent went through and systematically solved the problem. The agent ran a series of searches to gather the information it needed and then determined the final answer!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Agentic RAG\n",
|
||||
"\n",
|
||||
"Standard Retrieval Augmented Generation (RAG) runs a single vector search to obtain a context and builds a prompt with the context + input question. Agentic RAG is a more complex process that goes through multiple iterations. It can also utilize multiple databases to come to a final conclusion.\n",
|
||||
"\n",
|
||||
"The example below aggregates information from multiple sources and builds a report on a topic."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"\u001b[32;20;1m======== New task ========\u001b[0m\n",
|
||||
"\u001b[37;1m\n",
|
||||
"You're an expert researcher looking to write a paper on alien life.\n",
|
||||
"Search for websites, scientific papers and Wikipedia related to the topic.\n",
|
||||
"Write a report with summaries and references (with hyperlinks).\n",
|
||||
"Write the text as Markdown.\n",
|
||||
"\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: I will start by searching for Wikipedia articles related to alien life.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'wikipedia' with arguments: {'query': 'Alien life'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: The search results from Wikipedia look promising. I will now search for scientific papers related to alien life using the arxiv tool.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'arxiv' with arguments: {'query': 'alien life'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: The search results from arxiv look interesting. I will now search for websites related to alien life using the web_search tool.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'web_search' with arguments: {'query': 'alien life'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: Now that I have gathered information from Wikipedia, arxiv, and web search, I will write a report with summaries and references (with hyperlinks) in Markdown format.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'final_answer' with arguments: {'answer': \"## Alien Life: A Report on the Search for Extraterrestrial Life\\n\\n### Introduction\\n\\nThe search for alien life has been ongoing for centuries, with scientists and enthusiasts alike searching for signs of life beyond Earth. While no conclusive evidence has been found yet, there are many promising leads and discoveries that suggest the possibility of life existing elsewhere in the universe.\\n\\n### Wikipedia\\n\\nAccording to Wikipedia, extraterrestrial life refers to life that does not originate from Earth. The search for alien life is an active area of research, with scientists using various methods to detect signs of life, such as analyzing the atmospheres of exoplanets and searching for biosignatures.\\n\\n### Arxiv\\n\\nThe arxiv database has numerous papers on the search for alien life, including studies on the possibility of life existing on exoplanets and the use of astrobiology to detect signs of life. One paper, 'Extraterrestrial Life and Censorship,' discusses the development of the theory of cosmic life and the attempts to censor evidence incompatible with this theory.\\n\\n### Web Search\\n\\nA web search for alien life yielded numerous results, including articles from reputable sources such as NASA and Live Science. These articles discuss the latest discoveries and findings in the search for alien life, including the potential for life existing on exoplanets and the use of astrobiology to detect signs of life.\\n\\n### Conclusion\\n\\nWhile the search for alien life is an ongoing and active area of research, there are many promising leads and discoveries that suggest the possibility of life existing elsewhere in the universe. As scientists continue to search for signs of life, it is possible that we may one day find evidence of alien life.\\n\\n### References\\n\\n* [Extraterrestrial life - Wikipedia](https://en.wikipedia.org/wiki/Extraterrestrial_life)\\n* [What's the best evidence we've found for alien life?](https://www.livescience.com/space/extraterrestrial-life/whats-the-best-evidence-weve-found-for-alien-life)\\n* [NASA Research Gives Guideline for Future Alien Life Search](https://www.nasa.gov/universe/nasa-research-gives-guideline-for-future-alien-life-search/)\\n* [James Webb telescope sees potential signs of alien life in the atmosphere of a distant goldilocks water-world](https://www.livescience.com/space/exoplanets/james-webb-telescope-sees-potential-signs-of-alien-life-in-the-atmosphere-of-a-distant-goldilocks-water-world)\\n* [Alien life in Universe: Scientists say finding it is 'only a matter of time'](https://www.bbc.com/news/science-environment-66950930)\\n* [Extraterrestrial evidence: 10 incredible findings about aliens from 2020](https://www.livescience.com/alien-discoveries-2020.html)\\n* [How Scientists Could Tell the World if They Find Alien Life](https://www.scientificamerican.com/article/how-scientists-could-tell-the-world-if-they-find-alien-life/)\"}\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"## Alien Life: A Report on the Search for Extraterrestrial Life\n",
|
||||
"\n",
|
||||
"### Introduction\n",
|
||||
"\n",
|
||||
"The search for alien life has been ongoing for centuries, with scientists and enthusiasts alike searching for signs of life beyond Earth. While no conclusive evidence has been found yet, there are many promising leads and discoveries that suggest the possibility of life existing elsewhere in the universe.\n",
|
||||
"\n",
|
||||
"### Wikipedia\n",
|
||||
"\n",
|
||||
"According to Wikipedia, extraterrestrial life refers to life that does not originate from Earth. The search for alien life is an active area of research, with scientists using various methods to detect signs of life, such as analyzing the atmospheres of exoplanets and searching for biosignatures.\n",
|
||||
"\n",
|
||||
"### Arxiv\n",
|
||||
"\n",
|
||||
"The arxiv database has numerous papers on the search for alien life, including studies on the possibility of life existing on exoplanets and the use of astrobiology to detect signs of life. One paper, 'Extraterrestrial Life and Censorship,' discusses the development of the theory of cosmic life and the attempts to censor evidence incompatible with this theory.\n",
|
||||
"\n",
|
||||
"### Web Search\n",
|
||||
"\n",
|
||||
"A web search for alien life yielded numerous results, including articles from reputable sources such as NASA and Live Science. These articles discuss the latest discoveries and findings in the search for alien life, including the potential for life existing on exoplanets and the use of astrobiology to detect signs of life.\n",
|
||||
"\n",
|
||||
"### Conclusion\n",
|
||||
"\n",
|
||||
"While the search for alien life is an ongoing and active area of research, there are many promising leads and discoveries that suggest the possibility of life existing elsewhere in the universe. As scientists continue to search for signs of life, it is possible that we may one day find evidence of alien life.\n",
|
||||
"\n",
|
||||
"### References\n",
|
||||
"\n",
|
||||
"* [Extraterrestrial life - Wikipedia](https://en.wikipedia.org/wiki/Extraterrestrial_life)\n",
|
||||
"* [What's the best evidence we've found for alien life?](https://www.livescience.com/space/extraterrestrial-life/whats-the-best-evidence-weve-found-for-alien-life)\n",
|
||||
"* [NASA Research Gives Guideline for Future Alien Life Search](https://www.nasa.gov/universe/nasa-research-gives-guideline-for-future-alien-life-search/)\n",
|
||||
"* [James Webb telescope sees potential signs of alien life in the atmosphere of a distant goldilocks water-world](https://www.livescience.com/space/exoplanets/james-webb-telescope-sees-potential-signs-of-alien-life-in-the-atmosphere-of-a-distant-goldilocks-water-world)\n",
|
||||
"* [Alien life in Universe: Scientists say finding it is 'only a matter of time'](https://www.bbc.com/news/science-environment-66950930)\n",
|
||||
"* [Extraterrestrial evidence: 10 incredible findings about aliens from 2020](https://www.livescience.com/alien-discoveries-2020.html)\n",
|
||||
"* [How Scientists Could Tell the World if They Find Alien Life](https://www.scientificamerican.com/article/how-scientists-could-tell-the-world-if-they-find-alien-life/)"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from IPython.display import display, Markdown\n",
|
||||
"\n",
|
||||
"researcher = \"\"\"\n",
|
||||
"You're an expert researcher looking to write a paper on {topic}.\n",
|
||||
"Search for websites, scientific papers and Wikipedia related to the topic.\n",
|
||||
"Write a report with summaries and references (with hyperlinks).\n",
|
||||
"Write the text as Markdown.\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"display(Markdown(agent(researcher.format(topic=\"alien life\"))))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Let's unpack what just happened here. The Agent reviewed multiple sources and aggregated the references into a single report. This is quite powerful 💪\n",
|
||||
"\n",
|
||||
"Note that depending on the LLM used, errors can be seen as the Agent tries to get the function parameters right. This example is using Qwen3 4B. A more powerful LLM will likely lead to even better results. Remember that txtai supports a number of LLM frameworks (Hugging Face, llama.cpp and LiteLLM APIs)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Agent Teams\n",
|
||||
"\n",
|
||||
"Can agents also be tools? Yes!\n",
|
||||
"\n",
|
||||
"Next, we'll build a similar example but instead use an \"Agent Team\" to answer questions."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import Agent, LLM\n",
|
||||
"\n",
|
||||
"# Share the LLM instance across multiple agents\n",
|
||||
"llm = LLM(\"Qwen/Qwen3-4B-Instruct-2507\")\n",
|
||||
"\n",
|
||||
"websearcher = Agent(\n",
|
||||
" tools=[\"websearch\"],\n",
|
||||
" llm=llm,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"wikiman = Agent(\n",
|
||||
" tools=[{\n",
|
||||
" \"name\": \"wikipedia\",\n",
|
||||
" \"description\": \"Searches a Wikipedia database\",\n",
|
||||
" \"provider\": \"huggingface-hub\",\n",
|
||||
" \"container\": \"neuml/txtai-wikipedia\"\n",
|
||||
" }],\n",
|
||||
" llm=llm,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"researcher = Agent(\n",
|
||||
" tools=[{\n",
|
||||
" \"name\": \"arxiv\",\n",
|
||||
" \"description\": \"Searches a database of scientific papers\",\n",
|
||||
" \"provider\": \"huggingface-hub\",\n",
|
||||
" \"container\": \"neuml/txtai-arxiv\"\n",
|
||||
" }],\n",
|
||||
" llm=llm,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"agent = Agent(\n",
|
||||
" tools=[{\n",
|
||||
" \"name\": \"websearcher\",\n",
|
||||
" \"description\": \"I run web searches, there is no answer a web search can't solve!\",\n",
|
||||
" \"target\": websearcher\n",
|
||||
" }, {\n",
|
||||
" \"name\": \"wikiman\",\n",
|
||||
" \"description\": \"Wikipedia has all the answers, I search Wikipedia and answer questions\",\n",
|
||||
" \"target\": wikiman\n",
|
||||
" }, {\n",
|
||||
" \"name\": \"researcher\",\n",
|
||||
" \"description\": \"I'm a science guy. I search arXiv to get all my answers.\",\n",
|
||||
" \"target\": researcher\n",
|
||||
" }],\n",
|
||||
" llm=llm,\n",
|
||||
" max_iterations=10\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 20,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"The comprehensive report on fundamental concepts about Signal Processing is as follows:\n",
|
||||
"\n",
|
||||
"### Introduction\n",
|
||||
"\n",
|
||||
"Signal representation, filtering, and spectral analysis are fundamental concepts in signal processing. Signal representation involves encoding a signal in a mathematical or computational format, while filtering involves modifying the signal to remove unwanted parts or components. Spectral analysis involves analyzing the frequency content of a signal.\n",
|
||||
"\n",
|
||||
"### Signal Representation\n",
|
||||
"\n",
|
||||
"Signal representation involves encoding a signal in a mathematical or computational format. This can be done using various techniques, including Fourier analysis and wavelet analysis.\n",
|
||||
"\n",
|
||||
"### Filtering\n",
|
||||
"\n",
|
||||
"Filtering involves modifying the signal to remove unwanted parts or components. This can be done using various types of filters, including Butterworth filters, Chebyshev-I filters, and Elliptical filters.\n",
|
||||
"\n",
|
||||
"### Spectral Analysis\n",
|
||||
"\n",
|
||||
"Spectral analysis involves analyzing the frequency content of a signal. This can be done using various techniques, including Fourier analysis, wavelet analysis, and singular spectrum analysis.\n",
|
||||
"\n",
|
||||
"### Applications\n",
|
||||
"\n",
|
||||
"Butterworth filters are widely used in signal processing and audio processing for applications such as noise reduction and anti-aliasing filtering. They are also used in medical signal processing for filtering ECG and EEG signals.\n",
|
||||
"\n",
|
||||
"### Conclusion\n",
|
||||
"\n",
|
||||
"In conclusion, signal representation, filtering, and spectral analysis are fundamental concepts in signal processing. Butterworth filters are widely used in various applications, including audio processing and medical signal processing."
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"display(Markdown(agent(\"\"\"\n",
|
||||
"Work with your team and build a comprehensive report on fundamental concepts about Signal Processing.\n",
|
||||
"Write the output in Markdown.\n",
|
||||
"\"\"\", maxlength=16000)))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"📚 See the report above. It ran through similar logic as the first agent, except this time it ran with multiple agents! Note that the agent logging output is not included for brevity."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Agents as a service\n",
|
||||
"\n",
|
||||
"Agents are fully supported through txtai's application configuration via YAML. These services can be run standalone in Python or as a FastAPI service."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Writing config.yml\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"%%writefile config.yml\n",
|
||||
"\n",
|
||||
"agent:\n",
|
||||
" researcher:\n",
|
||||
" tools:\n",
|
||||
" - websearch\n",
|
||||
" - name: wikipedia\n",
|
||||
" description: Searches a Wikipedia database\n",
|
||||
" provider: huggingface-hub\n",
|
||||
" container: neuml/txtai-wikipedia\n",
|
||||
" - name: arxiv\n",
|
||||
" description: Searches a database of scientific papers\n",
|
||||
" provider: huggingface-hub\n",
|
||||
" container: neuml/txtai-arxiv\n",
|
||||
"\n",
|
||||
"llm:\n",
|
||||
" path: Qwen/Qwen3-4B-Instruct-2507"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!CONFIG=config.yml nohup uvicorn \"txtai.api:app\" &> api.log &\n",
|
||||
"!sleep 90"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'The Roman Empire was a vast and powerful state that existed from 27 BC to 476 AD, covering much of Europe, North Africa, and Western Asia. It was founded by Augustus Caesar and was ruled by a series of emperors, with its capital in Rome. The Roman Empire was known for its impressive architecture, engineering, and administrative achievements, including the construction of roads, bridges, and public buildings. It also had a significant impact on the development of law, governance, and culture in the ancient world.'"
|
||||
]
|
||||
},
|
||||
"execution_count": 15,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import requests\n",
|
||||
"\n",
|
||||
"requests.post(\n",
|
||||
" \"http://localhost:8000/agent\",\n",
|
||||
" headers={\"Content-Type\": \"application/json\"},\n",
|
||||
" json={\"name\": \"researcher\", \"text\": \"Tell me about the Roman Empire\", \"maxlength\": 160000}\n",
|
||||
").json()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"💥 Look at that! A full API service from a simple configuration file. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Vectorization with Model2Vec\n",
|
||||
"\n",
|
||||
"While the agent framework is the headline change, there is another major update - support for Model2Vec models. \n",
|
||||
"\n",
|
||||
"[Model2Vec](https://github.com/MinishLab/model2vec) is a technique to turn any sentence transformer into a really small static model, reducing model size by 15x and making the models up to 500x faster, with a small drop in performance."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"\"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\""
|
||||
]
|
||||
},
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"# Data to index\n",
|
||||
"data = [\n",
|
||||
" \"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",
|
||||
"\n",
|
||||
"# Create an embeddings\n",
|
||||
"embeddings = Embeddings(method=\"model2vec\", path=\"minishlab/M2V_base_output\")\n",
|
||||
"embeddings.index(data)\n",
|
||||
"\n",
|
||||
"uid = embeddings.search(\"climate change\")[0][0]\n",
|
||||
"data[uid]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "tvnMO1Eai6Gy"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook gave a quick overview of txtai 8.0. Updated documentation and more examples will be forthcoming. There is much to cover and much to build on!\n",
|
||||
"\n",
|
||||
"See the following links for more information.\n",
|
||||
"\n",
|
||||
"- [8.0 Release on GitHub](https://github.com/neuml/txtai/releases/tag/v8.0.0)\n",
|
||||
"- [Documentation site](https://neuml.github.io/txtai)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "T4",
|
||||
"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.20"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,501 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Granting autonomy to agents\n",
|
||||
"\n",
|
||||
"txtai 8.0 was recently released and added the ability to run agents. Agents automatically create workflows to answer multi-faceted user requests.\n",
|
||||
"\n",
|
||||
"Agents connect a series of tools with a reasoning engine (i.e. LLM). We're giving the agent a degree of latitude to go through it's own internal logic to address a user's request.\n",
|
||||
"\n",
|
||||
"This is a huge paradigm shift. We're talking about handing over control to a program and hoping it makes the right decisions itself. Perhaps there are some parallels to sending your kid to college - we hope we've raised them the right way to be able to make smart decisions 😂. \n",
|
||||
"\n",
|
||||
"This notebook will focus on examples that give agents autonomy to address requests. With this, we can start to the see the path ahead towards more and more automation of tasks."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[agent,graph]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Let's get creative\n",
|
||||
"\n",
|
||||
"In the first example, we'll define an agent that has access to the [txtai-wikipedia](https://huggingface.co/NeuML/txtai-wikipedia) embeddings database. Standard retrieval augmented generation (RAG) and vector search are typically designed for a single search. Agents enable a much more creative and iterative approach to search.\n",
|
||||
"\n",
|
||||
"This example won't ask the agent exactly what we're looking for. We'll ask the agent to tell us something interesting based on the agent's own ideas and research."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import Agent\n",
|
||||
"\n",
|
||||
"agent = Agent(\n",
|
||||
" tools=[{\n",
|
||||
" \"name\": \"wikipedia\",\n",
|
||||
" \"description\": \"Searches a Wikipedia database\",\n",
|
||||
" \"provider\": \"huggingface-hub\",\n",
|
||||
" \"container\": \"neuml/txtai-wikipedia\"\n",
|
||||
" }],\n",
|
||||
" llm=\"Qwen/Qwen3-4B-Instruct-2507\",\n",
|
||||
" max_iterations=10,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"\u001b[32;20;1m======== New task ========\u001b[0m\n",
|
||||
"\u001b[37;1m\n",
|
||||
"I'm bored 🥱. Think about 2-3 disparate topics and use those to search wikipedia to generate something fascinating.\n",
|
||||
"Write a report summarizing each article. Include a section with a list of article hyperlinks.\n",
|
||||
"Write the text as Markdown.\n",
|
||||
"\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: I will use Wikipedia to search for articles related to three disparate topics. The first topic will be the \"Great Barrier Reef\", the second topic will be \"Dark Matter\", and the third topic will be \"The Voynich Manuscript\".\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'wikipedia' with arguments: {'query': 'Great Barrier Reef'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: The results from the Wikipedia search show that the Great Barrier Reef is a vast and diverse ecosystem that supports a wide variety of life. However, it is also facing many threats, including climate change, pollution, and overfishing.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'wikipedia' with arguments: {'query': 'Dark Matter'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: The results from the Wikipedia search show that Dark Matter is a hypothetical form of matter that does not interact with light or other electromagnetic radiation, and is implied by gravitational effects that cannot be explained by general relativity.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'wikipedia' with arguments: {'query': 'The Voynich Manuscript'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: The results from the Wikipedia search show that the Voynich Manuscript is an illustrated codex that has been carbon-dated to the early 15th century, and its origins, authorship, and purpose are still debated.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'final_answer' with arguments: {'answer': \"### Report on Fascinating Topics\\n\\n#### The Great Barrier Reef\\n\\nThe Great Barrier Reef is the world's largest coral reef system, composed of over 2,900 individual reefs and 900 islands stretching for over 2,300 kilometers off the coast of Queensland, Australia. It is a vast and diverse ecosystem that supports a wide variety of life, including coral and other invertebrate species, bony fish, sharks, rays, marine mammals, marine turtles, sea snakes, and algae. However, it is also facing many threats, including climate change, pollution, and overfishing.\\n\\n#### Dark Matter\\n\\nDark Matter is a hypothetical form of matter that does not interact with light or other electromagnetic radiation. It is implied by gravitational effects that cannot be explained by general relativity, and is a topic of ongoing research in astrophysics and cosmology.\\n\\n#### The Voynich Manuscript\\n\\nThe Voynich Manuscript is an illustrated codex that has been carbon-dated to the early 15th century. Its origins, authorship, and purpose are still debated, and it is considered one of the most mysterious and enigmatic manuscripts in the world.\\n\\n### Article Hyperlinks\\n\\n* [Great Barrier Reef](https://en.wikipedia.org/wiki/Great_Barrier_Reef)\\n* [Dark Matter](https://en.wikipedia.org/wiki/Dark_matter)\\n* [Voynich Manuscript](https://en.wikipedia.org/wiki/Voynich_manuscript)\"}\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"### Report on Fascinating Topics\n",
|
||||
"\n",
|
||||
"#### The Great Barrier Reef\n",
|
||||
"\n",
|
||||
"The Great Barrier Reef is the world's largest coral reef system, composed of over 2,900 individual reefs and 900 islands stretching for over 2,300 kilometers off the coast of Queensland, Australia. It is a vast and diverse ecosystem that supports a wide variety of life, including coral and other invertebrate species, bony fish, sharks, rays, marine mammals, marine turtles, sea snakes, and algae. However, it is also facing many threats, including climate change, pollution, and overfishing.\n",
|
||||
"\n",
|
||||
"#### Dark Matter\n",
|
||||
"\n",
|
||||
"Dark Matter is a hypothetical form of matter that does not interact with light or other electromagnetic radiation. It is implied by gravitational effects that cannot be explained by general relativity, and is a topic of ongoing research in astrophysics and cosmology.\n",
|
||||
"\n",
|
||||
"#### The Voynich Manuscript\n",
|
||||
"\n",
|
||||
"The Voynich Manuscript is an illustrated codex that has been carbon-dated to the early 15th century. Its origins, authorship, and purpose are still debated, and it is considered one of the most mysterious and enigmatic manuscripts in the world.\n",
|
||||
"\n",
|
||||
"### Article Hyperlinks\n",
|
||||
"\n",
|
||||
"* [Great Barrier Reef](https://en.wikipedia.org/wiki/Great_Barrier_Reef)\n",
|
||||
"* [Dark Matter](https://en.wikipedia.org/wiki/Dark_matter)\n",
|
||||
"* [Voynich Manuscript](https://en.wikipedia.org/wiki/Voynich_manuscript)"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from IPython.display import display, Markdown\n",
|
||||
"\n",
|
||||
"answer = agent(\"\"\"\n",
|
||||
"I'm bored 🥱. Think about 2-3 disparate topics and use those to search wikipedia to generate something fascinating.\n",
|
||||
"Write a report summarizing each article. Include a section with a list of article hyperlinks.\n",
|
||||
"Write the text as Markdown.\n",
|
||||
"\"\"\", maxlength=16000)\n",
|
||||
"\n",
|
||||
"display(Markdown(answer))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"💥 Interesting indeed. The fundamental concept of search is we need to know what to look for. In this case, we didn't have that (i.e. we're bored 😀).\n",
|
||||
"\n",
|
||||
"Let's go to another example. This time we'll look at the [txtai-hfposts](https://huggingface.co/NeuML/txtai-hfposts) embeddings database. We'll ask the agent to research a specific topic then write a report about what was found.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"agent = Agent(\n",
|
||||
" tools=[{\n",
|
||||
" \"name\": \"hfposts\",\n",
|
||||
" \"description\": \"Searches a database of technical posts on Hugging Face\",\n",
|
||||
" \"provider\": \"huggingface-hub\",\n",
|
||||
" \"container\": \"neuml/txtai-hfposts\"\n",
|
||||
" }],\n",
|
||||
" llm=\"Qwen/Qwen3-4B-Instruct-2507\",\n",
|
||||
" max_iterations=10,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"\u001b[32;20;1m======== New task ========\u001b[0m\n",
|
||||
"\u001b[37;1m\n",
|
||||
"Read posts about medicine and write a report on what you learned.\n",
|
||||
"\n",
|
||||
"The report should be a Markdown table with the following columns.\n",
|
||||
" - Name\n",
|
||||
" - Description\n",
|
||||
" - Link to content\n",
|
||||
"\n",
|
||||
"Only include rows that have a valid web url.\n",
|
||||
"\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: First, I need to search for posts related to medicine using the `hfposts` tool.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'hfposts' with arguments: {'query': 'medicine'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: Now that I have the posts related to medicine, I need to filter out the ones that have a valid web URL. I will go through each post and check if it contains a URL.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'hfposts' with arguments: {'query': 'medicine url'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: Now that I have the posts with URLs, I need to extract the relevant information such as the name, description, and link to content. I will create a Markdown table with these details.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'hfposts' with arguments: {'query': 'medicine url'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: Now that I have the posts with URLs, I can extract the relevant information and create a Markdown table. I will format the data into the required structure.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'final_answer' with arguments: {'answer': '| Name | Description | Link to content |\\n| --- | --- | --- |\\n| OpenBioLLM Release | Introducing OpenBioLLM-Llama3-70B & 8B: The most capable openly available Medical-domain LLMs to date! | [https://huggingface.co/blog/aaditya/openbiollm](https://huggingface.co/blog/aaditya/openbiollm) |\\n| Last Week in Medical AI: Top Research Papers/Models (September 1 - September 7, 2024) | Outperforms industry giants like GPT-4, Gemini, Meditron-70B, Med-PaLM-1, and Med-PaLM-2 in the biomedical domain. | [https://x.com/OpenlifesciAI/status/1832476252260712788](https://x.com/OpenlifesciAI/status/1832476252260712788) |\\n| Last Week in Medical AI: Top Research Papers/Models (August 25 - August 31, 2024) | Includes MultiMed, a Multimodal Medical Benchmark, and A Foundation model for generating chest X-ray images. | [https://x.com/OpenlifesciAI/status/1829984701324448051](https://x.com/OpenlifesciAI/status/1829984701324448051) |\\n| Last Week in Medical AI: Top Research Papers/Models (October 5 - October 12, 2024) | Introduces MMedAgent: Learning to Use Medical Tools with Multi-modal Agent. | [https://youtu.be/OD3C5jirszw](https://youtu.be/OD3C5jirszw) |\\n| Last Week in Medical AI: Top Research Papers/Models (October 26 - November 2, 2024) | Google Presents MDAgents: An Adaptive Collaboration of LLMs for Medical Decision-Making. | [https://x.com/OpenlifesciAI/status/1852685220912464066](https://x.com/OpenlifesciAI/status/1852685220912464066) '}\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"| Name | Description | Link to content |\n",
|
||||
"| --- | --- | --- |\n",
|
||||
"| OpenBioLLM Release | Introducing OpenBioLLM-Llama3-70B & 8B: The most capable openly available Medical-domain LLMs to date! | [https://huggingface.co/blog/aaditya/openbiollm](https://huggingface.co/blog/aaditya/openbiollm) |\n",
|
||||
"| Last Week in Medical AI: Top Research Papers/Models (September 1 - September 7, 2024) | Outperforms industry giants like GPT-4, Gemini, Meditron-70B, Med-PaLM-1, and Med-PaLM-2 in the biomedical domain. | [https://x.com/OpenlifesciAI/status/1832476252260712788](https://x.com/OpenlifesciAI/status/1832476252260712788) |\n",
|
||||
"| Last Week in Medical AI: Top Research Papers/Models (August 25 - August 31, 2024) | Includes MultiMed, a Multimodal Medical Benchmark, and A Foundation model for generating chest X-ray images. | [https://x.com/OpenlifesciAI/status/1829984701324448051](https://x.com/OpenlifesciAI/status/1829984701324448051) |\n",
|
||||
"| Last Week in Medical AI: Top Research Papers/Models (October 5 - October 12, 2024) | Introduces MMedAgent: Learning to Use Medical Tools with Multi-modal Agent. | [https://youtu.be/OD3C5jirszw](https://youtu.be/OD3C5jirszw) |\n",
|
||||
"| Last Week in Medical AI: Top Research Papers/Models (October 26 - November 2, 2024) | Google Presents MDAgents: An Adaptive Collaboration of LLMs for Medical Decision-Making. | [https://x.com/OpenlifesciAI/status/1852685220912464066](https://x.com/OpenlifesciAI/status/1852685220912464066) "
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"answer = agent(\"\"\"\n",
|
||||
"Read posts about medicine and write a report on what you learned.\n",
|
||||
"\n",
|
||||
"The report should be a Markdown table with the following columns.\n",
|
||||
" - Name\n",
|
||||
" - Description\n",
|
||||
" - Link to content\n",
|
||||
"\n",
|
||||
"Only include rows that have a valid web url.\n",
|
||||
"\"\"\", maxlength=16000)\n",
|
||||
"\n",
|
||||
"display(Markdown(answer))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"🚀 Once again, very interesting! This time we asked the agent to go read about a topic and report back. The agent did that and left us with links to explore further."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Autonomous Embeddings\n",
|
||||
"\n",
|
||||
"For our last example, we're going to give an agent free rein to control an embeddings database.\n",
|
||||
"\n",
|
||||
"First, we will create an empty embeddings database and tell the agent how to add and search for data."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import Agent, Embeddings\n",
|
||||
"from txtai.pipeline import Textractor\n",
|
||||
"from txtai.workflow import Workflow, Task\n",
|
||||
"\n",
|
||||
"# Empty embeddings database\n",
|
||||
"embeddings = Embeddings(\n",
|
||||
" path=\"intfloat/e5-large\",\n",
|
||||
" instructions={\"query\": \"query: \", \"data\": \"passage: \"},\n",
|
||||
" content=True\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Textractor instance\n",
|
||||
"textractor = Textractor(sections=True, headers={\"user-agent\": \"Mozilla/5.0\"})\n",
|
||||
"\n",
|
||||
"def insert(elements):\n",
|
||||
" \"\"\"\n",
|
||||
" Inserts elements into the embeddings database.\n",
|
||||
"\n",
|
||||
" Args:\n",
|
||||
" elements: list of strings to insert\n",
|
||||
" \"\"\"\n",
|
||||
"\n",
|
||||
" def upsert(elements):\n",
|
||||
" embeddings.upsert(elements)\n",
|
||||
" return elements\n",
|
||||
"\n",
|
||||
" # Upsert workflow\n",
|
||||
" workflow = Workflow([Task(textractor), Task(upsert)])\n",
|
||||
" list(workflow(elements))\n",
|
||||
" return f\"{elements} inserted successfully\"\n",
|
||||
"\n",
|
||||
"agent = Agent(\n",
|
||||
" tools=[insert, embeddings.search, \"websearch\"],\n",
|
||||
" llm=\"Qwen/Qwen3-4B-Instruct-2507\",\n",
|
||||
" max_iterations=10\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Next, we'll define a workflow process. This creates the rules we want the agent to follow. Then we'll run two requests."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"\u001b[32;20;1m======== New task ========\u001b[0m\n",
|
||||
"\u001b[37;1m\n",
|
||||
"Run the following process:\n",
|
||||
" 1. Search your internal knowledge for txtai\n",
|
||||
" 2. If not found, find relevant urls and insert those as a list of strings ONLY. Then rerun the search for txtai.\n",
|
||||
" 3. Write a detailed report about txtai with Markdown sections covering the major topics. Include a section with hyperlink references.\n",
|
||||
"\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: First, I will attempt to search for \"txtai\" in my internal knowledge base.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'search' with arguments: {'query': 'txtai'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: Since the search for \"txtai\" did not yield any results, I will proceed to find relevant URLs and insert them as a list of strings.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'web_search' with arguments: {'query': 'txtai'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: Now that I have found relevant URLs, I will insert them into the database as a list of strings.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'insert' with arguments: {'elements': ['https://github.com/neuml/txtai', 'https://pypi.org/project/txtai/', 'https://neuml.github.io/txtai/', 'https://neuml.com/', 'https://medium.com/neuml/introducing-txtai-the-all-in-one-embeddings-database-c721f4ff91ad', 'https://neuml.github.io/txtai/embeddings/', 'https://neuml.github.io/txtai/examples/']}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: Now that the relevant URLs have been inserted, I will rerun the search for \"txtai\".\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'search' with arguments: {'query': 'txtai'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'final_answer' with arguments: {'answer': '# txtai Overview\\n\\ntxtai is an all-in-one open-source embeddings database designed for semantic search, LLM orchestration, and language model workflows. It combines vector indexes (both sparse and dense), graph networks, and relational databases to offer a comprehensive solution for various tasks.\\n\\n## Semantic Search\\n\\ntxtai supports semantic search, enabling users to find relevant information based on the meaning of the text rather than just keywords. This makes it suitable for applications such as recommendation systems, content retrieval, and more.\\n\\n## LLM Orchestration\\n\\ntxtai includes support for Large Language Models (LLMs) orchestration, which allows users to integrate and manage different LLMs within their workflows. This feature is particularly useful for tasks such as prompt chaining, summarization, transcription, translation, and more.\\n\\n## Embeddings Databases\\n\\ntxtai provides a unified interface for embedding, indexing, searching, and processing text and multimedia data. Users can create, index, and search embeddings using code snippets and examples provided in the documentation.\\n\\n## Hyperlink References\\n\\n- [txtai on GitHub](https://github.com/neuml/txtai)\\n- [PyPI - txtai](https://pypi.org/project/txtai/)\\n- [txtai Documentation](https://neuml.github.io/txtai/)\\n- [txtai Examples](https://neuml.github.io/txtai/examples/)\\n- [txtai Embeddings](https://neuml.github.io/txtai/embeddings/)'}\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"# txtai Overview\n",
|
||||
"\n",
|
||||
"txtai is an all-in-one open-source embeddings database designed for semantic search, LLM orchestration, and language model workflows. It combines vector indexes (both sparse and dense), graph networks, and relational databases to offer a comprehensive solution for various tasks.\n",
|
||||
"\n",
|
||||
"## Semantic Search\n",
|
||||
"\n",
|
||||
"txtai supports semantic search, enabling users to find relevant information based on the meaning of the text rather than just keywords. This makes it suitable for applications such as recommendation systems, content retrieval, and more.\n",
|
||||
"\n",
|
||||
"## LLM Orchestration\n",
|
||||
"\n",
|
||||
"txtai includes support for Large Language Models (LLMs) orchestration, which allows users to integrate and manage different LLMs within their workflows. This feature is particularly useful for tasks such as prompt chaining, summarization, transcription, translation, and more.\n",
|
||||
"\n",
|
||||
"## Embeddings Databases\n",
|
||||
"\n",
|
||||
"txtai provides a unified interface for embedding, indexing, searching, and processing text and multimedia data. Users can create, index, and search embeddings using code snippets and examples provided in the documentation.\n",
|
||||
"\n",
|
||||
"## Hyperlink References\n",
|
||||
"\n",
|
||||
"- [txtai on GitHub](https://github.com/neuml/txtai)\n",
|
||||
"- [PyPI - txtai](https://pypi.org/project/txtai/)\n",
|
||||
"- [txtai Documentation](https://neuml.github.io/txtai/)\n",
|
||||
"- [txtai Examples](https://neuml.github.io/txtai/examples/)\n",
|
||||
"- [txtai Embeddings](https://neuml.github.io/txtai/embeddings/)"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"prompt = \"\"\"\n",
|
||||
"Run the following process:\n",
|
||||
" 1. Search your internal knowledge for {topic}\n",
|
||||
" 2. If not found, find relevant urls and insert those as a list of strings ONLY. Then rerun the search for {topic}.\n",
|
||||
" 3. Write a detailed report about {topic} with Markdown sections covering the major topics. Include a section with hyperlink references.\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"answer = agent(prompt.format(topic=\"txtai\"))\n",
|
||||
"display(Markdown(answer))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"\u001b[32;20;1m======== New task ========\u001b[0m\n",
|
||||
"\u001b[37;1m\n",
|
||||
"Run the following process:\n",
|
||||
" 1. Search your internal knowledge for openscholar\n",
|
||||
" 2. If not found, find relevant urls and insert those as a list of strings ONLY. Then rerun the search for openscholar.\n",
|
||||
" 3. Write a detailed report about openscholar with Markdown sections covering the major topics. Include a section with hyperlink references.\n",
|
||||
"\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: First, I will search my internal knowledge for \"openscholar\".\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'search' with arguments: {'query': 'openscholar'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: Since the search did not return any relevant information about \"openscholar\", I will find relevant URLs and insert them as a list of strings. Then, I will rerun the search for \"openscholar\".\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'web_search' with arguments: {'query': 'openscholar'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: Now that I have found relevant URLs about \"openscholar\", I will insert these URLs into the database.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'insert' with arguments: {'elements': ['https://openscholar.allen.ai/', 'https://allenai.org/blog/openscholar', 'https://arxiv.org/abs/2411.14199', 'https://github.com/AkariAsai/OpenScholar', 'https://www.marktechpost.com/2024/11/22/the-allen-institute-for-ai-ai2-introduces-openscholar-an-open-ecosystem-for-literature-synthesis-featuring-advanced-datastores-and-expert-level-results/', 'https://venturebeat.com/ai/openscholar-the-open-source-a-i-thats-outperforming-gpt-4o-in-scientific-research/']}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: Now that the relevant URLs have been inserted, I will rerun the search for \"openscholar\" to gather more detailed information.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'search' with arguments: {'query': 'openscholar'}\u001b[0m\n",
|
||||
"\u001b[33;1m=== Agent thoughts:\u001b[0m\n",
|
||||
"\u001b[0mThought: The search has returned relevant information about OpenScholar. I will now compile a detailed report about OpenScholar with Markdown sections covering the major topics.\u001b[0m\n",
|
||||
"\u001b[33;1m>>> Calling tool: 'final_answer' with arguments: {'answer': \"# OpenScholar Report\\n\\n## Introduction\\nOpenScholar is an advanced AI system developed by the Allen Institute for AI and the University of Washington. It aims to revolutionize scientific research by processing millions of research papers and providing citation-backed answers to complex research questions.\\n\\n## Key Features\\n- **Data Processing**: Capable of handling 45 million research papers instantly.\\n- **Citation Accuracy**: Achieves high citation accuracy, challenging proprietary AI systems.\\n- **Customizable Tools**: Offers specialized tools and built-in features for creating impactful websites.\\n\\n## Research Impact\\n- **Data Utilization**: Integrates a vast corpus of scientific literature, making it the largest open-access corpus for scientific literature synthesis.\\n- **Expert-Level Results**: Provides answers that match or exceed those of human experts in various scientific domains.\\n\\n## Future Directions\\n- **Continued Development**: Ongoing efforts to improve the system's capabilities and expand its applications.\\n- **Community Engagement**: Plans for expanding the user base and fostering a community around OpenScholar.\\n\\n## References\\n- [OpenScholar](https://openscholar.allen.ai/)\\n- [VentureBeat Article](https://venturebeat.com/ai/openscholar-the-open-source-a-i-thats-outperforming-gpt-4o-in-scientific-research/)\\n- [Allen Institute for AI](https://allenai.org/blog/openscholar)\"}\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"# OpenScholar Report\n",
|
||||
"\n",
|
||||
"## Introduction\n",
|
||||
"OpenScholar is an advanced AI system developed by the Allen Institute for AI and the University of Washington. It aims to revolutionize scientific research by processing millions of research papers and providing citation-backed answers to complex research questions.\n",
|
||||
"\n",
|
||||
"## Key Features\n",
|
||||
"- **Data Processing**: Capable of handling 45 million research papers instantly.\n",
|
||||
"- **Citation Accuracy**: Achieves high citation accuracy, challenging proprietary AI systems.\n",
|
||||
"- **Customizable Tools**: Offers specialized tools and built-in features for creating impactful websites.\n",
|
||||
"\n",
|
||||
"## Research Impact\n",
|
||||
"- **Data Utilization**: Integrates a vast corpus of scientific literature, making it the largest open-access corpus for scientific literature synthesis.\n",
|
||||
"- **Expert-Level Results**: Provides answers that match or exceed those of human experts in various scientific domains.\n",
|
||||
"\n",
|
||||
"## Future Directions\n",
|
||||
"- **Continued Development**: Ongoing efforts to improve the system's capabilities and expand its applications.\n",
|
||||
"- **Community Engagement**: Plans for expanding the user base and fostering a community around OpenScholar.\n",
|
||||
"\n",
|
||||
"## References\n",
|
||||
"- [OpenScholar](https://openscholar.allen.ai/)\n",
|
||||
"- [VentureBeat Article](https://venturebeat.com/ai/openscholar-the-open-source-a-i-thats-outperforming-gpt-4o-in-scientific-research/)\n",
|
||||
"- [Allen Institute for AI](https://allenai.org/blog/openscholar)"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"answer = agent(prompt.format(topic=\"openscholar\"))\n",
|
||||
"display(Markdown(answer))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"🔥 Amazing.\n",
|
||||
"\n",
|
||||
"Remember, we started with an empty embeddings database. Then we gave basic instructions on how to use the available tools. From there, the agent autonomously operated and answered user requests. The agent also stored what it learned for future requests. This gave the agent it's own internal memory.\n",
|
||||
"\n",
|
||||
"Of course, we could program a process that implements this workflow. But think about the productivity gains we're opening up to so many more people. We're enabling people to control a process simply by pairing a set of tools with a description of what they want, in English.\n",
|
||||
"\n",
|
||||
"Exciting times!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook demonstrated ways to run agents in a more autonomous fashion. While the technology isn't perfect, we can certainly see the path ahead where new models will continue to do a better job. With the right agents and targeted tools, much can be done now though. \n",
|
||||
"\n",
|
||||
"Think about the differences between now and 6-12 months ago. Where will we be in another 6-12 months!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"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.20"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
@@ -0,0 +1,792 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "7btcrgth5z6f"
|
||||
},
|
||||
"source": [
|
||||
"# Getting started with LLM APIs\n",
|
||||
"\n",
|
||||
"_Thank you to [Igor Ribeiro Lima](https://github.com/igorlima) for writing this notebook!_\n",
|
||||
"\n",
|
||||
"This notebook demonstrates how to call different LLM APIs using just one set of common functions.\n",
|
||||
"\n",
|
||||
"While each LLM has its own configuration, here's a cool perk of using [`txtai`](https://neuml.github.io/txtai/) plus [`litellm`](https://docs.litellm.ai/docs/) working behind the scenes. The big win with `txtai` is that it helps standardize inputs and outputs across several LLMs. **This means we can seamlessly call any LLM API**.\n",
|
||||
"\n",
|
||||
"In this notebook, we use one set of common functions to call _**Gemini**_, _**VertexAI**_, _**Mistral**_, _**Cohere**_, _**AWS Bedrock**_, _**OpenAI**_, _**Claude by Anthropic**_, _**Groq**_, _and more_!\n",
|
||||
"\n",
|
||||
"Isn't that amazing? One method for many different LLM APIs - **that's the beauty of these libraries**!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "lhrk9y766Lw1"
|
||||
},
|
||||
"source": [
|
||||
"## Install dependencies\n",
|
||||
"\n",
|
||||
"This session guides you through all the essential dependencies you'll need to run the Python script for various LLM APIs.\n",
|
||||
"\n",
|
||||
"You can choose the dependencies based on your specific needs. Each dependency is carefully commented on in the code, so you'll know exactly which API it supports.\n",
|
||||
"\n",
|
||||
"The core dependencies you'll typically need are: `txtai` and `txtai[pipeline]`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "czzRItdyN-Zj"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"# The main point here is to follow the pattern of other notebooks.\n",
|
||||
"# For more information, kindly refer to the note in the bullet right below.\n",
|
||||
"\n",
|
||||
"# !pip install txtai==8.1.0\n",
|
||||
"# !pip install txtai[pipeline]\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline]\n",
|
||||
"\n",
|
||||
"# to use Vertex AI\n",
|
||||
"# https://github.com/BerriAI/litellm/issues/5483\n",
|
||||
"# !pip install google-cloud-aiplatform==1.75.0\n",
|
||||
"!pip install google-cloud-aiplatform\n",
|
||||
"\n",
|
||||
"# to use AWS Bedrock\n",
|
||||
"# https://docs.litellm.ai/docs/providers/bedrock\n",
|
||||
"# !pip install boto3==1.35.88\n",
|
||||
"!pip install boto3"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "evHD8Kw09Aa0"
|
||||
},
|
||||
"source": [
|
||||
"- _**A friendly reminder**: things can change unexpectedly in the ever-evolving world of coding. One day, your code works flawlessly; the next, it might throw a tantrum for no apparent reason. That's why it's essential to specify the dependencies' versions using the latest available versions when writing the code._\n",
|
||||
" - <sup><sub>Even recognizing the importance of version control, each version has one line above as a comment in the code. It serves as a guide to help you track dependencies and ensures the code runs smoothly across different environments, even if something goes awry.</sub></sup>\n",
|
||||
" - <sup><sub>While there is [a trade-off](https://github.com/neuml/txtai/pull/844#issuecomment-2564294186) in limiting code to a specific version, noting today's version as a comment should help you protect against potential issues with future updates.</sub></sup>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "wuCvBPCL-VQG"
|
||||
},
|
||||
"source": [
|
||||
"## LLM API Configuration\n",
|
||||
"\n",
|
||||
"This session is like a special Python dictionary that holds all the cool configurations for our LLM AI model. It includes details like environment variables, the model's name, text embedding parameters, and other fancy settings.\n",
|
||||
"\n",
|
||||
"One important key here is `IS_ENABLED`. When it's set to `True`, it's like giving the model a green light to shine! But if you ever feel like taking a break or don't need it for a while, you can easily set this key to `False`, and the model will chill out."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "jyahNyOrO2yp",
|
||||
"outputId": "8143c2b0-a6a9-429b-a66f-93f39e6885e4"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"config set!\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import os, getpass\n",
|
||||
"from txtai import LLM, Embeddings\n",
|
||||
"\n",
|
||||
"# https://neuml.github.io/txtai/install/\n",
|
||||
"# https://neuml.github.io/txtai/pipeline/text/llm/#example\n",
|
||||
"# https://neuml.github.io/txtai/embeddings/configuration/vectors/#method\n",
|
||||
"# https://docs.litellm.ai/docs/embedding/supported_embedding\n",
|
||||
"LLM_MODEL_CONFIG = {\n",
|
||||
" 'GEMINI': {\n",
|
||||
" 'IS_ENABLED': True,\n",
|
||||
" 'ENV_VAR': ['GEMINI_API_KEY'],\n",
|
||||
" # https://docs.litellm.ai/docs/providers/gemini\n",
|
||||
" # https://ai.google.dev/gemini-api/docs/models/gemini\n",
|
||||
" 'LLM_MODEL_NAME': 'gemini/gemini-pro',\n",
|
||||
" # https://github.com/BerriAI/litellm/tree/12c4e7e695edb07d403dd14fc768a736638bd3d1/litellm/llms/vertex_ai\n",
|
||||
" # https://github.com/BerriAI/litellm/blob/e19bb55e3b4c6a858b6e364302ebbf6633a51de5/model_prices_and_context_window.json#L2625\n",
|
||||
" 'TEXT_EMBEDDING_PATH': 'gemini/text-embedding-004'\n",
|
||||
" },\n",
|
||||
" 'COHERE': {\n",
|
||||
" 'IS_ENABLED': False,\n",
|
||||
" 'ENV_VAR': ['COHERE_API_KEY'],\n",
|
||||
" # https://docs.litellm.ai/docs/providers/cohere\n",
|
||||
" 'LLM_MODEL_NAME': 'command-light',\n",
|
||||
" 'TEXT_EMBEDDING_PATH': 'cohere/embed-english-v3.0'\n",
|
||||
" },\n",
|
||||
" 'AWS_BEDROCK': {\n",
|
||||
" 'IS_ENABLED': False,\n",
|
||||
" 'ENV_VAR': ['AWS_ACCESS_KEY_ID', 'AWS_SECRET_ACCESS_KEY', 'AWS_REGION_NAME'],\n",
|
||||
" # https://docs.litellm.ai/docs/providers/bedrock\n",
|
||||
" 'LLM_MODEL_NAME': 'bedrock/amazon.titan-text-lite-v1',\n",
|
||||
" 'TEXT_EMBEDDING_PATH': 'amazon.titan-embed-text-v1'\n",
|
||||
" },\n",
|
||||
" 'MISTRAL': {\n",
|
||||
" 'IS_ENABLED': False,\n",
|
||||
" 'ENV_VAR': ['MISTRAL_API_KEY'],\n",
|
||||
" # https://docs.litellm.ai/docs/providers/mistral\n",
|
||||
" 'LLM_MODEL_NAME': 'mistral/mistral-tiny',\n",
|
||||
" 'TEXT_EMBEDDING_PATH': 'mistral/mistral-embed'\n",
|
||||
" },\n",
|
||||
" 'VERTEXAI': {\n",
|
||||
" 'IS_ENABLED': False,\n",
|
||||
" 'ENV_VAR': ['GOOGLE_APPLICATION_CREDENTIALS', 'GOOGLE_VERTEX_PROJECT', 'GOOGLE_VERTEX_LOCATION'],\n",
|
||||
" 'ENV_VAR_SETUP': None,\n",
|
||||
" # https://docs.litellm.ai/docs/providers/vertex\n",
|
||||
" 'LLM_MODEL_NAME': 'vertex_ai/gemini-pro',\n",
|
||||
" 'TEXT_EMBEDDING_PATH': 'vertex_ai/text-embedding-004'\n",
|
||||
" },\n",
|
||||
" 'GROQ': {\n",
|
||||
" 'IS_ENABLED': False,\n",
|
||||
" 'ENV_VAR': ['GROQ_API_KEY'],\n",
|
||||
" # https://docs.litellm.ai/docs/providers/groq\n",
|
||||
" # https://console.groq.com/docs/models\n",
|
||||
" 'LLM_MODEL_NAME': 'groq/llama3-8b-8192',\n",
|
||||
" # the line below has been commented out for a reason.\n",
|
||||
" # to understand why, check out the Groq section right below.\n",
|
||||
" # https://groq.com/retrieval-augmented-generation-with-groq-api/\n",
|
||||
" # 'TEXT_EMBEDDING_PATH': 'jinaai/jina-embeddings-v2-base-en',\n",
|
||||
" },\n",
|
||||
" 'OPENAI': {\n",
|
||||
" 'IS_ENABLED': False,\n",
|
||||
" 'ENV_VAR': ['OPENAI_API_KEY'],\n",
|
||||
" # https://docs.litellm.ai/docs/providers/openai\n",
|
||||
" # https://platform.openai.com/docs/models\n",
|
||||
" 'LLM_MODEL_NAME': 'gpt-4o-mini-2024-07-18',\n",
|
||||
" 'TEXT_EMBEDDING_PATH': 'text-embedding-3-small',\n",
|
||||
" },\n",
|
||||
" 'CLAUDE': {\n",
|
||||
" 'IS_ENABLED': False,\n",
|
||||
" # You might need the Voyage API key, depending on the embedding model you choose. Read more below!\n",
|
||||
" # 'ENV_VAR': ['ANTHROPIC_API_KEY', 'VOYAGE_API_KEY'],\n",
|
||||
" 'ENV_VAR': ['ANTHROPIC_API_KEY'],\n",
|
||||
" # https://docs.litellm.ai/docs/providers/anthropic\n",
|
||||
" # https://docs.anthropic.com/en/docs/about-claude/models\n",
|
||||
" 'LLM_MODEL_NAME': 'anthropic/claude-3-5-haiku-20241022',\n",
|
||||
" # the line below has been commented out for a reason.\n",
|
||||
" # to understand why, check out the Claude Anthropic section right below.\n",
|
||||
" # https://docs.anthropic.com/en/docs/build-with-claude/embeddings\n",
|
||||
" # 'TEXT_EMBEDDING_PATH': 'voyage/voyage-01',\n",
|
||||
" },\n",
|
||||
" 'VOYAGE': {\n",
|
||||
" 'IS_ENABLED': False,\n",
|
||||
" # You might need the Anthropic API key, depending on the LLM you choose. Read more below!\n",
|
||||
" # 'ENV_VAR': ['ANTHROPIC_API_KEY', 'VOYAGE_API_KEY'],\n",
|
||||
" 'ENV_VAR': ['VOYAGE_API_KEY'],\n",
|
||||
" # the line below has been commented out for a reason.\n",
|
||||
" # to understand why, check out the Voyage AI section right below.\n",
|
||||
" # https://docs.litellm.ai/docs/providers/voyage\n",
|
||||
" # https://docs.voyageai.com/docs/introduction\n",
|
||||
" # 'LLM_MODEL_NAME': 'anthropic/claude-3-5-haiku-20241022',\n",
|
||||
" #\n",
|
||||
" # https://docs.voyageai.com/docs/embeddings\n",
|
||||
" 'TEXT_EMBEDDING_PATH': 'voyage/voyage-01',\n",
|
||||
" }\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"import litellm\n",
|
||||
"# # https://github.com/BerriAI/litellm/blob/11932d0576a073d83f38a418cbdf6b2d8d4ff46f/litellm/litellm_core_utils/get_llm_provider_logic.py#L322\n",
|
||||
"litellm.suppress_debug_info = True\n",
|
||||
"# https://docs.litellm.ai/docs/debugging/local_debugging#set-verbose\n",
|
||||
"litellm.set_verbose=False\n",
|
||||
"\n",
|
||||
"def customsetup():\n",
|
||||
" def vertexai():\n",
|
||||
" # https://docs.litellm.ai/docs/embedding/supported_embedding#usage---embedding\n",
|
||||
" # https://docs.litellm.ai/docs/providers/vertex\n",
|
||||
" litellm.vertex_project = os.environ['GOOGLE_VERTEX_PROJECT']\n",
|
||||
" litellm.vertex_location = os.environ['GOOGLE_VERTEX_LOCATION']\n",
|
||||
"\n",
|
||||
" LLM_MODEL_CONFIG['VERTEXAI']['ENV_VAR_SETUP'] = vertexai()\n",
|
||||
"\n",
|
||||
"customsetup()\n",
|
||||
"\n",
|
||||
"LLM_MODELS = {k: v for k, v in LLM_MODEL_CONFIG.items() if v['IS_ENABLED']}\n",
|
||||
"ENV_VARS = [v['ENV_VAR'] for k, v in LLM_MODELS.items()]\n",
|
||||
"ENV_VARS = [x for xs in ENV_VARS for x in xs] # flatten array\n",
|
||||
"\n",
|
||||
"print(\"config set!\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "pgQuDv4fZgHe"
|
||||
},
|
||||
"source": [
|
||||
"### Vertex AI"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "wG_2u_jKTyYg"
|
||||
},
|
||||
"source": [
|
||||
"When working with VertexAI in a Jupyter notebook, an environment variable called `GOOGLE_APPLICATION_CREDENTIALS` points to a credentials file. This variable is like giving VertexAI a secret map to find the JSON file containing your service account key.\n",
|
||||
"\n",
|
||||
"Imagine you've set `GOOGLE_APPLICATION_CREDENTIALS` to `application_default_credentials.json`. The service account key is stored in a file with that exact name. And guess what? This file must be in the same place you're working _(your current working directory)_.\n",
|
||||
"\n",
|
||||
"Here's how your directory might look:\n",
|
||||
"```\n",
|
||||
".\n",
|
||||
"├── ..\n",
|
||||
"├── sample_data/\n",
|
||||
"└── application_default_credentials.json\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"If you use Google Colab, your directory structure will differ slightly, but don't worry! It'll look something like this:\n",
|
||||
"\n",
|
||||
"```\n",
|
||||
".\n",
|
||||
"├── bin\n",
|
||||
"├── boot\n",
|
||||
"├── content/\n",
|
||||
"│ ├── sample_data/\n",
|
||||
"│ └── application_default_credentials.json\n",
|
||||
"├── datalab\n",
|
||||
"├── dev\n",
|
||||
"├── etc\n",
|
||||
"├── home\n",
|
||||
"├── lib\n",
|
||||
"├── lib32\n",
|
||||
"└── lib64\n",
|
||||
"```\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "0KsLGZFtrZVT"
|
||||
},
|
||||
"source": [
|
||||
"### Groq"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-4RUn2dUrcbZ"
|
||||
},
|
||||
"source": [
|
||||
"It's good to make the most of [free resources](https://github.com/BerriAI/litellm/issues/4922#issuecomment-2374234548) like this one! Of course, free options often come with some limitations, and this is no different. Currently, Groq doesn't offer an API for text embeddings.\n",
|
||||
"\n",
|
||||
"The Groq team highlights a pre-trained model that can be used locally. You can learn more about it on the [Groq Blog](https://groq.com/retrieval-augmented-generation-with-groq-api/), where they discuss the `jinaai/jina-embeddings-v2-base-en`, a model that can be hosted locally.\n",
|
||||
"\n",
|
||||
"This notebook primarily focuses on well-known online API models, so I've commented on the Groq line for text embedding. However, if you're curious, you can explore the model Groq recommends for local use by uncommenting it."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "I7a3d8EsrdBc"
|
||||
},
|
||||
"source": [
|
||||
"### Claude Anthropic"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "RHmrKCbZriz_"
|
||||
},
|
||||
"source": [
|
||||
"While Claude Anthropic isn't a free resource, they're well-known for their top-notch LLM outputs. The only tiny hiccup for this Notebook context is that they don't offer an API for text embeddings.\n",
|
||||
"\n",
|
||||
"But don't worry. In their documentation, they recommend checking out the Voyage AI embedding model. You can find all the deets right [here](https://docs.anthropic.com/en/docs/build-with-claude/embeddings).\n",
|
||||
"\n",
|
||||
"The Claude team highlights the Voyage AI embedding model on their documentation page. Voyage is an online resource, though the Claude Anthropic team does not host it.\n",
|
||||
"\n",
|
||||
"This notebook is about well-known online API models. So, I've commented on the Claude config line for text embedding. But if you're curious and want to explore the model Claude recommends, uncomment that line and explore away!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "IYpWxGBR2COD"
|
||||
},
|
||||
"source": [
|
||||
"### Voyage AI"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "j6vIcO0A2FjD"
|
||||
},
|
||||
"source": [
|
||||
"Voyage AI provides only [top-notch embeddings](https://docs.voyageai.com/docs/introduction). While it doesn't offer [its LLM model](https://docs.voyageai.com/docs/embeddings), that doesn't diminish its offerings. The cool thing? Their embeddings are versatile and can be integrated into any model you choose. So, feel free to experiment and have fun with them! That's the beauty of this notebook - it's all about exploration and learning."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "xZg16h4KHpE5"
|
||||
},
|
||||
"source": [
|
||||
"## Environment Variables\n",
|
||||
"\n",
|
||||
"This session has scripts to reset or set environment variables _(env vars)_. Using env vars is a great way to keep those sensitive API KEY values safe and sound, away from unwanted eyes."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Cqh5j0a4JcHg"
|
||||
},
|
||||
"source": [
|
||||
"### Script to Reset Environment Variables\n",
|
||||
"\n",
|
||||
"There are a bunch of reasons to reset the environment variable here. Let's go over a few:\n",
|
||||
"- **Switching API Keys:** Sometimes, we need to change to a different API key.\n",
|
||||
"- **Typos Happen:** We're all human, and typos can sneak in!\n",
|
||||
"- **Skipping Variables:** Sometimes we need to skip one or another, so we set the env var to empty.\n",
|
||||
"- **Fresh Start:** After tweaking a script or setup, it's always good to rerun with a fresh environment.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "VejHDLF1aKub"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"if 'MISTRAL_API_KEY' in os.environ:\n",
|
||||
" del os.environ['MISTRAL_API_KEY']"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "vzrNvU5XJpIm",
|
||||
"outputId": "05181b36-caa8-417d-fcd6-159e15118594"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"unset env var: GEMINI_API_KEY\n",
|
||||
"unset env var: COHERE_API_KEY\n",
|
||||
"unset env var: AWS_ACCESS_KEY_ID\n",
|
||||
"unset env var: AWS_SECRET_ACCESS_KEY\n",
|
||||
"unset env var: AWS_REGION_NAME\n",
|
||||
"unset env var: MISTRAL_API_KEY\n",
|
||||
"unset env var: GOOGLE_APPLICATION_CREDENTIALS\n",
|
||||
"unset env var: GOOGLE_VERTEX_PROJECT\n",
|
||||
"unset env var: GOOGLE_VERTEX_LOCATION\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"for ENV_VARS in [v['ENV_VAR'] for k, v in LLM_MODEL_CONFIG.items()]:\n",
|
||||
" for ENV_VAR in ENV_VARS:\n",
|
||||
" if ENV_VAR in os.environ:\n",
|
||||
" del os.environ[ENV_VAR]\n",
|
||||
" print(f'unset env var: {ENV_VAR}')\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "GqZceA9qJIdq"
|
||||
},
|
||||
"source": [
|
||||
"### Script to Set Environment Variables\n",
|
||||
"\n",
|
||||
"In this session, the script will stroll through the config dictionary to cherry-pick only the enabled LLM models.\n",
|
||||
"\n",
|
||||
"Once your favorites have been gathered, the script will prompt you to set all the necessary environment variables.\n",
|
||||
"\n",
|
||||
"The more LLM models you have enabled, the more prompts you'll see - but don't worry, each environment variable will prompt only once."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "zQmXuvbsHSwR",
|
||||
"outputId": "07e94a99-817b-4fae-e649-c6e0337bb1d1"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"enter GEMINI_API_KEY: ··········\n",
|
||||
"enter COHERE_API_KEY: ··········\n",
|
||||
"enter AWS_ACCESS_KEY_ID: ··········\n",
|
||||
"enter AWS_SECRET_ACCESS_KEY: ··········\n",
|
||||
"enter AWS_REGION_NAME: ··········\n",
|
||||
"enter MISTRAL_API_KEY: ··········\n",
|
||||
"enter GOOGLE_APPLICATION_CREDENTIALS: ··········\n",
|
||||
"enter GOOGLE_VERTEX_PROJECT: ··········\n",
|
||||
"enter GOOGLE_VERTEX_LOCATION: ··········\n",
|
||||
"enter GROQ_API_KEY: ··········\n",
|
||||
"enter OPENAI_API_KEY: ··········\n",
|
||||
"enter ANTHROPIC_API_KEY: ··········\n",
|
||||
"enter VOYAGE_API_KEY: ··········\n",
|
||||
"running custom env var setup for VERTEXAI...\n",
|
||||
"done setup for VERTEXAI\n",
|
||||
"all env var set!\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"for ENV_VAR in ENV_VARS:\n",
|
||||
" os.environ[ENV_VAR] = getpass.getpass(f\"enter {ENV_VAR}: \") if not ENV_VAR in os.environ else os.environ[ENV_VAR]\n",
|
||||
"\n",
|
||||
"LLM_MODELS_WITH_CUSTOM_SETUP = {k: v for k, v in LLM_MODELS.items() if 'ENV_VAR_SETUP' in v}\n",
|
||||
"for LLM_MODEL, LLM_CONFIG in LLM_MODELS_WITH_CUSTOM_SETUP.items():\n",
|
||||
" print(f\"running custom env var setup for {LLM_MODEL}...\")\n",
|
||||
" LLM_CONFIG['ENV_VAR_SETUP']()\n",
|
||||
" print(f\"done setup for {LLM_MODEL}\")\n",
|
||||
"\n",
|
||||
"print(\"all env var set!\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "tnTSFJZcJW5p"
|
||||
},
|
||||
"source": [
|
||||
"# Code Snippet\n",
|
||||
"\n",
|
||||
"You can run the code snippet after installing the necessary dependencies and setting up the required environment variables!\n",
|
||||
"\n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "oPq-sHZifHFZ"
|
||||
},
|
||||
"source": [
|
||||
"## Introduction\n",
|
||||
"\n",
|
||||
"This code snippet is designed to achieve only two tasks: _(i)_ run an LLM Pipeline; _(ii)_ text Embed using the LLM Model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "gKd9GB4UfJTt"
|
||||
},
|
||||
"source": [
|
||||
"The code snippet has two tasks:\n",
|
||||
"1. **Running an LLM Pipeline**:\n",
|
||||
" - To kickstart an LLM pipeline, you'll need a prompt and a model name. It's as simple as that!\n",
|
||||
"2. **Text Embedding using the LLM Model**:\n",
|
||||
" - To embed text with the LLM model, you'll need some text and the name of the text embedding model.\n",
|
||||
"\n",
|
||||
"In this script, we'll run an LLM pipeline for every given LLM and embed text using the provided LLM model.\n",
|
||||
"\n",
|
||||
"Here's the heart of the script:\n",
|
||||
"```python\n",
|
||||
"LLM_MODEL_NAME = \"model-name\"\n",
|
||||
"TEXT_EMBEDDING_PATH = \"text-embedding-name\"\n",
|
||||
"# A text prompt to send through the LLM pipeline\n",
|
||||
"LLM_PROMPT_INPUT = \"Where is one place you'd go in Washington, DC?\"\n",
|
||||
"# The embeddings dataset is versatile! It plays with lists, datasets, or even generators.\n",
|
||||
"EMBEDDING_DATA = [\n",
|
||||
" \"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",
|
||||
"]\n",
|
||||
"# LLM PIPELINE\n",
|
||||
"llm = LLM(LLM_MODEL_NAME, method=\"litellm\")\n",
|
||||
"print(llm([{\"role\": \"user\", \"content\": LLM_PROMPT_INPUT}]))\n",
|
||||
"# TEXT EMBEDDING\n",
|
||||
"embeddings = Embeddings(path=TEXT_EMBEDDING_PATH, method=\"litellm\")\n",
|
||||
"embeddings.index(EMBEDDING_DATA) # create an index for the text list\n",
|
||||
"for query in (\"feel good story\", \"climate change\"): # now, let's embark on an embeddings search for each query\n",
|
||||
" # extract the uid of the first result\n",
|
||||
" # search result format: (uid, score)\n",
|
||||
" uid = embeddings.search(query, 1)[0][0]\n",
|
||||
" # print the text\n",
|
||||
" print(\"%-20s %s\" % (query, EMBEDDING_DATA[uid]))\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"- <sup><sub>**friendly reminder**: _the library's author often [points out](https://github.com/neuml/txtai/pull/844#issuecomment-2563561232) that it's [not necessary](https://github.com/neuml/txtai/issues/843#issuecomment-2563244810) to explicitly pass the second argument `method='litellm'`. When you're learning something new, it's okay to avoid relying on shortcuts or \"magic\" until you're more comfortable. Once you understand the library better, you can start using these convenient features to your advantage. In this introduction, I'm intentionally including the second argument `method='litellm'` in the function. However, I'm choosing to leave it out in the Playground section_.</sub></sup>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "1I6o5f1ue_KX"
|
||||
},
|
||||
"source": [
|
||||
"## Playground\n",
|
||||
"\n",
|
||||
"Once you've installed the necessary dependencies and configured the environment variables, you can play with and explore the code snippet. Enjoy your coding journey! 😊\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "A3NetE3Sh0e0",
|
||||
"outputId": "d9be990a-1f57-4abc-f2be-86f8e8ef7c6c"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"--------------------------------------------------\n",
|
||||
"GEMINI\n",
|
||||
"The National Mall\n",
|
||||
"..................................................\n",
|
||||
"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",
|
||||
"public health story 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",
|
||||
"lucky Maine man wins $1M from $25 lottery ticket\n",
|
||||
"dishonest junk Make huge profits without work, earn up to $100,000 a day\n",
|
||||
"--------------------------------------------------\n",
|
||||
"COHERE\n",
|
||||
"As an AI chatbot, I do not have personal preferences or the ability to travel. However, I can suggest popular tourist destinations in Washington, DC, that many people enjoy visiting:\n",
|
||||
"\n",
|
||||
"1. The White House: The official residence and workplace of the President of the United States is a symbol of American democracy. Visitors can take a tour of the White House to learn about its history and see the beautiful architecture.\n",
|
||||
"\n",
|
||||
"2. National Mall: This iconic open park stretches from the United States Capitol to the Lincoln Memorial, featuring various monuments and memorials. Some notable landmarks include the Washington Monument, the Lincoln Memorial Reflecting Pool, the Vietnam Veterans Memorial, and the National World War II Memorial.\n",
|
||||
"\n",
|
||||
"3. Smithsonian Institution: The world's largest museum and research complex offers a wealth of knowledge and cultural experiences. It includes renowned museums such as the National Air and Space Museum, National Museum of Natural History, National Museum of African American History and Culture, and many more, all offering free admission.\n",
|
||||
"\n",
|
||||
"4. United States Capitol: A visit to the Capitol allows you to explore the seat of the United States Congress and learn about the legislative process. The Capitol also features impressive architecture and art, including the iconic Capitol Dome and the National Statuary Hall Collection.\n",
|
||||
"\n",
|
||||
"5. National Gallery of Art: Located on the National Mall, this renowned art museum houses a vast collection of paintings, sculptures, and other artworks from the Middle Ages to the present. It offers a chance to appreciate works by famous artists like Leonardo da Vinci, Rembrandt, and Claude Monet.\n",
|
||||
"\n",
|
||||
"These are just a few highlights, but Washington, DC, offers many more attractions, including historic sites, beautiful parks, vibrant neighborhoods, and cultural institutions, ensuring there's something for everyone to enjoy.\n",
|
||||
"..................................................\n",
|
||||
"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",
|
||||
"public health story 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",
|
||||
"lucky Maine man wins $1M from $25 lottery ticket\n",
|
||||
"dishonest junk US tops 5 million confirmed virus cases\n",
|
||||
"--------------------------------------------------\n",
|
||||
"AWS_BEDROCK\n",
|
||||
"I'd recommend the Smithsonian National Air and Space Museum. It is located in Washington, DC, and is the world's largest and most visited museum complex.\n",
|
||||
"..................................................\n",
|
||||
"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",
|
||||
"public health story 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",
|
||||
"lucky Maine man wins $1M from $25 lottery ticket\n",
|
||||
"dishonest junk Make huge profits without work, earn up to $100,000 a day\n",
|
||||
"--------------------------------------------------\n",
|
||||
"MISTRAL\n",
|
||||
"I would love to visit the Smithsonian National Museum of Natural History in Washington, DC. It's one of the most popular and largest museums in the world, with a vast array of exhibits showcasing natural history, including dinosaur fossils, mineral and gemstone collections, and marine life. I find the diversity and depth of knowledge presented in this museum fascinating. Additionally, it's free to the public, which makes it an accessible and enjoyable experience for everyone.\n",
|
||||
"..................................................\n",
|
||||
"Query Best Match\n",
|
||||
"..................................................\n",
|
||||
"feel good story The National Park Service warns against sacrificing slower friends in a bear attack\n",
|
||||
"climate change Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\n",
|
||||
"public health story The National Park Service warns against sacrificing slower friends in a bear attack\n",
|
||||
"war The National Park Service warns against sacrificing slower friends in a bear attack\n",
|
||||
"wildlife The National Park Service warns against sacrificing slower friends in a bear attack\n",
|
||||
"asia Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\n",
|
||||
"lucky Maine man wins $1M from $25 lottery ticket\n",
|
||||
"dishonest junk Make huge profits without work, earn up to $100,000 a day\n",
|
||||
"--------------------------------------------------\n",
|
||||
"VERTEXAI\n",
|
||||
"Washington, DC, is a fascinating city with a wealth of historical and cultural attractions. If I had the opportunity to visit, I would definitely make a stop at the Smithsonian National Air and Space Museum. This museum is home to a vast collection of aircraft and spacecraft, from the Wright brothers' first plane to the Apollo 11 command module. I am particularly interested in the history of space exploration, so I would be excited to see these iconic artifacts up close.\n",
|
||||
"\n",
|
||||
"In addition to the Air and Space Museum, there are many other places in Washington, DC, that I would like to visit. The National Mall is a beautiful park that is home to many of the city's most famous monuments, including the Washington Monument, the Lincoln Memorial, and the Vietnam Veterans Memorial. I would also like to visit the White House, the Capitol Building, and the Supreme Court. These buildings are all important symbols of American democracy, and I would be honored to have the opportunity to see them in person.\n",
|
||||
"\n",
|
||||
"Washington, DC, is a city with a rich history and a bright future. I am confident that it will continue to be a major center of culture and government for many years to come.\n",
|
||||
"..................................................\n",
|
||||
"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",
|
||||
"public health story 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",
|
||||
"lucky Maine man wins $1M from $25 lottery ticket\n",
|
||||
"dishonest junk Make huge profits without work, earn up to $100,000 a day\n",
|
||||
"--------------------------------------------------\n",
|
||||
"GROQ\n",
|
||||
"Washington, D.C. is a city with a plethora of amazing places to visit! But, if I had to pick just one place, I'd recommend the National Mall.\n",
|
||||
"\n",
|
||||
"The National Mall is a beautiful stretch of parkland that stretches from the Lincoln Memorial to the United States Capitol Building. It's home to some of the city's most iconic landmarks, including the Washington Monument, World War II Memorial, and Vietnam Veterans Memorial.\n",
|
||||
"\n",
|
||||
"You can take a leisurely stroll along the Mall, enjoy the scenic views of the city, and stop at one of the many museums or memorials along the way. The National Mall is also a great place to people-watch, with street performers and musicians adding to the lively atmosphere.\n",
|
||||
"\n",
|
||||
"Personally, I'd recommend starting at the Lincoln Memorial, where you can take in the stunning views of the Reflecting Pool and the Washington Monument. From there, you can walk to the World War II Memorial, where you can pay your respects to the millions of Americans who served in the war.\n",
|
||||
"\n",
|
||||
"Overall, the National Mall is a must-see destination in Washington, D.C. – and it's completely free!\n",
|
||||
"--------------------------------------------------\n",
|
||||
"OPENAI\n",
|
||||
"One iconic place to visit in Washington, DC, is the National Mall. It's home to many famous monuments and memorials, including the Lincoln Memorial, the Washington Monument, and the Vietnam Veterans Memorial. The expansive park offers a great opportunity to explore the history and culture of the nation, and it's a beautiful spot for walking, picnicking, and enjoying the iconic views.\n",
|
||||
"..................................................\n",
|
||||
"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",
|
||||
"public health story 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",
|
||||
"lucky Maine man wins $1M from $25 lottery ticket\n",
|
||||
"dishonest junk Make huge profits without work, earn up to $100,000 a day\n",
|
||||
"--------------------------------------------------\n",
|
||||
"CLAUDE\n",
|
||||
"One great place to visit in Washington, DC is the Smithsonian National Air and Space Museum. It's located on the National Mall and features fascinating exhibits about aviation and space exploration, including famous aircraft and spacecraft like the Wright Brothers' plane and the Apollo 11 command module. The museum is free to enter and is very popular with visitors of all ages.\n",
|
||||
"--------------------------------------------------\n",
|
||||
"VOYAGE\n",
|
||||
"..................................................\n",
|
||||
"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",
|
||||
"public health story Maine man wins $1M from $25 lottery ticket\n",
|
||||
"war The National Park Service warns against sacrificing slower friends in a bear attack\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",
|
||||
"lucky Maine man wins $1M from $25 lottery ticket\n",
|
||||
"dishonest junk Make huge profits without work, earn up to $100,000 a day\n",
|
||||
"--------------------------------------------------\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# A text prompt to run through the LLM pipeline\n",
|
||||
"# https://neuml.github.io/txtai/pipeline/text/llm/\n",
|
||||
"LLM_PROMPT_INPUT = \"Where is one place you'd go in Washington, DC?\"\n",
|
||||
"\n",
|
||||
"# The embeddings dataset is versatile! It plays with lists, datasets, or even generators.\n",
|
||||
"# https://neuml.github.io/txtai/embeddings/\n",
|
||||
"EMBEDDING_DATA = [\n",
|
||||
" \"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",
|
||||
"\n",
|
||||
"# https://neuml.github.io/txtai/pipeline/text/llm/\n",
|
||||
"def runllm(prompt=\"\", path=None):\n",
|
||||
" if path:\n",
|
||||
" # A quick note: you can skip specifying the `method` argument.\n",
|
||||
" # There's an autodetection logic designed to recognize it as a `litellm` model.\n",
|
||||
" # \n",
|
||||
" # llm = LLM(LLM_MODEL_NAME, method=\"litellm\")\n",
|
||||
" # \n",
|
||||
" llm = LLM(path)\n",
|
||||
"\n",
|
||||
" # OR: print(llm(LLM_PROMPT_INPUT, defaultrole=\"user\"))\n",
|
||||
" print(llm([{\"role\": \"user\", \"content\": prompt}]))\n",
|
||||
"\n",
|
||||
"# https://neuml.github.io/txtai/embeddings/\n",
|
||||
"def runembeddings(data=None, path=None):\n",
|
||||
" if path:\n",
|
||||
" embeddings = Embeddings(\n",
|
||||
" path=path,\n",
|
||||
" # a quick note: you can skip specifying the `method` argument - there is autodetection logic\n",
|
||||
" # method=\"litellm\"\n",
|
||||
" )\n",
|
||||
" # create an index for the list of text\n",
|
||||
" embeddings.index(data)\n",
|
||||
"\n",
|
||||
" print(\".\" * 50)\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\",\n",
|
||||
" \"public health story\", \"war\", \"wildlife\", \"asia\",\n",
|
||||
" \"lucky\", \"dishonest junk\"):\n",
|
||||
" # extract uid of first result\n",
|
||||
" # search result format: (uid, score)\n",
|
||||
" uid = embeddings.search(query, 1)[0][0]\n",
|
||||
" # print text\n",
|
||||
" print(\"%-20s %s\" % (query, data[uid]))\n",
|
||||
"\n",
|
||||
"# Let's LOOP THROUGH each enabled LLM and embedding model.\n",
|
||||
"for LLM_MODEL, LLM_CONFIG in LLM_MODELS.items():\n",
|
||||
" LLM_MODEL_NAME = LLM_CONFIG['LLM_MODEL_NAME'] if 'LLM_MODEL_NAME' in LLM_CONFIG else None\n",
|
||||
" TEXT_EMBEDDING_PATH = LLM_CONFIG['TEXT_EMBEDDING_PATH'] if 'TEXT_EMBEDDING_PATH' in LLM_CONFIG else None\n",
|
||||
" print(\"-\" * 50)\n",
|
||||
" print(LLM_MODEL)\n",
|
||||
"\n",
|
||||
" # https://neuml.github.io/txtai/pipeline/text/llm/\n",
|
||||
" runllm(prompt=LLM_PROMPT_INPUT, path=LLM_MODEL_NAME)\n",
|
||||
"\n",
|
||||
" # https://neuml.github.io/txtai/embeddings/\n",
|
||||
" runembeddings(data=EMBEDDING_DATA, path=TEXT_EMBEDDING_PATH)\n",
|
||||
"\n",
|
||||
"print(\"-\" * 50)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"collapsed_sections": [
|
||||
"lhrk9y766Lw1",
|
||||
"wuCvBPCL-VQG",
|
||||
"pgQuDv4fZgHe",
|
||||
"0KsLGZFtrZVT",
|
||||
"I7a3d8EsrdBc",
|
||||
"xZg16h4KHpE5",
|
||||
"Cqh5j0a4JcHg",
|
||||
"oPq-sHZifHFZ"
|
||||
],
|
||||
"provenance": [],
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1,275 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Chunking your data for RAG\n",
|
||||
"\n",
|
||||
"One of the major workflows in `txtai` is Retrieval Augmented Generation (RAG). Large Language Models (LLM) are built to generate coherent sounding text. While in many cases it is factually accurate, that is not what they're built to do. RAG steps in to help inject smaller pieces of knowledge into a LLM prompt and increase the overall accuracy of responses. The `R` in RAG is very important.\n",
|
||||
"\n",
|
||||
"This notebook will demonstrate how to extract, chunk and index text to support retrieval operations for RAG."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline-data]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Data chunking and indexing\n",
|
||||
"\n",
|
||||
"Let's dive right in and keep this example simple. The next section creates a [Textractor pipeline](https://neuml.github.io/txtai/pipeline/data/textractor/) and an [Embeddings database](https://neuml.github.io/txtai/embeddings/).\n",
|
||||
"\n",
|
||||
"The `Textractor` extracts chunks of text from files and the `Embeddings` takes those chunks and builds an index/database. We'll use a [late chunker](https://docs.chonkie.ai/chunkers/late-chunker) backed by [Chonkie](https://github.com/chonkie-inc/chonkie).\n",
|
||||
"\n",
|
||||
"Then, we'll build an indexing workflow that streams chunks from two files."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"from txtai.pipeline import Textractor\n",
|
||||
"\n",
|
||||
"# Text extraction pipeline with late chunking via Chonkie\n",
|
||||
"textractor = Textractor(chunker=\"late\")\n",
|
||||
"embeddings = Embeddings(content=True)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def stream():\n",
|
||||
" urls = [\"https://github.com/neuml/txtai\", \"https://arxiv.org/pdf/2005.11401\"]\n",
|
||||
" for x, url in enumerate(urls):\n",
|
||||
" chunks = textractor(url)\n",
|
||||
" \n",
|
||||
" # Add all chunks - use the same document id for each chunk\n",
|
||||
" for chunk in chunks:\n",
|
||||
" yield x, chunk\n",
|
||||
"\n",
|
||||
" # Add the document metadata with the same document id\n",
|
||||
" # Can be any metadata. Can also be the entire document.\n",
|
||||
" yield x, {\"url\": url}\n",
|
||||
"\n",
|
||||
"# Index the chunks and metadata\n",
|
||||
"embeddings.index(stream())"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"A key element of `txtai` that is commonly misunderstood is how to best store chunks of data and join them back to the main document. `txtai` allows re-using the same logical id multiple times.\n",
|
||||
"\n",
|
||||
"Behind the scenes, each chunk gets it's own unique index id. The backend database stores chunks in a table called `sections` and data in a table called `documents`. This has been the case as far back as `txtai` 4.0. `txtai` also has the ability to store associated binary data in a table called `objects`. It's important to note that each associated `document` or `object` is only stored once. \n",
|
||||
"\n",
|
||||
"To illustrate, let's look at the first 20 rows in the embeddings database created."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 63,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"{'indexid': 0, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': '**GitHub - neuml/txtai: 💡 All-in-one open-source embeddings database for semantic search, LLM orchestration and language model workflows**\\n\\n*💡 All-in-one open-source embeddings database for semantic search, LLM orchestration and language model workflows - neuml/txtai*\\n\\n\\n\\n**All-in-one embeddings database** \\ntxtai is an all-in-one embeddings database for semantic search, LLM orchestration and language model workflows.\\n\\nEmbeddings databases are a union of vector indexes (sparse and dense), graph networks and relational databases.\\n\\nThis foundation enables vector search and/or serves as a powerful knowledge source for large language model (LLM) applications.\\n\\nBuild autonomous agents, retrieval augmented generation (RAG) processes, multi-model workflows and more.\\n\\nSummary of txtai features:\\n\\n- 🔎 Vector search with SQL, object storage, topic modeling, graph analysis and multimodal indexing\\n- 📄 Create embeddings for text, documents, audio, images and video\\n- 💡 Pipelines powered by language models that run LLM prompts, question-answering, labeling, transcription, translation, summarization and more\\n- ↪️️ Workflows to join pipelines together and aggregate business logic. txtai processes can be simple microservices or multi-model workflows.\\n- 🤖 Agents that intelligently connect embeddings, pipelines, workflows and other agents together to autonomously solve complex problems\\n- ⚙️ Build with Python or YAML. API bindings available for [JavaScript](https://github.com/neuml/txtai.js) , [Java](https://github.com/neuml/txtai.java) , [Rust](https://github.com/neuml/txtai.rs) and [Go](https://github.com/neuml/txtai.go) .\\n- 🔋 Batteries included with defaults to get up and running fast\\n- ☁️ Run local or scale out with container orchestration\\ntxtai is built with Python 3.10+, [Hugging Face Transformers](https://github.com/huggingface/transformers) , [Sentence Transformers](https://github.'}\n",
|
||||
"{'indexid': 1, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi) . txtai is open-source under an Apache 2.0 license.\\n\\n*Interested in an easy and secure way to run hosted txtai applications? Then join the* [txtai.cloud](https://txtai.cloud) *preview to learn more.* \\n\\n## Why txtai?\\nNew vector databases, LLM frameworks and everything in between are sprouting up daily. Why build with txtai?\\n\\n- Up and running in minutes with [pip](https://neuml.github.io/txtai/install/) or [Docker](https://neuml.github.io/txtai/cloud/) \\n```\\n# Get started in a couple lines\\nimport txtai\\n\\nembeddings = txtai.Embeddings()\\nembeddings.index([\"Correct\", \"Not what we hoped\"])\\nembeddings.search(\"positive\", 1)\\n#[(0, 0.29862046241760254)]\\n```\\n\\n- Built-in API makes it easy to develop applications using your programming language of choice\\n```\\n# app.yml\\nembeddings:\\n path: sentence-transformers/all-MiniLM-L6-v2\\n```\\n\\n```\\nCONFIG=app.yml uvicorn \"txtai.api:app\"\\ncurl -X GET \"http://localhost:8000/search?query=positive\"\\n```\\n\\n- Run local - no need to ship data off to disparate remote services\\n- Work with micromodels all the way up to large language models (LLMs)\\n- Low footprint - install additional dependencies and scale up when needed\\n- [Learn by example](https://neuml.github.io/txtai/examples) - notebooks cover all available functionality\\n\\n## Use Cases\\nThe following sections introduce common txtai use cases. A comprehensive set of over 60 [example notebooks and applications](https://neuml.github.io/txtai/examples) are also available.\\n\\n\\n### Semantic Search\\nBuild semantic/similarity/vector/neural search applications.'}\n",
|
||||
"{'indexid': 2, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'Traditional search systems use keywords to find data. Semantic search has an understanding of natural language and identifies results that have the same meaning, not necessarily the same keywords.\\n\\nGet started with the following examples.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Introducing txtai](https://github.com/neuml/txtai/blob/master/examples/01_Introducing_txtai.ipynb) |Overview of the functionality provided by txtai||\\n|[Similarity search with images](https://github.com/neuml/txtai/blob/master/examples/13_Similarity_search_with_images.ipynb) |Embed images and text into the same space for search||\\n|[Build a QA database](https://github.com/neuml/txtai/blob/master/examples/34_Build_a_QA_database.ipynb) |Question matching with semantic search||\\n|[Semantic Graphs](https://github.com/neuml/txtai/blob/master/examples/38_Introducing_the_Semantic_Graph.ipynb) |Explore topics, data connectivity and run network analysis||\\n\\n### LLM Orchestration\\nAutonomous agents, retrieval augmented generation (RAG), chat with your data, pipelines and workflows that interface with large language models (LLMs).\\n\\nSee below to learn more.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Prompt templates and task chains](https://github.com/neuml/txtai/blob/master/examples/44_Prompt_templates_and_task_chains.ipynb) |Build model prompts and connect tasks together with workflows||\\n|[Integrate LLM frameworks](https://github.com/neuml/txtai/blob/master/examples/53_Integrate_LLM_Frameworks.ipynb) |Integrate llama.cpp, LiteLLM and custom generation frameworks||\\n|[Build knowledge graphs with LLMs](https://github.'}\n",
|
||||
"{'indexid': 3, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'com/neuml/txtai/blob/master/examples/57_Build_knowledge_graphs_with_LLM_driven_entity_extraction.ipynb) |Build knowledge graphs with LLM-driven entity extraction||\\n\\n#### Agents\\nAgents connect embeddings, pipelines, workflows and other agents together to autonomously solve complex problems.\\n\\ntxtai agents are built on top of the Transformers Agent framework. This supports all LLMs txtai supports (Hugging Face, llama.cpp, OpenAI / Claude / AWS Bedrock via LiteLLM).\\n\\nSee the link below to learn more.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Analyzing Hugging Face Posts with Graphs and Agents](https://github.com/neuml/txtai/blob/master/examples/68_Analyzing_Hugging_Face_Posts_with_Graphs_and_Agents.ipynb) |Explore a rich dataset with Graph Analysis and Agents||\\n|[Granting autonomy to agents](https://github.com/neuml/txtai/blob/master/examples/69_Granting_autonomy_to_agents.ipynb) |Agents that iteratively solve problems as they see fit||\\n|[Analyzing LinkedIn Company Posts with Graphs and Agents](https://github.com/neuml/txtai/blob/master/examples/71_Analyzing_LinkedIn_Company_Posts_with_Graphs_and_Agents.ipynb) |Exploring how to improve social media engagement with AI||\\n\\n#### Retrieval augmented generation\\nRetrieval augmented generation (RAG) reduces the risk of LLM hallucinations by constraining the output with a knowledge base as context. RAG is commonly used to \"chat with your data\".\\n\\nA novel feature of txtai is that it can provide both an answer and source citation.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Build RAG pipelines with txtai](https://github.'}\n",
|
||||
"{'indexid': 4, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'com/neuml/txtai/blob/master/examples/52_Build_RAG_pipelines_with_txtai.ipynb) |Guide on retrieval augmented generation including how to create citations||\\n|[How RAG with txtai works](https://github.com/neuml/txtai/blob/master/examples/63_How_RAG_with_txtai_works.ipynb) |Create RAG processes, API services and Docker instances||\\n|[Advanced RAG with graph path traversal](https://github.com/neuml/txtai/blob/master/examples/58_Advanced_RAG_with_graph_path_traversal.ipynb) |Graph path traversal to collect complex sets of data for advanced RAG||\\n|[Speech to Speech RAG](https://github.com/neuml/txtai/blob/master/examples/65_Speech_to_Speech_RAG.ipynb) |Full cycle speech to speech workflow with RAG||\\n\\n### Language Model Workflows\\nLanguage model workflows, also known as semantic workflows, connect language models together to build intelligent applications.\\n\\nWhile LLMs are powerful, there are plenty of smaller, more specialized models that work better and faster for specific tasks. This includes models for extractive question-answering, automatic summarization, text-to-speech, transcription and translation.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Run pipeline workflows](https://github.com/neuml/txtai/blob/master/examples/14_Run_pipeline_workflows.ipynb) |Simple yet powerful constructs to efficiently process data||\\n|[Building abstractive text summaries](https://github.com/neuml/txtai/blob/master/examples/09_Building_abstractive_text_summaries.ipynb) |Run abstractive text summarization||\\n|[Transcribe audio to text](https://github.com/neuml/txtai/blob/master/examples/11_Transcribe_audio_to_text.'}\n",
|
||||
"{'indexid': 5, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'ipynb) |Convert audio files to text||\\n|[Translate text between languages](https://github.com/neuml/txtai/blob/master/examples/12_Translate_text_between_languages.ipynb) |Streamline machine translation and language detection||\\n\\n## Installation\\nThe easiest way to install is via pip and PyPI\\n\\n```\\npip install txtai\\n```\\n\\nPython 3.10+ is supported. Using a Python [virtual environment](https://docs.python.org/3/library/venv.html) is recommended.\\n\\nSee the detailed [install instructions](https://neuml.github.io/txtai/install) for more information covering [optional dependencies](https://neuml.github.io/txtai/install/#optional-dependencies) , [environment specific prerequisites](https://neuml.github.io/txtai/install/#environment-specific-prerequisites) , [installing from source](https://neuml.github.io/txtai/install/#install-from-source) , [conda support](https://neuml.github.io/txtai/install/#conda) and how to [run with containers](https://neuml.github.io/txtai/cloud) .\\n\\n\\n## Model guide\\nSee the table below for the current recommended models. These models all allow commercial use and offer a blend of speed and performance.\\n\\n|Component|Model(s)|\\n|---|---|\\n|[Embeddings](https://neuml.github.io/txtai/embeddings) |[all-MiniLM-L6-v2](https://hf.co/sentence-transformers/all-MiniLM-L6-v2) |\\n|[Image Captions](https://neuml.github.io/txtai/pipeline/image/caption) |[BLIP](https://hf.co/Salesforce/blip-image-captioning-base) |'}\n",
|
||||
"{'indexid': 6, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': '|[Labels - Zero Shot](https://neuml.github.io/txtai/pipeline/text/labels) |[BART-Large-MNLI](https://hf.co/facebook/bart-large) |\\n|[Labels - Fixed](https://neuml.github.io/txtai/pipeline/text/labels) |Fine-tune with [training pipeline](https://neuml.github.io/txtai/pipeline/train/trainer) |\\n|[Large Language Model (LLM)](https://neuml.github.io/txtai/pipeline/text/llm) |[Llama 3.1 Instruct](https://hf.co/meta-llama/Llama-3.1-8B-Instruct) |\\n|[Summarization](https://neuml.github.io/txtai/pipeline/text/summary) |[DistilBART](https://hf.co/sshleifer/distilbart-cnn-12-6) |\\n|[Text-to-Speech](https://neuml.github.io/txtai/pipeline/audio/texttospeech) |[ESPnet JETS](https://hf.co/NeuML/ljspeech-jets-onnx) |\\n|[Transcription](https://neuml.github.io/txtai/pipeline/audio/transcription) |[Whisper](https://hf.co/openai/whisper-base) |\\n|[Translation](https://neuml.github.io/txtai/pipeline/text/translation) |[OPUS Model Series](https://hf.co/Helsinki-NLP) |\\nModels can be loaded as either a path from the Hugging Face Hub or a local directory. Model paths are optional, defaults are loaded when not specified. For tasks with no recommended model, txtai uses the default models as shown in the Hugging Face Tasks guide.\\n\\nSee the following links to learn more.\\n\\n\\n## Powered by txtai\\nThe following applications are powered by txtai.'}\n",
|
||||
"{'indexid': 7, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': \"|Application|Description|\\n|---|---|\\n|[rag](https://github.com/neuml/rag) |Retrieval Augmented Generation (RAG) application|\\n|[ragdata](https://github.com/neuml/ragdata) |Build knowledge bases for RAG|\\n|[paperai](https://github.com/neuml/paperai) |Semantic search and workflows for medical/scientific papers|\\n|[annotateai](https://github.com/neuml/annotateai) |Automatically annotate papers with LLMs|\\nIn addition to this list, there are also many other [open-source projects](https://github.com/neuml/txtai/network/dependents) , [published research](https://scholar.google.com/scholar?q=txtai&hl=en&as_ylo=2022) and closed proprietary/commercial projects that have built on txtai in production.\\n\\n\\n## Further Reading\\n- [Tutorial series on Hashnode](https://neuml.hashnode.dev/series/txtai-tutorial) | [dev.to](https://dev.to/neuml/tutorial-series-on-txtai-ibg) \\n- [What's new in txtai 8.0](https://medium.com/neuml/whats-new-in-txtai-8-0-2d7d0ab4506b) | [7.0](https://medium.com/neuml/whats-new-in-txtai-7-0-855ad6a55440) | [6.0](https://medium.com/neuml/whats-new-in-txtai-6-0-7d93eeedf804) | [5.0](https://medium.com/neuml/whats-new-in-txtai-5-0-e5c75a13b101) | [4.0](https://medium.\"}\n",
|
||||
"{'indexid': 8, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'com/neuml/whats-new-in-txtai-4-0-bbc3a65c3d1c) \\n- [Getting started with semantic search](https://medium.com/neuml/getting-started-with-semantic-search-a9fd9d8a48cf) | [workflows](https://medium.com/neuml/getting-started-with-semantic-workflows-2fefda6165d9) | [rag](https://medium.com/neuml/getting-started-with-rag-9a0cca75f748) \\n\\n## Documentation\\n[Full documentation on txtai](https://neuml.github.io/txtai) including configuration settings for embeddings, pipelines, workflows, API and a FAQ with common questions/issues is available.\\n\\n\\n## Contributing\\nFor those who would like to contribute to txtai, please see [this guide](https://github.com/neuml/.github/blob/master/CONTRIBUTING.md) .'}\n",
|
||||
"{'indexid': 9, 'id': '1', 'url': 'https://arxiv.org/pdf/2005.11401', 'text': 'Retrieval-Augmented Generation for\\nKnowledge-Intensive NLP Tasks\\n\\nPatrick Lewis†‡, Ethan Perez?,\\n\\nAleksandra Piktus†, Fabio Petroni†, Vladimir Karpukhin†, Naman Goyal†, Heinrich Küttler†,\\n\\nMike Lewis†, Wen-tau Yih†, Tim Rocktäschel†‡, Sebastian Riedel†‡, Douwe Kiela†\\n\\n†Facebook AI Research; ‡University College London; ?New York University;\\nplewis@fb.com\\n\\nAbstract\\n\\nLarge pre-trained language models have been shown to store factual knowledge\\nin their parameters, and achieve state-of-the-art results when fine-tuned on down-\\nstream NLP tasks. However, their ability to access and precisely manipulate knowl-\\nedge is still limited, and hence on knowledge-intensive tasks, their performance\\nlags behind task-specific architectures. Additionally, providing provenance for their\\ndecisions and updating their world knowledge remain open research problems. Pre-\\ntrained models with a differentiable access mechanism to explicit non-parametric\\nmemory have so far been only investigated for extractive downstream tasks. We\\nexplore a general-purpose fine-tuning recipe for retrieval-augmented generation\\n(RAG) — models which combine pre-trained parametric and non-parametric mem-\\nory for language generation. We introduce RAG models where the parametric\\nmemory is a pre-trained seq2seq model and the non-parametric memory is a dense\\nvector index of Wikipedia, accessed with a pre-trained neural retriever. We com-\\npare two RAG formulations, one which conditions on the same retrieved passages\\nacross the whole generated sequence, and another which can use different passages\\nper token. We fine-tune and evaluate our models on a wide range of knowledge-\\nintensive NLP tasks and set the state of the art on three open domain QA tasks,\\noutperforming parametric seq2seq models and task-specific retrieve-and-extract\\narchitectures. For language generation tasks, we find that RAG models generate\\nmore specific, diverse and factual language than a state-of-the-art parametric-only\\nseq2seq baseline.\\n\\n1 Introduction\\n\\nPre-trained neural language models have been shown to learn a substantial amount of in-depth knowl-\\nedge from data [47]. They can do so without any access to an external memory, as a parameterized\\nimplicit knowledge base [51, 52].'}\n",
|
||||
"{'indexid': 10, 'id': '1', 'url': 'https://arxiv.org/pdf/2005.11401', 'text': 'While this development is exciting, such models do have down-\\nsides: They cannot easily expand or revise their memory, can’t straightforwardly provide insight into\\ntheir predictions, and may produce “hallucinations” [38]. Hybrid models that combine parametric\\nmemory with non-parametric (i.e., retrieval-based) memories [20, 26, 48] can address some of these\\nissues because knowledge can be directly revised and expanded, and accessed knowledge can be\\ninspected and interpreted. REALM [20] and ORQA [31], two recently introduced models that\\ncombine masked language models [8] with a differentiable retriever, have shown promising results,\\n\\nar\\nX\\n\\niv\\n:2\\n\\n00\\n5.\\n\\n11\\n40\\n\\n1v\\n4 \\n\\n [\\ncs\\n\\n.C\\nL\\n\\n] \\n 1\\n\\n2 \\nA\\n\\npr\\n 2\\n\\n02\\n1\\n\\n\\nThe\\tDivine\\nComedy\\t(x) q\\n\\nQuery\\nEncoder\\n\\nq(x)\\n\\nMIPS pθ\\n\\nGenerator\\xa0pθ\\n(Parametric)\\n\\nMargin-\\nalize\\n\\nThis\\t14th\\tcentury\\twork\\nis\\tdivided\\tinto\\t3\\nsections:\\t\"Inferno\",\\n\"Purgatorio\"\\t&\\n\"Paradiso\"\\t\\t\\t\\t\\t\\t\\t\\t\\t(y)\\n\\nEnd-to-End Backprop through q and\\xa0pθ\\n\\nBarack\\tObama\\twas\\nborn\\tin\\tHawaii.(x)\\n\\nFact Verification: Fact Query\\n\\nsupports\\t(y)\\n\\nQuestion Generation\\n\\nFact Verification:\\nLabel Generation\\n\\nDocument\\nIndex\\n\\nDefine\\t\"middle\\tear\"(x)\\n\\nQuestion Answering:\\nQuestion Query\\n\\nThe\\tmiddle\\tear\\tincludes\\nthe\\ttympanic\\tcavity\\tand\\nthe\\tthree\\tossicles.\\t\\t(y)\\n\\nQuestion Answering:\\nAnswer GenerationRetriever pη\\n\\n(Non-Parametric)\\nz4\\n\\nz3\\nz2\\n\\nz1\\n\\nd(z)\\n\\nJeopardy Question\\nGeneration:\\n\\nAnswer Query\\n\\nFigure 1: Overview of our approach. We combine a pre-trained retriever (Query Encoder + Document\\nIndex) with a pre-trained seq2seq model (Generator) and fine-tune end-to-end. For query x, we use\\nMaximum Inner Product Search (MIPS) to find the top-K documents zi. For final prediction y, we\\ntreat z as a latent variable and marginalize over seq2seq predictions given different documents.\\n\\nbut have only explored open-domain extractive question answering. Here, we bring hybrid parametric\\nand non-parametric memory to the “workhorse of NLP,” i.e. sequence-to-sequence (seq2seq) models.\\n\\nWe endow pre-trained, parametric-memory generation models with a non-parametric memory through'}\n",
|
||||
"{'indexid': 11, 'id': '1', 'url': 'https://arxiv.org/pdf/2005.11401', 'text': 'a general-purpose fine-tuning approach which we refer to as retrieval-augmented generation (RAG).\\nWe build RAG models where the parametric memory is a pre-trained seq2seq transformer, and the\\nnon-parametric memory is a dense vector index of Wikipedia, accessed with a pre-trained neural\\nretriever. We combine these components in a probabilistic model trained end-to-end (Fig. 1). The\\nretriever (Dense Passage Retriever [26], henceforth DPR) provides latent documents conditioned on\\nthe input, and the seq2seq model (BART [32]) then conditions on these latent documents together with\\nthe input to generate the output. We marginalize the latent documents with a top-K approximation,\\neither on a per-output basis (assuming the same document is responsible for all tokens) or a per-token\\nbasis (where different documents are responsible for different tokens). Like T5 [51] or BART, RAG\\ncan be fine-tuned on any seq2seq task, whereby both the generator and retriever are jointly learned.\\n\\nThere has been extensive previous work proposing architectures to enrich systems with non-parametric\\nmemory which are trained from scratch for specific tasks, e.g. memory networks [64, 55], stack-\\naugmented networks [25] and memory layers [30]. In contrast, we explore a setting where both\\nparametric and non-parametric memory components are pre-trained and pre-loaded with extensive\\nknowledge. Crucially, by using pre-trained access mechanisms, the ability to access knowledge is\\npresent without additional training.\\n\\nOur results highlight the benefits of combining parametric and non-parametric memory with genera-\\ntion for knowledge-intensive tasks—tasks that humans could not reasonably be expected to perform\\nwithout access to an external knowledge source. Our RAG models achieve state-of-the-art results\\non open Natural Questions [29], WebQuestions [3] and CuratedTrec [2] and strongly outperform\\nrecent approaches that use specialised pre-training objectives on TriviaQA [24]. Despite these being'}\n",
|
||||
"{'indexid': 12, 'id': '1', 'url': 'https://arxiv.org/pdf/2005.11401', 'text': 'extractive tasks, we find that unconstrained generation outperforms previous extractive approaches.\\nFor knowledge-intensive generation, we experiment with MS-MARCO [1] and Jeopardy question\\ngeneration, and we find that our models generate responses that are more factual, specific, and\\ndiverse than a BART baseline. For FEVER [56] fact verification, we achieve results within 4.3% of\\nstate-of-the-art pipeline models which use strong retrieval supervision. Finally, we demonstrate that\\nthe non-parametric memory can be replaced to update the models’ knowledge as the world changes.1\\n\\n2 Methods\\n\\nWe explore RAG models, which use the input sequence x to retrieve text documents z and use them\\nas additional context when generating the target sequence y. As shown in Figure 1, our models\\nleverage two components: (i) a retriever pη(z|x) with parameters η that returns (top-K truncated)\\ndistributions over text passages given a query x and (ii) a generator pθ(yi|x, z, y1:i−1) parametrized\\n\\n1Code to run experiments with RAG has been open-sourced as part of the HuggingFace Transform-\\ners Library [66] and can be found at https://github.com/huggingface/transformers/blob/master/\\nexamples/rag/. An interactive demo of RAG models can be found at https://huggingface.co/rag/\\n\\n2\\n\\n[https://github.com/huggingface/transformers/blob/master/examples/rag/](https://github.com/huggingface/transformers/blob/master/examples/rag/) \\n[https://github.com/huggingface/transformers/blob/master/examples/rag/](https://github.com/huggingface/transformers/blob/master/examples/rag/) \\n[https://huggingface.co/rag/](https://huggingface.co/rag/) \\n\\nby θ that generates a current token based on a context of the previous i− 1 tokens y1:i−1, the original\\ninput x and a retrieved passage z.\\n\\nTo train the retriever and generator end-to-end, we treat the retrieved document as a latent variable.\\nWe propose two models that marginalize over the latent documents in different ways to produce a'}\n",
|
||||
"{'indexid': 13, 'id': '1', 'url': 'https://arxiv.org/pdf/2005.11401', 'text': 'distribution over generated text. In one approach, RAG-Sequence, the model uses the same document\\nto predict each target token. The second approach, RAG-Token, can predict each target token based\\non a different document. In the following, we formally introduce both models and then describe the\\npη and pθ components, as well as the training and decoding procedure.\\n\\n2.1 Models\\n\\nRAG-Sequence Model The RAG-Sequence model uses the same retrieved document to generate\\nthe complete sequence. Technically, it treats the retrieved document as a single latent variable that\\nis marginalized to get the seq2seq probability p(y|x) via a top-K approximation. Concretely, the\\ntop K documents are retrieved using the retriever, and the generator produces the output sequence\\nprobability for each document, which are then marginalized,\\n\\npRAG-Sequence(y|x) ≈\\n∑\\n\\nz∈top-k(p(·|x))\\n\\npη(z|x)pθ(y|x, z) =\\n∑\\n\\nz∈top-k(p(·|x))\\n\\npη(z|x)\\nN∏\\ni\\n\\npθ(yi|x, z, y1:i−1)\\n\\nRAG-Token Model In the RAG-Token model we can draw a different latent document for each\\ntarget token and marginalize accordingly. This allows the generator to choose content from several\\ndocuments when producing an answer. Concretely, the top K documents are retrieved using the\\nretriever, and then the generator produces a distribution for the next output token for each document,\\nbefore marginalizing, and repeating the process with the following output token, Formally, we define:\\n\\npRAG-Token(y|x) ≈\\nN∏\\ni\\n\\n∑\\nz∈top-k(p(·|x))\\n\\npη(z|x)pθ(yi|x, z, y1:i−1)\\n\\nFinally, we note that RAG can be used for sequence classification tasks by considering the target class\\nas a target sequence of length one, in which case RAG-Sequence and RAG-Token are equivalent.\\n\\n2.2 Retriever: DPR\\n\\nThe retrieval component pη(z|x) is based on DPR [26]. DPR follows a bi-encoder architecture:\\n\\npη(z|x) ∝ exp\\n(\\nd(z)>q(x)\\n\\n)\\nd(z) = BERTd(z), q(x) = BERTq(x)'}\n",
|
||||
"{'indexid': 14, 'id': '1', 'url': 'https://arxiv.org/pdf/2005.11401', 'text': 'where d(z) is a dense representation of a document produced by a BERTBASE document encoder [8],\\nand q(x) a query representation produced by a query encoder, also based on BERTBASE. Calculating\\ntop-k(pη(·|x)), the list of k documents z with highest prior probability pη(z|x), is a Maximum Inner\\nProduct Search (MIPS) problem, which can be approximately solved in sub-linear time [23]. We use\\na pre-trained bi-encoder from DPR to initialize our retriever and to build the document index. This\\nretriever was trained to retrieve documents which contain answers to TriviaQA [24] questions and\\nNatural Questions [29]. We refer to the document index as the non-parametric memory.\\n\\n2.3 Generator: BART\\n\\nThe generator component pθ(yi|x, z, y1:i−1) could be modelled using any encoder-decoder. We use\\nBART-large [32], a pre-trained seq2seq transformer [58] with 400M parameters. To combine the input\\nx with the retrieved content z when generating from BART, we simply concatenate them. BART was\\npre-trained using a denoising objective and a variety of different noising functions. It has obtained\\nstate-of-the-art results on a diverse set of generation tasks and outperforms comparably-sized T5\\nmodels [32]. We refer to the BART generator parameters θ as the parametric memory henceforth.\\n\\n2.4 Training\\n\\nWe jointly train the retriever and generator components without any direct supervision on what\\ndocument should be retrieved. Given a fine-tuning training corpus of input/output pairs (xj , yj), we\\n\\n3\\n\\n\\nminimize the negative marginal log-likelihood of each target,\\n∑\\nj − log p(yj |xj) using stochastic\\n\\ngradient descent with Adam [28]. Updating the document encoder BERTd during training is costly as\\nit requires the document index to be periodically updated as REALM does during pre-training [20].'}\n",
|
||||
"{'indexid': 15, 'id': '1', 'url': 'https://arxiv.org/pdf/2005.11401', 'text': 'We do not find this step necessary for strong performance, and keep the document encoder (and\\nindex) fixed, only fine-tuning the query encoder BERTq and the BART generator.\\n\\n2.5 Decoding\\n\\nAt test time, RAG-Sequence and RAG-Token require different ways to approximate argmaxy p(y|x).\\n\\nRAG-Token The RAG-Token model can be seen as a standard, autoregressive seq2seq genera-\\ntor with transition probability: p′θ(yi|x, y1:i−1) =\\n\\n∑\\nz∈top-k(p(·|x)) pη(zi|x)pθ(yi|x, zi, y1:i−1) To\\n\\ndecode, we can plug p′θ(yi|x, y1:i−1) into a standard beam decoder.\\n\\nRAG-Sequence For RAG-Sequence, the likelihood p(y|x) does not break into a conventional per-\\ntoken likelihood, hence we cannot solve it with a single beam search. Instead, we run beam search for\\neach document z, scoring each hypothesis using pθ(yi|x, z, y1:i−1). This yields a set of hypotheses\\nY , some of which may not have appeared in the beams of all documents. To estimate the probability\\nof an hypothesis y we run an additional forward pass for each document z for which y does not\\nappear in the beam, multiply generator probability with pη(z|x) and then sum the probabilities across\\nbeams for the marginals. We refer to this decoding procedure as “Thorough Decoding.” For longer\\noutput sequences, |Y | can become large, requiring many forward passes. For more efficient decoding,\\nwe can make a further approximation that pθ(y|x, zi) ≈ 0 where y was not generated during beam\\nsearch from x, zi. This avoids the need to run additional forward passes once the candidate set Y has\\nbeen generated. We refer to this decoding procedure as “Fast Decoding.”\\n\\n3 Experiments\\n\\nWe experiment with RAG in a wide range of knowledge-intensive tasks. For all experiments, we use\\na single Wikipedia dump for our non-parametric knowledge source. Following Lee et al. [31] and\\nKarpukhin et al. [26], we use the December 2018 dump. Each Wikipedia article is split into disjoint'}\n",
|
||||
"{'indexid': 16, 'id': '1', 'url': 'https://arxiv.org/pdf/2005.11401', 'text': '100-word chunks, to make a total of 21M documents. We use the document encoder to compute an\\nembedding for each document, and build a single MIPS index using FAISS [23] with a Hierarchical\\nNavigable Small World approximation for fast retrieval [37]. During training, we retrieve the top\\nk documents for each query. We consider k ∈ {5, 10} for training and set k for test time using dev\\ndata. We now discuss experimental details for each task.\\n\\n3.1 Open-domain Question Answering\\n\\nOpen-domain question answering (QA) is an important real-world application and common testbed\\nfor knowledge-intensive tasks [20]. We treat questions and answers as input-output text pairs (x, y)\\nand train RAG by directly minimizing the negative log-likelihood of answers. We compare RAG to\\nthe popular extractive QA paradigm [5, 7, 31, 26], where answers are extracted spans from retrieved\\ndocuments, relying primarily on non-parametric knowledge. We also compare to “Closed-Book\\nQA” approaches [52], which, like RAG, generate answers, but which do not exploit retrieval, instead\\nrelying purely on parametric knowledge. We consider four popular open-domain QA datasets: Natural\\nQuestions (NQ) [29], TriviaQA (TQA) [24]. WebQuestions (WQ) [3] and CuratedTrec (CT) [2]. As\\nCT and WQ are small, we follow DPR [26] by initializing CT and WQ models with our NQ RAG\\nmodel. We use the same train/dev/test splits as prior work [31, 26] and report Exact Match (EM)\\nscores. For TQA, to compare with T5 [52], we also evaluate on the TQA Wiki test set.\\n\\n3.2 Abstractive Question Answering\\n\\nRAG models can go beyond simple extractive QA and answer questions with free-form, abstractive\\ntext generation. To test RAG’s natural language generation (NLG) in a knowledge-intensive setting,\\nwe use the MSMARCO NLG task v2.1 [43]. The task consists of questions, ten gold passages'}\n",
|
||||
"{'indexid': 17, 'id': '1', 'url': 'https://arxiv.org/pdf/2005.11401', 'text': 'retrieved from a search engine for each question, and a full sentence answer annotated from the\\nretrieved passages. We do not use the supplied passages, only the questions and answers, to treat\\n\\n4\\n\\n\\nMSMARCO as an open-domain abstractive QA task. MSMARCO has some questions that cannot be\\nanswered in a way that matches the reference answer without access to the gold passages, such as\\n“What is the weather in Volcano, CA?” so performance will be lower without using gold passages.\\nWe also note that some MSMARCO questions cannot be answered using Wikipedia alone. Here,\\nRAG can rely on parametric knowledge to generate reasonable responses.\\n\\n3.3 Jeopardy Question Generation\\n\\nTo evaluate RAG’s generation abilities in a non-QA setting, we study open-domain question gen-\\neration. Rather than use questions from standard open-domain QA tasks, which typically consist\\nof short, simple questions, we propose the more demanding task of generating Jeopardy questions.\\nJeopardy is an unusual format that consists of trying to guess an entity from a fact about that entity.\\nFor example, “The World Cup” is the answer to the question “In 1986 Mexico scored as the first\\ncountry to host this international sports competition twice.” As Jeopardy questions are precise,\\nfactual statements, generating Jeopardy questions conditioned on their answer entities constitutes a\\nchallenging knowledge-intensive generation task.\\n\\nWe use the splits from SearchQA [10], with 100K train, 14K dev, and 27K test examples. As\\nthis is a new task, we train a BART model for comparison. Following [67], we evaluate using the\\nSQuAD-tuned Q-BLEU-1 metric [42]. Q-BLEU is a variant of BLEU with a higher weight for\\nmatching entities and has higher correlation with human judgment for question generation than\\nstandard metrics. We also perform two human evaluations, one to assess generation factuality, and\\none for specificity. We define factuality as whether a statement can be corroborated by trusted external'}\n",
|
||||
"{'indexid': 18, 'id': '1', 'url': 'https://arxiv.org/pdf/2005.11401', 'text': 'sources, and specificity as high mutual dependence between the input and output [33]. We follow\\nbest practice and use pairwise comparative evaluation [34]. Evaluators are shown an answer and two\\ngenerated questions, one from BART and one from RAG. They are then asked to pick one of four\\noptions—quuestion A is better, question B is better, both are good, or neither is good.\\n\\n3.4 Fact Verification\\n\\nFEVER [56] requires classifying whether a natural language claim is supported or refuted by\\nWikipedia, or whether there is not enough information to decide. The task requires retrieving\\nevidence from Wikipedia relating to the claim and then reasoning over this evidence to classify\\nwhether the claim is true, false, or unverifiable from Wikipedia alone. FEVER is a retrieval problem\\ncoupled with an challenging entailment reasoning task. It also provides an appropriate testbed for\\nexploring the RAG models’ ability to handle classification rather than generation. We map FEVER\\nclass labels (supports, refutes, or not enough info) to single output tokens and directly train with\\nclaim-class pairs. Crucially, unlike most other approaches to FEVER, we do not use supervision on\\nretrieved evidence. In many real-world applications, retrieval supervision signals aren’t available, and\\nmodels that do not require such supervision will be applicable to a wider range of tasks. We explore\\ntwo variants: the standard 3-way classification task (supports/refutes/not enough info) and the 2-way\\n(supports/refutes) task studied in Thorne and Vlachos [57]. In both cases we report label accuracy.\\n\\n4 Results\\n\\n4.1 Open-domain Question Answering\\n\\nTable 1 shows results for RAG along with state-of-the-art models. On all four open-domain QA\\ntasks, RAG sets a new state of the art (only on the T5-comparable split for TQA). RAG combines\\nthe generation flexibility of the “closed-book” (parametric only) approaches and the performance of\\n\"open-book\" retrieval-based approaches. Unlike REALM and T5+SSM, RAG enjoys strong results\\nwithout expensive, specialized “salient span masking” pre-training [20]. It is worth noting that RAG’s\\nretriever is initialized using DPR’s retriever, which uses retrieval supervision on Natural Questions\\nand TriviaQA. RAG compares favourably to the DPR QA system, which uses a BERT-based “cross-'}\n",
|
||||
"{'indexid': 19, 'id': '1', 'url': 'https://arxiv.org/pdf/2005.11401', 'text': 'encoder” to re-rank documents, along with an extractive reader. RAG demonstrates that neither a\\nre-ranker nor extractive reader is necessary for state-of-the-art performance.\\n\\nThere are several advantages to generating answers even when it is possible to extract them. Docu-\\nments with clues about the answer but do not contain the answer verbatim can still contribute towards\\na correct answer being generated, which is not possible with standard extractive approaches, leading\\n\\n5\\n\\n\\nTable 1: Open-Domain QA Test Scores. For TQA,\\nleft column uses the standard test set for Open-\\nDomain QA, right column uses the TQA-Wiki\\ntest set. See Appendix D for further details.\\n\\nModel NQ TQA WQ CT\\n\\nClosed\\nBook\\n\\nT5-11B [52] 34.5 - /50.1 37.4 -\\nT5-11B+SSM[52] 36.6 - /60.5 44.7 -\\n\\nOpen\\nBook\\n\\nREALM [20] 40.4 - / - 40.7 46.8\\nDPR [26] 41.5 57.9/ - 41.1 50.6\\n\\nRAG-Token 44.1 55.2/66.1 45.5 50.0\\nRAG-Seq. 44.5 56.8/68.0 45.2 52.2\\n\\nTable 2: Generation and classification Test Scores.\\nMS-MARCO SotA is [4], FEVER-3 is [68] and\\nFEVER-2 is [57] *Uses gold context/evidence.\\nBest model without gold access underlined.\\n\\nModel Jeopardy MSMARCO FVR3 FVR2\\nB-1 QB-1 R-L B-1 Label Acc.\\n\\nSotA - - 49.8* 49.9* 76.8 92.2*\\n\\nBART 15.1 19.7 38.2 41.6 64.0 81.1\\n\\nRAG-Tok. 17.3 22.2 40.1 41.5 72.5 89.5RAG-Seq. 14.7 21.4 40.8 44.2\\n\\nto more effective marginalization over documents. Furthermore, RAG can generate correct answers\\neven when the correct answer is not in any retrieved document, achieving 11.8% accuracy in such\\ncases for NQ, where an extractive model would score 0%.\\n\\n4.2 Abstractive Question Answering'}\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"for x in embeddings.search(\"SELECT indexid, id, url, text from txtai\", 20):\n",
|
||||
" print(x)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Note that the id/metadata is the same but the indexid and chunk text change with each row."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Retrieval\n",
|
||||
"\n",
|
||||
"Last thing here is to illustrate a couple retrieval operations. LLMs are great at generating answers when we properly bound the context. See the two examples below."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 61,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"including less of emphasis on lightly editing a retrieved item, but on aggregating content from several\n",
|
||||
"pieces of retrieved content, as well as learning latent retrieval, and retrieving evidence documents\n",
|
||||
"rather than related training pairs. This said, RAG techniques may work well in these settings, and\n",
|
||||
"could represent promising future work.\n",
|
||||
"\n",
|
||||
"6 Discussion\n",
|
||||
"\n",
|
||||
"In this work, we presented hybrid generation models with access to parametric and non-parametric\n",
|
||||
"memory. We showed that our RAG models obtain state of the art results on open-domain QA. We\n",
|
||||
"found that people prefer RAG’s generation over purely parametric BART, finding RAG more factual\n",
|
||||
"and specific. We conducted an thorough investigation of the learned retrieval component, validating\n",
|
||||
"its effectiveness, and we illustrated how the retrieval index can be hot-swapped to update the model\n",
|
||||
"without requiring any retraining. In future work, it may be fruitful to investigate if the two components\n",
|
||||
"can be jointly pre-trained from scratch, either with a denoising objective similar to BART or some\n",
|
||||
"another objective. Our work opens up new research directions on how parametric and non-parametric\n",
|
||||
"memories interact and how to most effectively combine them, showing promise in being applied to a\n",
|
||||
"wide variety of NLP tasks.\n",
|
||||
"\n",
|
||||
"9\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"Broader Impact\n",
|
||||
"\n",
|
||||
"This work offers several positive societal benefits over previous work: the fact that it is more\n",
|
||||
"strongly grounded in real factual knowledge (in this case Wikipedia) makes it “hallucinate” less\n",
|
||||
"with generations that are more factual, and offers more control and interpretability. RAG could be\n",
|
||||
"employed in a wide variety of scenarios with direct benefit to society, for example by endowing it\n",
|
||||
"with a medical index and asking it open-domain questions on that topic, or by helping people be more\n",
|
||||
"effective at their jobs.\n",
|
||||
"\n",
|
||||
"With these advantages also come potential downsides: Wikipedia, or any potential external knowledge\n",
|
||||
"source, will probably never be entirely factual and completely devoid of bias. Since RAG can be\n",
|
||||
"employed as a language model, similar concerns as for GPT-2 [50] are valid here, although arguably\n",
|
||||
"to a lesser extent, including that it might be used to generate abuse, faked or misleading content in\n",
|
||||
"the news or on social media; to impersonate others; or to automate the production of spam/phishing\n",
|
||||
"content [54]. Advanced language models may also lead to the automation of various jobs in the\n",
|
||||
"coming decades [16].\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"print(embeddings.search(\"What is it called when LLM generation is bounded with factually correct data?\", 1)[0][\"text\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 62,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Traditional search systems use keywords to find data. Semantic search has an understanding of natural language and identifies results that have the same meaning, not necessarily the same keywords.\n",
|
||||
"\n",
|
||||
"Get started with the following examples.\n",
|
||||
"\n",
|
||||
"|Notebook|Description||\n",
|
||||
"|---|---|---|\n",
|
||||
"|[Introducing txtai](https://github.com/neuml/txtai/blob/master/examples/01_Introducing_txtai.ipynb) |Overview of the functionality provided by txtai||\n",
|
||||
"|[Similarity search with images](https://github.com/neuml/txtai/blob/master/examples/13_Similarity_search_with_images.ipynb) |Embed images and text into the same space for search||\n",
|
||||
"|[Build a QA database](https://github.com/neuml/txtai/blob/master/examples/34_Build_a_QA_database.ipynb) |Question matching with semantic search||\n",
|
||||
"|[Semantic Graphs](https://github.com/neuml/txtai/blob/master/examples/38_Introducing_the_Semantic_Graph.ipynb) |Explore topics, data connectivity and run network analysis||\n",
|
||||
"\n",
|
||||
"### LLM Orchestration\n",
|
||||
"Autonomous agents, retrieval augmented generation (RAG), chat with your data, pipelines and workflows that interface with large language models (LLMs).\n",
|
||||
"\n",
|
||||
"See below to learn more.\n",
|
||||
"\n",
|
||||
"|Notebook|Description||\n",
|
||||
"|---|---|---|\n",
|
||||
"|[Prompt templates and task chains](https://github.com/neuml/txtai/blob/master/examples/44_Prompt_templates_and_task_chains.ipynb) |Build model prompts and connect tasks together with workflows||\n",
|
||||
"|[Integrate LLM frameworks](https://github.com/neuml/txtai/blob/master/examples/53_Integrate_LLM_Frameworks.ipynb) |Integrate llama.cpp, LiteLLM and custom generation frameworks||\n",
|
||||
"|[Build knowledge graphs with LLMs](https://github.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"print(embeddings.search(\"Tell me about semantic search\", 1)[0][\"text\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Note how both answers give more than enough information for a LLM to answer the question."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook covered how to build a retrieval system for RAG with txtai. Chunking and retrieval are key pieces of a RAG system, arguably the most important. With the commoditization of LLMs, it's going to be more and more important on how data is presented to LLMs. When given concise information, LLMs can take it from there!\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"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": 2
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,454 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "666da745",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Medical RAG Research with txtai\n",
|
||||
"\n",
|
||||
"Large Language Models (LLMs) have captured the public's attention with their impressive capabilities. The Generative AI era has reached a fever pitch with some predicting the coming rise of superintelligence.\n",
|
||||
"\n",
|
||||
"LLMs are far from perfect though and we're still a ways away from true AI. One big challenge is with hallucinations. Hallucinations is the term for when an LLM generates output that is factually incorrect. The alarming part of this is that on a cursory glance, it actually sounds like factual content. The default behavior of LLMs is to produce plausible answers even when no plausible answer exists. LLMs are not great at saying I don't know.\n",
|
||||
"\n",
|
||||
"Retrieval Augmented Generation (RAG) helps reduce the risk of hallucinations by limiting the context in which a LLM can generate answers. This is typically done with a search query that hydrates a prompt with a relevant context. RAG has been one of the most practical use cases of the Generative AI era.\n",
|
||||
"\n",
|
||||
"This notebook will demonstrate how to build a Medical RAG Research process with txtai."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a4c7cc15",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "a246bb3a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "fc5c27ed",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Medical Dataset\n",
|
||||
"\n",
|
||||
"For this example, we'll use a [PubMed subset of article metadata for H5N1](https://huggingface.co/datasets/NeuML/pubmed-h5n1). This dataset was created using [`paperetl`](https://github.com/neuml/paperetl), an open-source library for parsing medical and scientific papers.\n",
|
||||
"\n",
|
||||
"[PubMed](https://pubmed.ncbi.nlm.nih.gov/) has over 38 million article abstracts as of June 2025. `paperetl` supports loading the full dataset with all 38 million articles or just a smaller subset. The dataset link above has more details on how this can be changed for different codes and keywords. This link also has information on how the article abstracts can be loaded in addition to the metadata."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 80,
|
||||
"id": "5dd4145d",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"ds = load_dataset(\"neuml/pubmed-h5n1\", split=\"train\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "2d3b1cd8",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Next, we'll build a `txtai` embeddings index with the articles. We'll use a vector embeddings model that specializes in vectorizing medical papers: [PubMedBERT Embeddings](https://huggingface.co/NeuML/pubmedbert-base-embeddings). "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 81,
|
||||
"id": "04e69829",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"embeddings = Embeddings(path=\"neuml/pubmedbert-base-embeddings\", content=True, columns={\"text\": \"title\"})\n",
|
||||
"embeddings.index(x for x in ds if x[\"title\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 82,
|
||||
"id": "a801391e",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"7865"
|
||||
]
|
||||
},
|
||||
"execution_count": 82,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings.count()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "907898a4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# RAG Pipeline\n",
|
||||
"\n",
|
||||
"There are a number of [prior examples](https://neuml.github.io/txtai/examples/#llm) on how to run RAG with `txtai`. The [RAG pipeline](https://neuml.github.io/txtai/pipeline/text/rag/) takes two main parameters, an embeddings database and an LLM. The embeddings database is the one just created above. For this example, we'll use a [simple local LLM with 600M parameters](https://huggingface.co/Qwen/Qwen3-0.6B).\n",
|
||||
"\n",
|
||||
"Substitute your own embeddings database to change the knowledge base. `txtai` supports running local LLMs via [transformers](https://github.com/huggingface/transformers) or [llama.cpp](https://github.com/abetlen/llama-cpp-python). It also supports a wide variety of LLMs via [LiteLLM](https://github.com/BerriAI/litellm). For example, setting the 2nd RAG pipeline parameter below to `gpt-4o` along with the appropriate environment variables with access keys switches to a hosted LLM. See [this documentation page](https://neuml.github.io/txtai/pipeline/text/llm/) for more on this."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 83,
|
||||
"id": "cd05fc41",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Device set to use cuda:0\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import RAG\n",
|
||||
"\n",
|
||||
"# Prompt templates\n",
|
||||
"system = \"You are a friendly medical assistant that answers questions\"\n",
|
||||
"template = \"\"\"\n",
|
||||
"Answer the following question using the provided context.\n",
|
||||
"\n",
|
||||
"Question:\n",
|
||||
"{question}\n",
|
||||
"\n",
|
||||
"Context:\n",
|
||||
"{context}\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"# Create RAG pipeline\n",
|
||||
"rag = RAG(embeddings, \"Qwen/Qwen3-0.6B\", system=system, template=template, output=\"flatten\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "98800220",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# RAG Queries\n",
|
||||
"\n",
|
||||
"Now that the pipeline is setup, let's run a query."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 85,
|
||||
"id": "4d93ac4b",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"<think>\n",
|
||||
"Okay, let's see. The user is asking about H5N1. The context provided starts with \"Why tell me now?\" and then goes into facts about H5N1. The first sentence mentions that people and healthcare providers are weighing in on pandemic messages. Then it says H5N1 is avian influenza, a potential pandemic.\n",
|
||||
"\n",
|
||||
"Wait, but the user's question is about H5N1. The context doesn't go into specifics about what H5N1 is, but it does state that it's avian influenza. So I need to make sure I answer based on that. The answer should be concise, maybe mention that H5N1 is avian flu and it's a potential pandemic. Also, note that people are weighing in on messages. But I need to check if there's any more information. The context ends there. So the answer should be straightforward.\n",
|
||||
"</think>\n",
|
||||
"\n",
|
||||
"H5N1 influenza viruses are a type of avian influenza, a potential pandemic influenza virus that could cause widespread illness and death. While the context highlights the importance of public health and preparedness, it does not provide more specific details about its characteristics or risks.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"print(rag(\"Tell me about H5N1\"))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "39ab9946",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Notice that this LLM outputs a thinking or reasoning section then the answer.\n",
|
||||
"\n",
|
||||
"Let's review the context to validate this answer is derived from the knowledge base."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 79,
|
||||
"id": "53de0f96",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '16775537',\n",
|
||||
" 'text': '\"Why tell me now?\" the public and healthcare providers weigh in on pandemic influenza messages.',\n",
|
||||
" 'score': 0.7156285643577576},\n",
|
||||
" {'id': '22308474',\n",
|
||||
" 'text': 'H5N1 influenza viruses: facts, not fear.',\n",
|
||||
" 'score': 0.658343493938446},\n",
|
||||
" {'id': '16440117',\n",
|
||||
" 'text': 'Avian influenza--a pandemic waiting to happen?',\n",
|
||||
" 'score': 0.5827972888946533},\n",
|
||||
" {'id': '20667302',\n",
|
||||
" 'text': 'The influenza A(H5N1) epidemic at six and a half years: 500 notified human cases and more to come.',\n",
|
||||
" 'score': 0.5593500137329102},\n",
|
||||
" {'id': '18936262',\n",
|
||||
" 'text': 'What Australians know and believe about bird flu: results of a population telephone survey.',\n",
|
||||
" 'score': 0.5568690299987793},\n",
|
||||
" {'id': '30349811',\n",
|
||||
" 'text': 'Back to the Future: Lessons Learned From the 1918 Influenza Pandemic.',\n",
|
||||
" 'score': 0.5540266036987305},\n",
|
||||
" {'id': '17276785',\n",
|
||||
" 'text': 'Pandemic influenza: what infection control professionals should know.',\n",
|
||||
" 'score': 0.5519200563430786},\n",
|
||||
" {'id': '16681227',\n",
|
||||
" 'text': 'A pandemic flu: not if, but when. SARS was the wake-up call we slept through.',\n",
|
||||
" 'score': 0.5518345832824707},\n",
|
||||
" {'id': '22402712',\n",
|
||||
" 'text': 'Ferretting out the facts behind the H5N1 controversy.',\n",
|
||||
" 'score': 0.5508109331130981},\n",
|
||||
" {'id': '25546511',\n",
|
||||
" 'text': \"One-way trip: influenza virus' adaptation to gallinaceous poultry may limit its pandemic potential.\",\n",
|
||||
" 'score': 0.5494509339332581}]"
|
||||
]
|
||||
},
|
||||
"execution_count": 79,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings.search(\"Tell me about H5N1\", limit=10)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "025e6092",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The answer is doing a good job being based on the context above. Also keep in mind this is a small 600M parameter model, which is even more impressive.\n",
|
||||
"\n",
|
||||
"Let's try another query."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 86,
|
||||
"id": "1899047a",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"<think>\n",
|
||||
"Okay, let's see. The user is asking about the locations that have had H5N1 outbreaks, and the provided context mentions a few places: Indonesia and Bangladesh. The context also has a title about a decade of avian influenza in Bangladesh and mentions \"H5N1.\" \n",
|
||||
"\n",
|
||||
"Wait, the user's question is in English, so I need to make sure I'm interpreting the context correctly. The context includes two sentences: one about a decade in Bangladesh and another about H5N1. The user is probably looking for specific locations where H5N1 has been reported. \n",
|
||||
"\n",
|
||||
"Looking at the context again, it says \"Human avian influenza in Indonesia\" and \"A Decade of Avian Influenza in Bangladesh: Where Are We Now? Are we ready for pandemic influenza H5N1?\" So the outbreaks are in Indonesia and Bangladesh. \n",
|
||||
"\n",
|
||||
"I should confirm that there are no other mentions of other locations. The context doesn't provide more information beyond those two countries. Therefore, the answer should list Indonesia and Bangladesh as the locations with H5N1 outbreaks.\n",
|
||||
"</think>\n",
|
||||
"\n",
|
||||
"The locations with H5N1 outbreaks are Indonesia and Bangladesh.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"print(rag(\"What locations have had H5N1 outbreaks?\"))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 87,
|
||||
"id": "39d4efa6",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '21706937',\n",
|
||||
" 'text': 'Human avian influenza in Indonesia: are they really clustered?',\n",
|
||||
" 'score': 0.6269429326057434},\n",
|
||||
" {'id': '31514405',\n",
|
||||
" 'text': 'A Decade of Avian Influenza in Bangladesh: Where Are We Now?',\n",
|
||||
" 'score': 0.5972536206245422},\n",
|
||||
" {'id': '15889987',\n",
|
||||
" 'text': 'Are we ready for pandemic influenza H5N1?',\n",
|
||||
" 'score': 0.5863772630691528},\n",
|
||||
" {'id': '17717543',\n",
|
||||
" 'text': 'Commentary: From scarcity to abundance: pandemic vaccines and other agents for \"have not\" countries.',\n",
|
||||
" 'score': 0.5844159126281738},\n",
|
||||
" {'id': '22491771',\n",
|
||||
" 'text': 'Two years after pandemic influenza A/2009/H1N1: what have we learned?',\n",
|
||||
" 'score': 0.5812581777572632},\n",
|
||||
" {'id': '39666804',\n",
|
||||
" 'text': \"Why hasn't the bird flu pandemic started?\",\n",
|
||||
" 'score': 0.5738048553466797},\n",
|
||||
" {'id': '23402131',\n",
|
||||
" 'text': 'Where do avian influenza viruses meet in the Americas?',\n",
|
||||
" 'score': 0.5638074278831482},\n",
|
||||
" {'id': '20667302',\n",
|
||||
" 'text': 'The influenza A(H5N1) epidemic at six and a half years: 500 notified human cases and more to come.',\n",
|
||||
" 'score': 0.560465395450592},\n",
|
||||
" {'id': '17338983',\n",
|
||||
" 'text': 'Human avian influenza: how ready are we?',\n",
|
||||
" 'score': 0.555113673210144},\n",
|
||||
" {'id': '24518630',\n",
|
||||
" 'text': 'Recognizing true H5N1 infections in humans during confirmed outbreaks.',\n",
|
||||
" 'score': 0.5501888990402222}]"
|
||||
]
|
||||
},
|
||||
"execution_count": 87,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"embeddings.search(\"What locations have had H5N1 outbreaks?\", limit=10)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "efa8e59e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Once again the answer is based on the context which mentions the two countries in the answer. The context also discusses the Americas but it doesn't have as strong of language connecting H5N1 outbreaks to the location."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "18e04661",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Add citations\n",
|
||||
"\n",
|
||||
"The last item we'll cover is citations. One of the most important aspects of a RAG process is being able to ensure the answer is based on reality. There are a number of ways to do this but in this example, we'll ask the LLM to perform this step."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 96,
|
||||
"id": "ce10ff39",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Device set to use cuda:0\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# Prompt templates\n",
|
||||
"system = \"You are a friendly medical assistant that answers questions\"\n",
|
||||
"template = \"\"\"\n",
|
||||
"Answer the following question using the provided context.\n",
|
||||
"\n",
|
||||
"After the answer, write a citation section with ALL the original article ids used for the answer.\n",
|
||||
"\n",
|
||||
"Question:\n",
|
||||
"{question}\n",
|
||||
"\n",
|
||||
"Context:\n",
|
||||
"{context}\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"def context(question):\n",
|
||||
" context = []\n",
|
||||
" for x in embeddings.search(question, limit=10):\n",
|
||||
" context.append(f\"ARTICLE ID: {x['id']}, TEXT: {x['text']}\")\n",
|
||||
"\n",
|
||||
" return context\n",
|
||||
"\n",
|
||||
"# Create RAG pipeline\n",
|
||||
"rag = RAG(embeddings, \"Qwen/Qwen3-0.6B\", system=system, template=template, output=\"flatten\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "f79bdc9a",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"H5N1 is a type of avian influenza virus. \n",
|
||||
"\n",
|
||||
"**Citation Section:** \n",
|
||||
"- ARTICLE ID: 22010536, TEXT: Is avian influenza virus A(H5N1) a real threat to human health?\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"question = \"What is H5N1?\"\n",
|
||||
"print(rag(question, context(question), maxlength=2048, stripthink=True))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "15539c5d",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"As expected, the answer adds a citation section. Also note that the RAG pipeline stripped the thinking section from the result."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "3bdf84b5",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook covered how to build a Medical RAG Research process with `txtai`. It also covered how to modify this logic to add in your own knowledge base or use a more sophisticated LLM.\n",
|
||||
"\n",
|
||||
"With an important space such as the medical domain, it's vital to ensure that answers are derived from reliable knowledge. This notebook shows how to add that reliability via RAG. But as with anything in an important domain, there should be a human in the loop and answers shouldn't be blindly relied upon. "
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"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.10.18"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,309 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "e3wdiK5fGUoZ"
|
||||
},
|
||||
"source": [
|
||||
"# 💡 What's new in txtai 9.0\n",
|
||||
"\n",
|
||||
"The 9.0 release adds first class support for sparse vector models (i.e. [SPLADE](https://en.wikipedia.org/wiki/Learned_sparse_retrieval)), late interaction models (i.e. [ColBERT](https://huggingface.co/colbert-ir/colbertv2.0)), fixed dimensional encoding (i.e. [MUVERA](https://arxiv.org/abs/2405.19504)) and reranking pipelines ✨ \n",
|
||||
"\n",
|
||||
"The embeddings framework was overhauled to seamlessly support both sparse and dense vector models. Previously, sparse vector support was limited to keyword/term indexes. Now learned sparse retrieval models such as SPLADE are supported. These models can help improve the accuracy of retrieval/search operations, which also improves RAG and Agents.\n",
|
||||
"\n",
|
||||
"Support for late interaction models, such as ColBERT, were also added to the embeddings framework. Unlike traditional vector models that pool outputs into single vector outputs, late interaction models produce multiple vectors. These models are paired with the MUVERA algorithm to transform multiple vectors into fixed dimensional single vectors for search.\n",
|
||||
"\n",
|
||||
"LLMs are quickly converging to produce similar outputs for similar inputs and becoming standard commodities. The retrieval or context layer makes or breaks projects. This is known as putting the R in RAG!\n",
|
||||
"\n",
|
||||
"**Standard upgrade disclaimer below**\n",
|
||||
"\n",
|
||||
"While everything is backwards compatible, it's prudent to backup production indexes before upgrading and test before deploying."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "p8BbfjrhH-V2"
|
||||
},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"id": "-OXsTQgaGQPM"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[ann,vectors]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "n4EXrtcYIIYE"
|
||||
},
|
||||
"source": [
|
||||
"# Sparse vector indexes\n",
|
||||
"\n",
|
||||
"The first major change added with this release is `learned sparse retrieval` (aka sparse vector indexes) models. This effort was multi-faceted in that it required both changes to how vectors were generated as well as how they are stored.\n",
|
||||
"\n",
|
||||
"`txtai` uses approximate nearest neighbor (ANN) search for it's vector search operations. The default library is [Faiss](https://github.com/facebookresearch/faiss). There is support for other libraries but in all cases the existing ANN backends only supported dense (i.e. NumPy) vectors.\n",
|
||||
"\n",
|
||||
"There aren't many options out there for sparse ANN search that supports `txtai` requirements, so IVFSparse was introduced. IVFSparse is an Inverted file (IVF) index with flat vector file storage and sparse array support. There is also support for storing sparse vectors in Postgres via [pgvector](https://github.com/pgvector/pgvector).\n",
|
||||
"\n",
|
||||
"Let's see it in action.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"id": "hzPD8_cQJNtN"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '0',\n",
|
||||
" 'text': 'US tops 5 million confirmed virus cases',\n",
|
||||
" 'score': 0.019873601198196412},\n",
|
||||
" {'id': '1',\n",
|
||||
" 'text': \"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\",\n",
|
||||
" 'score': 0.018737798929214476}]"
|
||||
]
|
||||
},
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"# Works with a list, dataset or generator\n",
|
||||
"data = [\n",
|
||||
" \"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",
|
||||
"\n",
|
||||
"# Create an embeddings\n",
|
||||
"embeddings = Embeddings(sparse=True, content=True)\n",
|
||||
"embeddings.index(data)\n",
|
||||
"embeddings.search(\"North America\", 10)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Late interaction models\n",
|
||||
"\n",
|
||||
"Late interaction models encode data into multi-vector outputs. In other words, multiple input tokens map to multiple output vectors. Then at search time, the maximum similarity algorithm is used to find the best matches between the corpus and a query. This algorithm has achieved excellent results on retrieval benchmarks such as [MTEB](https://github.com/embeddings-benchmark/mteb).\n",
|
||||
"\n",
|
||||
"The downside of this approach is that it produces multiple vectors as opposed a single vector for each input. For example, if a text element tokenizes to many input tokens, there will be many output vectors vs a single one as with standard pooled vector approaches.\n",
|
||||
"\n",
|
||||
"Starting with the 9.0 release, late interaction models are supported with embeddings instances. Late interaction vectors will be transformed into fixed dimensional vectors using the MUVERA algorithm. See below. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '0',\n",
|
||||
" 'text': 'US tops 5 million confirmed virus cases',\n",
|
||||
" 'score': 0.04216160625219345},\n",
|
||||
" {'id': '1',\n",
|
||||
" 'text': \"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\",\n",
|
||||
" 'score': 0.029944246634840965},\n",
|
||||
" {'id': '3',\n",
|
||||
" 'text': 'The National Park Service warns against sacrificing slower friends in a bear attack',\n",
|
||||
" 'score': 0.015931561589241028}]"
|
||||
]
|
||||
},
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"# Works with a list, dataset or generator\n",
|
||||
"data = [\n",
|
||||
" \"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",
|
||||
"\n",
|
||||
"# Create an embeddings\n",
|
||||
"embeddings = Embeddings(path=\"colbert-ir/colbertv2.0\", content=True)\n",
|
||||
"embeddings.index(data)\n",
|
||||
"embeddings.search(\"North America\", 10)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Reranking pipeline\n",
|
||||
"\n",
|
||||
"Another major new component in this release is the Reranker pipeline. This pipeline takes an embeddings instance, a similarity instance and uses the similarity instance to rerank outputs. This is a key component of the MUVERA paper - using the standard vector index to retrieve candidates then reranking the outputs using the late interaction model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': '1',\n",
|
||||
" 'text': \"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg\",\n",
|
||||
" 'score': 0.3324427008628845},\n",
|
||||
" {'id': '0',\n",
|
||||
" 'text': 'US tops 5 million confirmed virus cases',\n",
|
||||
" 'score': 0.24423550069332123},\n",
|
||||
" {'id': '3',\n",
|
||||
" 'text': 'The National Park Service warns against sacrificing slower friends in a bear attack',\n",
|
||||
" 'score': 0.16353240609169006}]"
|
||||
]
|
||||
},
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"from txtai.pipeline import Reranker, Similarity\n",
|
||||
"\n",
|
||||
"# Works with a list, dataset or generator\n",
|
||||
"data = [\n",
|
||||
" \"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",
|
||||
"\n",
|
||||
"# Create an embeddings\n",
|
||||
"embeddings = Embeddings(path=\"colbert-ir/colbertv2.0\", content=True)\n",
|
||||
"embeddings.index(data)\n",
|
||||
"\n",
|
||||
"similarity = Similarity(path=\"colbert-ir/colbertv2.0\", lateencode=True)\n",
|
||||
"\n",
|
||||
"ranker = Reranker(embeddings, similarity)\n",
|
||||
"ranker(\"North America\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Notice that while the outputs are the same, the scoring and order is different.\n",
|
||||
"\n",
|
||||
"Let's try a more interesting example."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[{'id': 'ChatGPT',\n",
|
||||
" 'text': 'ChatGPT is a generative artificial intelligence chatbot developed by OpenAI and released on November 30, 2022. It uses large language models (LLMs) such as GPT-4o as well as other multimodal models to create human-like responses in text, speech, and images. It has access to features such as searching the web, using apps, and running programs. It is credited with accelerating the AI boom, an ongoing period of rapid investment in and public attention to the field of artificial intelligence (AI). Some observers have raised concern about the potential of ChatGPT and similar programs to displace human intelligence, enable plagiarism, or fuel misinformation.',\n",
|
||||
" 'score': 0.6639302968978882},\n",
|
||||
" {'id': 'ChatGPT Search',\n",
|
||||
" 'text': 'ChatGPT Search (originally SearchGPT) is a search engine developed by OpenAI. It combines traditional search engine features with generative pretrained transformers (GPT) to generate responses, including citations to external websites.',\n",
|
||||
" 'score': 0.6477508544921875},\n",
|
||||
" {'id': 'ChatGPT in education',\n",
|
||||
" 'text': 'The usage of ChatGPT in education has sparked considerable debate and exploration. ChatGPT is a chatbot based on large language models (LLMs) that was released by OpenAI in November 2022.',\n",
|
||||
" 'score': 0.5918337106704712}]"
|
||||
]
|
||||
},
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"from txtai.pipeline import Reranker, Similarity\n",
|
||||
"\n",
|
||||
"# Create an embeddings\n",
|
||||
"embeddings = Embeddings()\n",
|
||||
"embeddings.load(provider=\"huggingface-hub\", container=\"neuml/txtai-wikipedia\")\n",
|
||||
"\n",
|
||||
"similarity = Similarity(path=\"colbert-ir/colbertv2.0\", lateencode=True)\n",
|
||||
"\n",
|
||||
"ranker = Reranker(embeddings, similarity)\n",
|
||||
"ranker(\"Tell me about ChatGPT\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "tvnMO1Eai6Gy"
|
||||
},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook gave a quick overview of txtai 9.0. Updated documentation and more examples will be forthcoming. There is much to cover and much to build on!\n",
|
||||
"\n",
|
||||
"See the following links for more information.\n",
|
||||
"\n",
|
||||
"- [9.0 Release on GitHub](https://github.com/neuml/txtai/releases/tag/v9.0.0)\n",
|
||||
"- [Documentation site](https://neuml.github.io/txtai)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"gpuType": "T4",
|
||||
"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.10.18"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,482 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "64be0343",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Accessing Low Level Vector APIs\n",
|
||||
"\n",
|
||||
"The primary interface to build vector databases with `txtai` is through [Embeddings instances](https://neuml.github.io/txtai/embeddings/). `txtai` also supports accessing all of it's features through lower level APIs. \n",
|
||||
"\n",
|
||||
"Let's dive in.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e8198312",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "3f2d0792",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[ann] gguf"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "46740a21",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Load a dataset\n",
|
||||
"\n",
|
||||
"We'll use a [subset](https://huggingface.co/datasets/m-a-p/FineFineWeb-test) of the [FineFineWeb dataset](https://huggingface.co/datasets/m-a-p/FineFineWeb). This dataset is a domain-labeled version of the general purpose [FineWeb dataset](https://huggingface.co/datasets/HuggingFaceFW/fineweb). "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "3c13e554",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"\n",
|
||||
"ds = load_dataset(\"m-a-p/FineFineWeb-test\", split=\"train\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ae49af96",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Building an Embeddings database\n",
|
||||
"\n",
|
||||
"Before going into the low-level API, let's recap how we build an Embeddings database."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 18,
|
||||
"id": "2de3bf4c",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"0.6012564897537231 The National Aeronautics and Space Administration (NASA) is the United States’ civil space program. \n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"embeddings = Embeddings()\n",
|
||||
"embeddings.index(ds[\"text\"][:10000])\n",
|
||||
"for uid, score in embeddings.search(\"nasa\", 1):\n",
|
||||
" print(score, ds[uid][\"text\"][:100])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d6f196cb",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This simple example abstracts the heavy lifting behind the `Embeddings` interface. Behind the scenes, it defaults to vectorizing text using [all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2). Vectors are stored in a [Faiss index](https://github.com/facebookresearch/faiss).\n",
|
||||
"\n",
|
||||
"The first 10K records are vectorized and stored in the vector index. Then at query time, the query is vectorized and a vector similarity search is run.\n",
|
||||
"\n",
|
||||
"While the `Embeddings` interface is convenient, it's also possible to access lower level APIs. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ecd576e8",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Vectors Interface\n",
|
||||
"\n",
|
||||
"First, let's vectorize our data using the low level APIs. We'll use the default Hugging Face vectorizer available in `txtai`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"id": "03c85684",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai.ann import ANNFactory\n",
|
||||
"from txtai.vectors import VectorsFactory\n",
|
||||
"\n",
|
||||
"vectors = VectorsFactory.create({\"path\": \"sentence-transformers/all-MiniLM-L6-v2\"})\n",
|
||||
"data = vectors.vectorize(ds[\"text\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "36fac81a",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# ANN Interface\n",
|
||||
"\n",
|
||||
"Now that we have a NumPy array of vectors, let's store them in an Approximate Neighest Neighbor (ANN) backend. Recall earlier, we used the default Faiss interface. For this example, we're going to use the [PyTorch ANN](https://neuml.github.io/txtai/embeddings/configuration/ann/#torch). This will allow us to use new features that are available as of `txtai` 9.1."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"id": "34ebdc85",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"ann = ANNFactory.create({\n",
|
||||
" \"backend\": \"torch\",\n",
|
||||
" \"torch\": {\n",
|
||||
" \"safetensors\": True,\n",
|
||||
" }\n",
|
||||
"})\n",
|
||||
"ann.index(data)\n",
|
||||
"ann.save(\"vectors.safetensors\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "98a26c9a",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This ANN builds a Torch tensor with the vectors and stores them in a [Safetensors](https://github.com/huggingface/safetensors) file.\n",
|
||||
"\n",
|
||||
"The code below shows how the file is simply a standard Safetensors file."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"id": "bcaa4832",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"data (1411868, 384)\n",
|
||||
"Memory = 2068.17 MB\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from safetensors import safe_open\n",
|
||||
"\n",
|
||||
"def tensorinfo():\n",
|
||||
" memory = 0\n",
|
||||
" with safe_open(\"vectors.safetensors\", framework=\"np\") as f:\n",
|
||||
" for key in f.keys():\n",
|
||||
" array = f.get_tensor(key)\n",
|
||||
" print(key, array.shape)\n",
|
||||
" memory += array.nbytes\n",
|
||||
"\n",
|
||||
" print(f\"Memory = {memory / 1024 / 1024:.2f} MB\")\n",
|
||||
"\n",
|
||||
"tensorinfo()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d6525b1a",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Vector search\n",
|
||||
"\n",
|
||||
"Now let's show how these low-level APIs can be used to implement vector search."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"id": "09b81958",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"The answer to your question, that how many miles is it from earth to mars, is very easy to know. Because of huge satellites which are being sent to\n",
|
||||
"mars in search of life from many countries, we have discovered a lot about mars. According to experts, earth and mars reaches to their closest points\n",
|
||||
"in every 26 months. This situation is considered as opposition of mars as the location of sun and mars in totally opposite to each other in relation\n",
|
||||
"to earth. When this opposition takes place, the planet is visible with a red tint in the sky from earth. And this also gives mars a name, i.e. the red\n",
|
||||
"planet. Mars is also the fourth planet from sun, which is located between Jupiter and earth. Its distance from sun is not only opposite but is also\n",
|
||||
"much further away, than that of the earth and sun. The distance between the sun and mars is said to be 140 million miles. Mars can reach about 128\n",
|
||||
"million miles closer to the sun whereas it can even travel around 154 million miles away from it. The assumed distance between mars and earth is said\n",
|
||||
"to be between 40 to 225 million miles. The distance between these two planets keeps on changing throughout the year because of the elliptical path in\n",
|
||||
"which all the planets rotate. As the distance between mars, sun and earth is so much high, it takes a Martian year, for mars to go around the sun. The\n",
|
||||
"Martian period includes a time of around 687 earth days. This means that, it takes more than 2 years for the mars to reach its initial rotation point.\n",
|
||||
"If we talk about one Martian day, it is the total time which is taken by a planet to spin around once. This day usually lasts longer than our regular\n",
|
||||
"earth days. So this was the actual reason which states the distance between earth and mars. \n",
|
||||
"\n",
|
||||
" 0.7060051560401917\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import textwrap\n",
|
||||
"\n",
|
||||
"def search(text):\n",
|
||||
" result = ann.search(vectors.vectorize([text]), 1)\n",
|
||||
" index, score = result[0][0]\n",
|
||||
" print(textwrap.fill(ds[index][\"text\"], width=150), \"\\n\\n\", score)\n",
|
||||
"\n",
|
||||
"search(\"How far is earth from mars?\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "8adf64db",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Torch 4-bit quantization\n",
|
||||
"\n",
|
||||
"`txtai` 9.1 adds a new feature: 4-bit vector quantization. This means that instead of using 32-bit floats for each vector dimension, this method uses 4 bits. This reduces memory usage to ~12-13% of the original size."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 12,
|
||||
"id": "cbe5af13",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"absmax (8471208,)\n",
|
||||
"code (16,)\n",
|
||||
"data (271078656, 1)\n",
|
||||
"shape (2,)\n",
|
||||
"Memory = 290.84 MB\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"ann = ANNFactory.create({\n",
|
||||
" \"backend\": \"torch\",\n",
|
||||
" \"torch\": {\n",
|
||||
" \"safetensors\": True,\n",
|
||||
" \"quantize\": {\n",
|
||||
" \"type\": \"nf4\"\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
"})\n",
|
||||
"ann.index(data)\n",
|
||||
"ann.save(\"vectors.safetensors\")\n",
|
||||
"\n",
|
||||
"tensorinfo()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ca7b1e0e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Note how the unquantized vectors took 2068.17 MB and this only takes 290.84 MB! With quantization and ever growing GPUs, this opens the possibility of pinning your entire vector database in GPU memory!\n",
|
||||
"\n",
|
||||
"For example, let's extrapolate this dataset to 100M rows.\n",
|
||||
"\n",
|
||||
"```\n",
|
||||
"(290.84 MB / 1,411,868) * 100,000,000 = 20,599.7 MB\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"An entire 100M row dataset could fit into a single RTX 3090 / 4090 consumer GPU!\n",
|
||||
"\n",
|
||||
"Let's confirm search still works the same."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 13,
|
||||
"id": "a02578dd",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"The answer to your question, that how many miles is it from earth to mars, is very easy to know. Because of huge satellites which are being sent to\n",
|
||||
"mars in search of life from many countries, we have discovered a lot about mars. According to experts, earth and mars reaches to their closest points\n",
|
||||
"in every 26 months. This situation is considered as opposition of mars as the location of sun and mars in totally opposite to each other in relation\n",
|
||||
"to earth. When this opposition takes place, the planet is visible with a red tint in the sky from earth. And this also gives mars a name, i.e. the red\n",
|
||||
"planet. Mars is also the fourth planet from sun, which is located between Jupiter and earth. Its distance from sun is not only opposite but is also\n",
|
||||
"much further away, than that of the earth and sun. The distance between the sun and mars is said to be 140 million miles. Mars can reach about 128\n",
|
||||
"million miles closer to the sun whereas it can even travel around 154 million miles away from it. The assumed distance between mars and earth is said\n",
|
||||
"to be between 40 to 225 million miles. The distance between these two planets keeps on changing throughout the year because of the elliptical path in\n",
|
||||
"which all the planets rotate. As the distance between mars, sun and earth is so much high, it takes a Martian year, for mars to go around the sun. The\n",
|
||||
"Martian period includes a time of around 687 earth days. This means that, it takes more than 2 years for the mars to reach its initial rotation point.\n",
|
||||
"If we talk about one Martian day, it is the total time which is taken by a planet to spin around once. This day usually lasts longer than our regular\n",
|
||||
"earth days. So this was the actual reason which states the distance between earth and mars. \n",
|
||||
"\n",
|
||||
" 0.6982609033584595\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"search(\"How far is earth from mars?\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f0e62670",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Same result. Note the score is slightly different but this is expected."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "89abb301",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# GGUF Support\n",
|
||||
"\n",
|
||||
"`txtai` 9.1 also adds support for [GGML](https://github.com/ggml-org/ggml) / [GGUF](https://huggingface.co/docs/hub/en/gguf) popularized by [llama.cpp](https://github.com/ggml-org/llama.cpp)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 19,
|
||||
"id": "e04498fc",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"ann = ANNFactory.create({\n",
|
||||
" \"backend\": \"ggml\",\n",
|
||||
" \"ggml\": {\n",
|
||||
" \"quantize\": \"Q4_0\"\n",
|
||||
" }\n",
|
||||
"})\n",
|
||||
"ann.index(data)\n",
|
||||
"ann.save(\"vectors.gguf\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7bea0f58",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Now let's check out the generated file using the [gguf](https://github.com/ggml-org/llama.cpp/tree/master/gguf-py) package provided by llama.cpp."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 20,
|
||||
"id": "b0e6f85e",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Tensor Name | Shape | Size | Quantization\n",
|
||||
"--------------------------------------------------------------------------------\n",
|
||||
"data | 384x1411868 | 258.52 MB | Q4_0\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from gguf.gguf_reader import GGUFReader\n",
|
||||
"\n",
|
||||
"reader = GGUFReader(\"vectors.gguf\")\n",
|
||||
"\n",
|
||||
"# List all tensors\n",
|
||||
"info = \"{:<30} | {:<15} | {:<12} | {}\"\n",
|
||||
"print(info.format(\"Tensor Name\", \"Shape\", \"Size\", \"Quantization\"))\n",
|
||||
"print(\"-\" * 80)\n",
|
||||
"for tensor in reader.tensors:\n",
|
||||
" shape = \"x\".join(map(str, tensor.shape))\n",
|
||||
" size = f\"{tensor.n_elements / 2 / 1024 / 1024:.2f} MB\"\n",
|
||||
" quantization = tensor.tensor_type.name\n",
|
||||
" print(info.format(tensor.name, shape, size, quantization))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7be79bf7",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"And search like we did with Torch."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 18,
|
||||
"id": "450dc024",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"The answer to your question, that how many miles is it from earth to mars, is very easy to know. Because of huge satellites which are being sent to\n",
|
||||
"mars in search of life from many countries, we have discovered a lot about mars. According to experts, earth and mars reaches to their closest points\n",
|
||||
"in every 26 months. This situation is considered as opposition of mars as the location of sun and mars in totally opposite to each other in relation\n",
|
||||
"to earth. When this opposition takes place, the planet is visible with a red tint in the sky from earth. And this also gives mars a name, i.e. the red\n",
|
||||
"planet. Mars is also the fourth planet from sun, which is located between Jupiter and earth. Its distance from sun is not only opposite but is also\n",
|
||||
"much further away, than that of the earth and sun. The distance between the sun and mars is said to be 140 million miles. Mars can reach about 128\n",
|
||||
"million miles closer to the sun whereas it can even travel around 154 million miles away from it. The assumed distance between mars and earth is said\n",
|
||||
"to be between 40 to 225 million miles. The distance between these two planets keeps on changing throughout the year because of the elliptical path in\n",
|
||||
"which all the planets rotate. As the distance between mars, sun and earth is so much high, it takes a Martian year, for mars to go around the sun. The\n",
|
||||
"Martian period includes a time of around 687 earth days. This means that, it takes more than 2 years for the mars to reach its initial rotation point.\n",
|
||||
"If we talk about one Martian day, it is the total time which is taken by a planet to spin around once. This day usually lasts longer than our regular\n",
|
||||
"earth days. So this was the actual reason which states the distance between earth and mars. \n",
|
||||
"\n",
|
||||
" 0.7043964862823486\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"search(\"How far is earth from mars?\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "af1a344c",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"While the `Embeddings` interface is the preferred way to build vector databases with `txtai`, it's entirely possible to also build with the low level APIs!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"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.10.19"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "4c8ab3b0",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# RAG is more than Vector Search\n",
|
||||
"\n",
|
||||
"Retrieval Augmented Generation (RAG) is often associated with vector search. And while that is a primary use case, any search will do.\n",
|
||||
"\n",
|
||||
"- ✅ Vector Search\n",
|
||||
"- ✅ Web Search\n",
|
||||
"- ✅ SQL Query\n",
|
||||
"\n",
|
||||
"This notebook will go over a few RAG examples covering different retrieval methods. These examples require txtai 9.3+."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ab8a77d1",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "21bebb46",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline-data]\n",
|
||||
"\n",
|
||||
"# Download example SQL database\n",
|
||||
"!wget https://huggingface.co/NeuML/txtai-wikipedia-slim/resolve/main/documents"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1ece2b09",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# RAG with Late Interaction\n",
|
||||
"\n",
|
||||
"The first example will cover RAG with ColBERT / Late Interaction retrieval. TxtAI 9.0 added support for [MUVERA](https://arxiv.org/abs/2405.19504) and [ColBERT](https://arxiv.org/abs/2112.01488) multi-vector ranking. \n",
|
||||
"\n",
|
||||
"We'll build a pipeline that reads the ColBERT v2 paper, extracts the text into sections and builds an index with a ColBERT model. Then we'll wrap that as a [Reranker pipeline](https://neuml.github.io/txtai/pipeline/text/reranker/) using the same ColBERT model. Finally a RAG pipeline will utilize this for retrieval.\n",
|
||||
"\n",
|
||||
"_Note: This uses the custom [ColBERT Muvera Nano](https://huggingface.co/NeuML/colbert-muvera-nano) model which is only 970K parameters! That's right thousands. It's surprisingly effective._"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"id": "dae2e6dc",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"This paper introduces ColBERTv2, a neural information retrieval model that enhances the quality and efficiency of late interaction by combining an aggressive residual compression mechanism with a denoised supervision strategy, achieving state-of-the-art performance across diverse benchmarks while reducing the model's space footprint by 6–10× compared to previous methods.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import Embeddings, RAG, Textractor\n",
|
||||
"from txtai.pipeline import Reranker, Similarity\n",
|
||||
"\n",
|
||||
"# Get text from ColBERT v2 paper\n",
|
||||
"textractor = Textractor(sections=True, backend=\"docling\")\n",
|
||||
"data = textractor(\"https://arxiv.org/pdf/2112.01488\")\n",
|
||||
"\n",
|
||||
"# MUVERA fixed dimensional encodings\n",
|
||||
"embeddings = Embeddings(content=True, path=\"neuml/colbert-muvera-nano\", vectors={\"trust_remote_code\": True})\n",
|
||||
"embeddings.index(data)\n",
|
||||
"\n",
|
||||
"# Re-rank using same late interaction model\n",
|
||||
"reranker = Reranker(embeddings, Similarity(\"neuml/colbert-muvera-nano\", lateencode=True, vectors={\"trust_remote_code\": True}))\n",
|
||||
"\n",
|
||||
"template = \"\"\"\n",
|
||||
" Answer the following question using the provided context.\n",
|
||||
"\n",
|
||||
" Question:\n",
|
||||
" {question}\n",
|
||||
"\n",
|
||||
" Context:\n",
|
||||
" {context}\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"# RAG with late interaction models\n",
|
||||
"rag = RAG(reranker, \"Qwen/Qwen3-4B-Instruct-2507\", template=template, output=\"flatten\")\n",
|
||||
"print(rag(\"Write a sentence abstract about this paper\", maxlength=2048))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b31e46de",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# RAG with a Web Search\n",
|
||||
"\n",
|
||||
"Next we'll run a RAG pipeline using a web search as the retrieval method."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"id": "b61b684b",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Artificial intelligence (AI) is the capability of computational systems to perform tasks typically associated with human intelligence, such as learning, reasoning, problem-solving, perception, and decision-making. It involves technologies like machine learning, deep learning, and natural language processing, and enables machines to simulate human-like learning, comprehension, problem solving, decision-making, creativity, and autonomy.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from smolagents import WebSearchTool\n",
|
||||
"\n",
|
||||
"tool = WebSearchTool()\n",
|
||||
"\n",
|
||||
"def websearch(queries, limit):\n",
|
||||
" results = []\n",
|
||||
" for query in queries:\n",
|
||||
" result = [\n",
|
||||
" {\"id\": i, \"text\": f'{x[\"title\"]} {x[\"description\"]}', \"score\": 1.0} for i, x in enumerate(tool.search(query))\n",
|
||||
" ]\n",
|
||||
" results.append(result[:limit])\n",
|
||||
"\n",
|
||||
" return results\n",
|
||||
"\n",
|
||||
"# RAG with a websearch\n",
|
||||
"rag = RAG(websearch, \"Qwen/Qwen3-4B-Instruct-2507\", template=template, output=\"flatten\")\n",
|
||||
"print(rag(\"What is AI?\", maxlength=2048))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "45c8a096",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# RAG with a SQL Query\n",
|
||||
"\n",
|
||||
"The last example we'll cover is running RAG with a SQL query. We'll use the SQL database that's a component of the [txtai-wikipedia-slim](https://huggingface.co/NeuML/txtai-wikipedia-slim) embeddings database.\n",
|
||||
"\n",
|
||||
"Since this is just a database with Wikipedia abstracts, we'll need a way to build a SQL query from a search query. For that we'll use an LLM to extract a keyword to use in a `LIKE` clause.\n",
|
||||
"\n",
|
||||
"Given that the LLM used was released in August 2025, let's ask it a question that can only be accurated answered with external data. `Who won the 2025 World Series?` which ended in November."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "c7ff2a35",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"In the 2025 World Series, the Los Angeles Dodgers defeated the Toronto Blue Jays in seven games to win the championship. The series took place from October 24 to November 1 (ending early on November 2, Toronto time). Dodgers pitcher Yoshinobu Yamamoto was named the World Series MVP. The series was televised by Fox in the United States and by Sportsnet in Canada.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import sqlite3\n",
|
||||
"\n",
|
||||
"from txtai import LLM\n",
|
||||
"\n",
|
||||
"def keyword(query):\n",
|
||||
" return llm(f\"\"\"\n",
|
||||
" Extract a keyword for this search query: {query}.\n",
|
||||
" Return only text with no other formatting or explanation.\n",
|
||||
" \"\"\")\n",
|
||||
"\n",
|
||||
"def sqlsearch(queries, limit):\n",
|
||||
" results = []\n",
|
||||
" sql = \"SELECT id, text FROM sections WHERE id LIKE ? LIMIT ?\"\n",
|
||||
"\n",
|
||||
" for query in queries:\n",
|
||||
" # Extract a keyword for this search\n",
|
||||
" query = keyword(query)\n",
|
||||
"\n",
|
||||
" # Run the SQL Query\n",
|
||||
" results.append([\n",
|
||||
" {\"id\": uid, \"text\": text, \"score\": 1.0}\n",
|
||||
" for uid, text in cursor.execute(sql, [f\"%{query}%\", limit])\n",
|
||||
" ])\n",
|
||||
"\n",
|
||||
" return results\n",
|
||||
"\n",
|
||||
"# Load the database\n",
|
||||
"cursor = sqlite3.connect(\"documents\")\n",
|
||||
"\n",
|
||||
"# Load the LLM\n",
|
||||
"llm = LLM(\"Qwen/Qwen3-4B-Instruct-2507\")\n",
|
||||
"\n",
|
||||
"# RAG with a SQL query\n",
|
||||
"rag = RAG(sqlsearch, llm, template=template, output=\"flatten\")\n",
|
||||
"print(rag(\"Tell me what happened in the 2025 World Series\", maxlength=2048))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a9f78076",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"And as we see, this answer is using the SQL database!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "97b1e282",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook showed that RAG is about much more than vector search. With txtai 9.3+, any callable method is now supported for retrieval. Enjoy!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"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.10.19"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,360 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c65bdb85",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Distilling Knowledge into Tiny LLMs\n",
|
||||
"\n",
|
||||
"Large Language Models (LLMs) are the magic behind AI. These massive billion and trillion parameter models have been shown to generalize well when trained on enough data.\n",
|
||||
"\n",
|
||||
"A big problem is that they are hard to run and expensive. So many just call LLMs through APIs such as OpenAI or Claude. Additionally, in many instances, developers spend a lot of time with complex prompt logic hoping to cover all the edge cases and believe they need a model that's large enough to handle all the rules.\n",
|
||||
"\n",
|
||||
"If you truly want control over your business processes, running a local model is a better choice. And the good news is that it doesn't have to be a giant and expensive multi-billion parameter model. We can finetune LLMs to handle our specific business logic, which helps us take control and limit prompt complexity. \n",
|
||||
"\n",
|
||||
"This notebook will show how we can distill knowledge into tiny LLMs."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "279ecd4b",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "4f7bca3a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline-train] datasets"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c7b1b4be",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# The LLM\n",
|
||||
"\n",
|
||||
"We'll use a [600M parameter Qwen3 model](https://hf.co/qwen/qwen3-0.6b) for this example. Our target task will be translating user requests into linux commands."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"id": "e1d95dd8",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import LLM\n",
|
||||
"\n",
|
||||
"llm = LLM(\"Qwen/Qwen3-0.6B\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "8f276cfc",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Let's try one with the base model as it is."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"id": "686c3983",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'ps -e'"
|
||||
]
|
||||
},
|
||||
"execution_count": 11,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"llm(\"\"\"\n",
|
||||
"Translate the following request into a linux command. Only print the command.\n",
|
||||
"\n",
|
||||
"Find number of logged in users\n",
|
||||
"\"\"\", maxlength=1024)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "dcd92414",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"As we can see, the model actually has a good understanding and at least prints a command. But in this case it's not correct. Let's get to fine-tuning!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "73666982",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Finetuning the LLM with knowledge\n",
|
||||
"\n",
|
||||
"Yes, 600M parameters is small and we can't possibly expect it to do well with everything. But the good news is that we can distill knowledge into this tiny LLM and make it better. We'll use this [linux commands dataset](https://huggingface.co/datasets/mecha-org/linux-command-dataset) from the Hugging Face Hub. We'll also use this [training pipeline from txtai](https://neuml.github.io/txtai/pipeline/train/trainer).\n",
|
||||
"\n",
|
||||
"First, we'll create the training dataset. We'll use the same prompt strategy from above.\n",
|
||||
"\n",
|
||||
"```python\n",
|
||||
"\"\"\"\n",
|
||||
"Translate the following request into a linux command. Only print the command.\n",
|
||||
"\n",
|
||||
"{user request}\n",
|
||||
"\"\"\"\n",
|
||||
"```"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"id": "43d6b563",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"from transformers import AutoTokenizer\n",
|
||||
"\n",
|
||||
"# LLM path\n",
|
||||
"path = \"Qwen/Qwen3-0.6B\"\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(path)\n",
|
||||
"\n",
|
||||
"# Load the training dataset\n",
|
||||
"dataset = load_dataset(\"mecha-org/linux-command-dataset\", split=\"train\")\n",
|
||||
"\n",
|
||||
"def prompt(row):\n",
|
||||
" text = tokenizer.apply_chat_template([\n",
|
||||
" {\"role\": \"system\", \"content\": \"Translate the following request into a linux command. Only print the command.\"},\n",
|
||||
" {\"role\": \"user\", \"content\": row[\"input\"]},\n",
|
||||
" {\"role\": \"assistant\", \"content\": row[\"output\"]}\n",
|
||||
" ], tokenize=False, enable_thinking=False)\n",
|
||||
"\n",
|
||||
" return {\"text\": text}\n",
|
||||
"\n",
|
||||
"# Map to training prompts\n",
|
||||
"train = dataset.map(prompt, remove_columns=[\"input\", \"output\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"id": "7f71dce0",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"\n",
|
||||
" <div>\n",
|
||||
" \n",
|
||||
" <progress value='210' max='210' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
|
||||
" [210/210 01:12, Epoch 1/1]\n",
|
||||
" </div>\n",
|
||||
" <table border=\"1\" class=\"dataframe\">\n",
|
||||
" <thead>\n",
|
||||
" <tr style=\"text-align: left;\">\n",
|
||||
" <th>Step</th>\n",
|
||||
" <th>Training Loss</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" <tr>\n",
|
||||
" <td>50</td>\n",
|
||||
" <td>0.625300</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <td>100</td>\n",
|
||||
" <td>0.490200</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <td>150</td>\n",
|
||||
" <td>0.403300</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <td>200</td>\n",
|
||||
" <td>0.391800</td>\n",
|
||||
" </tr>\n",
|
||||
" </tbody>\n",
|
||||
"</table><p>"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.HTML object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai.pipeline import HFTrainer\n",
|
||||
"\n",
|
||||
"# Load the training pipeline\n",
|
||||
"trainer = HFTrainer()\n",
|
||||
"\n",
|
||||
"# Train the model\n",
|
||||
"# Set output_dir to save, trained in memory for this example\n",
|
||||
"model = trainer(\n",
|
||||
" \"Qwen/Qwen3-0.6B\",\n",
|
||||
" train,\n",
|
||||
" task=\"language-generation\",\n",
|
||||
" maxlength=512,\n",
|
||||
" bf16=True,\n",
|
||||
" per_device_train_batch_size=4,\n",
|
||||
" num_train_epochs=1,\n",
|
||||
" logging_steps=50,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "68b5bc4e",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'who | wc -l'"
|
||||
]
|
||||
},
|
||||
"execution_count": 17,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from txtai import LLM\n",
|
||||
"\n",
|
||||
"llm = LLM(model)\n",
|
||||
"\n",
|
||||
"llm([\n",
|
||||
" {\"role\": \"system\", \"content\": \"Translate the following request into a linux command. Only print the command.\"},\n",
|
||||
" {\"role\": \"user\", \"content\": \"Find number of logged in users\"}\n",
|
||||
"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 14,
|
||||
"id": "7427186d",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'ls ~/'"
|
||||
]
|
||||
},
|
||||
"execution_count": 14,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"llm([\n",
|
||||
" {\"role\": \"system\", \"content\": \"Translate the following request into a linux command. Only print the command.\"},\n",
|
||||
" {\"role\": \"user\", \"content\": \"List the files in my home directory\"}\n",
|
||||
"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 16,
|
||||
"id": "a141d63b",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'zip -r data.zip data'"
|
||||
]
|
||||
},
|
||||
"execution_count": 16,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"llm([\n",
|
||||
" {\"role\": \"system\", \"content\": \"Translate the following request into a linux command. Only print the command.\"},\n",
|
||||
" {\"role\": \"user\", \"content\": \"Zip the data directory with all it's contents\"}\n",
|
||||
"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "fc7d0fda",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"It even works well without the system prompt."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 20,
|
||||
"id": "a7d0c671",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'du -sh ~'"
|
||||
]
|
||||
},
|
||||
"execution_count": 20,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"llm(\"Calculate the total amount of disk space used for my home directory. Only print the total.\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "48cfcacf",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook demonstrated how it's very straightforward to distill knowledge into LLMs with `txtai`. Don't always go for the giant LLM, spend a little time finetuning a tiny LLM, it is well worth it!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"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.10.19"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,377 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f95d2373",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# OpenCode as a txtai LLM\n",
|
||||
"\n",
|
||||
"[OpenCode](https://github.com/anomalyco/opencode) is an open source AI coding agent. It's rapidly growing in popularity as an open alternative to [Claude Code](https://code.claude.com/docs/en/overview). \n",
|
||||
"\n",
|
||||
"OpenCode shows the power of permissive open source. It's rapidly undergoing mass adoption and as of January 2026 sits at `81K+` GitHub ⭐'s. OpenCode supports running as a [local server](https://opencode.ai/docs/server/) via `opencode serve`.\n",
|
||||
"\n",
|
||||
"This notebook will explore how TxtAI integrates this as a LLM pipeline."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a8719b92",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "e00c85bc",
|
||||
"metadata": {
|
||||
"vscode": {
|
||||
"languageId": "shellscript"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[api,pipeline-llm]\n",
|
||||
"\n",
|
||||
"# Install OpenCode if not already installed and run serve\n",
|
||||
"!curl -fsSL https://opencode.ai/install | bash\n",
|
||||
"!opencode serve &"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "73ffe60e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Create a OpenCode LLM\n",
|
||||
"\n",
|
||||
"Now that everything is installed, let's test it out!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"id": "0e35a304",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"```python\n",
|
||||
"print(\"Hello, World!\")\n",
|
||||
"```"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from IPython.display import Markdown, display\n",
|
||||
"\n",
|
||||
"from txtai import LLM\n",
|
||||
"\n",
|
||||
"def generate(request):\n",
|
||||
" display(Markdown(llm(request)))\n",
|
||||
"\n",
|
||||
"# Connect to OpenCode server and use default LLM provider\n",
|
||||
"llm = LLM(\"opencode\")\n",
|
||||
"\n",
|
||||
"generate(\"Show a Python Hello World example\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "55dd0fad",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"OK, now let's try this for a more advanced use case."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"id": "b3080841",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"```python\n",
|
||||
"from txtai import Embeddings, LLM\n",
|
||||
"\n",
|
||||
"# Create embeddings model\n",
|
||||
"embeddings = Embeddings(path=\"sentence-transformers/all-MiniLM-L6-v2\")\n",
|
||||
"\n",
|
||||
"# Create LLM\n",
|
||||
"llm = LLM(\"huggingface/Qwen/Qwen2.5-1.5B-Instruct\")\n",
|
||||
"\n",
|
||||
"# Sample documents\n",
|
||||
"documents = [\n",
|
||||
" \"Python is a programming language created by Guido van Rossum\",\n",
|
||||
" \"Machine learning is a subset of artificial intelligence\",\n",
|
||||
" \"Deep learning uses neural networks with multiple layers\"\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"# Build the index\n",
|
||||
"embeddings.index([(i, doc) for i, doc in enumerate(documents)])\n",
|
||||
"\n",
|
||||
"# RAG function\n",
|
||||
"def rag_query(question):\n",
|
||||
" # Search for relevant documents\n",
|
||||
" results = embeddings.search(question, 2)\n",
|
||||
" \n",
|
||||
" # Build context from retrieved documents\n",
|
||||
" context = \" \".join([documents[result[0]] for result in results])\n",
|
||||
" \n",
|
||||
" # Generate response using LLM with context\n",
|
||||
" prompt = f\"Context: {context}\\n\\nQuestion: {question}\\n\\nAnswer:\"\n",
|
||||
" return llm(prompt)\n",
|
||||
"\n",
|
||||
"# Example usage\n",
|
||||
"question = \"What is Python?\"\n",
|
||||
"answer = rag_query(question)\n",
|
||||
"print(f\"Question: {question}\")\n",
|
||||
"print(f\"Answer: {answer}\")\n",
|
||||
"```"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"generate(\"Show a simple and basic TxtAI RAG example\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b44e0309",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Interesting! It's not limited to Python though."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"id": "0a2ec1b4",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"```assembly\n",
|
||||
"section .data\n",
|
||||
" hello db 'Hello, World!', 10 ; String with newline\n",
|
||||
" hello_len equ $ - hello ; Length of string\n",
|
||||
"\n",
|
||||
"section .text\n",
|
||||
" global _start\n",
|
||||
"\n",
|
||||
"_start:\n",
|
||||
" ; sys_write system call\n",
|
||||
" mov eax, 4 ; sys_write\n",
|
||||
" mov ebx, 1 ; stdout\n",
|
||||
" mov ecx, hello ; message\n",
|
||||
" mov edx, hello_len ; message length\n",
|
||||
" int 0x80 ; kernel interrupt\n",
|
||||
"\n",
|
||||
" ; sys_exit system call\n",
|
||||
" mov eax, 1 ; sys_exit\n",
|
||||
" mov ebx, 0 ; exit code 0\n",
|
||||
" int 0x80 ; kernel interrupt\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"Compile and run:\n",
|
||||
"```bash\n",
|
||||
"nasm -f elf32 hello.asm -o hello.o\n",
|
||||
"ld -m elf_i386 hello.o -o hello\n",
|
||||
"./hello\n",
|
||||
"```"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"generate(\"Show a x86 assembly hello world\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"id": "4e28d5c4",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"```python\n",
|
||||
"a = None\n",
|
||||
"if a is not None:\n",
|
||||
" a.startswith(\"hello\")\n",
|
||||
"else:\n",
|
||||
" print(\"a is None, cannot call startswith\")\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"**What's wrong:** The variable `a` is `None`, which is a `NoneType` object. The `startswith()` method only exists for strings, not for `NoneType`. You need to check if `a` is not `None` before calling string methods on it."
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"generate(\"\"\"\n",
|
||||
"Fix the following error in the code below and explain what's wrong.\n",
|
||||
"\n",
|
||||
"Error:\n",
|
||||
"AttributeError: 'NoneType' object has no attribute 'startswith'\n",
|
||||
"\n",
|
||||
"Code:\n",
|
||||
"a = None\n",
|
||||
"a.startswith(\"hello\")\n",
|
||||
"\"\"\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "156122f0",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# OpenAI-compatible Endpoint\n",
|
||||
"\n",
|
||||
"Now that OpenCode is integrated into the TxtAI ecosystem, it opens a lot of different opportunities. Let's say we want to host a local OpenAI compatible endpoint, no problem!\n",
|
||||
"\n",
|
||||
"Write the following file.\n",
|
||||
"\n",
|
||||
"## config.yml\n",
|
||||
"\n",
|
||||
"```yaml\n",
|
||||
"# Enable OpenAI compat endpoint\n",
|
||||
"openai: True\n",
|
||||
"\n",
|
||||
"llm:\n",
|
||||
" path: opencode\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"and start the following process.\n",
|
||||
"\n",
|
||||
"```\n",
|
||||
"CONFIG=config.yml uvicorn \"txtai.api:app\"\n",
|
||||
"```"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 16,
|
||||
"id": "33057f59",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"```python\n",
|
||||
"def factorial(n):\n",
|
||||
" if n < 0:\n",
|
||||
" raise ValueError(\"Factorial is not defined for negative numbers\")\n",
|
||||
" if n == 0 or n == 1:\n",
|
||||
" return 1\n",
|
||||
" result = 1\n",
|
||||
" for i in range(2, n + 1):\n",
|
||||
" result *= i\n",
|
||||
" return result\n",
|
||||
"\n",
|
||||
"# Example usage:\n",
|
||||
"# print(factorial(5)) # Output: 120\n",
|
||||
"```"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from IPython.display import Markdown, display\n",
|
||||
"\n",
|
||||
"from openai import OpenAI\n",
|
||||
"\n",
|
||||
"client = OpenAI(\n",
|
||||
" base_url=\"http://localhost:8000/v1\",\n",
|
||||
" api_key=\"api-key\",\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"response = client.chat.completions.create(\n",
|
||||
" messages=[{\n",
|
||||
" \"role\": \"user\",\n",
|
||||
" \"content\": \"Show a Python factorial function\",\n",
|
||||
" }],\n",
|
||||
" model=\"opencode\",\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"display(Markdown((response.choices[0].message.content)))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "58a9169b",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Just like before we get an answer! This time via the OpenAI client. Fun times."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "47f9dbbd",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This notebook showed how to connect TxtAI and OpenCode. This will lead to some interesting integrations. One such example is [ncoder](https://github.com/neuml/ncoder), an AI coding agent that integrates with Jupyter Notebooks. Stay tuned for more!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"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.10.19"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,358 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "5a13697b",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Agentic College Search\n",
|
||||
"\n",
|
||||
"This example will demonstrate how to use `txtai` agents to identify list of strong engineering colleges with baseball programs 🎓⚙️⚾.\n",
|
||||
"\n",
|
||||
"We'll setup a default [agents.md](https://github.com/agentsmd/agents.md) file and then give an LLM access to the web.\n",
|
||||
"\n",
|
||||
"Let's get started!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "face0f9e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "f5da3842",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[agent]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f9a6a500",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Define an `agents.md` file\n",
|
||||
"\n",
|
||||
"Next, we'll define an agents.md file. This file defines default behavior for our college search agent and is inserted into the system prompt of every agent run."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "4e0c500d",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%writefile agents.md\n",
|
||||
"# College Researcher\n",
|
||||
"\n",
|
||||
"You are a college research agent. Your job is to help find the best matching schools given a set of criteria.\n",
|
||||
"\n",
|
||||
"## Academics\n",
|
||||
"\n",
|
||||
"Work to match the request with a set of schools that match the academic requirements\n",
|
||||
"\n",
|
||||
"## Athletics\n",
|
||||
"\n",
|
||||
"Some requests may be interested in collegiate athletics. Make sure to balance that with other parts of the request."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "83dee7de",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Identify the Top 20 schools\n",
|
||||
"\n",
|
||||
"Next, let's run a prompt to identify the best 20 colleges given a set of criteria.\n",
|
||||
"\n",
|
||||
"First, we'll setup the scaffolding code to create and run an agent. We'll use a [Qwen3 4B non-thinking LLM](https://huggingface.co/Qwen/Qwen3-4B-Instruct-2507) as the agent's model. We'll add the `websearch` and `webview` tools to the agent along with the `agents.md` file previously created.\n",
|
||||
"\n",
|
||||
"Additionally, we'll add a sliding window of the last 5 responses as \"agent memory\". This will help create a rolling dialogue. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "410f2585",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import Agent\n",
|
||||
"from IPython.display import display, Markdown\n",
|
||||
"\n",
|
||||
"def run(query, reset=False):\n",
|
||||
" answer = agent(query, maxlength=60000, reset=reset)\n",
|
||||
" display(Markdown(answer))\n",
|
||||
"\n",
|
||||
"agent = Agent(model=\"Qwen/Qwen3-4B-Instruct-2507\", tools=[\"websearch\", \"webview\"], memory=5, instructions=\"agents.md\", max_steps=50)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "53aa302b",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Now we can run the search. To come up with our list of colleges, a series of web searches were executed by the agent. It's like if you had the superpower 🚀 to kick off 10+ web searches with slight variations each time as you go. It's a rapid fire set of operations to collect information for the LLM for analysis.\n",
|
||||
"\n",
|
||||
"_Note: The output below has the `verbosity_level` set to 0 for brevity. If you run this again, you'll get the full details with each step._"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"id": "f2d1ef2e",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #d4b702; text-decoration-color: #d4b702\">╭──────────────────────────────────────────────────── </span><span style=\"color: #d4b702; text-decoration-color: #d4b702; font-weight: bold\">New run</span><span style=\"color: #d4b702; text-decoration-color: #d4b702\"> ────────────────────────────────────────────────────╮</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">List 20 universities located in the Mid-Atlantic or Northeast that would be good candidates to consider with </span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">top 50 engineering programs and an active baseball team. Visit websites covering the top engineering schools in</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">the Mid-Atlantic and Northeast. Add a rationale for each of the picks in a few words.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">╰─ PipelineModel - Qwen/Qwen3-4B-Instruct-2507 ───────────────────────────────────────────────────────────────────╯</span>\n",
|
||||
"</pre>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"\u001b[38;2;212;183;2m╭─\u001b[0m\u001b[38;2;212;183;2m───────────────────────────────────────────────────\u001b[0m\u001b[38;2;212;183;2m \u001b[0m\u001b[1;38;2;212;183;2mNew run\u001b[0m\u001b[38;2;212;183;2m \u001b[0m\u001b[38;2;212;183;2m───────────────────────────────────────────────────\u001b[0m\u001b[38;2;212;183;2m─╮\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mList 20 universities located in the Mid-Atlantic or Northeast that would be good candidates to consider with \u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mtop 50 engineering programs and an active baseball team. Visit websites covering the top engineering schools in\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mthe Mid-Atlantic and Northeast. Add a rationale for each of the picks in a few words.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m╰─\u001b[0m\u001b[38;2;212;183;2m PipelineModel - Qwen/Qwen3-4B-Instruct-2507 \u001b[0m\u001b[38;2;212;183;2m──────────────────────────────────────────────────────────────────\u001b[0m\u001b[38;2;212;183;2m─╯\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"Here are 20 universities in the Mid-Atlantic or Northeast with top 50 engineering programs and active baseball teams, along with a brief rationale for each:\n",
|
||||
"\n",
|
||||
"1. **Massachusetts Institute of Technology (MIT)** – Top-ranked engineering program in the Northeast with a strong research focus and active baseball team.\n",
|
||||
"2. **Princeton University** – Renowned for engineering excellence and competitive baseball in the Ivy League.\n",
|
||||
"3. **Columbia University** – Strong engineering programs and a consistently active baseball team.\n",
|
||||
"4. **University of Pennsylvania** – Top-tier engineering and a well-established baseball program in the Ivy League.\n",
|
||||
"5. **Carnegie Mellon University** – Leading in computer and engineering fields with a competitive baseball team.\n",
|
||||
"6. **Cornell University** – Excellent engineering programs and a strong baseball presence.\n",
|
||||
"7. **Northeastern University** – Strong engineering programs and active baseball team.\n",
|
||||
"8. **New York University** – Top engineering programs with a competitive baseball team.\n",
|
||||
"9. **Rensselaer Polytechnic Institute (RPI)** – Strong engineering focus and active baseball in the Liberty League.\n",
|
||||
"10. **Rutgers University** – Top engineering programs and a well-established baseball team.\n",
|
||||
"11. **Lehigh University** – Strong engineering and active baseball in the Patriot League.\n",
|
||||
"12. **University of Connecticut** – Top engineering programs and a competitive baseball team.\n",
|
||||
"13. **University of Massachusetts Amherst** – Solid engineering programs and active baseball.\n",
|
||||
"14. **University of Massachusetts Lowell** – Strong engineering and active baseball in the America East Conference.\n",
|
||||
"15. **Syracuse University** – Top engineering programs and a competitive baseball team.\n",
|
||||
"16. **Boston University** – Strong engineering and active baseball.\n",
|
||||
"17. **Dartmouth College** – Excellent engineering programs and competitive baseball.\n",
|
||||
"18. **Brown University** – Top engineering and active baseball.\n",
|
||||
"19. **Bucknell University** – Solid engineering and active baseball.\n",
|
||||
"20. **University of Rochester** – Strong engineering and competitive baseball.\n",
|
||||
"\n",
|
||||
"These institutions combine top-tier engineering programs with active collegiate baseball teams, making them excellent candidates for students seeking both academic excellence and athletic engagement."
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"run((\n",
|
||||
" \"List 20 universities located in the Mid-Atlantic or Northeast that would be good candidates to consider \"\n",
|
||||
" \"with top 50 engineering programs and an active baseball team. \"\n",
|
||||
" \"Visit websites covering the top engineering schools in the Mid-Atlantic and Northeast. \"\n",
|
||||
" \"Add a rationale for each of the picks in a few words. \"\n",
|
||||
"))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a7533137",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Interesting! Looks like quite an impressive list of colleges. This output is the product of taking the results of 10+ web searches and analyzing the results with an LLM. A very practical and realistic use of \"Agentic AI\". This is a great list of strong engineering schools with baseball programs 🎓⚙️⚾.\n",
|
||||
"\n",
|
||||
"Let's narrow it down to the top 5 engineering schools."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 105,
|
||||
"id": "5e261825",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #d4b702; text-decoration-color: #d4b702\">╭──────────────────────────────────────────────────── </span><span style=\"color: #d4b702; text-decoration-color: #d4b702; font-weight: bold\">New run</span><span style=\"color: #d4b702; text-decoration-color: #d4b702\"> ────────────────────────────────────────────────────╮</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">Read the list below and pick the Top 5 engineering programs.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">Use the following conversation history to help answer the question above.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">User: List 20 universities located in the Mid-Atlantic or Northeast that would be good candidates to consider </span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">with top 50 engineering programs and an active baseball team. Visit websites covering the top engineering </span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">schools in the Mid-Atlantic and Northeast. Add a rationale for each of the picks in a few words. </span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">Assistant: Here are 20 universities in the Mid-Atlantic or Northeast with top 50 engineering programs and </span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">active baseball teams, along with a brief rationale for each:</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">1. **Massachusetts Institute of Technology (MIT)** – Top-ranked engineering program in the Northeast with a </span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">strong research focus and active baseball team.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">2. **Princeton University** – Renowned for engineering excellence and competitive baseball in the Ivy League.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">3. **Columbia University** – Strong engineering programs and a consistently active baseball team.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">4. **University of Pennsylvania** – Top-tier engineering and a well-established baseball program in the Ivy </span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">League.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">5. **Carnegie Mellon University** – Leading in computer and engineering fields with a competitive baseball </span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">team.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">6. **Cornell University** – Excellent engineering programs and a strong baseball presence.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">7. **Northeastern University** – Strong engineering programs and active baseball team.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">8. **New York University** – Top engineering programs with a competitive baseball team.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">9. **Rensselaer Polytechnic Institute (RPI)** – Strong engineering focus and active baseball in the Liberty </span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">League.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">10. **Rutgers University** – Top engineering programs and a well-established baseball team.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">11. **Lehigh University** – Strong engineering and active baseball in the Patriot League.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">12. **University of Connecticut** – Top engineering programs and a competitive baseball team.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">13. **University of Massachusetts Amherst** – Solid engineering programs and active baseball.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">14. **University of Massachusetts Lowell** – Strong engineering and active baseball in the America East </span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">Conference.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">15. **Syracuse University** – Top engineering programs and a competitive baseball team.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">16. **Boston University** – Strong engineering and active baseball.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">17. **Dartmouth College** – Excellent engineering programs and competitive baseball.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">18. **Brown University** – Top engineering and active baseball.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">19. **Bucknell University** – Solid engineering and active baseball.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">20. **University of Rochester** – Strong engineering and competitive baseball.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">These institutions combine top-tier engineering programs with active collegiate baseball teams, making them </span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">excellent candidates for students seeking both academic excellence and athletic engagement.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"font-weight: bold\">If the history is irrelevant, forget it and use other tools to answer the question.</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span> <span style=\"color: #d4b702; text-decoration-color: #d4b702\">│</span>\n",
|
||||
"<span style=\"color: #d4b702; text-decoration-color: #d4b702\">╰─ PipelineModel - Qwen/Qwen3-4B-Instruct-2507 ───────────────────────────────────────────────────────────────────╯</span>\n",
|
||||
"</pre>\n"
|
||||
],
|
||||
"text/plain": [
|
||||
"\u001b[38;2;212;183;2m╭─\u001b[0m\u001b[38;2;212;183;2m───────────────────────────────────────────────────\u001b[0m\u001b[38;2;212;183;2m \u001b[0m\u001b[1;38;2;212;183;2mNew run\u001b[0m\u001b[38;2;212;183;2m \u001b[0m\u001b[38;2;212;183;2m───────────────────────────────────────────────────\u001b[0m\u001b[38;2;212;183;2m─╮\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mRead the list below and pick the Top 5 engineering programs.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mUse the following conversation history to help answer the question above.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mUser: List 20 universities located in the Mid-Atlantic or Northeast that would be good candidates to consider \u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mwith top 50 engineering programs and an active baseball team. Visit websites covering the top engineering \u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mschools in the Mid-Atlantic and Northeast. Add a rationale for each of the picks in a few words. \u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mAssistant: Here are 20 universities in the Mid-Atlantic or Northeast with top 50 engineering programs and \u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mactive baseball teams, along with a brief rationale for each:\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m1. **Massachusetts Institute of Technology (MIT)** – Top-ranked engineering program in the Northeast with a \u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mstrong research focus and active baseball team.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m2. **Princeton University** – Renowned for engineering excellence and competitive baseball in the Ivy League.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m3. **Columbia University** – Strong engineering programs and a consistently active baseball team.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m4. **University of Pennsylvania** – Top-tier engineering and a well-established baseball program in the Ivy \u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mLeague.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m5. **Carnegie Mellon University** – Leading in computer and engineering fields with a competitive baseball \u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mteam.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m6. **Cornell University** – Excellent engineering programs and a strong baseball presence.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m7. **Northeastern University** – Strong engineering programs and active baseball team.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m8. **New York University** – Top engineering programs with a competitive baseball team.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m9. **Rensselaer Polytechnic Institute (RPI)** – Strong engineering focus and active baseball in the Liberty \u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mLeague.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m10. **Rutgers University** – Top engineering programs and a well-established baseball team.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m11. **Lehigh University** – Strong engineering and active baseball in the Patriot League.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m12. **University of Connecticut** – Top engineering programs and a competitive baseball team.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m13. **University of Massachusetts Amherst** – Solid engineering programs and active baseball.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m14. **University of Massachusetts Lowell** – Strong engineering and active baseball in the America East \u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mConference.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m15. **Syracuse University** – Top engineering programs and a competitive baseball team.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m16. **Boston University** – Strong engineering and active baseball.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m17. **Dartmouth College** – Excellent engineering programs and competitive baseball.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m18. **Brown University** – Top engineering and active baseball.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m19. **Bucknell University** – Solid engineering and active baseball.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1m20. **University of Rochester** – Strong engineering and competitive baseball.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mThese institutions combine top-tier engineering programs with active collegiate baseball teams, making them \u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mexcellent candidates for students seeking both academic excellence and athletic engagement.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[1mIf the history is irrelevant, forget it and use other tools to answer the question.\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m│\u001b[0m \u001b[38;2;212;183;2m│\u001b[0m\n",
|
||||
"\u001b[38;2;212;183;2m╰─\u001b[0m\u001b[38;2;212;183;2m PipelineModel - Qwen/Qwen3-4B-Instruct-2507 \u001b[0m\u001b[38;2;212;183;2m──────────────────────────────────────────────────────────────────\u001b[0m\u001b[38;2;212;183;2m─╯\u001b[0m\n"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"The top 5 engineering programs from the list of universities in the Mid-Atlantic or Northeast with active baseball teams are:\n",
|
||||
"\n",
|
||||
"1. **Massachusetts Institute of Technology (MIT)** – Top-ranked engineering program in the Northeast with a strong research focus and active baseball team.\n",
|
||||
"2. **Princeton University** – Renowned for engineering excellence and competitive baseball in the Ivy League.\n",
|
||||
"3. **Carnegie Mellon University** – Leading in computer and engineering fields with a competitive baseball team.\n",
|
||||
"4. **Cornell University** – Excellent engineering programs and a strong baseball presence.\n",
|
||||
"5. **University of Pennsylvania** – Top-tier engineering and a well-established baseball program in the Ivy League.\n",
|
||||
"\n",
|
||||
"These schools stand out due to their exceptional engineering programs and active collegiate baseball teams, offering a balanced academic and athletic experience."
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"run(\"Read the list below and pick the Top 5 engineering programs.\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e6c27f16",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This example showed how to run an Agentic Loop that identified the best set of colleges to research further. It's an agent with memory and can be used to build an ongoing dialogue with an agent. Give it a try!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"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.10.19"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,410 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "8d8cb3ad",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# TxtAI got skills\n",
|
||||
"\n",
|
||||
"This example will demonstrate how to use `txtai` agents with [`skill.md`](https://agentskills.io/specification) files.\n",
|
||||
"\n",
|
||||
"We'll setup a `skill.md` file with details on how to use TxtAI and run a series of agent requests.\n",
|
||||
"\n",
|
||||
"Let's get started!"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "00a5a0ea",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "0ff7aa19",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[agent]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "98bcd5c3",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Define a `skill.md` file\n",
|
||||
"\n",
|
||||
"Next, we'll create our `skill.md` file. This file has examples on how to build embeddings databases, how to use re-ranker pipelines, RAG pipelines and more.\n",
|
||||
"\n",
|
||||
"The upside of a `skill.md` file vs an `agents.md` file is that it can be dynamically added to the agent context. The description helps the agent decide if the `skill` is necessary given the request. Think of it like a dynamic knowledge base that's easy to modify."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "49179643",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%writefile skill.md\n",
|
||||
"---\n",
|
||||
"name: txtai\n",
|
||||
"description: Examples on how to build txtai embeddings databases, txtai RAG pipelines, txtai reranker pipelines and txtai translation pipelines\n",
|
||||
"---\n",
|
||||
"\n",
|
||||
"# Build an embeddings database\n",
|
||||
"\n",
|
||||
"```python\n",
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"# Create embeddings model, backed by sentence-transformers & transformers\n",
|
||||
"embeddings = Embeddings(path=\"sentence-transformers/nli-mpnet-base-v2\")\n",
|
||||
"\n",
|
||||
"data = [\n",
|
||||
" \"US tops 5 million confirmed virus cases\",\n",
|
||||
" \"Canada's last fully intact ice shelf has suddenly collapsed, \" +\n",
|
||||
" \"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 \" +\n",
|
||||
" \"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",
|
||||
"\n",
|
||||
"# Index the list of text\n",
|
||||
"embeddings.index(data)\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"# Search an embeddings database\n",
|
||||
"\n",
|
||||
"```python\n",
|
||||
"embeddings.search(\"Search query\")\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"# Build a RAG pipeline\n",
|
||||
"\n",
|
||||
"```python\n",
|
||||
"from txtai import Embeddings, RAG\n",
|
||||
"\n",
|
||||
"# Input data\n",
|
||||
"data = [\n",
|
||||
" \"US tops 5 million confirmed virus cases\",\n",
|
||||
" \"Canada's last fully intact ice shelf has suddenly collapsed, \" +\n",
|
||||
" \"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 \" +\n",
|
||||
" \"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",
|
||||
"\n",
|
||||
"# Build embeddings index\n",
|
||||
"embeddings = Embeddings(content=True)\n",
|
||||
"embeddings.index(data)\n",
|
||||
"\n",
|
||||
"# Create the RAG pipeline\n",
|
||||
"rag = RAG(embeddings, \"Qwen/Qwen3-0.6B\", template=\"\"\"\n",
|
||||
" Answer the following question using the provided context.\n",
|
||||
"\n",
|
||||
" Question:\n",
|
||||
" {question}\n",
|
||||
"\n",
|
||||
" Context:\n",
|
||||
" {context}\n",
|
||||
"\"\"\")\n",
|
||||
"\n",
|
||||
"# Run RAG pipeline\n",
|
||||
"rag(\"What was won?\")\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"# Translate text from English into French\n",
|
||||
"\n",
|
||||
"```python\n",
|
||||
"from txtai.pipeline import Translation\n",
|
||||
"\n",
|
||||
"# Create and run pipeline\n",
|
||||
"translate = Translation()\n",
|
||||
"translate(\"This is a test translation\", \"fr\")\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"# Re-ranker pipeline\n",
|
||||
"\n",
|
||||
"```python\n",
|
||||
"from txtai import Embeddings\n",
|
||||
"from txtai.pipeline import Reranker, Similarity\n",
|
||||
"\n",
|
||||
"# Embeddings instance\n",
|
||||
"embeddings = Embeddings()\n",
|
||||
"embeddings.load(provider=\"huggingface-hub\", container=\"neuml/txtai-wikipedia\")\n",
|
||||
"\n",
|
||||
"# Similarity instance\n",
|
||||
"similarity = Similarity(path=\"colbert-ir/colbertv2.0\", lateencode=True)\n",
|
||||
"\n",
|
||||
"# Reranking pipeline\n",
|
||||
"reranker = Reranker(embeddings, similarity)\n",
|
||||
"reranker(\"Tell me about AI\")\n",
|
||||
"```"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "4ed26969",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Query for TxtAI questions\n",
|
||||
"\n",
|
||||
"Now, let's try this out and see if the LLM is smart enough to use the defined skill vs. going out on the web.\n",
|
||||
"\n",
|
||||
"Let's setup the scaffolding code to create and run an agent. We'll use a [Qwen3 4B non-thinking LLM](https://huggingface.co/Qwen/Qwen3-4B-Instruct-2507) as the agent's model. We'll add the `websearch` and `webview` tools to the agent along with the `skill.md` file previously created.\n",
|
||||
"\n",
|
||||
"Additionally, we'll add a sliding window of the last 2 responses as \"agent memory\". This will help create a rolling dialogue. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "b6a1aaf1",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from txtai import Agent\n",
|
||||
"from IPython.display import display, Markdown\n",
|
||||
"\n",
|
||||
"def run(query, reset=False):\n",
|
||||
" answer = agent(query, maxlength=50000, reset=reset)\n",
|
||||
" display(Markdown(answer))\n",
|
||||
"\n",
|
||||
"agent = Agent(\n",
|
||||
" model=\"Qwen/Qwen3-4B-Instruct-2507\",\n",
|
||||
" tools=[\"websearch\", \"webview\", \"skill.md\"],\n",
|
||||
" memory=2,\n",
|
||||
" verbosity_level=0\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "19302f54",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"First, we'll ask how to build a TxtAI embeddings database."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"id": "121d0b8a",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"To create a txtai embeddings program that indexes data, use the following Python code:\n",
|
||||
"\n",
|
||||
"```python\n",
|
||||
"from txtai import Embeddings\n",
|
||||
"\n",
|
||||
"# Create embeddings model using sentence-transformers\n",
|
||||
"embeddings = Embeddings(path=\"sentence-transformers/nli-mpnet-base-v2\")\n",
|
||||
"\n",
|
||||
"# Sample data to index\n",
|
||||
"data = [\n",
|
||||
" \"US tops 5 million confirmed virus cases\",\n",
|
||||
" \"Canada's last fully intact ice shelf has suddenly collapsed, \" +\n",
|
||||
" \"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 \" +\n",
|
||||
" \"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",
|
||||
"\n",
|
||||
"# Index the data\n",
|
||||
"embeddings.index(data)\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"This program initializes an embeddings database using the `sentence-transformers/nli-mpnet-base-v2` model and indexes a list of text data. The indexed data can later be searched or used in other applications like retrieval, RAG, or translation."
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"run(\"Write a txtai embeddings program that indexes data\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f723e07f",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The Agent pulled the correct section from the `skill.md` file. Now, let's look for an example on how to use a re-ranker pipeline."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "221a590c",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"```python\n",
|
||||
"from txtai import Embeddings\n",
|
||||
"from txtai.pipeline import Reranker, Similarity\n",
|
||||
"\n",
|
||||
"# Embeddings instance\n",
|
||||
"embeddings = Embeddings()\n",
|
||||
"embeddings.load(provider=\"huggingface-hub\", container=\"neuml/txtai-wikipedia\")\n",
|
||||
"\n",
|
||||
"# Similarity instance\n",
|
||||
"similarity = Similarity(path=\"colbert-ir/colbertv2.0\", lateencode=True)\n",
|
||||
"\n",
|
||||
"# Reranking pipeline\n",
|
||||
"reranker = Reranker(embeddings, similarity)\n",
|
||||
"reranker(\"Tell me about AI\")\n",
|
||||
"```"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"run(\"Write a txtai re-ranker pipeline\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "4af3a6b6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Great! This also works. But remember we have access to an LLM here. It doesn't have to blindly just pull the text. Let's ask it to modify the last example."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"id": "c9f1eed4",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"```python\n",
|
||||
"from txtai import Embeddings\n",
|
||||
"from txtai.pipeline import Reranker, Similarity\n",
|
||||
"\n",
|
||||
"# Embeddings instance\n",
|
||||
"embeddings = Embeddings()\n",
|
||||
"embeddings.load(provider=\"huggingface-hub\", container=\"neuml/txtai-wikipedia\")\n",
|
||||
"\n",
|
||||
"# Similarity instance with a different reranker model and lateencode disabled\n",
|
||||
"similarity = Similarity(path=\"BAAI/bge-reranker-base\", lateencode=False)\n",
|
||||
"\n",
|
||||
"# Reranking pipeline\n",
|
||||
"reranker = Reranker(embeddings, similarity)\n",
|
||||
"reranker(\"Tell me about AI\")\n",
|
||||
"```"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"run(\"Update the similarity path to use another reranker model. Disable lateencode.\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "4a3bb21e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Notice it just edited the code with a different reranker model and even added a comment to note this change.\n",
|
||||
"\n",
|
||||
"We can also clear the rolling dialogue and start fresh. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"id": "ad53f7a8",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/markdown": [
|
||||
"```python\n",
|
||||
"from txtai.pipeline import Translation\n",
|
||||
"\n",
|
||||
"# Create and run pipeline for English to Spanish translation\n",
|
||||
"translate = Translation()\n",
|
||||
"translated_text = translate(\"This is a test translation\", \"es\")\n",
|
||||
"print(translated_text)\n",
|
||||
"```"
|
||||
],
|
||||
"text/plain": [
|
||||
"<IPython.core.display.Markdown object>"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"run(\"Write a txtai program that translate text from English to Spanish\", reset=True)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "007daee0",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This example shows how to add a `skill.md` file to `txtai` agents. Go give it a try!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"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.10.19"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,337 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "40f124c0",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Progressive Distillation\n",
|
||||
"\n",
|
||||
"Now that almost everyone has thought about or is actively integrating AI workflows into their projects, some might ask is this all worth the cost? Many think the current economics of the AI space don't scale and that there will be upward price movement. Others still might not be comfortable with sending their data to remote services for processing. Then there is the crowd that wants to deploy models in small spaces with limited compute.\n",
|
||||
"\n",
|
||||
"Are there ways we can deploy small models locally and run at a lower cost? Yes with [Knowledge Distillation](https://en.wikipedia.org/wiki/Knowledge_distillation). Knowledge distillation can get a bad rap due to it's questionable use in training some Large Language Models (LLMs). But it's a perfectly valid way to transfer performance from a larger model to a smaller one. Especially when both models are yours and/or open. \n",
|
||||
"\n",
|
||||
"This notebook will explore progressive distillation which is a technique to incrementally transfer knowledge from a series of larger teacher models into a smaller student."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f4aad867",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Install dependencies\n",
|
||||
"\n",
|
||||
"Install `txtai` and all dependencies."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"id": "16e2c760",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline-train] datasets"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e3a23082",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Setup the Training Pipeline\n",
|
||||
"\n",
|
||||
"The first step we need to do is setup up the training pipeline. We'll use the Hugging Face Training framework to build a series of models.\n",
|
||||
"\n",
|
||||
"The following code establishes a `train` method, `test` method and loads the classification training data."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "ea7ee576",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from datasets import load_dataset\n",
|
||||
"from transformers import AutoModelForSequenceClassification, AutoTokenizer\n",
|
||||
"from txtai.pipeline import HFTrainer, Labels\n",
|
||||
"\n",
|
||||
"def train(teacher, student, distillation, **kwargs):\n",
|
||||
" trainer = HFTrainer()\n",
|
||||
"\n",
|
||||
" model = AutoModelForSequenceClassification.from_pretrained(student, trust_remote_code=True)\n",
|
||||
" tokenizer = AutoTokenizer.from_pretrained(student, trust_remote_code=True)\n",
|
||||
"\n",
|
||||
" return trainer(\n",
|
||||
" (model, tokenizer),\n",
|
||||
" ds[\"train\"],\n",
|
||||
" columns=(\"sentence\", \"label\"),\n",
|
||||
" maxlength=maxlength,\n",
|
||||
" teacher=teacher,\n",
|
||||
" distillation=distillation,\n",
|
||||
" **kwargs\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
"def test(model):\n",
|
||||
" labels = Labels(model, dynamic=False, trust_remote_code=True)\n",
|
||||
"\n",
|
||||
" # Determine accuracy on test set\n",
|
||||
" results = [row[\"label\"] == labels(row[\"sentence\"], max_length=maxlength)[0][0] for row in ds[\"validation\"]]\n",
|
||||
" print(sum(results) / len(ds[\"validation\"]))\n",
|
||||
"\n",
|
||||
"# Hugging Face dataset\n",
|
||||
"ds = load_dataset(\"nyu-mll/glue\", \"sst2\")\n",
|
||||
"maxlength = 128"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "732fad7d",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Progressive Distillation\n",
|
||||
"\n",
|
||||
"We're going to build a [bert-hash-femto](https://huggingface.co/NeuML/bert-hash-femto) classifier which is a extremely small 250K parameter model. This model was pretrained using the same recipe as [BERT](https://arxiv.org/abs/1810.04805). The paper [Well-Read Students Learn Better: On the Importance of Pre-training Compact Models](https://arxiv.org/abs/1908.08962) established that small models perform better with Knowledge Distillation tasks when they are pretrained. [This article](https://huggingface.co/blog/NeuML/bert-hash-embeddings) is also a good source for more information on the topic.\n",
|
||||
"\n",
|
||||
"The next series of steps will do the following:\n",
|
||||
"\n",
|
||||
"- Train a [BERT Mini](https://huggingface.co/google/bert_uncased_L-4_H-256_A-4) model using a [BERT Large](https://huggingface.co/assemblyai/bert-large-uncased-sst2) teacher\n",
|
||||
"- Train a [BERT Tiny](https://huggingface.co/google/bert_uncased_L-2_H-128_A-2) model using the BERT Mini teacher above\n",
|
||||
"- Train a [BERT Hash Femto](https://huggingface.co/NeuML/bert-hash-femto) model using the BERT Tiny teacher above\n",
|
||||
"\n",
|
||||
"This path goes from a 11M parameter model -> 4.4M parameter model -> 250K parameter model\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "43baccd6",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[transformers] \u001b[1mBertForSequenceClassification LOAD REPORT\u001b[0m from: google/bert_uncased_L-4_H-256_A-4\n",
|
||||
"Key | Status | \n",
|
||||
"-------------------------------------------+------------+-\n",
|
||||
"cls.predictions.transform.dense.weight | UNEXPECTED | \n",
|
||||
"cls.predictions.decoder.bias | UNEXPECTED | \n",
|
||||
"cls.predictions.transform.LayerNorm.bias | UNEXPECTED | \n",
|
||||
"cls.predictions.decoder.weight | UNEXPECTED | \n",
|
||||
"cls.seq_relationship.bias | UNEXPECTED | \n",
|
||||
"cls.seq_relationship.weight | UNEXPECTED | \n",
|
||||
"cls.predictions.transform.dense.bias | UNEXPECTED | \n",
|
||||
"cls.predictions.bias | UNEXPECTED | \n",
|
||||
"cls.predictions.transform.LayerNorm.weight | UNEXPECTED | \n",
|
||||
"classifier.weight | MISSING | \n",
|
||||
"classifier.bias | MISSING | \n",
|
||||
"\n",
|
||||
"Notes:\n",
|
||||
"- UNEXPECTED:\tcan be ignored when loading from different task/architecture; not ok if you expect identical arch.\n",
|
||||
"- MISSING:\tthose params were newly initialized because missing from the checkpoint. Consider training on your downstream task.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"0.875\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"test(train(\n",
|
||||
" \"assemblyai/bert-large-uncased-sst2\",\n",
|
||||
" \"google/bert_uncased_L-4_H-256_A-4\",\n",
|
||||
" (1.0, 1.0),\n",
|
||||
" learning_rate=1e-4,\n",
|
||||
" num_train_epochs=5,\n",
|
||||
" per_device_train_batch_size=32,\n",
|
||||
" output_dir=\"bert-mini-sst2\"\n",
|
||||
"))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"id": "9ba3158e",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[transformers] \u001b[1mBertForSequenceClassification LOAD REPORT\u001b[0m from: google/bert_uncased_L-2_H-128_A-2\n",
|
||||
"Key | Status | \n",
|
||||
"-------------------------------------------+------------+-\n",
|
||||
"cls.predictions.transform.dense.weight | UNEXPECTED | \n",
|
||||
"cls.predictions.transform.LayerNorm.bias | UNEXPECTED | \n",
|
||||
"cls.seq_relationship.weight | UNEXPECTED | \n",
|
||||
"cls.seq_relationship.bias | UNEXPECTED | \n",
|
||||
"cls.predictions.transform.dense.bias | UNEXPECTED | \n",
|
||||
"cls.predictions.bias | UNEXPECTED | \n",
|
||||
"cls.predictions.transform.LayerNorm.weight | UNEXPECTED | \n",
|
||||
"classifier.weight | MISSING | \n",
|
||||
"classifier.bias | MISSING | \n",
|
||||
"\n",
|
||||
"Notes:\n",
|
||||
"- UNEXPECTED:\tcan be ignored when loading from different task/architecture; not ok if you expect identical arch.\n",
|
||||
"- MISSING:\tthose params were newly initialized because missing from the checkpoint. Consider training on your downstream task.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"0.8325688073394495\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"test(train(\n",
|
||||
" \"bert-mini-sst2\",\n",
|
||||
" \"google/bert_uncased_L-2_H-128_A-2\",\n",
|
||||
" (1.0, 1.0),\n",
|
||||
" learning_rate=1e-4,\n",
|
||||
" num_train_epochs=5,\n",
|
||||
" per_device_train_batch_size=32,\n",
|
||||
" output_dir=\"bert-tiny-sst2\"\n",
|
||||
"))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"id": "e98d950e",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[transformers] \u001b[1mBertHashForSequenceClassification LOAD REPORT\u001b[0m from: neuml/bert-hash-femto\n",
|
||||
"Key | Status | \n",
|
||||
"-------------------------+---------+-\n",
|
||||
"bert.pooler.dense.bias | MISSING | \n",
|
||||
"bert.pooler.dense.weight | MISSING | \n",
|
||||
"classifier.bias | MISSING | \n",
|
||||
"classifier.weight | MISSING | \n",
|
||||
"\n",
|
||||
"Notes:\n",
|
||||
"- MISSING:\tthose params were newly initialized because missing from the checkpoint. Consider training on your downstream task.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"0.8084862385321101\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"test(train(\n",
|
||||
" \"bert-tiny-sst2\",\n",
|
||||
" \"neuml/bert-hash-femto\",\n",
|
||||
" (1.0, 1.0),\n",
|
||||
" learning_rate=3e-4,\n",
|
||||
" num_train_epochs=5,\n",
|
||||
" per_device_train_batch_size=32,\n",
|
||||
" output_dir=\"bert-femto-sst2\"\n",
|
||||
"))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "5f1c94e7",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Let's look at the performance. The 11M parameter model registered an accuracy score of `0.8750` on the SST2 dev set. Then the 4M parameter model scored `0.8326` and finally the tiny femto model scored `0.8085`. As we can see each score got progressively worse but each model has less capability. Let's train a femto model directly to compare."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"id": "b89d8292",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[transformers] \u001b[1mBertHashForSequenceClassification LOAD REPORT\u001b[0m from: neuml/bert-hash-femto\n",
|
||||
"Key | Status | \n",
|
||||
"-------------------------+---------+-\n",
|
||||
"bert.pooler.dense.bias | MISSING | \n",
|
||||
"bert.pooler.dense.weight | MISSING | \n",
|
||||
"classifier.bias | MISSING | \n",
|
||||
"classifier.weight | MISSING | \n",
|
||||
"\n",
|
||||
"Notes:\n",
|
||||
"- MISSING:\tthose params were newly initialized because missing from the checkpoint. Consider training on your downstream task.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"0.801605504587156\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"test(train(\n",
|
||||
" None,\n",
|
||||
" \"neuml/bert-hash-femto\",\n",
|
||||
" None,\n",
|
||||
" learning_rate=3e-4,\n",
|
||||
" num_train_epochs=5,\n",
|
||||
" per_device_train_batch_size=32,\n",
|
||||
" output_dir=\"bert-femto-sst2\"\n",
|
||||
"))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "4864b8d2",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This scored `0.8016` a sizable shift down from the progressively distilled model. Now keep in mind the femto model only has `250K` parameters but we gave it a sizable accuracy boost within it's capabilities."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e4e04226",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Wrapping up\n",
|
||||
"\n",
|
||||
"This example showed how progressive distillation can boost overall model performance, especially for tiny models. Incrementally compressing knowledge into a series of smaller subsets enabled the final model to learn more efficiently than directly training the model on the dataset without distillation. Put progressive distillation in the toolkit when working with tiny models!"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"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.10.20"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
"""
|
||||
Agent Quick Start
|
||||
Easy to use way to get started with AI Agents.
|
||||
|
||||
TxtAI has many example notebooks covering everything the framework provides
|
||||
Examples: https://neuml.github.io/txtai/examples
|
||||
|
||||
Install TxtAI
|
||||
pip install txtai[agent]
|
||||
"""
|
||||
|
||||
# pylint: disable=C0103
|
||||
from datetime import datetime
|
||||
from txtai import Agent
|
||||
|
||||
# Step 1: Define your Embeddings database
|
||||
#
|
||||
# Replace provider/container with a path to a local Embeddings database
|
||||
# See RAG Quickstart for an example of building your own custom database
|
||||
embeddings = {
|
||||
"name": "wikipedia",
|
||||
"description": "Searches a Wikipedia database",
|
||||
# "path": "path to your embeddings database"
|
||||
"provider": "huggingface-hub",
|
||||
"container": "neuml/txtai-wikipedia",
|
||||
}
|
||||
|
||||
|
||||
# Step 2: Define other tools
|
||||
#
|
||||
# Add any Python function. Just need to describe it.
|
||||
def today() -> str:
|
||||
"""
|
||||
Gets the current date and time
|
||||
|
||||
Returns:
|
||||
current date and time
|
||||
"""
|
||||
|
||||
return datetime.today().isoformat()
|
||||
|
||||
|
||||
# Step 3: Create a list of available tools
|
||||
#
|
||||
# Combine defined tools with default tools
|
||||
tools = [
|
||||
embeddings, # Embeddings database with YOUR data
|
||||
today, # Python function
|
||||
"websearch", # Runs a websearch using default engine
|
||||
"webview", # Loads a web page
|
||||
]
|
||||
|
||||
# Step 4: Set LLM configuration
|
||||
#
|
||||
# LLM APIs
|
||||
# model = "gpt-5.1"
|
||||
# model = "claude-opus-4-5-20251101"
|
||||
# model = "gemini/gemini-3-pro-preview"
|
||||
#
|
||||
# Local LLMs
|
||||
# model = "ollama/gpt-oss
|
||||
# model = "openai/gpt-oss-20b"
|
||||
# model = "unsloth/gpt-oss-20b-GGUF/gpt-oss-20b-Q4_K_M.gguf"
|
||||
#
|
||||
# Pass multiple options as a dictionary
|
||||
# model = {
|
||||
# "path": "unsloth/Qwen3-30B-A3B-Instruct-2507-GGUF/Qwen3-30B-A3B-Instruct-2507-Q4_K_M.gguf",
|
||||
# "n_ctx": 25000
|
||||
# }
|
||||
model = "Qwen/Qwen3-4B-Instruct-2507"
|
||||
|
||||
# Step 4: Create an Agent
|
||||
#
|
||||
# Set LLM, tools and other configuration
|
||||
# See this for more options: https://huggingface.co/docs/smolagents/reference/agents#agents
|
||||
agent = Agent(model=model, tools=tools, max_steps=10)
|
||||
|
||||
print(agent("Tell me about the Roman Empire"))
|
||||
print(agent("What is the current date?"))
|
||||
print(agent("Get the 5 top news stories for today", maxlength=25000))
|
||||
@@ -0,0 +1,66 @@
|
||||
"""
|
||||
Application that builds a summary of an article.
|
||||
|
||||
Requires streamlit to be installed.
|
||||
pip install streamlit
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import streamlit as st
|
||||
|
||||
from txtai.pipeline import Summary, Textractor
|
||||
from txtai.workflow import UrlTask, Task, Workflow
|
||||
|
||||
|
||||
class Application:
|
||||
"""
|
||||
Main application.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Creates a new application.
|
||||
"""
|
||||
|
||||
textract = Textractor(paragraphs=True, minlength=100, join=True)
|
||||
summary = Summary("sshleifer/distilbart-cnn-12-6")
|
||||
|
||||
self.workflow = Workflow([UrlTask(textract), Task(summary)])
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
Runs a Streamlit application.
|
||||
"""
|
||||
|
||||
st.title("Article Summary")
|
||||
st.markdown("This application builds a summary of an article.")
|
||||
|
||||
url = st.text_input("URL")
|
||||
if url:
|
||||
# Run workflow and get summary
|
||||
summary = list(self.workflow([url]))[0]
|
||||
|
||||
# Write results
|
||||
st.write(summary)
|
||||
st.markdown("*Source: " + url + "*")
|
||||
|
||||
|
||||
@st.cache(allow_output_mutation=True)
|
||||
def create():
|
||||
"""
|
||||
Creates and caches a Streamlit application.
|
||||
|
||||
Returns:
|
||||
Application
|
||||
"""
|
||||
|
||||
return Application()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
# Create and run application
|
||||
app = create()
|
||||
app.run()
|
||||
@@ -0,0 +1,712 @@
|
||||
"""
|
||||
Baseball statistics application with txtai and Streamlit.
|
||||
|
||||
Install txtai and streamlit (>= 1.23) to run:
|
||||
pip install txtai streamlit
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
|
||||
import altair as alt
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import streamlit as st
|
||||
|
||||
from txtai import Embeddings
|
||||
|
||||
|
||||
class Stats:
|
||||
"""
|
||||
Base stats class. Contains methods for loading, indexing and searching baseball stats.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Creates a new Stats instance.
|
||||
"""
|
||||
|
||||
# Load columns
|
||||
self.columns = self.loadcolumns()
|
||||
|
||||
# Load stats data
|
||||
self.stats = self.load()
|
||||
|
||||
# Load names
|
||||
self.names = self.loadnames()
|
||||
|
||||
# Build index
|
||||
self.vectors, self.data, self.maxyear, self.embeddings = self.index()
|
||||
|
||||
def loadcolumns(self):
|
||||
"""
|
||||
Returns a list of data columns.
|
||||
|
||||
Returns:
|
||||
list of columns
|
||||
"""
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
def load(self):
|
||||
"""
|
||||
Loads and returns raw stats.
|
||||
|
||||
Returns:
|
||||
stats
|
||||
"""
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
def metric(self):
|
||||
"""
|
||||
Primary metric column.
|
||||
|
||||
Returns:
|
||||
metric column name
|
||||
"""
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
def vector(self, row):
|
||||
"""
|
||||
Build a vector for input row.
|
||||
|
||||
Args:
|
||||
row: input row
|
||||
|
||||
Returns:
|
||||
row vector
|
||||
"""
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
def loadnames(self):
|
||||
"""
|
||||
Loads a name - player id dictionary.
|
||||
|
||||
Returns:
|
||||
{player name: player id}
|
||||
"""
|
||||
|
||||
# Get unique names
|
||||
names = {}
|
||||
rows = self.stats.sort_values(by=self.metric(), ascending=False)[["nameFirst", "nameLast", "playerID"]].drop_duplicates().reset_index()
|
||||
for x, row in rows.iterrows():
|
||||
# Name key
|
||||
key = f"{row['nameFirst']} {row['nameLast']}"
|
||||
key += f" ({row['playerID']})" if key in names else ""
|
||||
|
||||
if key not in names:
|
||||
# Scale scores of top n players
|
||||
exponent = 2 if ((len(rows) - x) / len(rows)) >= 0.95 else 1
|
||||
|
||||
# score = num seasons ^ exponent
|
||||
score = math.pow(len(self.stats[self.stats["playerID"] == row["playerID"]]), exponent)
|
||||
|
||||
# Save name key - values pair
|
||||
names[key] = (row["playerID"], score)
|
||||
|
||||
return names
|
||||
|
||||
def index(self):
|
||||
"""
|
||||
Builds an embeddings index to stats data. Returns vectors, input data and embeddings index.
|
||||
|
||||
Returns:
|
||||
vectors, data, embeddings
|
||||
"""
|
||||
|
||||
# Build data dictionary
|
||||
vectors = {f'{row["yearID"]}{row["playerID"]}': self.transform(row) for _, row in self.stats.iterrows()}
|
||||
data = {f'{row["yearID"]}{row["playerID"]}': dict(row) for _, row in self.stats.iterrows()}
|
||||
maxyear = max(row["yearID"] for _, row in self.stats.iterrows())
|
||||
|
||||
embeddings = Embeddings({"transform": Stats.transform})
|
||||
embeddings.index((uid, vectors[uid], None) for uid in vectors)
|
||||
|
||||
return vectors, data, maxyear, embeddings
|
||||
|
||||
def metrics(self, name):
|
||||
"""
|
||||
Looks up a player's active years, best statistical year and key metrics.
|
||||
|
||||
Args:
|
||||
name: player name
|
||||
|
||||
Returns:
|
||||
active, best, metrics
|
||||
"""
|
||||
|
||||
if name in self.names:
|
||||
# Get player stats
|
||||
stats = self.stats[self.stats["playerID"] == self.names[name][0]]
|
||||
|
||||
# Build key metrics
|
||||
metrics = stats[["yearID", self.metric()]]
|
||||
|
||||
# Get best year, sort by primary metric
|
||||
best = int(stats.sort_values(by=self.metric(), ascending=False)["yearID"].iloc[0])
|
||||
|
||||
# Get years active, best year, along with metric trends
|
||||
return metrics["yearID"].tolist(), best, metrics
|
||||
|
||||
return range(1871, datetime.datetime.today().year), 1950, None
|
||||
|
||||
def search(self, name=None, year=None, window=None, row=None, limit=10):
|
||||
"""
|
||||
Runs an embeddings search. This method takes either a player-year or stats row as input.
|
||||
|
||||
Args:
|
||||
name: player name to search
|
||||
year: year to search
|
||||
window: limit to window recent seasons
|
||||
row: row of stats to search
|
||||
limit: max results to return
|
||||
|
||||
Returns:
|
||||
list of results
|
||||
"""
|
||||
|
||||
if row:
|
||||
query = self.vector(row)
|
||||
else:
|
||||
# Lookup player key and build vector id
|
||||
name = self.names.get(name)
|
||||
query = f"{year}{name[0] if name else name}"
|
||||
query = self.vectors.get(query)
|
||||
|
||||
results, ids = [], set()
|
||||
if query is not None:
|
||||
candidates = limit * 100 if window else limit * 5
|
||||
for uid, _ in self.embeddings.search(query, candidates):
|
||||
# Only add unique players
|
||||
if uid[4:] not in ids:
|
||||
result = self.data[uid].copy()
|
||||
|
||||
# Add first player if this is a player comparison. Limit results to window, if necessary
|
||||
if (not ids and not row) or not window or result["yearID"] > self.maxyear - window:
|
||||
result["link"] = f'https://www.baseball-reference.com/players/{result["nameLast"].lower()[0]}/{result["bbrefID"]}.shtml'
|
||||
results.append(result)
|
||||
ids.add(uid[4:])
|
||||
|
||||
if len(ids) >= limit:
|
||||
break
|
||||
|
||||
return results
|
||||
|
||||
def transform(self, row):
|
||||
"""
|
||||
Transforms a stats row into a vector.
|
||||
|
||||
Args:
|
||||
row: stats row
|
||||
|
||||
Returns:
|
||||
vector
|
||||
"""
|
||||
|
||||
if isinstance(row, np.ndarray):
|
||||
return row
|
||||
|
||||
return np.array([0.0 if not row[x] or np.isnan(row[x]) else row[x] for x in self.columns])
|
||||
|
||||
|
||||
class Batting(Stats):
|
||||
"""
|
||||
Batting stats.
|
||||
"""
|
||||
|
||||
def loadcolumns(self):
|
||||
return [
|
||||
"birthMonth",
|
||||
"yearID",
|
||||
"age",
|
||||
"height",
|
||||
"weight",
|
||||
"G",
|
||||
"AB",
|
||||
"R",
|
||||
"H",
|
||||
"1B",
|
||||
"2B",
|
||||
"3B",
|
||||
"HR",
|
||||
"RBI",
|
||||
"SB",
|
||||
"CS",
|
||||
"BB",
|
||||
"SO",
|
||||
"IBB",
|
||||
"HBP",
|
||||
"SH",
|
||||
"SF",
|
||||
"GIDP",
|
||||
"POS",
|
||||
"AVG",
|
||||
"OBP",
|
||||
"TB",
|
||||
"SLG",
|
||||
"OPS",
|
||||
"OPS+",
|
||||
]
|
||||
|
||||
def load(self):
|
||||
# Retrieve raw data
|
||||
players = pd.read_csv("https://hf.co/datasets/neuml/baseballdata/resolve/main/People.csv")
|
||||
batting = pd.read_csv("https://hf.co/datasets/neuml/baseballdata/resolve/main/Batting.csv")
|
||||
fielding = pd.read_csv("https://hf.co/datasets/neuml/baseballdata/resolve/main/Fielding.csv")
|
||||
|
||||
# Merge player data in
|
||||
batting = pd.merge(players, batting, how="inner", on=["playerID"])
|
||||
|
||||
# Require player to have at least 350 plate appearances.
|
||||
batting = batting[((batting["AB"] + batting["BB"]) >= 350) & (batting["stint"] == 1)]
|
||||
|
||||
# Derive primary player positions
|
||||
positions = self.positions(fielding)
|
||||
|
||||
# Calculated columns
|
||||
batting["age"] = batting["yearID"] - batting["birthYear"]
|
||||
batting["POS"] = batting.apply(lambda row: self.position(positions, row), axis=1)
|
||||
batting["AVG"] = batting["H"] / batting["AB"]
|
||||
batting["OBP"] = (batting["H"] + batting["BB"]) / (batting["AB"] + batting["BB"])
|
||||
batting["1B"] = batting["H"] - batting["2B"] - batting["3B"] - batting["HR"]
|
||||
batting["TB"] = batting["1B"] + 2 * batting["2B"] + 3 * batting["3B"] + 4 * batting["HR"]
|
||||
batting["SLG"] = batting["TB"] / batting["AB"]
|
||||
batting["OPS"] = batting["OBP"] + batting["SLG"]
|
||||
batting["OPS+"] = 100 + (batting["OPS"] - batting["OPS"].mean()) * 100
|
||||
|
||||
return batting
|
||||
|
||||
def metric(self):
|
||||
return "OPS+"
|
||||
|
||||
def vector(self, row):
|
||||
row["TB"] = row["1B"] + 2 * row["2B"] + 3 * row["3B"] + 4 * row["HR"]
|
||||
row["AVG"] = row["H"] / row["AB"]
|
||||
row["OBP"] = (row["H"] + row["BB"]) / (row["AB"] + row["BB"])
|
||||
row["SLG"] = row["TB"] / row["AB"]
|
||||
row["OPS"] = row["OBP"] + row["SLG"]
|
||||
row["OPS+"] = 100 + (row["OPS"] - self.stats["OPS"].mean()) * 100
|
||||
|
||||
return self.transform(row)
|
||||
|
||||
def positions(self, fielding):
|
||||
"""
|
||||
Derives primary positions for players.
|
||||
|
||||
Args:
|
||||
fielding: fielding data
|
||||
|
||||
Returns:
|
||||
{player id: (position, number of games)}
|
||||
"""
|
||||
|
||||
positions = {}
|
||||
for _, row in fielding.iterrows():
|
||||
uid = f'{row["yearID"]}{row["playerID"]}'
|
||||
position = row["POS"] if row["POS"] else 0
|
||||
if position == "P":
|
||||
position = 1
|
||||
elif position == "C":
|
||||
position = 2
|
||||
elif position == "1B":
|
||||
position = 3
|
||||
elif position == "2B":
|
||||
position = 4
|
||||
elif position == "3B":
|
||||
position = 5
|
||||
elif position == "SS":
|
||||
position = 6
|
||||
elif position == "OF":
|
||||
position = 7
|
||||
|
||||
# Save position if not set or player played more at this position
|
||||
if uid not in positions or positions[uid][1] < row["G"]:
|
||||
positions[uid] = (position, row["G"])
|
||||
|
||||
return positions
|
||||
|
||||
def position(self, positions, row):
|
||||
"""
|
||||
Looks up primary position for player row.
|
||||
|
||||
Arg:
|
||||
positions: all player positions
|
||||
row: player row
|
||||
|
||||
Returns:
|
||||
primary player positions
|
||||
"""
|
||||
|
||||
uid = f'{row["yearID"]}{row["playerID"]}'
|
||||
return positions[uid][0] if uid in positions else 0
|
||||
|
||||
|
||||
class Pitching(Stats):
|
||||
"""
|
||||
Pitching stats.
|
||||
"""
|
||||
|
||||
def loadcolumns(self):
|
||||
return [
|
||||
"birthMonth",
|
||||
"yearID",
|
||||
"age",
|
||||
"height",
|
||||
"weight",
|
||||
"W",
|
||||
"L",
|
||||
"G",
|
||||
"GS",
|
||||
"CG",
|
||||
"SHO",
|
||||
"SV",
|
||||
"IPouts",
|
||||
"H",
|
||||
"ER",
|
||||
"HR",
|
||||
"BB",
|
||||
"SO",
|
||||
"BAOpp",
|
||||
"ERA",
|
||||
"IBB",
|
||||
"WP",
|
||||
"HBP",
|
||||
"BK",
|
||||
"BFP",
|
||||
"GF",
|
||||
"R",
|
||||
"SH",
|
||||
"SF",
|
||||
"GIDP",
|
||||
"WHIP",
|
||||
"WADJ",
|
||||
]
|
||||
|
||||
def load(self):
|
||||
# Retrieve raw data
|
||||
players = pd.read_csv("https://hf.co/datasets/neuml/baseballdata/resolve/main/People.csv")
|
||||
pitching = pd.read_csv("https://hf.co/datasets/neuml/baseballdata/resolve/main/Pitching.csv")
|
||||
|
||||
# Merge player data in
|
||||
pitching = pd.merge(players, pitching, how="inner", on=["playerID"])
|
||||
|
||||
# Require player to have 20 appearances
|
||||
pitching = pitching[(pitching["G"] >= 20) & (pitching["stint"] == 1)]
|
||||
|
||||
# Calculated columns
|
||||
pitching["age"] = pitching["yearID"] - pitching["birthYear"]
|
||||
pitching["WHIP"] = (pitching["BB"] + pitching["H"]) / (pitching["IPouts"] / 3)
|
||||
pitching["WADJ"] = (pitching["W"] + pitching["SV"]) / (pitching["ERA"] + pitching["WHIP"])
|
||||
|
||||
return pitching
|
||||
|
||||
def metric(self):
|
||||
return "WADJ"
|
||||
|
||||
def vector(self, row):
|
||||
row["WHIP"] = (row["BB"] + row["H"]) / (row["IPouts"] / 3) if row["IPouts"] else None
|
||||
row["WADJ"] = (row["W"] + row["SV"]) / (row["ERA"] + row["WHIP"]) if row["ERA"] and row["WHIP"] else None
|
||||
|
||||
return self.transform(row)
|
||||
|
||||
|
||||
class Application:
|
||||
"""
|
||||
Main application.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Creates a new application.
|
||||
"""
|
||||
|
||||
# Batting stats
|
||||
self.batting = Batting()
|
||||
|
||||
# Pitching stats
|
||||
self.pitching = Pitching()
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
Runs a Streamlit application.
|
||||
"""
|
||||
|
||||
st.set_page_config(layout="wide", page_title="Baseball Stats", page_icon="⚾")
|
||||
st.title("⚾ Baseball Stats")
|
||||
st.markdown(
|
||||
"""
|
||||
This application finds the best matching players using vector search with [txtai](https://github.com/neuml/txtai).
|
||||
Raw data is from the [Lahman Baseball Database](https://sabr.org/lahman-database/). Read [this
|
||||
article](https://medium.com/neuml/explore-baseball-history-with-vector-search-5778d98d6846) for more details.
|
||||
"""
|
||||
)
|
||||
|
||||
player, search = st.tabs(["Player", "Search"])
|
||||
|
||||
# Player tab
|
||||
with player:
|
||||
self.player()
|
||||
|
||||
# Search
|
||||
with search:
|
||||
self.search()
|
||||
|
||||
def player(self):
|
||||
"""
|
||||
Player tab.
|
||||
"""
|
||||
|
||||
st.markdown("Match by player-season. Each player search defaults to the best season sorted by OPS or Wins Adjusted.")
|
||||
|
||||
# Get parameters
|
||||
params = self.params()
|
||||
|
||||
# Category and stats
|
||||
category = self.category(params.get("category"), "category")
|
||||
stats = self.batting if category == "Batting" else self.pitching
|
||||
|
||||
# Player name
|
||||
name = self.name(stats.names, params.get("name"))
|
||||
|
||||
# Limit player-year comparisons using this window
|
||||
window = self.window(params.get("window"), "window")
|
||||
|
||||
# Player metrics
|
||||
active, best, metrics = stats.metrics(name)
|
||||
|
||||
# Player year
|
||||
year = self.year(active, params.get("year"), best)
|
||||
|
||||
# Display metrics chart
|
||||
if len(active) > 1:
|
||||
self.chart(category, metrics)
|
||||
|
||||
# Run search
|
||||
results = stats.search(name, year, window)
|
||||
|
||||
# Display results
|
||||
self.table(results, ["link", "nameFirst", "nameLast", "teamID"] + stats.columns[1:])
|
||||
|
||||
# Save query parameters
|
||||
st.query_params.clear()
|
||||
for key, value in [("category", category), ("name", name), ("window", window), ("year", year)]:
|
||||
if value:
|
||||
st.query_params[key] = value
|
||||
|
||||
def search(self):
|
||||
"""
|
||||
Stats search tab.
|
||||
"""
|
||||
|
||||
st.markdown("Find players with similar statistics.")
|
||||
|
||||
stats, category = None, self.category("Batting", "searchcategory")
|
||||
|
||||
# Limit player-year comparisons using this window
|
||||
window = self.window(None, "searchwindow")
|
||||
|
||||
with st.form("search"):
|
||||
if category == "Batting":
|
||||
stats, columns = self.batting, self.batting.columns[:-6]
|
||||
elif category == "Pitching":
|
||||
stats, columns = self.pitching, self.pitching.columns[:-2]
|
||||
|
||||
# Enter stats with data editor
|
||||
inputs = st.data_editor(pd.DataFrame([dict((column, None) for column in columns)]), hide_index=True).astype(float)
|
||||
|
||||
submitted = st.form_submit_button("Search")
|
||||
if submitted:
|
||||
# Run search
|
||||
results = stats.search(window=window, row=inputs.to_dict(orient="records")[0])
|
||||
|
||||
# Display table
|
||||
self.table(results, ["link", "nameFirst", "nameLast", "teamID"] + stats.columns[1:])
|
||||
|
||||
def params(self):
|
||||
"""
|
||||
Get application parameters. This method combines URL parameters with session parameters.
|
||||
|
||||
Returns:
|
||||
parameters
|
||||
"""
|
||||
|
||||
# Get query parameters
|
||||
params = {x: st.query_params.get(x) for x in ["category", "name", "window", "year"]}
|
||||
|
||||
# Sync parameters with session state
|
||||
if all(x in st.session_state for x in ["category", "name", "year"]):
|
||||
# Copy session year if category and name are unchanged
|
||||
params["year"] = str(st.session_state["year"]) if all(params.get(x) == st.session_state[x] for x in ["category", "name"]) else None
|
||||
|
||||
# Copy category, name and window from session state
|
||||
params["category"] = st.session_state["category"]
|
||||
params["name"] = st.session_state["name"]
|
||||
params["window"] = st.session_state["window"]
|
||||
|
||||
return params
|
||||
|
||||
def category(self, category, key):
|
||||
"""
|
||||
Builds category input widget.
|
||||
|
||||
Args:
|
||||
category: category parameter
|
||||
key: widget key
|
||||
|
||||
Returns:
|
||||
category component
|
||||
"""
|
||||
|
||||
# List of stat categories
|
||||
categories = ["Batting", "Pitching"]
|
||||
|
||||
# Get category parameter, default if not available or valid
|
||||
default = categories.index(category) if category and category in categories else 0
|
||||
|
||||
# Radio box component
|
||||
return st.radio("Stat", categories, index=default, horizontal=True, key=key)
|
||||
|
||||
def window(self, window, key):
|
||||
"""
|
||||
Limit results to last N seasons.
|
||||
|
||||
Args:
|
||||
window: limit to window seasons
|
||||
key: widget key
|
||||
|
||||
Returns:
|
||||
window component
|
||||
"""
|
||||
|
||||
# Get window size
|
||||
window = st.text_input("Limit to last N seasons", value=window, key=key)
|
||||
|
||||
# Clear invalid input
|
||||
if window and (not window.isnumeric() or int(window) < 1):
|
||||
st.error("Window must be a number greater or equal to 1")
|
||||
return None
|
||||
|
||||
# Convert to int
|
||||
return int(window) if window else window
|
||||
|
||||
def name(self, names, name):
|
||||
"""
|
||||
Builds name input widget.
|
||||
|
||||
Args:
|
||||
names: list of all allowable names
|
||||
|
||||
Returns:
|
||||
name component
|
||||
"""
|
||||
|
||||
# Get name parameter, default to random weighted value if not valid
|
||||
name = name if name and name in names else random.choices(list(names.keys()), weights=[names[x][1] for x in names])[0]
|
||||
|
||||
# Sort names for display
|
||||
names = sorted(names)
|
||||
|
||||
# Select box component
|
||||
return st.selectbox("Name", names, names.index(name), key="name")
|
||||
|
||||
def year(self, years, year, best):
|
||||
"""
|
||||
Builds year input widget.
|
||||
|
||||
Args:
|
||||
years: active years for a player
|
||||
year: year parameter
|
||||
best: default to best year if year is invalid
|
||||
|
||||
Returns:
|
||||
year component
|
||||
"""
|
||||
|
||||
# Get year parameter, default if not available or valid
|
||||
year = int(year) if year and year.isdigit() and int(year) in years else best
|
||||
|
||||
# Slider component
|
||||
return int(st.select_slider("Year", years, year, key="year") if len(years) > 1 else years[0])
|
||||
|
||||
def chart(self, category, metrics):
|
||||
"""
|
||||
Displays a metric chart.
|
||||
|
||||
Args:
|
||||
category: Batting or Pitching
|
||||
metrics: player metrics to plot
|
||||
"""
|
||||
|
||||
# Key metric
|
||||
metric = self.batting.metric() if category == "Batting" else self.pitching.metric()
|
||||
|
||||
# Cast year to string
|
||||
metrics["yearID"] = metrics["yearID"].astype(str)
|
||||
|
||||
# Metric over years
|
||||
chart = (
|
||||
alt.Chart(metrics)
|
||||
.mark_line(interpolate="monotone", point=True, strokeWidth=2.5, opacity=0.75)
|
||||
.encode(x=alt.X("yearID", title=""), y=alt.Y(metric, scale=alt.Scale(zero=False)))
|
||||
)
|
||||
|
||||
# Create metric median rule line
|
||||
rule = alt.Chart(metrics).mark_rule(color="gray", strokeDash=[3, 5], opacity=0.5).encode(y=f"median({metric})")
|
||||
|
||||
# Layered chart configuration
|
||||
chart = (chart + rule).encode(y=alt.Y(title=metric)).properties(height=200).configure_axis(grid=False)
|
||||
|
||||
# Draw chart
|
||||
st.altair_chart(chart + rule, theme="streamlit", width="stretch")
|
||||
|
||||
def table(self, results, columns):
|
||||
"""
|
||||
Displays a list of results as a table.
|
||||
|
||||
Args:
|
||||
results: list of results
|
||||
columns: column names
|
||||
"""
|
||||
|
||||
if results:
|
||||
st.dataframe(
|
||||
results,
|
||||
column_order=columns,
|
||||
column_config={
|
||||
"link": st.column_config.LinkColumn("Link", width="small", display_text=":material/open_in_new:"),
|
||||
"yearID": st.column_config.NumberColumn("Year", format="%d"),
|
||||
"nameFirst": "First",
|
||||
"nameLast": "Last",
|
||||
"teamID": "Team",
|
||||
"age": "Age",
|
||||
"weight": "Weight",
|
||||
"height": "Height",
|
||||
},
|
||||
)
|
||||
else:
|
||||
st.write("Player-Year not found")
|
||||
|
||||
|
||||
@st.cache_resource(show_spinner=False)
|
||||
def create():
|
||||
"""
|
||||
Creates and caches a Streamlit application.
|
||||
|
||||
Returns:
|
||||
Application
|
||||
"""
|
||||
|
||||
return Application()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
# Create and run application
|
||||
app = create()
|
||||
app.run()
|
||||
@@ -0,0 +1,731 @@
|
||||
"""
|
||||
Runs benchmark evaluations with the BEIR dataset.
|
||||
|
||||
Install txtai and the following dependencies to run:
|
||||
pip install txtai pytrec_eval rank-bm25 bm25s elasticsearch psutil
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import os
|
||||
import pickle
|
||||
import sqlite3
|
||||
import time
|
||||
|
||||
import psutil
|
||||
import yaml
|
||||
|
||||
import numpy as np
|
||||
|
||||
from bm25s import BM25 as BM25Sparse
|
||||
from elasticsearch import Elasticsearch
|
||||
from elasticsearch.helpers import bulk
|
||||
from pytrec_eval import RelevanceEvaluator
|
||||
from rank_bm25 import BM25Okapi
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
from txtai.pipeline import LLM, RAG, Similarity, Tokenizer
|
||||
from txtai.scoring import ScoringFactory
|
||||
|
||||
|
||||
class Index:
|
||||
"""
|
||||
Base index definition. Defines methods to index and search a dataset.
|
||||
"""
|
||||
|
||||
def __init__(self, path, config, output, refresh):
|
||||
"""
|
||||
Creates a new index.
|
||||
|
||||
Args:
|
||||
path: path to dataset
|
||||
config: path to config file
|
||||
output: path to store index
|
||||
refresh: overwrites existing index if True, otherwise existing index is loaded
|
||||
"""
|
||||
|
||||
self.path = path
|
||||
self.config = config
|
||||
self.output = output
|
||||
self.refresh = refresh
|
||||
|
||||
# Build and save index
|
||||
self.backend = self.index()
|
||||
|
||||
def __call__(self, limit, filterscores=True):
|
||||
"""
|
||||
Main evaluation logic. Loads an index, runs the dataset queries and returns the results.
|
||||
|
||||
Args:
|
||||
limit: maximum results
|
||||
filterscores: if exact matches should be filtered out
|
||||
|
||||
Returns:
|
||||
search results
|
||||
"""
|
||||
|
||||
uids, queries = self.load()
|
||||
|
||||
# Run queries in batches
|
||||
offset, results = 0, {}
|
||||
for batch in self.batch(queries, 256):
|
||||
for i, r in enumerate(self.search(batch, limit + 1)):
|
||||
# Get result as list of (id, score) tuples
|
||||
r = list(r)
|
||||
r = [(x["id"], x["score"]) for x in r] if r and isinstance(r[0], dict) else r
|
||||
|
||||
if filterscores:
|
||||
r = [(uid, score) for uid, score in r if uid != uids[offset + i]][:limit]
|
||||
|
||||
results[uids[offset + i]] = dict(r)
|
||||
|
||||
# Increment offset
|
||||
offset += len(batch)
|
||||
|
||||
return results
|
||||
|
||||
def search(self, queries, limit):
|
||||
"""
|
||||
Runs a search for a set of queries.
|
||||
|
||||
Args:
|
||||
queries: list of queries to run
|
||||
limit: maximum results
|
||||
|
||||
Returns:
|
||||
search results
|
||||
"""
|
||||
|
||||
return self.backend.batchsearch(queries, limit)
|
||||
|
||||
def index(self):
|
||||
"""
|
||||
Indexes a dataset.
|
||||
"""
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
def rows(self):
|
||||
"""
|
||||
Iterates over the dataset yielding a row at a time for indexing.
|
||||
"""
|
||||
|
||||
# Data file
|
||||
path = f"{self.path}/corpus.jsonl"
|
||||
|
||||
# Get total count
|
||||
with open(path, encoding="utf-8") as f:
|
||||
total = sum(1 for _ in f)
|
||||
|
||||
# Yield data
|
||||
with open(path, encoding="utf-8") as f:
|
||||
for line in tqdm(f, total=total):
|
||||
row = json.loads(line)
|
||||
text = f'{row["title"]}. {row["text"]}' if row.get("title") else row["text"]
|
||||
if text:
|
||||
yield (row["_id"], text, None)
|
||||
|
||||
def load(self):
|
||||
"""
|
||||
Loads queries for the dataset. Returns a list of expected result ids and input queries.
|
||||
|
||||
Returns:
|
||||
(result ids, input queries)
|
||||
"""
|
||||
|
||||
with open(f"{self.path}/queries.jsonl", encoding="utf-8") as f:
|
||||
data = [json.loads(query) for query in f]
|
||||
uids, queries = [x["_id"] for x in data], [x["text"] for x in data]
|
||||
|
||||
return uids, queries
|
||||
|
||||
def batch(self, data, size):
|
||||
"""
|
||||
Splits data into equal sized batches.
|
||||
|
||||
Args:
|
||||
data: input data
|
||||
size: batch size
|
||||
|
||||
Returns:
|
||||
data split into equal size batches
|
||||
"""
|
||||
|
||||
return [data[x : x + size] for x in range(0, len(data), size)]
|
||||
|
||||
def readconfig(self, key, default):
|
||||
"""
|
||||
Reads configuration from a config file. Returns default configuration
|
||||
if config file is not found or config key isn't present.
|
||||
|
||||
Args:
|
||||
key: configuration key to lookup
|
||||
default: default configuration
|
||||
|
||||
Returns:
|
||||
config if found, otherwise returns default config
|
||||
"""
|
||||
|
||||
if self.config and os.path.exists(self.config):
|
||||
# Read configuration
|
||||
with open(self.config, "r", encoding="utf-8") as f:
|
||||
# Check for config
|
||||
config = yaml.safe_load(f)
|
||||
if key in config:
|
||||
return config[key]
|
||||
|
||||
return default
|
||||
|
||||
|
||||
class Embed(Index):
|
||||
"""
|
||||
Embeddings index using txtai.
|
||||
"""
|
||||
|
||||
def index(self):
|
||||
if os.path.exists(self.output) and not self.refresh:
|
||||
embeddings = Embeddings()
|
||||
embeddings.load(self.output)
|
||||
else:
|
||||
# Read configuration
|
||||
config = self.readconfig("embeddings", {"batch": 8192, "encodebatch": 128, "faiss": {"quantize": True, "sample": 0.05}})
|
||||
|
||||
# Build index
|
||||
embeddings = Embeddings(config)
|
||||
embeddings.index(self.rows())
|
||||
embeddings.save(self.output)
|
||||
|
||||
return embeddings
|
||||
|
||||
|
||||
class Hybrid(Index):
|
||||
"""
|
||||
Hybrid embeddings + BM25 index using txtai.
|
||||
"""
|
||||
|
||||
def index(self):
|
||||
if os.path.exists(self.output) and not self.refresh:
|
||||
embeddings = Embeddings()
|
||||
embeddings.load(self.output)
|
||||
else:
|
||||
# Read configuration
|
||||
config = self.readconfig(
|
||||
"hybrid",
|
||||
{
|
||||
"batch": 8192,
|
||||
"encodebatch": 128,
|
||||
"faiss": {"quantize": True, "sample": 0.05},
|
||||
"scoring": {"method": "bm25", "terms": True, "normalize": True},
|
||||
},
|
||||
)
|
||||
|
||||
# Build index
|
||||
embeddings = Embeddings(config)
|
||||
embeddings.index(self.rows())
|
||||
embeddings.save(self.output)
|
||||
|
||||
return embeddings
|
||||
|
||||
|
||||
class RetrievalAugmentedGeneration(Embed):
|
||||
"""
|
||||
Retrieval augmented generation (RAG) using txtai.
|
||||
"""
|
||||
|
||||
def __init__(self, path, config, output, refresh):
|
||||
# Parent logic
|
||||
super().__init__(path, config, output, refresh)
|
||||
|
||||
# Read LLM configuration
|
||||
llm = self.readconfig("llm", {})
|
||||
|
||||
# Read RAG configuration
|
||||
rag = self.readconfig("rag", {})
|
||||
|
||||
# Load RAG pipeline
|
||||
self.rag = RAG(self.backend, LLM(**llm), output="reference", **rag)
|
||||
|
||||
def search(self, queries, limit):
|
||||
# Set context window size to limit and run
|
||||
self.rag.context = limit
|
||||
return [[(x["reference"], 1)] for x in self.rag(queries, maxlength=4096)]
|
||||
|
||||
|
||||
class Score(Index):
|
||||
"""
|
||||
BM25 index using txtai.
|
||||
"""
|
||||
|
||||
def index(self):
|
||||
# Read configuration
|
||||
config = self.readconfig("scoring", {"method": "bm25", "terms": True})
|
||||
|
||||
# Create scoring instance
|
||||
scoring = ScoringFactory.create(config)
|
||||
|
||||
output = os.path.join(self.output, "scoring")
|
||||
if os.path.exists(output) and not self.refresh:
|
||||
scoring.load(output)
|
||||
else:
|
||||
scoring.index(self.rows())
|
||||
scoring.save(output)
|
||||
|
||||
return scoring
|
||||
|
||||
|
||||
class Similar(Index):
|
||||
"""
|
||||
Search data using a similarity pipeline.
|
||||
"""
|
||||
|
||||
def index(self):
|
||||
# Load similarity pipeline
|
||||
model = Similarity(**self.readconfig("similar", {}))
|
||||
|
||||
# Get datasets
|
||||
data = list(self.rows())
|
||||
ids = [x[0] for x in data]
|
||||
texts = [x[1] for x in data]
|
||||
|
||||
# Encode texts
|
||||
data = model.encode(texts, "data")
|
||||
|
||||
return (ids, data, model)
|
||||
|
||||
def search(self, queries, limit):
|
||||
# Unpack backend
|
||||
ids, data, model = self.backend
|
||||
|
||||
# Run model inference
|
||||
results = []
|
||||
for result in model(queries, data, limit=limit):
|
||||
results.append([(ids[x], score) for x, score in result])
|
||||
|
||||
return results
|
||||
|
||||
|
||||
class Rerank(Embed):
|
||||
"""
|
||||
Embeddings index using txtai combined with a similarity pipeline
|
||||
"""
|
||||
|
||||
def index(self):
|
||||
# Build embeddings index
|
||||
embeddings = super().index()
|
||||
|
||||
# Combine similar pipeline with embeddings
|
||||
model = Similar(self.path, self.config, self.output, self.refresh)
|
||||
return model.index() + (embeddings,)
|
||||
|
||||
def search(self, queries, limit):
|
||||
# Unpack backend
|
||||
ids, data, model, embeddings = self.backend
|
||||
|
||||
# Run initial query
|
||||
indices = []
|
||||
for r in embeddings.batchsearch(queries, limit * 10):
|
||||
indices.append({x: ids.index(uid) for x, (uid, _) in enumerate(r)})
|
||||
|
||||
# Run model inference
|
||||
results = []
|
||||
for x, query in enumerate(queries):
|
||||
queue = data[list(indices[x].values())]
|
||||
if len(queue) > 0:
|
||||
result = model(query, queue, limit=limit)
|
||||
results.append([(ids[indices[x][i]], score) for i, score in result])
|
||||
|
||||
return results
|
||||
|
||||
|
||||
class RankBM25(Index):
|
||||
"""
|
||||
BM25 index using rank-bm25.
|
||||
"""
|
||||
|
||||
def search(self, queries, limit):
|
||||
ids, backend = self.backend
|
||||
tokenizer, results = Tokenizer(), []
|
||||
for query in queries:
|
||||
scores = backend.get_scores(tokenizer(query))
|
||||
topn = np.argsort(scores)[::-1][:limit]
|
||||
results.append([(ids[x], scores[x]) for x in topn])
|
||||
|
||||
return results
|
||||
|
||||
def index(self):
|
||||
output = os.path.join(self.output, "rank")
|
||||
if os.path.exists(output) and not self.refresh:
|
||||
with open(output, "rb") as f:
|
||||
ids, model = pickle.load(f)
|
||||
else:
|
||||
# Tokenize data
|
||||
tokenizer, data = Tokenizer(), []
|
||||
for uid, text, _ in self.rows():
|
||||
data.append((uid, tokenizer(text)))
|
||||
|
||||
ids = [uid for uid, _ in data]
|
||||
model = BM25Okapi([text for _, text in data])
|
||||
|
||||
# Save model
|
||||
with open(output, "wb") as out:
|
||||
pickle.dump(model, out)
|
||||
|
||||
return ids, model
|
||||
|
||||
|
||||
class BM25S(Index):
|
||||
"""
|
||||
BM25 as implemented by bm25s
|
||||
"""
|
||||
|
||||
def __init__(self, path, config, output, refresh):
|
||||
# Corpus ids
|
||||
self.ids = None
|
||||
|
||||
# Parent logic
|
||||
super().__init__(path, config, output, refresh)
|
||||
|
||||
def search(self, queries, limit):
|
||||
tokenizer = Tokenizer()
|
||||
results, scores = self.backend.retrieve([tokenizer(x) for x in queries], corpus=self.ids, k=limit)
|
||||
|
||||
# List of queries => list of matches (id, score)
|
||||
x = []
|
||||
for a, b in zip(results, scores):
|
||||
x.append([(str(c), float(d)) for c, d in zip(a, b)])
|
||||
|
||||
return x
|
||||
|
||||
def index(self):
|
||||
tokenizer = Tokenizer()
|
||||
ids, texts = [], []
|
||||
|
||||
for uid, text, _ in self.rows():
|
||||
ids.append(uid)
|
||||
texts.append(text)
|
||||
|
||||
self.ids = ids
|
||||
|
||||
if os.path.exists(self.output) and not self.refresh:
|
||||
model = BM25Sparse.load(self.output)
|
||||
else:
|
||||
model = BM25Sparse(method="lucene", k1=1.2, b=0.75)
|
||||
model.index([tokenizer(x) for x in texts], leave_progress=False)
|
||||
model.save(self.output)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
class SQLiteFTS(Index):
|
||||
"""
|
||||
BM25 index using SQLite's FTS extension.
|
||||
"""
|
||||
|
||||
def search(self, queries, limit):
|
||||
tokenizer, results = Tokenizer(), []
|
||||
for query in queries:
|
||||
query = tokenizer(query)
|
||||
query = " OR ".join([f'"{q}"' for q in query])
|
||||
|
||||
self.backend.execute(
|
||||
f"SELECT id, bm25(textindex) * -1 score FROM textindex WHERE text MATCH ? ORDER BY bm25(textindex) LIMIT {limit}", [query]
|
||||
)
|
||||
|
||||
results.append(list(self.backend))
|
||||
|
||||
return results
|
||||
|
||||
def index(self):
|
||||
if os.path.exists(self.output) and not self.refresh:
|
||||
# Load existing database
|
||||
connection = sqlite3.connect(self.output)
|
||||
else:
|
||||
# Delete existing database
|
||||
if os.path.exists(self.output):
|
||||
os.remove(self.output)
|
||||
|
||||
# Create new database
|
||||
connection = sqlite3.connect(self.output)
|
||||
|
||||
# Tokenize data
|
||||
tokenizer, data = Tokenizer(), []
|
||||
for uid, text, _ in self.rows():
|
||||
data.append((uid, " ".join(tokenizer(text))))
|
||||
|
||||
# Create table
|
||||
connection.execute("CREATE VIRTUAL TABLE textindex using fts5(id, text)")
|
||||
|
||||
# Load data and build index
|
||||
connection.executemany("INSERT INTO textindex VALUES (?, ?)", data)
|
||||
|
||||
connection.commit()
|
||||
|
||||
return connection.cursor()
|
||||
|
||||
|
||||
class Elastic(Index):
|
||||
"""
|
||||
BM25 index using Elasticsearch.
|
||||
"""
|
||||
|
||||
def search(self, queries, limit):
|
||||
# Generate bulk queries
|
||||
request = []
|
||||
for query in queries:
|
||||
req_head = {"index": "textindex", "search_type": "dfs_query_then_fetch"}
|
||||
req_body = {
|
||||
"_source": False,
|
||||
"query": {"multi_match": {"query": query, "type": "best_fields", "fields": ["text"], "tie_breaker": 0.5}},
|
||||
"size": limit,
|
||||
}
|
||||
request.extend([req_head, req_body])
|
||||
|
||||
# Run ES query
|
||||
response = self.backend.msearch(body=request, request_timeout=600)
|
||||
|
||||
# Read responses
|
||||
results = []
|
||||
for resp in response["responses"]:
|
||||
result = resp["hits"]["hits"]
|
||||
results.append([(r["_id"], r["_score"]) for r in result])
|
||||
|
||||
return results
|
||||
|
||||
def index(self):
|
||||
es = Elasticsearch("http://localhost:9200")
|
||||
|
||||
# Delete existing index
|
||||
# pylint: disable=W0702
|
||||
try:
|
||||
es.indices.delete(index="textindex")
|
||||
except:
|
||||
pass
|
||||
|
||||
bulk(es, ({"_index": "textindex", "_id": uid, "text": text} for uid, text, _ in self.rows()))
|
||||
es.indices.refresh(index="textindex")
|
||||
|
||||
return es
|
||||
|
||||
|
||||
def relevance(path):
|
||||
"""
|
||||
Loads relevance data for evaluation.
|
||||
|
||||
Args:
|
||||
path: path to dataset test file
|
||||
|
||||
Returns:
|
||||
relevance data
|
||||
"""
|
||||
|
||||
rel = {}
|
||||
with open(f"{path}/qrels/test.tsv", encoding="utf-8") as f:
|
||||
reader = csv.reader(f, delimiter="\t", quoting=csv.QUOTE_MINIMAL)
|
||||
next(reader)
|
||||
|
||||
for row in reader:
|
||||
queryid, corpusid, score = row[0], row[1], int(row[2])
|
||||
if queryid not in rel:
|
||||
rel[queryid] = {corpusid: score}
|
||||
else:
|
||||
rel[queryid][corpusid] = score
|
||||
|
||||
return rel
|
||||
|
||||
|
||||
def create(method, path, config, output, refresh):
|
||||
"""
|
||||
Creates a new index.
|
||||
|
||||
Args:
|
||||
method: indexing method
|
||||
path: path to dataset
|
||||
config: path to config file
|
||||
output: path to store index
|
||||
refresh: overwrites existing index if True, otherwise existing index is loaded
|
||||
|
||||
Returns:
|
||||
Index
|
||||
"""
|
||||
|
||||
if method == "hybrid":
|
||||
return Hybrid(path, config, output, refresh)
|
||||
if method == "rag":
|
||||
return RetrievalAugmentedGeneration(path, config, output, refresh)
|
||||
if method == "scoring":
|
||||
return Score(path, config, output, refresh)
|
||||
if method == "rank":
|
||||
return RankBM25(path, config, output, refresh)
|
||||
if method == "bm25s":
|
||||
return BM25S(path, config, output, refresh)
|
||||
if method == "sqlite":
|
||||
return SQLiteFTS(path, config, output, refresh)
|
||||
if method == "es":
|
||||
return Elastic(path, config, output, refresh)
|
||||
if method == "similar":
|
||||
return Similar(path, config, output, refresh)
|
||||
if method == "rerank":
|
||||
return Rerank(path, config, output, refresh)
|
||||
|
||||
# Default
|
||||
return Embed(path, config, output, refresh)
|
||||
|
||||
|
||||
def compute(results):
|
||||
"""
|
||||
Computes metrics using the results from an evaluation run.
|
||||
|
||||
Args:
|
||||
results: evaluation results
|
||||
|
||||
Returns:
|
||||
metrics
|
||||
"""
|
||||
|
||||
metrics = {}
|
||||
for r in results:
|
||||
for metric in results[r]:
|
||||
if metric not in metrics:
|
||||
metrics[metric] = []
|
||||
|
||||
metrics[metric].append(results[r][metric])
|
||||
|
||||
return {metric: round(np.mean(values), 5) for metric, values in metrics.items()}
|
||||
|
||||
|
||||
def evaluate(methods, path, args):
|
||||
"""
|
||||
Runs an evaluation.
|
||||
|
||||
Args:
|
||||
methods: list of indexing methods to test
|
||||
path: path to dataset
|
||||
args: command line arguments
|
||||
|
||||
Returns:
|
||||
{calculated performance metrics}
|
||||
"""
|
||||
|
||||
print(f"------ {os.path.basename(path)} ------")
|
||||
|
||||
# Performance stats
|
||||
performance = {}
|
||||
|
||||
# Calculate stats for each model type
|
||||
topk = args.topk
|
||||
evaluator = RelevanceEvaluator(relevance(path), {f"ndcg_cut.{topk}", f"map_cut.{topk}", f"recall.{topk}", f"P.{topk}"})
|
||||
for method in methods:
|
||||
# Stats for this source
|
||||
stats = {}
|
||||
performance[method] = stats
|
||||
|
||||
# Create index and get results
|
||||
start = time.time()
|
||||
output = args.output if args.output else f"{path}/{method}"
|
||||
index = create(method, path, args.config, output, args.refresh)
|
||||
|
||||
# Add indexing metrics
|
||||
stats["index"] = round(time.time() - start, 2)
|
||||
stats["memory"] = int(psutil.Process().memory_info().rss / (1024 * 1024))
|
||||
stats["disk"] = int(sum(d.stat().st_size for d in os.scandir(output) if d.is_file()) / 1024) if os.path.isdir(output) else 0
|
||||
|
||||
print("INDEX TIME =", time.time() - start)
|
||||
print(f"MEMORY USAGE = {stats['memory']} MB")
|
||||
print(f"DISK USAGE = {stats['disk']} KB")
|
||||
|
||||
start = time.time()
|
||||
results = index(topk)
|
||||
|
||||
# Add search metrics
|
||||
stats["search"] = round(time.time() - start, 2)
|
||||
print("SEARCH TIME =", time.time() - start)
|
||||
|
||||
# Calculate stats
|
||||
metrics = compute(evaluator.evaluate(results))
|
||||
|
||||
# Add accuracy metrics
|
||||
for stat in [f"ndcg_cut_{topk}", f"map_cut_{topk}", f"recall_{topk}", f"P_{topk}"]:
|
||||
stats[stat] = metrics[stat]
|
||||
|
||||
# Print model stats
|
||||
print(f"------ {method} ------")
|
||||
print(f"NDCG@{topk} =", metrics[f"ndcg_cut_{topk}"])
|
||||
print(f"MAP@{topk} =", metrics[f"map_cut_{topk}"])
|
||||
print(f"Recall@{topk} =", metrics[f"recall_{topk}"])
|
||||
print(f"P@{topk} =", metrics[f"P_{topk}"])
|
||||
|
||||
print()
|
||||
return performance
|
||||
|
||||
|
||||
def benchmarks(args):
|
||||
"""
|
||||
Main benchmark execution method.
|
||||
|
||||
Args:
|
||||
args: command line arguments
|
||||
"""
|
||||
|
||||
# Directory where BEIR datasets are stored
|
||||
directory = args.directory if args.directory else "beir"
|
||||
|
||||
if args.sources and args.methods:
|
||||
sources, methods = args.sources.split(","), args.methods.split(",")
|
||||
mode = "a"
|
||||
else:
|
||||
# Default sources and methods
|
||||
sources = [
|
||||
"trec-covid",
|
||||
"nfcorpus",
|
||||
"nq",
|
||||
"hotpotqa",
|
||||
"fiqa",
|
||||
"arguana",
|
||||
"webis-touche2020",
|
||||
"quora",
|
||||
"dbpedia-entity",
|
||||
"scidocs",
|
||||
"fever",
|
||||
"climate-fever",
|
||||
"scifact",
|
||||
]
|
||||
methods = ["embed", "hybrid", "rag", "scoring", "rank", "bm25s", "sqlite", "es", "similar", "rerank"]
|
||||
mode = "w"
|
||||
|
||||
# Run and save benchmarks
|
||||
with open("benchmarks.json", mode, encoding="utf-8") as f:
|
||||
for source in sources:
|
||||
# Run evaluations
|
||||
results = evaluate(methods, f"{directory}/{source}", args)
|
||||
|
||||
# Save as JSON lines output
|
||||
for method, stats in results.items():
|
||||
stats["source"] = source
|
||||
stats["method"] = method
|
||||
stats["name"] = args.name if args.name else method
|
||||
|
||||
json.dump(stats, f)
|
||||
f.write("\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Command line parser
|
||||
parser = argparse.ArgumentParser(description="Benchmarks")
|
||||
parser.add_argument("-c", "--config", help="path to config file", metavar="CONFIG")
|
||||
parser.add_argument("-d", "--directory", help="root directory path with datasets", metavar="DIRECTORY")
|
||||
parser.add_argument("-m", "--methods", help="comma separated list of methods", metavar="METHODS")
|
||||
parser.add_argument("-n", "--name", help="name to assign to this run, defaults to method name", metavar="NAME")
|
||||
parser.add_argument("-o", "--output", help="index output directory path", metavar="OUTPUT")
|
||||
parser.add_argument(
|
||||
"-r",
|
||||
"--refresh",
|
||||
help="refreshes index if set, otherwise uses existing index if available",
|
||||
action="store_true",
|
||||
)
|
||||
parser.add_argument("-s", "--sources", help="comma separated list of data sources", metavar="SOURCES")
|
||||
parser.add_argument("-t", "--topk", help="top k results to use for the evaluation", metavar="TOPK", type=int, default=10)
|
||||
|
||||
# Calculate benchmarks
|
||||
benchmarks(parser.parse_args())
|
||||
@@ -0,0 +1,186 @@
|
||||
"""
|
||||
Search application using Open Library book data. Requires the following steps to be run:
|
||||
|
||||
Install Streamlit
|
||||
pip install streamlit
|
||||
|
||||
Download and prepare data
|
||||
mkdir openlibrary && cd openlibrary
|
||||
wget -O works.txt.gz https://openlibrary.org/data/ol_dump_works_latest.txt.gz
|
||||
gunzip works.txt.gz
|
||||
grep "\"description\":" works.txt > filtered.txt
|
||||
|
||||
Build index
|
||||
python books.py openlibrary
|
||||
|
||||
Run application
|
||||
streamlit run books.py openlibrary
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
|
||||
import pandas as pd
|
||||
import streamlit as st
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
|
||||
class Application:
|
||||
"""
|
||||
Main application.
|
||||
"""
|
||||
|
||||
def __init__(self, path):
|
||||
"""
|
||||
Creates a new application.
|
||||
|
||||
Args:
|
||||
path: root path to data
|
||||
"""
|
||||
|
||||
self.path = path
|
||||
self.dbpath = os.path.join(self.path, "books")
|
||||
|
||||
def rows(self, index):
|
||||
"""
|
||||
Iterates over dataset yielding each row.
|
||||
|
||||
Args:
|
||||
index: yields rows for embeddings indexing if True, otherwise yields database rows
|
||||
"""
|
||||
|
||||
with open(os.path.join(self.path, "filtered.txt"), encoding="utf-8") as infile:
|
||||
for x, row in enumerate(infile):
|
||||
if x % 1000 == 0:
|
||||
print(f"Processed {x} rows", end="\r")
|
||||
|
||||
row = row.split("\t")
|
||||
uid, data = row[1], json.loads(row[4])
|
||||
|
||||
description = data["description"]
|
||||
if isinstance(description, dict):
|
||||
description = description["value"]
|
||||
|
||||
if "title" in data:
|
||||
if index:
|
||||
yield (uid, data["title"] + ". " + description, None)
|
||||
else:
|
||||
cover = f"{data['covers'][0]}" if "covers" in data and data["covers"] else None
|
||||
yield (uid, data["title"], description, cover)
|
||||
|
||||
def database(self):
|
||||
"""
|
||||
Builds a SQLite database.
|
||||
"""
|
||||
|
||||
# Database file path
|
||||
dbfile = os.path.join(self.dbpath, "books.sqlite")
|
||||
|
||||
# Delete existing file
|
||||
if os.path.exists(dbfile):
|
||||
os.remove(dbfile)
|
||||
|
||||
# Create output database
|
||||
db = sqlite3.connect(dbfile)
|
||||
|
||||
# Create database cursor
|
||||
cur = db.cursor()
|
||||
|
||||
cur.execute("CREATE TABLE books (Id TEXT PRIMARY KEY, Title TEXT, Description TEXT, Cover TEXT)")
|
||||
|
||||
for uid, title, description, cover in self.rows(False):
|
||||
cur.execute("INSERT INTO books (Id, Title, Description, Cover) VALUES (?, ?, ?, ?)", (uid, title, description, cover))
|
||||
|
||||
# Finish and close database
|
||||
db.commit()
|
||||
db.close()
|
||||
|
||||
def build(self):
|
||||
"""
|
||||
Builds an embeddings index and database.
|
||||
"""
|
||||
|
||||
# Build embeddings index
|
||||
embeddings = Embeddings({"path": "sentence-transformers/msmarco-distilbert-base-v4"})
|
||||
embeddings.index(self.rows(True))
|
||||
embeddings.save(self.dbpath)
|
||||
|
||||
# Build SQLite DB
|
||||
self.database()
|
||||
|
||||
@st.cache(allow_output_mutation=True)
|
||||
def load(self):
|
||||
"""
|
||||
Loads and caches embeddings index.
|
||||
|
||||
Returns:
|
||||
embeddings index
|
||||
"""
|
||||
|
||||
embeddings = Embeddings()
|
||||
embeddings.load(self.dbpath)
|
||||
|
||||
return embeddings
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
Runs a Streamlit application.
|
||||
"""
|
||||
|
||||
# Build embeddings index
|
||||
embeddings = self.load()
|
||||
|
||||
db = sqlite3.connect(os.path.join(self.dbpath, "books.sqlite"))
|
||||
cur = db.cursor()
|
||||
|
||||
st.title("Book search")
|
||||
|
||||
st.markdown(
|
||||
"This application builds a local txtai index using book data from [openlibrary.org](https://openlibrary.org). "
|
||||
+ "Links to the Open Library pages and covers are shown in the application."
|
||||
)
|
||||
|
||||
query = st.text_input("Search query:")
|
||||
if query:
|
||||
ids = [uid for uid, score in embeddings.search(query, 10) if score >= 0.5]
|
||||
|
||||
results = []
|
||||
for uid in ids:
|
||||
cur.execute("SELECT Title, Description, Cover FROM books WHERE Id=?", (uid,))
|
||||
result = cur.fetchone()
|
||||
|
||||
if result:
|
||||
# Build cover image
|
||||
cover = (
|
||||
f"<img src='http://covers.openlibrary.org/b/id/{result[2]}-M.jpg'/>"
|
||||
if result[2]
|
||||
else "<img src='http://openlibrary.org/images/icons/avatar_book-lg.png'/>"
|
||||
)
|
||||
|
||||
# Append book link
|
||||
cover = f"<a target='_blank' href='https://openlibrary.org/{uid}'>{cover}</a>"
|
||||
title = f"<a target='_blank' href='https://openlibrary.org/{uid}'>{result[0]}</a>"
|
||||
|
||||
results.append({"Cover": cover, "Title": title, "Description": result[1]})
|
||||
|
||||
st.write(pd.DataFrame(results).to_html(escape=False, index=False), unsafe_allow_html=True)
|
||||
|
||||
db.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
# Application is used both to index and search
|
||||
app = Application(sys.argv[1])
|
||||
|
||||
# pylint: disable=W0212
|
||||
if st._is_running_with_streamlit:
|
||||
# Run application using existing index/db
|
||||
app.run()
|
||||
else:
|
||||
# Not running through streamlit, build database/index
|
||||
app.build()
|
||||
@@ -0,0 +1,105 @@
|
||||
"""
|
||||
Builds a similarity index for a directory of images
|
||||
|
||||
Requires streamlit to be installed.
|
||||
pip install streamlit
|
||||
"""
|
||||
|
||||
import glob
|
||||
import os
|
||||
import sys
|
||||
|
||||
import streamlit as st
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
|
||||
class Application:
|
||||
"""
|
||||
Main application
|
||||
"""
|
||||
|
||||
def __init__(self, directory):
|
||||
"""
|
||||
Creates a new application.
|
||||
|
||||
Args:
|
||||
directory: directory of images
|
||||
"""
|
||||
|
||||
self.embeddings = self.build(directory)
|
||||
|
||||
def build(self, directory):
|
||||
"""
|
||||
Builds an image embeddings index.
|
||||
|
||||
Args:
|
||||
directory: directory with images
|
||||
|
||||
Returns:
|
||||
Embeddings index
|
||||
"""
|
||||
|
||||
embeddings = Embeddings({"method": "sentence-transformers", "path": "clip-ViT-B-32"})
|
||||
embeddings.index(self.images(directory))
|
||||
|
||||
# Update model to support multilingual queries
|
||||
embeddings.config["path"] = "sentence-transformers/clip-ViT-B-32-multilingual-v1"
|
||||
embeddings.model = embeddings.loadvectors()
|
||||
|
||||
return embeddings
|
||||
|
||||
def images(self, directory):
|
||||
"""
|
||||
Generator that loops over each image in a directory.
|
||||
|
||||
Args:
|
||||
directory: directory with images
|
||||
"""
|
||||
|
||||
for path in glob.glob(directory + "/*jpg") + glob.glob(directory + "/*png"):
|
||||
yield (path, Image.open(path), None)
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
Runs a Streamlit application.
|
||||
"""
|
||||
|
||||
st.title("Image search")
|
||||
|
||||
st.markdown("This application shows how images and text can be embedded into the same space to support similarity search. ")
|
||||
st.markdown(
|
||||
"[sentence-transformers](https://github.com/UKPLab/sentence-transformers/tree/master/examples/applications/image-search) "
|
||||
+ "recently added support for the [OpenAI CLIP model](https://github.com/openai/CLIP). This model embeds text and images into "
|
||||
+ "the same space, enabling image similarity search. txtai can directly utilize these models."
|
||||
)
|
||||
|
||||
query = st.text_input("Search query:")
|
||||
if query:
|
||||
index, _ = self.embeddings.search(query, 1)[0]
|
||||
st.image(Image.open(index))
|
||||
|
||||
|
||||
@st.cache(allow_output_mutation=True)
|
||||
def create(directory):
|
||||
"""
|
||||
Creates and caches a Streamlit application.
|
||||
|
||||
Args:
|
||||
directory: directory of images to index
|
||||
|
||||
Returns:
|
||||
Application
|
||||
"""
|
||||
|
||||
return Application(directory)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
# Create and run application
|
||||
app = create(sys.argv[1])
|
||||
app.run()
|
||||
@@ -0,0 +1,70 @@
|
||||
"""
|
||||
RAG Quick Start
|
||||
Easy to use way to get started with RAG using YOUR data
|
||||
|
||||
For a complete application see this: https://github.com/neuml/rag
|
||||
|
||||
TxtAI has many example notebooks covering everything the framework provides
|
||||
Examples: https://neuml.github.io/txtai/examples
|
||||
|
||||
Install TxtAI
|
||||
pip install txtai[pipeline-data]
|
||||
"""
|
||||
|
||||
# pylint: disable=C0103
|
||||
import os
|
||||
|
||||
from txtai import Embeddings, RAG
|
||||
from txtai.pipeline import Textractor
|
||||
|
||||
# Step 1: Collect files from local directory
|
||||
#
|
||||
# Defaults to "data". Set to whereever your files are.
|
||||
path = "data"
|
||||
files = [os.path.join(path, f) for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]
|
||||
|
||||
# Step 2: Text Extraction / Chunking
|
||||
#
|
||||
# Using section based chunking here. More complex options available such as semantic chunking, iterative chunking etc.
|
||||
# Documentation: https://neuml.github.io/txtai/pipeline/data/textractor
|
||||
# Supports Chonkie chunking as well: https://docs.chonkie.ai/oss/chunkers/overview
|
||||
textractor = Textractor(backend="docling", sections=True)
|
||||
chunks = []
|
||||
for f in files:
|
||||
for chunk in textractor(f):
|
||||
chunks.append((f, chunk))
|
||||
|
||||
# Step 3: Build an embeddings database
|
||||
#
|
||||
# The `path` parameter sets the vector embeddings model. Supports Hugging Face models, llama.cpp, Ollama, vLLM and more.
|
||||
# Documentation: https://neuml.github.io/txtai/embeddings/
|
||||
embeddings = Embeddings(content=True, path="Qwen/Qwen3-Embedding-0.6B", maxlength=2048)
|
||||
embeddings.index(chunks)
|
||||
|
||||
# Step 4: Create RAG pipeline
|
||||
#
|
||||
# Combines an embeddings database and an LLM.
|
||||
# Supports Hugging Face models, llama.cpp, Ollama, vLLM and more
|
||||
# Documentation: https://neuml.github.io/txtai/pipeline/text/rag
|
||||
|
||||
# User prompt template
|
||||
template = """
|
||||
Answer the following question using the provided context.
|
||||
|
||||
Question:
|
||||
{question}
|
||||
|
||||
Context:
|
||||
{context}
|
||||
"""
|
||||
|
||||
rag = RAG(
|
||||
embeddings,
|
||||
"Qwen/Qwen3-0.6B",
|
||||
system="You are a friendly assistant",
|
||||
template=template,
|
||||
output="flatten",
|
||||
)
|
||||
|
||||
question = "Summarize the main advancements made by BERT"
|
||||
print(rag(question, maxlength=2048, stripthink=True))
|
||||
@@ -0,0 +1,73 @@
|
||||
"""
|
||||
Basic similarity search example. Used in the original txtai demo.
|
||||
|
||||
Requires streamlit to be installed.
|
||||
pip install streamlit
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import streamlit as st
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
|
||||
class Application:
|
||||
"""
|
||||
Main application.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Creates a new application.
|
||||
"""
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
self.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
Runs a Streamlit application.
|
||||
"""
|
||||
|
||||
st.title("Similarity Search")
|
||||
st.markdown("This application runs a basic similarity search that identifies the best matching row for a query.")
|
||||
|
||||
data = [
|
||||
"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",
|
||||
]
|
||||
|
||||
data = st.text_area("Data", value="\n".join(data))
|
||||
query = st.text_input("Query")
|
||||
|
||||
data = data.split("\n")
|
||||
|
||||
if query:
|
||||
# Get index of best section that best matches query
|
||||
uid = self.embeddings.similarity(query, data)[0][0]
|
||||
st.write(data[uid])
|
||||
|
||||
|
||||
@st.cache(allow_output_mutation=True)
|
||||
def create():
|
||||
"""
|
||||
Creates and caches a Streamlit application.
|
||||
|
||||
Returns:
|
||||
Application
|
||||
"""
|
||||
|
||||
return Application()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
# Create and run application
|
||||
app = create()
|
||||
app.run()
|
||||
@@ -0,0 +1,73 @@
|
||||
"""
|
||||
Application that queries Wikipedia API and summarizes the top result.
|
||||
|
||||
Requires streamlit to be installed.
|
||||
pip install streamlit
|
||||
"""
|
||||
|
||||
import os
|
||||
import urllib.parse
|
||||
|
||||
import requests
|
||||
import streamlit as st
|
||||
|
||||
from txtai.pipeline import Summary
|
||||
|
||||
|
||||
class Application:
|
||||
"""
|
||||
Main application.
|
||||
"""
|
||||
|
||||
SEARCH_TEMPLATE = "https://en.wikipedia.org/w/api.php?action=opensearch&search=%s&limit=1&namespace=0&format=json"
|
||||
CONTENT_TEMPLATE = "https://en.wikipedia.org/w/api.php?format=json&action=query&prop=extracts&exintro&explaintext&redirects=1&titles=%s"
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Creates a new application.
|
||||
"""
|
||||
|
||||
self.summary = Summary("sshleifer/distilbart-cnn-12-6")
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
Runs a Streamlit application.
|
||||
"""
|
||||
|
||||
st.title("Wikipedia")
|
||||
st.markdown("This application queries the Wikipedia API and summarizes the top result.")
|
||||
|
||||
query = st.text_input("Query")
|
||||
|
||||
if query:
|
||||
query = urllib.parse.quote_plus(query)
|
||||
data = requests.get(Application.SEARCH_TEMPLATE % query).json()
|
||||
if data and data[1]:
|
||||
page = urllib.parse.quote_plus(data[1][0])
|
||||
content = requests.get(Application.CONTENT_TEMPLATE % page).json()
|
||||
content = list(content["query"]["pages"].values())[0]["extract"]
|
||||
|
||||
st.write(self.summary(content))
|
||||
st.markdown("*Source: " + data[3][0] + "*")
|
||||
else:
|
||||
st.markdown("*No results found*")
|
||||
|
||||
|
||||
@st.cache(allow_output_mutation=True)
|
||||
def create():
|
||||
"""
|
||||
Creates and caches a Streamlit application.
|
||||
|
||||
Returns:
|
||||
Application
|
||||
"""
|
||||
|
||||
return Application()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
# Create and run application
|
||||
app = create()
|
||||
app.run()
|
||||
@@ -0,0 +1,59 @@
|
||||
"""
|
||||
Workflow Quick Start
|
||||
Easy to use way to get started with deterministic workflows.
|
||||
|
||||
TxtAI has many example notebooks covering everything the framework provides
|
||||
Examples: https://neuml.github.io/txtai/examples
|
||||
|
||||
Install TxtAI
|
||||
pip install txtai[pipeline-data]
|
||||
"""
|
||||
|
||||
from txtai import LLM, Workflow
|
||||
from txtai.pipeline import Summary, Textractor, Translation
|
||||
from txtai.workflow import Task
|
||||
|
||||
# Step 1: Define available pipelines
|
||||
textractor = Textractor(backend="docling", headers={"user-agent": "Mozilla/5.0"})
|
||||
summary = Summary()
|
||||
translate = Translation()
|
||||
|
||||
# Step 2: Define workflow tasks
|
||||
workflow = Workflow([Task(textractor), Task(summary), Task(lambda inputs: [translate(x, "fr") for x in inputs])])
|
||||
|
||||
# Step 3: Run the workflow
|
||||
print(list(workflow(["https://neuml.com"])))
|
||||
|
||||
# Each component above is a single model that specializes in a task
|
||||
# LLMs can also be used to accomplish the same tasks
|
||||
|
||||
|
||||
# pylint: disable=E0102,C0116
|
||||
def summary(text):
|
||||
return f"""
|
||||
Summarize the following text in 40 words or less.
|
||||
|
||||
{text}
|
||||
"""
|
||||
|
||||
|
||||
def translate(text, language):
|
||||
return f"""
|
||||
Translate the following text to {language}.
|
||||
|
||||
{text}
|
||||
"""
|
||||
|
||||
|
||||
textractor = Textractor(backend="docling", headers={"user-agent": "Mozilla/5.0"})
|
||||
llm = LLM("Qwen/Qwen3-4B-Instruct-2507")
|
||||
|
||||
workflow = Workflow(
|
||||
[
|
||||
Task(textractor),
|
||||
Task(lambda inputs: llm([summary(x) for x in inputs], maxlength=25000, defaultrole="user")),
|
||||
Task(lambda inputs: llm([translate(x, "fr") for x in inputs], maxlength=25000, defaultrole="user")),
|
||||
]
|
||||
)
|
||||
|
||||
print(list(workflow(["https://neuml.com"])))
|
||||
@@ -0,0 +1,745 @@
|
||||
"""
|
||||
Build txtai workflows.
|
||||
|
||||
Requires streamlit to be installed.
|
||||
pip install streamlit
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import copy
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
|
||||
import uvicorn
|
||||
import yaml
|
||||
|
||||
import pandas as pd
|
||||
import streamlit as st
|
||||
|
||||
import txtai.api.application
|
||||
import txtai.app
|
||||
|
||||
|
||||
class Server(uvicorn.Server):
|
||||
"""
|
||||
Threaded uvicorn server used to bring up an API service.
|
||||
"""
|
||||
|
||||
def __init__(self, application=None, host="127.0.0.1", port=8000, log_level="info"):
|
||||
"""
|
||||
Initialize server configuration.
|
||||
"""
|
||||
|
||||
config = uvicorn.Config(application, host=host, port=port, log_level=log_level)
|
||||
super().__init__(config)
|
||||
|
||||
def install_signal_handlers(self):
|
||||
"""
|
||||
Signal handlers no-op.
|
||||
"""
|
||||
|
||||
@contextlib.contextmanager
|
||||
def service(self):
|
||||
"""
|
||||
Runs threaded server service.
|
||||
"""
|
||||
|
||||
# pylint: disable=W0201
|
||||
thread = threading.Thread(target=self.run)
|
||||
thread.start()
|
||||
try:
|
||||
while not self.started:
|
||||
time.sleep(1e-3)
|
||||
yield
|
||||
|
||||
finally:
|
||||
self.should_exit = True
|
||||
thread.join()
|
||||
|
||||
|
||||
class Process:
|
||||
"""
|
||||
Container for an active Workflow process instance.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@st.cache_resource(show_spinner=False)
|
||||
def get(name, config):
|
||||
"""
|
||||
Lookup or creates a new workflow process instance.
|
||||
|
||||
Args:
|
||||
name: workflow name
|
||||
config: application configuration
|
||||
|
||||
Returns:
|
||||
Process
|
||||
"""
|
||||
|
||||
process = Process()
|
||||
|
||||
# Build workflow
|
||||
with st.spinner("Building workflow...."):
|
||||
process.build(name, config)
|
||||
|
||||
return process
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Creates a new Process.
|
||||
"""
|
||||
|
||||
# Application handle
|
||||
self.application = None
|
||||
|
||||
# Workflow name
|
||||
self.name = None
|
||||
|
||||
# Workflow data
|
||||
self.data = None
|
||||
|
||||
def build(self, name, config):
|
||||
"""
|
||||
Builds an application.
|
||||
|
||||
Args:
|
||||
name: workflow name
|
||||
config: application configuration
|
||||
"""
|
||||
|
||||
# Create application
|
||||
self.application = txtai.app.Application(config)
|
||||
|
||||
# Workflow name
|
||||
self.name = name
|
||||
|
||||
def run(self, data):
|
||||
"""
|
||||
Runs a workflow using data as input.
|
||||
|
||||
Args:
|
||||
data: input data
|
||||
"""
|
||||
|
||||
if data and self.application:
|
||||
# Build tuples for embedding index
|
||||
if self.application.embeddings:
|
||||
data = [(x, element, None) for x, element in enumerate(data)]
|
||||
|
||||
# Process workflow
|
||||
with st.spinner("Running workflow...."):
|
||||
results = []
|
||||
for result in self.application.workflow(self.name, data):
|
||||
# Store result
|
||||
results.append(result)
|
||||
|
||||
# Write result if this isn't an indexing workflow
|
||||
if not self.application.embeddings:
|
||||
st.write(result)
|
||||
|
||||
# Store workflow results
|
||||
self.data = results
|
||||
|
||||
def search(self, query):
|
||||
"""
|
||||
Runs a search.
|
||||
|
||||
Args:
|
||||
query: input query
|
||||
"""
|
||||
|
||||
if self.application and query:
|
||||
st.markdown(
|
||||
"""
|
||||
<style>
|
||||
table td:nth-child(1) {
|
||||
display: none
|
||||
}
|
||||
table th:nth-child(1) {
|
||||
display: none
|
||||
}
|
||||
table {text-align: left !important}
|
||||
</style>
|
||||
""",
|
||||
unsafe_allow_html=True,
|
||||
)
|
||||
|
||||
results = []
|
||||
for result in self.application.search(query, 5):
|
||||
# Text is only present when content is stored
|
||||
if "text" not in result:
|
||||
uid, score = result["id"], result["score"]
|
||||
results.append({"text": self.find(uid), "score": f"{score:.2}"})
|
||||
else:
|
||||
if "id" in result and "text" in result:
|
||||
result["text"] = self.content(result.pop("id"), result["text"])
|
||||
if "score" in result and result["score"]:
|
||||
result["score"] = f'{result["score"]:.2}'
|
||||
|
||||
results.append(result)
|
||||
|
||||
df = pd.DataFrame(results)
|
||||
st.write(df.to_html(escape=False), unsafe_allow_html=True)
|
||||
|
||||
def find(self, key):
|
||||
"""
|
||||
Lookup record from cached data by uid key.
|
||||
|
||||
Args:
|
||||
key: id to search for
|
||||
|
||||
Returns:
|
||||
text for matching id
|
||||
"""
|
||||
|
||||
# Lookup text by id
|
||||
text = [text for uid, text, _ in self.data if uid == key][0]
|
||||
return self.content(key, text)
|
||||
|
||||
def content(self, uid, text):
|
||||
"""
|
||||
Builds a content reference for uid and text.
|
||||
|
||||
Args:
|
||||
uid: record id
|
||||
text: record text
|
||||
|
||||
Returns:
|
||||
content
|
||||
"""
|
||||
|
||||
if uid and isinstance(uid, str) and uid.lower().startswith("http"):
|
||||
return f"<a href='{uid}' rel='noopener noreferrer' target='blank'>{text}</a>"
|
||||
|
||||
return text
|
||||
|
||||
|
||||
class Application:
|
||||
"""
|
||||
Main application.
|
||||
"""
|
||||
|
||||
def load(self, components):
|
||||
"""
|
||||
Load an existing workflow file.
|
||||
|
||||
Args:
|
||||
components: list of components to load
|
||||
|
||||
Returns:
|
||||
(names of components loaded, workflow config, file changed)
|
||||
"""
|
||||
|
||||
workflow = st.file_uploader("Load workflow", type=["yml"])
|
||||
if workflow:
|
||||
# Detect file upload change
|
||||
upload = workflow.name != self.state("path")
|
||||
st.session_state["path"] = workflow.name
|
||||
|
||||
workflow = yaml.safe_load(workflow)
|
||||
|
||||
st.markdown("---")
|
||||
|
||||
# Get tasks for first workflow
|
||||
tasks = list(workflow["workflow"].values())[0]["tasks"]
|
||||
selected = []
|
||||
|
||||
for task in tasks:
|
||||
name = task.get("action", task.get("task"))
|
||||
if name in components:
|
||||
selected.append(name)
|
||||
elif name in ["index", "upsert"]:
|
||||
selected.append("embeddings")
|
||||
|
||||
return (selected, workflow, upload)
|
||||
|
||||
return (None, None, None)
|
||||
|
||||
def state(self, key):
|
||||
"""
|
||||
Lookup a session state variable.
|
||||
|
||||
Args:
|
||||
key: variable key
|
||||
|
||||
Returns:
|
||||
variable value
|
||||
"""
|
||||
|
||||
if key in st.session_state:
|
||||
return st.session_state[key]
|
||||
|
||||
return None
|
||||
|
||||
def appsetting(self, workflow, name):
|
||||
"""
|
||||
Looks up an application configuration setting.
|
||||
|
||||
Args:
|
||||
workflow: workflow configuration
|
||||
name: setting name
|
||||
|
||||
Returns:
|
||||
app setting value
|
||||
"""
|
||||
|
||||
if workflow:
|
||||
config = workflow.get("app")
|
||||
if config:
|
||||
return config.get(name)
|
||||
|
||||
return None
|
||||
|
||||
def setting(self, config, name, default=None):
|
||||
"""
|
||||
Looks up a component configuration setting.
|
||||
|
||||
Args:
|
||||
config: component configuration
|
||||
name: setting name
|
||||
default: default setting value
|
||||
|
||||
Returns:
|
||||
setting value
|
||||
"""
|
||||
|
||||
return config.get(name, default) if config else default
|
||||
|
||||
def text(self, label, component, config, name, default=None):
|
||||
"""
|
||||
Create a new text input field.
|
||||
|
||||
Args:
|
||||
label: field label
|
||||
component: component name
|
||||
config: component configuration
|
||||
name: setting name
|
||||
default: default setting value
|
||||
|
||||
Returns:
|
||||
text input field value
|
||||
"""
|
||||
|
||||
default = self.setting(config, name, default)
|
||||
if not default:
|
||||
default = ""
|
||||
elif isinstance(default, list):
|
||||
default = ",".join(default)
|
||||
elif isinstance(default, dict):
|
||||
default = ",".join(default.keys())
|
||||
|
||||
return st.text_input(label, value=default, key=component + name)
|
||||
|
||||
def number(self, label, component, config, name, default=None):
|
||||
"""
|
||||
Creates a new numeric input field.
|
||||
|
||||
Args:
|
||||
label: field label
|
||||
component: component name
|
||||
config: component configuration
|
||||
name: setting name
|
||||
default: default setting value
|
||||
|
||||
Returns:
|
||||
numeric value
|
||||
"""
|
||||
|
||||
value = self.text(label, component, config, name, default)
|
||||
return int(value) if value else None
|
||||
|
||||
def boolean(self, label, component, config, name, default=False):
|
||||
"""
|
||||
Creates a new checkbox field.
|
||||
|
||||
Args:
|
||||
label: field label
|
||||
component: component name
|
||||
config: component configuration
|
||||
name: setting name
|
||||
default: default setting value
|
||||
|
||||
Returns:
|
||||
boolean value
|
||||
"""
|
||||
|
||||
default = self.setting(config, name, default)
|
||||
return st.checkbox(label, value=default, key=component + name)
|
||||
|
||||
def select(self, label, component, config, name, options, default=0):
|
||||
"""
|
||||
Creates a new select box field.
|
||||
|
||||
Args:
|
||||
label: field label
|
||||
component: component name
|
||||
config: component configuration
|
||||
name: setting name
|
||||
options: list of dropdown options
|
||||
default: default setting value
|
||||
|
||||
Returns:
|
||||
boolean value
|
||||
"""
|
||||
|
||||
index = self.setting(config, name)
|
||||
index = [x for x, option in enumerate(options) if option == default]
|
||||
|
||||
# Derive default index
|
||||
default = index[0] if index else default
|
||||
|
||||
return st.selectbox(label, options, index=default, key=component + name)
|
||||
|
||||
def split(self, text):
|
||||
"""
|
||||
Splits text on commas and returns a list.
|
||||
|
||||
Args:
|
||||
text: input text
|
||||
|
||||
Returns:
|
||||
list
|
||||
"""
|
||||
|
||||
return [x.strip() for x in text.split(",")]
|
||||
|
||||
def options(self, component, workflow, index):
|
||||
"""
|
||||
Extracts component settings into a component configuration dict.
|
||||
|
||||
Args:
|
||||
component: component type
|
||||
workflow: existing workflow, can be None
|
||||
index: task index
|
||||
|
||||
Returns:
|
||||
dict with component settings
|
||||
"""
|
||||
|
||||
# pylint: disable=R0912, R0915
|
||||
options = {"type": component}
|
||||
|
||||
st.markdown("---")
|
||||
|
||||
# Lookup component configuration
|
||||
# - Runtime components have config defined within tasks
|
||||
# - Pipeline components have config defined at workflow root
|
||||
config = None
|
||||
if workflow:
|
||||
if component in ["service", "translation"]:
|
||||
# Service config is found in tasks section
|
||||
tasks = list(workflow["workflow"].values())[0]["tasks"]
|
||||
tasks = [task for task in tasks if task.get("task") == component or task.get("action") == component]
|
||||
if tasks:
|
||||
config = tasks[0]
|
||||
else:
|
||||
config = workflow.get(component)
|
||||
|
||||
if component == "embeddings":
|
||||
st.markdown(f"**{index + 1}.) Embeddings Index** \n*Index workflow output*")
|
||||
options["index"] = self.text("Embeddings storage path", component, config, "index")
|
||||
options["path"] = self.text("Embeddings model path", component, config, "path", "sentence-transformers/nli-mpnet-base-v2")
|
||||
options["upsert"] = self.boolean("Upsert", component, config, "upsert")
|
||||
options["content"] = self.boolean("Content", component, config, "content")
|
||||
|
||||
elif component in ("segmentation", "textractor"):
|
||||
if component == "segmentation":
|
||||
st.markdown(f"**{index + 1}.) Segment** \n*Split text into semantic units*")
|
||||
else:
|
||||
st.markdown(f"**{index + 1}.) Textract** \n*Extract text from documents*")
|
||||
|
||||
options["sentences"] = self.boolean("Split sentences", component, config, "sentences")
|
||||
options["lines"] = self.boolean("Split lines", component, config, "lines")
|
||||
options["paragraphs"] = self.boolean("Split paragraphs", component, config, "paragraphs")
|
||||
options["join"] = self.boolean("Join tokenized", component, config, "join")
|
||||
options["minlength"] = self.number("Min section length", component, config, "minlength")
|
||||
|
||||
elif component == "service":
|
||||
st.markdown(f"**{index + 1}.) Service** \n*Extract data from an API*")
|
||||
options["url"] = self.text("URL", component, config, "url")
|
||||
options["method"] = self.select("Method", component, config, "method", ["get", "post"], 0)
|
||||
options["params"] = self.text("URL parameters", component, config, "params")
|
||||
options["batch"] = self.boolean("Run as batch", component, config, "batch", True)
|
||||
options["extract"] = self.text("Subsection(s) to extract", component, config, "extract")
|
||||
|
||||
if options["params"]:
|
||||
options["params"] = {key: None for key in self.split(options["params"])}
|
||||
if options["extract"]:
|
||||
options["extract"] = self.split(options["extract"])
|
||||
|
||||
elif component == "summary":
|
||||
st.markdown(f"**{index + 1}.) Summary** \n*Abstractive text summarization*")
|
||||
options["path"] = self.text("Model", component, config, "path", "sshleifer/distilbart-cnn-12-6")
|
||||
options["minlength"] = self.number("Min length", component, config, "minlength")
|
||||
options["maxlength"] = self.number("Max length", component, config, "maxlength")
|
||||
|
||||
elif component == "tabular":
|
||||
st.markdown(f"**{index + 1}.) Tabular** \n*Split tabular data into rows and columns*")
|
||||
options["idcolumn"] = self.text("Id columns", component, config, "idcolumn")
|
||||
options["textcolumns"] = self.text("Text columns", component, config, "textcolumns")
|
||||
options["content"] = self.text("Content", component, config, "content")
|
||||
|
||||
if options["textcolumns"]:
|
||||
options["textcolumns"] = self.split(options["textcolumns"])
|
||||
|
||||
if options["content"]:
|
||||
options["content"] = self.split(options["content"])
|
||||
if len(options["content"]) == 1 and options["content"][0] == "1":
|
||||
options["content"] = options["content"][0]
|
||||
|
||||
elif component == "transcription":
|
||||
st.markdown(f"**{index + 1}.) Transcribe** \n*Transcribe audio to text*")
|
||||
options["path"] = self.text("Model", component, config, "path", "facebook/wav2vec2-base-960h")
|
||||
|
||||
elif component == "translation":
|
||||
st.markdown(f"**{index + 1}.) Translate** \n*Machine translation*")
|
||||
options["target"] = self.text("Target language code", component, config, "args", "en")
|
||||
|
||||
return options
|
||||
|
||||
def config(self, components):
|
||||
"""
|
||||
Builds configuration for components
|
||||
|
||||
Args:
|
||||
components: list of components to add to configuration
|
||||
|
||||
Returns:
|
||||
(workflow name, configuration)
|
||||
"""
|
||||
|
||||
data = {}
|
||||
tasks = []
|
||||
name = None
|
||||
|
||||
for component in components:
|
||||
component = dict(component)
|
||||
name = wtype = component.pop("type")
|
||||
|
||||
if wtype == "embeddings":
|
||||
index = component.pop("index")
|
||||
upsert = component.pop("upsert")
|
||||
|
||||
data[wtype] = component
|
||||
data["writable"] = True
|
||||
|
||||
if index:
|
||||
data["path"] = index
|
||||
|
||||
name = "index"
|
||||
tasks.append({"action": "upsert" if upsert else "index"})
|
||||
|
||||
elif wtype == "segmentation":
|
||||
data[wtype] = component
|
||||
tasks.append({"action": wtype})
|
||||
|
||||
elif wtype == "service":
|
||||
config = {**component}
|
||||
config["task"] = wtype
|
||||
tasks.append(config)
|
||||
|
||||
elif wtype == "summary":
|
||||
data[wtype] = {"path": component.pop("path")}
|
||||
tasks.append({"action": wtype})
|
||||
|
||||
elif wtype == "tabular":
|
||||
data[wtype] = component
|
||||
tasks.append({"action": wtype})
|
||||
|
||||
elif wtype == "textractor":
|
||||
data[wtype] = component
|
||||
tasks.append({"action": wtype, "task": "url"})
|
||||
|
||||
elif wtype == "transcription":
|
||||
data[wtype] = {"path": component.pop("path")}
|
||||
tasks.append({"action": wtype, "task": "url"})
|
||||
|
||||
elif wtype == "translation":
|
||||
data[wtype] = {}
|
||||
tasks.append({"action": wtype, "args": list(component.values())})
|
||||
|
||||
# Add in workflow
|
||||
data["workflow"] = {name: {"tasks": tasks}}
|
||||
|
||||
# Return workflow name and application configuration
|
||||
return (name, data)
|
||||
|
||||
def api(self, config):
|
||||
"""
|
||||
Starts an internal uvicorn server to host an API service for the current workflow.
|
||||
|
||||
Args:
|
||||
config: workflow configuration as YAML string
|
||||
"""
|
||||
|
||||
# Generate workflow file
|
||||
workflow = os.path.join(tempfile.gettempdir(), "workflow.yml")
|
||||
with open(workflow, "w", encoding="utf-8") as f:
|
||||
f.write(config)
|
||||
|
||||
os.environ["CONFIG"] = workflow
|
||||
txtai.api.application.start()
|
||||
server = Server(txtai.api.application.app)
|
||||
with server.service():
|
||||
uid = 0
|
||||
while True:
|
||||
stop = st.empty()
|
||||
click = stop.button("stop", key=uid)
|
||||
if not click:
|
||||
time.sleep(5)
|
||||
uid += 1
|
||||
stop.empty()
|
||||
|
||||
def inputs(self, selected, workflow):
|
||||
"""
|
||||
Generate process input fields.
|
||||
|
||||
Args:
|
||||
selected: list of selected components
|
||||
workflow: workflow configuration
|
||||
|
||||
Returns:
|
||||
True if inputs changed, False otherwise
|
||||
"""
|
||||
|
||||
change, query = False, None
|
||||
with st.expander("Data", expanded="embeddings" not in selected):
|
||||
default = self.appsetting(workflow, "data")
|
||||
default = default if default else ""
|
||||
|
||||
data = st.text_area("Input", height=10, value=default)
|
||||
|
||||
if selected and data and data != self.state("data"):
|
||||
change = True
|
||||
|
||||
# Save data and workflow state
|
||||
st.session_state["data"] = data
|
||||
|
||||
if "embeddings" in selected:
|
||||
default = self.appsetting(workflow, "query")
|
||||
default = default if default else ""
|
||||
|
||||
# Set query and limit
|
||||
query = st.text_input("Query", value=default)
|
||||
|
||||
if selected and query and query != self.state("query"):
|
||||
change = True
|
||||
|
||||
# Save query state
|
||||
st.session_state["query"] = query
|
||||
|
||||
return change or self.state("api") or self.state("download")
|
||||
|
||||
def data(self):
|
||||
"""
|
||||
Gets input data.
|
||||
|
||||
Returns:
|
||||
input data
|
||||
"""
|
||||
|
||||
data = self.state("data")
|
||||
|
||||
# Split on newlines if urls detected, allows a list of urls to be processed
|
||||
if re.match(r"^(http|https|file):\/\/", data):
|
||||
return [x for x in data.split("\n") if x]
|
||||
|
||||
return [data]
|
||||
|
||||
def process(self, components, index):
|
||||
"""
|
||||
Processes the current application action.
|
||||
|
||||
Args:
|
||||
components: workflow components
|
||||
index: True if this is an indexing workflow
|
||||
"""
|
||||
|
||||
# Generate application configuration
|
||||
name, config = self.config(components)
|
||||
|
||||
# Get workflow process
|
||||
process = Process.get(name, copy.deepcopy(config))
|
||||
|
||||
# Run workflow process
|
||||
process.run(self.data())
|
||||
|
||||
# Run search
|
||||
if index:
|
||||
process.search(self.state("query"))
|
||||
|
||||
return name, config
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
Runs Streamlit application.
|
||||
"""
|
||||
|
||||
build = False
|
||||
with st.sidebar:
|
||||
st.image("https://github.com/neuml/txtai/raw/master/logo.png", width=256)
|
||||
st.markdown("# Workflow builder \n*Build and apply workflows to data* ")
|
||||
st.markdown("---")
|
||||
|
||||
# Component configuration
|
||||
labels = {"segmentation": "segment", "textractor": "textract", "transcription": "transcribe", "translation": "translate"}
|
||||
components = ["embeddings", "segmentation", "service", "summary", "tabular", "textractor", "transcription", "translation"]
|
||||
|
||||
selected, workflow, upload = self.load(components)
|
||||
selected = st.multiselect("Select components", components, default=selected, format_func=lambda text: labels.get(text, text))
|
||||
|
||||
if selected:
|
||||
st.markdown(
|
||||
"""
|
||||
<style>
|
||||
[data-testid="stForm"] {
|
||||
border: 0;
|
||||
padding: 0;
|
||||
}
|
||||
</style>
|
||||
""",
|
||||
unsafe_allow_html=True,
|
||||
)
|
||||
|
||||
with st.form("workflow"):
|
||||
# Get selected options
|
||||
components = [self.options(component, workflow, x) for x, component in enumerate(selected)]
|
||||
st.markdown("---")
|
||||
|
||||
# Build or re-build workflow when build button clicked or new workflow loaded
|
||||
build = st.form_submit_button("Build", help="Build the workflow and run within this application")
|
||||
|
||||
# Generate input fields
|
||||
inputs = self.inputs(selected, workflow)
|
||||
|
||||
# Only execute if build button clicked, new workflow uploaded or inputs changed
|
||||
if build or upload or inputs:
|
||||
# Process current action
|
||||
name, config = self.process(components, "embeddings" in selected)
|
||||
|
||||
with st.sidebar:
|
||||
with st.expander("Other Actions", expanded=True):
|
||||
col1, col2 = st.columns(2)
|
||||
|
||||
# Add state information to configuration and export to YAML string
|
||||
config = config.copy()
|
||||
config.update({"app": {"data": self.state("data"), "query": self.state("query")}})
|
||||
config = yaml.dump(config)
|
||||
|
||||
api = col1.button("API", key="api", help="Start an API instance within this application")
|
||||
if api:
|
||||
with st.spinner(f"Running workflow '{name}' via API service, click stop to terminate"):
|
||||
self.api(config)
|
||||
|
||||
col2.download_button("Export", config, file_name="workflow.yml", key="download", help="Export the API workflow as YAML")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
# Create and run application
|
||||
app = Application()
|
||||
app.run()
|
||||
Reference in New Issue
Block a user