neuml/txtai
1
0
Fork 0
mirror of https://github.com/neuml/txtai.git synced 2026-10-11 21:39:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

Replace fasttext and pymagnitude with staticvectors. Closes #857. Closes #858. Closes #859

Browse Source
This commit is contained in:
davidmezzetti
2025-01-27 11:17:51 -05:00
parent 12aa955b77
commit f79d8062ae
13 changed files with 270 additions and 458 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/embeddings/configuration/vectors.md
-7
View File
@@ -56,13 +56,6 @@ Builds embeddings using a [Model2Vec](https://github.com/MinishLab/model2vec) mo
Builds embeddings using a word embeddings model and static vectors. While Transformers models are preferred in most cases, this method can be useful for low resource and historical languages where there isn't much linguistic data available.
#### storevectors
```yaml
storevectors: boolean
```
Enables copying of a vectors model set in path into the embeddings models output directory on save. This option enables a fully encapsulated index with no external file dependencies.
#### pca
```yaml
pca: int
examples/01_Introducing_txtai.ipynb
+1 -1
View File
@@ -80,7 +80,7 @@
"!pip install git+https://github.com/neuml/txtai#egg=txtai[graph]\n",
"\n",
"# Install translation pipeline dependencies for later examples\n",
"!pip install sentencepiece sacremoses fasttext"
"!pip install sentencepiece sacremoses staticvectors"
],
"execution_count": 1,
"outputs": []
examples/03_Build_an_Embeddings_index_from_a_data_source.ipynb
+181 -257
View File
@@ -1,15 +1,4 @@
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"provenance": []
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
}
},
"cells": [
{
"cell_type": "markdown",
@@ -21,7 +10,7 @@
"\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 a Embedding index backed by word embeddings and run queries against the generated Embeddings index.\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."
]
@@ -34,20 +23,20 @@
"source": [
"# Install dependencies\n",
"\n",
"Install `txtai` and all dependencies. Since this notebook is building word vectors, we need to install the similarity extras package."
"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#egg=txtai[similarity]"
],
"execution_count": null,
"outputs": []
"!pip install git+https://github.com/neuml/txtai"
]
},
{
"cell_type": "markdown",
@@ -64,93 +53,16 @@
},
{
"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"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "bzdaJiZYBIHE"
},
"source": [
"# Build Word Vectors\n",
"\n",
"This example will build a search system backed by word embeddings. Word embeddings provide a good tradeoff of performance to functionality for a similarity search system. With that being said, transformer models are making great progress in scaling performance down to smaller models and are the preferred vector backend for txtai in most cases.\n",
"\n",
"For this notebook, we'll build our own custom embeddings for demo purposes. A number of pre-trained word embedding models are available:\n",
"\n",
" - [General language models from pymagnitude](https://github.com/plasticityai/magnitude)\n",
" - [CORD-19 fastText](https://www.kaggle.com/davidmezzetti/cord19-fasttext-vectors)"
]
},
{
"cell_type": "code",
"metadata": {
"id": "fJcn-CAH-u3K",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "ab5a4340-13fc-4b47-b881-9388e2534ece"
},
"source": [
"import os\n",
"import sqlite3\n",
"import tempfile\n",
"\n",
"from txtai.pipeline import Tokenizer\n",
"from txtai.vectors import WordVectors\n",
"\n",
"print(\"Streaming tokens to temporary file\")\n",
"\n",
"# Stream tokens to temp working file\n",
"with tempfile.NamedTemporaryFile(mode=\"w\", suffix=\".txt\", delete=False) as output:\n",
" # Save file path\n",
" tokens = output.name\n",
"\n",
" db = sqlite3.connect(\"articles.sqlite\")\n",
" cur = db.cursor()\n",
" cur.execute(\"SELECT Text from sections\")\n",
"\n",
" for row in cur:\n",
" output.write(\" \".join(Tokenizer.tokenize(row[0])) + \"\\n\")\n",
"\n",
" # Free database resources\n",
" db.close()\n",
"\n",
"# Build word vectors model - 300 dimensions, 3 min occurrences\n",
"WordVectors.build(tokens, 300, 3, \"cord19-300d\")\n",
"\n",
"# Remove temporary tokens file\n",
"os.remove(tokens)\n",
"\n",
"# Show files\n",
"!ls -l"
],
"execution_count": null,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Streaming tokens to temporary file\n",
"Building 300 dimension model\n",
"Converting vectors to magnitude format\n",
"total 78948\n",
"-rw-r--r-- 1 root root 8065024 Dec 7 15:48 articles.sqlite\n",
"-rw-r--r-- 1 root root 24145920 Dec 27 17:42 cord19-300d.magnitude\n",
"-rw-r--r-- 1 root root 48625387 Dec 27 17:41 cord19-300d.txt\n",
"drwxr-xr-x 1 root root 4096 Dec 3 14:33 sample_data\n"
]
}
]
},
{
@@ -161,25 +73,34 @@
"source": [
"# Build an embeddings index\n",
"\n",
"The following steps build an embeddings index using the word vector model just created. This model builds a BM25 + fastText index. BM25 is used to build a weighted average of the word embeddings for a section. More information on this method can be found in this [Medium article](https://towardsdatascience.com/building-a-sentence-embedding-index-with-fasttext-and-bm25-f07e7148d240?gi=79da927aa10)."
"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": {
"id": "5PrrxGRPGHqX",
"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.embeddings import Embeddings\n",
"from txtai.pipeline import Tokenizer\n",
"from txtai import Embeddings\n",
"\n",
"def stream():\n",
" # Connection to database file\n",
@@ -196,46 +117,24 @@
"\n",
" # Only process certain document sections\n",
" if not name or not re.search(r\"background|(?<!.*?results.*?)discussion|introduction|reference\", name.lower()):\n",
" # Tokenize text\n",
" tokens = Tokenizer.tokenize(text)\n",
"\n",
" document = (uid, tokens, None)\n",
" document = (uid, text, None)\n",
"\n",
" count += 1\n",
" if count % 1000 == 0:\n",
" print(\"Streamed %d documents\" % (count), end=\"\\r\")\n",
"\n",
" # Skip documents with no tokens parsed\n",
" if tokens:\n",
" yield document\n",
" yield document\n",
"\n",
" print(\"Iterated over %d total rows\" % (count))\n",
"\n",
" # Free database resources\n",
" db.close()\n",
"\n",
"# BM25 + fastText vectors\n",
"embeddings = Embeddings({\"path\": \"cord19-300d.magnitude\",\n",
" \"scoring\": \"bm25\",\n",
" \"pca\": 3})\n",
"\n",
"# Build scoring index if scoring method provided\n",
"if embeddings.config.get(\"scoring\"):\n",
" embeddings.score(stream())\n",
"# Create embeddings index \n",
"embeddings = Embeddings(path=\"neuml/pubmedbert-base-embeddings\")\n",
"\n",
"# Build embeddings index\n",
"embeddings.index(stream())\n"
],
"execution_count": null,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Iterated over 21499 total rows\n",
"Iterated over 21499 total rows\n"
]
}
]
},
{
@@ -251,14 +150,70 @@
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {
"id": "CRbDhvvDKEl-",
"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",
@@ -291,62 +246,6 @@
"# file.write(htmlData)\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>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>Work-related and Personal Factors Associated with Mental Well-being during COVID-19 Response: A Survey of Health Care and Other Workers</td>\n",
" <td>2020-06-11 00:00:00</td>\n",
" <td>http://medrxiv.org/cgi/content/short/2020.06.09.20126722v1?rss=1</td>\n",
" <td>Poor family supportive behaviors by supervisors were also associated with these outcomes [1.40 (1.21 - 1.62), 1.69 (1.48 - 1.92), 1.54 (1.44 - 1.64)].</td>\n",
" </tr>\n",
" <tr>\n",
" <td>No evidence that androgen regulation of pulmonary TMPRSS2 explains sex-discordant COVID-19 outcomes</td>\n",
" <td>2020-04-21 00:00:00</td>\n",
" <td>https://doi.org/10.1101/2020.04.21.051201</td>\n",
" <td>In addition to male sex, smoking is a risk factor for COVID-19 susceptibility and poor clinical outcomes .</td>\n",
" </tr>\n",
" <tr>\n",
" <td>Current status of potential therapeutic candidates for the COVID-19 crisis</td>\n",
" <td>2020-04-22 00:00:00</td>\n",
" <td>https://doi.org/10.1016/j.bbi.2020.04.046</td>\n",
" <td>There was no difference on 28-day mortality between heparin users and nonusers.</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",
" </tbody>\n",
"</table>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {}
}
]
},
{
@@ -362,29 +261,98 @@
},
{
"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\")"
],
"execution_count": null,
"outputs": []
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {
"id": "19fmKawThs6d",
"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",
@@ -415,74 +383,6 @@
"\n",
"df = pd.DataFrame(results, columns=[\"Title\", \"Published\", \"Reference\", \"Match\", \"Risk Factors\", \"Locations\"])\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>Work-related and Personal Factors Associated with Mental Well-being during COVID-19 Response: A Survey of Health Care and Other Workers</td>\n",
" <td>2020-06-11 00:00:00</td>\n",
" <td>http://medrxiv.org/cgi/content/short/2020.06.09.20126722v1?rss=1</td>\n",
" <td>Poor family supportive behaviors by supervisors were also associated with these outcomes [1.40 (1.21 - 1.62), 1.69 (1.48 - 1.92), 1.54 (1.44 - 1.64)].</td>\n",
" <td>Poor family supportive behaviors</td>\n",
" <td>None</td>\n",
" </tr>\n",
" <tr>\n",
" <td>No evidence that androgen regulation of pulmonary TMPRSS2 explains sex-discordant COVID-19 outcomes</td>\n",
" <td>2020-04-21 00:00:00</td>\n",
" <td>https://doi.org/10.1101/2020.04.21.051201</td>\n",
" <td>In addition to male sex, smoking is a risk factor for COVID-19 susceptibility and poor clinical outcomes .</td>\n",
" <td>Higher morbidity and mortality</td>\n",
" <td>None</td>\n",
" </tr>\n",
" <tr>\n",
" <td>Current status of potential therapeutic candidates for the COVID-19 crisis</td>\n",
" <td>2020-04-22 00:00:00</td>\n",
" <td>https://doi.org/10.1016/j.bbi.2020.04.046</td>\n",
" <td>There was no difference on 28-day mortality between heparin users and nonusers.</td>\n",
" <td>elicited strong inflammatory responses are favorable or detrimental</td>\n",
" <td>None</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>sex (male), age (≥60), and severe pneumonia</td>\n",
" <td>None</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {}
}
]
},
{
@@ -494,5 +394,29 @@
"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
}
examples/24_Whats_new_in_txtai_4_0.ipynb
+1 -1
View File
@@ -429,7 +429,7 @@
"source": [
"# External vector models\n",
"\n",
"txtai supports generating vectors with [Hugging Face Transformers](https://github.com/huggingface/transformers), [PyTorch](https://github.com/pytorch/pytorch), [ONNX](https://github.com/microsoft/onnxruntime) and [Word Vector](https://github.com/neuml/magnitude) models.\n",
"txtai supports generating vectors with [Hugging Face Transformers](https://github.com/huggingface/transformers), [PyTorch](https://github.com/pytorch/pytorch), [ONNX](https://github.com/microsoft/onnxruntime) and [Word Vector](https://github.com/neuml/staticvectors) models.\n",
"\n",
"This release adds support for pre-computed vectors using external models. External models may be an API, custom library and/or another way to vectorize data. This adds flexibility given the high computation cost in building embeddings vectors. Embeddings generation could be outsourced or consolidated to a group of servers with GPUs, leaving index servers to run on lower resourced machines. \n",
"\n",
setup.py
+2 -3
View File
@@ -69,7 +69,7 @@ extras["pipeline-image"] = ["imagehash>=4.2.1", "pillow>=7.1.2", "timm>=0.4.12"]
extras["pipeline-llm"] = ["litellm>=1.37.16", "llama-cpp-python>=0.2.75"]
extras["pipeline-text"] = ["fasttext-wheel>=0.9.2", "sentencepiece>=0.1.91"]
extras["pipeline-text"] = ["sentencepiece>=0.1.91", "staticvectors>=0.2.0"]
extras["pipeline-train"] = [
"accelerate>=0.26.0",
@@ -93,14 +93,13 @@ extras["pipeline"] = (
extras["scoring"] = ["sqlalchemy>=2.0.20"]
extras["vectors"] = [
"fasttext-wheel>=0.9.2",
"litellm>=1.37.16",
"llama-cpp-python>=0.2.75",
"model2vec>=0.3.0",
"pymagnitude-lite>=0.1.43",
"scikit-learn>=0.23.1",
"sentence-transformers>=2.2.0",
"skops>=0.9.0",
"staticvectors>=0.2.0",
]
extras["workflow"] = [
src/python/txtai/embeddings/base.py
-7
View File
@@ -4,7 +4,6 @@ Embeddings module
import json
import os
import shutil
import tempfile
import numpy as np
@@ -618,12 +617,6 @@ class Embeddings:
# Create output directory, if necessary
os.makedirs(path, exist_ok=True)
# Copy vectors model
if self.config.get("storevectors"):
shutil.copyfile(self.config["path"], os.path.join(path, os.path.basename(self.config["path"])))
self.config["path"] = os.path.basename(self.config["path"])
# Save index configuration
Configuration().save(self.config, path)
src/python/txtai/embeddings/index/configuration.py
-4
View File
@@ -41,10 +41,6 @@ class Configuration:
# Add format parameter
config["format"] = "json" if jsonconfig else "pickle"
# Build full path to embedding vectors file
if config.get("storevectors"):
config["path"] = os.path.join(path, config["path"])
return config
def save(self, config, path):
src/python/txtai/pipeline/text/translation.py
+9 -35
View File
@@ -2,17 +2,14 @@
Translation module
"""
import os
# Conditional import
try:
import fasttext
from staticvectors import StaticVectors
FASTTEXT = True
STATICVECTORS = True
except ImportError:
FASTTEXT = False
STATICVECTORS = False
from huggingface_hub import hf_hub_download
from huggingface_hub.hf_api import HfApi
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
@@ -26,7 +23,7 @@ class Translation(HFModel):
"""
# Default language detection model
DEFAULT_LANG_DETECT = "julien-c/fasttext-language-id/lid.176.ftz"
DEFAULT_LANG_DETECT = "neuml/language-id-quantized"
def __init__(self, path=None, quantize=False, gpu=True, batch=64, langdetect=None, findmodels=True):
"""
@@ -137,7 +134,7 @@ class Translation(HFModel):
def defaultdetect(self, texts):
"""
Default fasttext language detection model.
Default language detection model.
Args:
texts: list of text
@@ -147,43 +144,20 @@ class Translation(HFModel):
"""
if not self.detector:
if not FASTTEXT:
if not STATICVECTORS:
raise ImportError('Language detection is not available - install "pipeline" extra to enable')
# Suppress unnecessary warning
fasttext.FastText.eprint = lambda x: None
# Get model path
path = self.langdetect if self.langdetect else Translation.DEFAULT_LANG_DETECT
# Load language detection model
path = path if os.path.exists(path) else self.download(path)
self.detector = fasttext.load_model(path)
self.detector = StaticVectors(path)
# Transform texts to format expected by language detection model
texts = [x.lower().replace("\n", " ").replace("\r\n", " ") for x in texts]
return [x[0].split("__")[-1] for x in self.detector.predict(texts)[0]]
def download(self, path):
"""
Downloads path from the Hugging Face Hub.
Args:
path: full model path
Returns:
local cached model path
"""
# Split into parts
parts = path.split("/")
# Calculate repo id split
repo = 2 if len(parts) > 2 else 1
# Download and cache file
return hf_hub_download(repo_id="/".join(parts[:repo]), filename="/".join(parts[repo:]))
# Detect languages
return [x[0][0] for x in self.detector.predict(texts)]
def translate(self, texts, source, target, showmodels=False):
"""
src/python/txtai/vectors/factory.py
+1 -1
View File
@@ -91,7 +91,7 @@ class VectorsFactory:
method = "llama.cpp"
elif Model2Vec.ismodel(path):
method = "model2vec"
elif WordVectors.isdatabase(path):
elif WordVectors.ismodel(path):
method = "words"
else:
method = "transformers"
src/python/txtai/vectors/words.py
+56 -79
View File
@@ -2,23 +2,24 @@
Word Vectors module
"""
import json
import logging
import os
import tempfile
from errno import ENOENT
from multiprocessing import Pool
import numpy as np
# Conditionally import Word Vector libraries as they aren't installed by default
try:
import fasttext
from pymagnitude import converter, Magnitude
from transformers.utils import cached_file
WORDS = True
# Conditional import
try:
from staticvectors import Database, StaticVectors
STATICVECTORS = True
except ImportError:
WORDS = False
STATICVECTORS = False
from ..pipeline import Tokenizer
@@ -66,24 +67,59 @@ class WordVectors(Vectors):
Builds vectors using weighted word embeddings.
"""
@staticmethod
def ismodel(path):
"""
Checks if path is a WordVectors model.
Args:
path: input path
Returns:
True if this is a WordVectors model, False otherwise
"""
# Check if this is a SQLite database
if WordVectors.isdatabase(path):
return True
try:
# Download file and parse JSON
path = cached_file(path_or_repo_id=path, filename="config.json")
if path:
with open(path, encoding="utf-8") as f:
config = json.load(f)
return config.get("model_type") == "staticvectors"
# Ignore this error - invalid repo or directory
except OSError:
pass
return False
@staticmethod
def isdatabase(path):
"""
Checks if this is a SQLite database file which is the file format used for word vectors databases.
Args:
path: path to check
Returns:
True if this is a SQLite database
"""
return isinstance(path, str) and Database.isdatabase(path)
def __init__(self, config, scoring, models):
# Check before parent constructor since it calls loadmodel
if not WORDS:
# Raise error if trying to create Word Vectors without vectors extra
raise ImportError(
'Word vector models are not available - install "vectors" extra to enable. Otherwise, specify '
+ 'method="transformers" to use transformer backed models'
)
if not STATICVECTORS:
raise ImportError('staticvectors is not available - install "vectors" extra to enable')
super().__init__(config, scoring, models)
def loadmodel(self, path):
# Ensure that vector path exists
if not path or not os.path.isfile(path):
raise IOError(ENOENT, "Vector model file not found", path)
# Load magnitude model. If this is a training run (uninitialized config), block until vectors are fully loaded
return Magnitude(path, case_insensitive=True, blocking=not self.initialized)
return StaticVectors(path)
def encode(self, data):
# Iterate over each data element, tokenize (if necessary) and build an aggregated embeddings vector
@@ -156,67 +192,8 @@ class WordVectors(Vectors):
word vectors array
"""
return self.model.query(tokens)
return self.model.embeddings(tokens)
def tokens(self, data):
# Skip tokenization rules
return data
@staticmethod
def isdatabase(path):
"""
Checks if this is a SQLite database file which is the file format used for word vectors databases.
Args:
path: path to check
Returns:
True if this is a SQLite database
"""
if isinstance(path, str) and os.path.isfile(path) and os.path.getsize(path) >= 100:
# Read 100 byte SQLite header
with open(path, "rb") as f:
header = f.read(100)
# Check for SQLite header
return header.startswith(b"SQLite format 3\000")
return False
@staticmethod
def build(data, size, mincount, path):
"""
Builds fastText vectors from a file.
Args:
data: path to input data file
size: number of vector dimensions
mincount: minimum number of occurrences required to register a token
path: path to output file
"""
# Train on data file using largest dimension size
model = fasttext.train_unsupervised(data, dim=size, minCount=mincount)
# Output file path
logger.info("Building %d dimension model", size)
# Output vectors in vec/txt format
with open(path + ".txt", "w", encoding="utf-8") as output:
words = model.get_words()
output.write(f"{len(words)} {model.get_dimension()}\n")
for word in words:
# Skip end of line token
if word != "</s>":
vector = model.get_word_vector(word)
data = ""
for v in vector:
data += " " + str(v)
output.write(word + data + "\n")
# Build magnitude vectors database
logger.info("Converting vectors to magnitude format")
converter.convert(path + ".txt", path + ".magnitude", subword=True)
test/python/testembeddings.py
+2 -25
View File
@@ -13,7 +13,6 @@ import numpy as np
from txtai.embeddings import Embeddings, Reducer
from txtai.serialize import SerializeFactory
from txtai.vectors import WordVectors
# pylint: disable=R0904
@@ -541,27 +540,11 @@ class TestEmbeddings(unittest.TestCase):
# Mock CPU count
cpucount.return_value = 1
# Initialize model path
path = os.path.join(tempfile.gettempdir(), "model")
os.makedirs(path, exist_ok=True)
# Build tokens file
with tempfile.NamedTemporaryFile(mode="w", delete=False) as output:
tokens = output.name
for x in self.data:
output.write(x + "\n")
# Word vectors path
vectors = os.path.join(path, "test-10d")
# Build word vectors, if they don't already exist
WordVectors.build(tokens, 10, 1, vectors)
# Create dataset
data = [(x, row.split(), None) for x, row in enumerate(self.data)]
# Create embeddings model, backed by word vectors
embeddings = Embeddings({"path": vectors + ".magnitude", "storevectors": True, "scoring": "bm25", "pca": 3, "quantize": True})
embeddings = Embeddings({"path": "neuml/glove-6B-quantized", "scoring": "bm25", "pca": 3, "quantize": True})
# Call scoring and index methods
embeddings.score(data)
@@ -589,17 +572,11 @@ class TestEmbeddings(unittest.TestCase):
# Mock CPU count
cpucount.return_value = 1
# Initialize model path
path = os.path.join(tempfile.gettempdir(), "model")
# Word vectors path
vectors = os.path.join(path, "test-10d")
# Create dataset
data = [(x, row.split(), None) for x, row in enumerate(self.data)]
# Create embeddings model, backed by word vectors
embeddings = Embeddings({"path": vectors + ".magnitude", "storevectors": True, "scoring": "bm25", "pca": 3})
embeddings = Embeddings({"path": "neuml/glove-6B/model.sqlite", "scoring": "bm25", "pca": 3})
# Call scoring and index methods
embeddings.score(data)
test/python/testoptional.py
+1 -1
View File
@@ -27,7 +27,6 @@ class TestOptional(unittest.TestCase):
"docling.document_converter",
"duckdb",
"fastapi",
"fasttext",
"grand-cypher",
"grand-graph",
"hnswlib",
@@ -53,6 +52,7 @@ class TestOptional(unittest.TestCase):
"soundfile",
"sqlalchemy",
"sqlite_vec",
"staticvectors",
"tika",
"ttstokenizer",
"xmltodict",
test/python/testvectors/testwordvectors.py
+16 -37
View File
@@ -10,7 +10,7 @@ from unittest.mock import patch
import numpy as np
from txtai.vectors import WordVectors, VectorsFactory
from txtai.vectors import VectorsFactory
class TestWordVectors(unittest.TestCase):
@@ -21,32 +21,11 @@ class TestWordVectors(unittest.TestCase):
@classmethod
def setUpClass(cls):
"""
Test a WordVectors build.
Sets the pretrained model to use
"""
# Word vectors path
path = os.path.join(tempfile.gettempdir(), "vectors")
# Save model path
cls.path = path + ".magnitude"
# Build word vectors
WordVectors.build("README.md", 10, 3, path)
def testBlocking(self):
"""
Test blocking load of vector model
"""
config = {"path": self.path}
model = VectorsFactory.create(config, None)
self.assertFalse(model.initialized)
config["dimensions"] = 10
model = VectorsFactory.create(config, None)
self.assertTrue(model.initialized)
# Test with pretrained glove quantized vectors
cls.path = "neuml/glove-6B-quantized"
@patch("os.cpu_count")
def testIndex(self, cpucount):
@@ -65,13 +44,13 @@ class TestWordVectors(unittest.TestCase):
ids, dimension, batches, stream = model.index(documents)
self.assertEqual(len(ids), 1000)
self.assertEqual(dimension, 10)
self.assertEqual(dimension, 300)
self.assertEqual(batches, 1000)
self.assertIsNotNone(os.path.exists(stream))
# Test shape of serialized embeddings
with open(stream, "rb") as queue:
self.assertEqual(np.load(queue).shape, (1, 10))
self.assertEqual(np.load(queue).shape, (1, 300))
@patch("os.cpu_count")
def testIndexBatch(self, cpucount):
@@ -90,14 +69,14 @@ class TestWordVectors(unittest.TestCase):
ids, dimension, batches, stream = model.index(documents, 512)
self.assertEqual(len(ids), 1000)
self.assertEqual(dimension, 10)
self.assertEqual(dimension, 300)
self.assertEqual(batches, 2)
self.assertIsNotNone(os.path.exists(stream))
# Test shape of serialized embeddings
with open(stream, "rb") as queue:
self.assertEqual(np.load(queue).shape, (512, 10))
self.assertEqual(np.load(queue).shape, (488, 10))
self.assertEqual(np.load(queue).shape, (512, 300))
self.assertEqual(np.load(queue).shape, (488, 300))
def testIndexSerial(self):
"""
@@ -112,13 +91,13 @@ class TestWordVectors(unittest.TestCase):
ids, dimension, batches, stream = model.index(documents)
self.assertEqual(len(ids), 1000)
self.assertEqual(dimension, 10)
self.assertEqual(dimension, 300)
self.assertEqual(batches, 1000)
self.assertIsNotNone(os.path.exists(stream))
# Test shape of serialized embeddings
with open(stream, "rb") as queue:
self.assertEqual(np.load(queue).shape, (1, 10))
self.assertEqual(np.load(queue).shape, (1, 300))
def testIndexSerialBatch(self):
"""
@@ -133,14 +112,14 @@ class TestWordVectors(unittest.TestCase):
ids, dimension, batches, stream = model.index(documents, 512)
self.assertEqual(len(ids), 1000)
self.assertEqual(dimension, 10)
self.assertEqual(dimension, 300)
self.assertEqual(batches, 2)
self.assertIsNotNone(os.path.exists(stream))
# Test shape of serialized embeddings
with open(stream, "rb") as queue:
self.assertEqual(np.load(queue).shape, (512, 10))
self.assertEqual(np.load(queue).shape, (488, 10))
self.assertEqual(np.load(queue).shape, (512, 300))
self.assertEqual(np.load(queue).shape, (488, 300))
def testLookup(self):
"""
@@ -148,7 +127,7 @@ class TestWordVectors(unittest.TestCase):
"""
model = VectorsFactory.create({"path": self.path}, None)
self.assertEqual(model.lookup(["txtai", "embeddings", "sentence"]).shape, (3, 10))
self.assertEqual(model.lookup(["txtai", "embeddings", "sentence"]).shape, (3, 300))
def testNoExist(self):
"""
@@ -165,4 +144,4 @@ class TestWordVectors(unittest.TestCase):
"""
model = VectorsFactory.create({"path": self.path}, None)
self.assertEqual(len(model.transform((None, ["txtai"], None))), 10)
self.assertEqual(len(model.transform((None, ["txtai"], None))), 300)