Add query translation logic in embeddings searches, closes #266

src/python/txtai/embeddings/__init__.py

@@ -7,4 +7,7 @@ from .base import Embeddings
 from .cloud import Cloud
 from .documents import Documents
 from .explain import Explain
+from .query import Query
 from .reducer import Reducer
+from .search import Search
+from .transform import Transform

src/python/txtai/embeddings/base.py

@@ -18,6 +18,7 @@ from ..vectors import VectorsFactory
 from .archive import Archive
 from .explain import Explain
 from .reducer import Reducer
+from .query import Query
 from .search import Search
 from .transform import Action, Transform
 
@@ -55,6 +56,9 @@ class Embeddings:
         # Document database
         self.database = None
 
+        # Query model
+        self.query = None
+
         # Index archive
         self.archive = None
 
@@ -429,6 +433,9 @@ class Embeddings:
         # Sentence vectors model - transforms data to embeddings vectors
         self.model = self.loadvectors()
 
+        # Query model
+        self.query = self.loadquery()
+
         # Document database - stores document content
         self.database = self.createdatabase()
         if self.database:
@@ -484,7 +491,7 @@ class Embeddings:
         Closes this embeddings index and frees all resources.
         """
 
-        self.config, self.reducer, self.scoring, self.model, self.ann, self.archive = None, None, None, None, None, None
+        self.config, self.reducer, self.scoring, self.model, self.ann, self.query, self.archive = None, None, None, None, None, None, None
 
         # Close database connection if open
         if self.database:
@@ -528,6 +535,9 @@ class Embeddings:
         # Sentence vectors model - transforms data to embeddings vectors
         self.model = self.loadvectors() if self.config else None
 
+        # Query model
+        self.query = self.loadquery() if self.config else None
+
     def loadvectors(self):
         """
         Loads a vector model set in config.
@@ -538,6 +548,19 @@ class Embeddings:
 
         return VectorsFactory.create(self.config, self.scoring)
 
+    def loadquery(self):
+        """
+        Loads a query model set in config.
+
+        Returns:
+            query model
+        """
+
+        if "query" in self.config:
+            return Query(**self.config["query"])
+
+        return None
+
     def checkarchive(self, path):
         """
         Checks if path is an archive file.

src/python/txtai/embeddings/query.py

@@ -0,0 +1,69 @@
+"""
+Query module
+"""
+
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, T5ForConditionalGeneration
+
+
+class Query:
+    """
+    Query translation model.
+    """
+
+    def __init__(self, path, prefix=None, maxlength=512):
+        """
+        Creates a query translation model.
+
+        Args:
+            path: path to query model
+            prefix: text prefix
+            maxlength: max sequence length to generate
+        """
+
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(path)
+
+        # Default prefix if not provided for T5 models
+        if not prefix and isinstance(self.model, T5ForConditionalGeneration):
+            prefix = "translate English to SQL: "
+
+        self.prefix = prefix
+        self.maxlength = maxlength
+
+    def __call__(self, query):
+        """
+        Runs query translation model.
+
+        Args:
+            query: input query
+
+        Returns:
+            transformed query
+        """
+
+        # Add prefix, if necessary
+        if self.prefix:
+            query = f"{self.prefix}{query}"
+
+        # Tokenize and generate text using model
+        features = self.tokenizer([query], return_tensors="pt")
+        output = self.model.generate(input_ids=features["input_ids"], attention_mask=features["attention_mask"], max_length=self.maxlength)
+
+        # Decode tokens to text
+        result = self.tokenizer.decode(output[0], skip_special_tokens=True)
+
+        # Clean and return generated text
+        return self.clean(result)
+
+    def clean(self, text):
+        """
+        Applies a series of rules to clean generated text.
+
+        Args:
+            text: input text
+
+        Returns:
+            clean text
+        """
+
+        return text.replace("$=", "<=")

src/python/txtai/embeddings/search.py

@@ -2,8 +2,13 @@
 Search module
 """
 
+import logging
+
 import numpy as np
 
+# Logging configuration
+logger = logging.getLogger(__name__)
+
 
 class Search:
     """
@@ -25,6 +30,7 @@ class Search:
         self.ann = embeddings.ann
         self.database = embeddings.database
         self.transform = embeddings.transform
+        self.query = embeddings.query
 
     def __call__(self, queries, limit):
         """
@@ -90,7 +96,7 @@ class Search:
         """
 
         # Parse queries
-        queries = [self.database.parse(query) for query in queries]
+        queries = self.parse(queries)
 
         # Override limit with query limit, if applicable
         limit = max(limit, self.limit(queries))
@@ -125,6 +131,37 @@ class Search:
 
         return results
 
+    def parse(self, queries):
+        """
+        Parses a list of database queries.
+
+        Args:
+            queries: list of queries
+
+        Returns:
+            parsed queries
+        """
+
+        # Parsed queries
+        parsed = []
+
+        for query in queries:
+            # Parse query
+            parse = self.database.parse(query)
+
+            # Transform query if SQL not parsed and reparse
+            if self.query and "select" not in parse:
+                # Generate query
+                query = self.query(query)
+                logging.debug(query)
+
+                # Reparse query
+                parse = self.database.parse(query)
+
+            parsed.append(parse)
+
+        return parsed
+
     def limit(self, queries):
         """
         Parses the largest LIMIT clause from queries.

test/python/testdatabase/testembeddings.py

@@ -330,6 +330,21 @@ class TestEmbeddings(unittest.TestCase):
         self.assertRaises(NotImplementedError, database.embed, None, None)
         self.assertRaises(NotImplementedError, database.query, None, None)
 
+    def testQueryModel(self):
+        """
+        Test index
+        """
+
+        embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": True, "query": {"path": "NeuML/t5-small-txtsql"}})
+
+        # Create an index for the list of text
+        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
+
+        # Search for best match
+        result = embeddings.search("feel good story with lottery in text", 1)[0]
+
+        self.assertEqual(result["text"], self.data[4])
+
     def testReindex(self):
         """
         Test reindex