diff --git a/docs/embeddings/configuration/vectors.md b/docs/embeddings/configuration/vectors.md index d71ba055..519322b6 100644 --- a/docs/embeddings/configuration/vectors.md +++ b/docs/embeddings/configuration/vectors.md @@ -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 diff --git a/examples/01_Introducing_txtai.ipynb b/examples/01_Introducing_txtai.ipynb index c7e13177..45b80b21 100644 --- a/examples/01_Introducing_txtai.ipynb +++ b/examples/01_Introducing_txtai.ipynb @@ -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": [] diff --git a/examples/03_Build_an_Embeddings_index_from_a_data_source.ipynb b/examples/03_Build_an_Embeddings_index_from_a_data_source.ipynb index 393ebe43..155fc99a 100644 --- a/examples/03_Build_an_Embeddings_index_from_a_data_source.ipynb +++ b/examples/03_Build_an_Embeddings_index_from_a_data_source.ipynb @@ -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|(?\n", + " \n", + " \n", + " Title\n", + " Published\n", + " Reference\n", + " Match\n", + " \n", + " \n", + " \n", + " \n", + " Management of osteoarthritis during COVID‐19 pandemic\n", + " 2020-05-21 00:00:00\n", + " https://doi.org/10.1002/cpt.1910\n", + " Indeed, risk factors are sex, obesity, genetic factors and mechanical factors (3) .\n", + " \n", + " \n", + " Does apolipoprotein E genotype predict COVID-19 severity?\n", + " 2020-04-27 00:00:00\n", + " https://doi.org/10.1093/qjmed/hcaa142\n", + " Risk factors associated with subsequent death include older age, hypertension, diabetes, ischemic heart disease, obesity and chronic lung disease; however, sometimes there are no obvious risk factors .\n", + " \n", + " \n", + " Prevalence and Impact of Myocardial Injury in Patients Hospitalized with COVID-19 Infection\n", + " 2020-04-24 00:00:00\n", + " http://medrxiv.org/cgi/content/short/2020.04.20.20072702v1?rss=1\n", + " This risk was consistent across patients stratified by history of CVD, risk factors but no CVD, and neither CVD nor risk factors.\n", + " \n", + " \n", + " COVID-19 and associations with frailty and multimorbidity: a prospective analysis of UK Biobank participants\n", + " 2020-07-23 00:00:00\n", + " https://www.ncbi.nlm.nih.gov/pubmed/32705587/\n", + " BACKGROUND: Frailty and multimorbidity have been suggested as risk factors for severe COVID-19 disease.\n", + " \n", + " \n", + " Risk Stratification for Healthcare workers during the CoViD-19 Pandemic; using demographics, co-morbid disease and clinical domain in order to assign clinical duties\n", + " 2020-05-09 00:00:00\n", + " http://medrxiv.org/cgi/content/short/2020.05.05.20091967v1?rss=1\n", + " Vascular disease, diabetes and chronic pulmonary disease further increased risk.\n", + " \n", + " \n", + "" + ], + "text/plain": [ + "" + ] + }, + "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": [ - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
TitlePublishedReferenceMatch
Management of osteoarthritis during COVID‐19 pandemic2020-05-21 00:00:00https://doi.org/10.1002/cpt.1910Indeed, risk factors are sex, obesity, genetic factors and mechanical factors (3) .
Work-related and Personal Factors Associated with Mental Well-being during COVID-19 Response: A Survey of Health Care and Other Workers2020-06-11 00:00:00http://medrxiv.org/cgi/content/short/2020.06.09.20126722v1?rss=1Poor 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)].
No evidence that androgen regulation of pulmonary TMPRSS2 explains sex-discordant COVID-19 outcomes2020-04-21 00:00:00https://doi.org/10.1101/2020.04.21.051201In addition to male sex, smoking is a risk factor for COVID-19 susceptibility and poor clinical outcomes .
Current status of potential therapeutic candidates for the COVID-19 crisis2020-04-22 00:00:00https://doi.org/10.1016/j.bbi.2020.04.046There was no difference on 28-day mortality between heparin users and nonusers.
COVID-19: what has been learned and to be learned about the novel coronavirus disease2020-03-15 00:00:00https://doi.org/10.7150/ijbs.45134• Three major risk factors for COVID-19 were sex (male), age (≥60), and severe pneumonia.
" - ], - "text/plain": [ - "" - ] - }, - "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": [ + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
TitlePublishedReferenceMatchRisk FactorsLocations
Management of osteoarthritis during COVID‐19 pandemic2020-05-21 00:00:00https://doi.org/10.1002/cpt.1910Indeed, risk factors are sex, obesity, genetic factors and mechanical factors (3) .sex, obesity, genetic factors and mechanical factorshospitals and clinics
Does apolipoprotein E genotype predict COVID-19 severity?2020-04-27 00:00:00https://doi.org/10.1093/qjmed/hcaa142Risk 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 .NoneNone
Prevalence and Impact of Myocardial Injury in Patients Hospitalized with COVID-19 Infection2020-04-24 00:00:00http://medrxiv.org/cgi/content/short/2020.04.20.20072702v1?rss=1This risk was consistent across patients stratified by history of CVD, risk factors but no CVD, and neither CVD nor risk factors.neither CVD nor risk factorsMount Sinai Health System
COVID-19 and associations with frailty and multimorbidity: a prospective analysis of UK Biobank participants2020-07-23 00:00:00https://www.ncbi.nlm.nih.gov/pubmed/32705587/BACKGROUND: Frailty and multimorbidity have been suggested as risk factors for severe COVID-19 disease.Frailty and multimorbidity213 countries and territories
Risk Stratification for Healthcare workers during the CoViD-19 Pandemic; using demographics, co-morbid disease and clinical domain in order to assign clinical duties2020-05-09 00:00:00http://medrxiv.org/cgi/content/short/2020.05.05.20091967v1?rss=1Vascular disease, diabetes and chronic pulmonary disease further increased risk.Vascular disease, diabetes and chronic pulmonary diseaseNone
" + ], + "text/plain": [ + "" + ] + }, + "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": [ - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
TitlePublishedReferenceMatchRisk FactorsLocations
Management of osteoarthritis during COVID‐19 pandemic2020-05-21 00:00:00https://doi.org/10.1002/cpt.1910Indeed, risk factors are sex, obesity, genetic factors and mechanical factors (3) .sex, obesity, genetic factors and mechanical factorsNone
Work-related and Personal Factors Associated with Mental Well-being during COVID-19 Response: A Survey of Health Care and Other Workers2020-06-11 00:00:00http://medrxiv.org/cgi/content/short/2020.06.09.20126722v1?rss=1Poor 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)].Poor family supportive behaviorsNone
No evidence that androgen regulation of pulmonary TMPRSS2 explains sex-discordant COVID-19 outcomes2020-04-21 00:00:00https://doi.org/10.1101/2020.04.21.051201In addition to male sex, smoking is a risk factor for COVID-19 susceptibility and poor clinical outcomes .Higher morbidity and mortalityNone
Current status of potential therapeutic candidates for the COVID-19 crisis2020-04-22 00:00:00https://doi.org/10.1016/j.bbi.2020.04.046There was no difference on 28-day mortality between heparin users and nonusers.elicited strong inflammatory responses are favorable or detrimentalNone
COVID-19: what has been learned and to be learned about the novel coronavirus disease2020-03-15 00:00:00https://doi.org/10.7150/ijbs.45134• Three major risk factors for COVID-19 were sex (male), age (≥60), and severe pneumonia.sex (male), age (≥60), and severe pneumoniaNone
" - ], - "text/plain": [ - "" - ] - }, - "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." ] } - ] -} \ No newline at end of file + ], + "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 +} diff --git a/examples/24_Whats_new_in_txtai_4_0.ipynb b/examples/24_Whats_new_in_txtai_4_0.ipynb index ea2b4d43..c9e39bc7 100644 --- a/examples/24_Whats_new_in_txtai_4_0.ipynb +++ b/examples/24_Whats_new_in_txtai_4_0.ipynb @@ -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", diff --git a/setup.py b/setup.py index 5eb1f518..34157424 100644 --- a/setup.py +++ b/setup.py @@ -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"] = [ diff --git a/src/python/txtai/embeddings/base.py b/src/python/txtai/embeddings/base.py index 268cd8e0..3df5d0da 100644 --- a/src/python/txtai/embeddings/base.py +++ b/src/python/txtai/embeddings/base.py @@ -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) diff --git a/src/python/txtai/embeddings/index/configuration.py b/src/python/txtai/embeddings/index/configuration.py index 500dff07..bf8f317d 100644 --- a/src/python/txtai/embeddings/index/configuration.py +++ b/src/python/txtai/embeddings/index/configuration.py @@ -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): diff --git a/src/python/txtai/pipeline/text/translation.py b/src/python/txtai/pipeline/text/translation.py index d15e0a58..41060fbd 100644 --- a/src/python/txtai/pipeline/text/translation.py +++ b/src/python/txtai/pipeline/text/translation.py @@ -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): """ diff --git a/src/python/txtai/vectors/factory.py b/src/python/txtai/vectors/factory.py index 3d291d8d..5ef28fe3 100644 --- a/src/python/txtai/vectors/factory.py +++ b/src/python/txtai/vectors/factory.py @@ -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" diff --git a/src/python/txtai/vectors/words.py b/src/python/txtai/vectors/words.py index a0be2eaf..cc14b970 100644 --- a/src/python/txtai/vectors/words.py +++ b/src/python/txtai/vectors/words.py @@ -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 != "": - 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) diff --git a/test/python/testembeddings.py b/test/python/testembeddings.py index 97d92c65..86b07a9c 100644 --- a/test/python/testembeddings.py +++ b/test/python/testembeddings.py @@ -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) diff --git a/test/python/testoptional.py b/test/python/testoptional.py index ce48252f..74ebfc2d 100644 --- a/test/python/testoptional.py +++ b/test/python/testoptional.py @@ -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", diff --git a/test/python/testvectors/testwordvectors.py b/test/python/testvectors/testwordvectors.py index 9130ad91..74d82c01 100644 --- a/test/python/testvectors/testwordvectors.py +++ b/test/python/testvectors/testwordvectors.py @@ -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)