Add semantic search explainability, closes #248

src/python/txtai/api/routers/embeddings.py

@@ -120,6 +120,38 @@ def count():
     return application.get().count()
 
 
+@router.post("/explain")
+def explain(query: str = Body(...), texts: List[str] = Body(default=None), limit: int = Body(default=None)):
+    """
+    Explains the importance of each input token in text for a query.
+
+    Args:
+        query: query text
+        texts: list of text
+
+    Returns:
+        list of dict where a higher scores represents higher importance relative to the query
+    """
+
+    return application.get().explain(query, texts, limit)
+
+
+@router.post("/batchexplain")
+def batchexplain(queries: List[str] = Body(...), texts: List[str] = Body(default=None), limit: int = Body(default=None)):
+    """
+    Explains the importance of each input token in text for a query.
+
+    Args:
+        query: query text
+        texts: list of text
+
+    Returns:
+        list of dict where a higher scores represents higher importance relative to the query
+    """
+
+    return application.get().batchexplain(queries, texts, limit)
+
+
 @router.get("/transform")
 def transform(text: str):
     """

src/python/txtai/app/base.py

@@ -462,6 +462,42 @@ class Application:
 
         return None
 
+    def explain(self, query, texts=None, limit=10):
+        """
+        Explains the importance of each input token in text for a query.
+
+        Args:
+            query: query text
+            texts: optional list of text, otherwise runs search query
+            limit: optional limit if texts is None
+
+        Returns:
+            list of dict per input text where a higher token scores represents higher importance relative to the query
+        """
+
+        if self.embeddings:
+            return self.embeddings.explain(query, texts, limit)
+
+        return None
+
+    def batchexplain(self, queries, texts=None, limit=10):
+        """
+        Explains the importance of each input token in text for a list of queries.
+
+        Args:
+            query: queries text
+            texts: optional list of text, otherwise runs search queries
+            limit: optional limit if texts is None
+
+        Returns:
+            list of dict per input text per query where a higher token scores represents higher importance relative to the query
+        """
+
+        if self.embeddings:
+            return self.embeddings.batchexplain(queries, texts, limit)
+
+        return None
+
     def transform(self, text):
         """
         Transforms text into embeddings arrays.

src/python/txtai/embeddings/__init__.py

@@ -6,4 +6,5 @@ from .archive import Archive
 from .base import Embeddings
 from .cloud import Cloud
 from .documents import Documents
+from .explain import Explain
 from .reducer import Reducer

src/python/txtai/embeddings/base.py

@@ -16,6 +16,7 @@ from ..scoring import ScoringFactory
 from ..vectors import VectorsFactory
 
 from .archive import Archive
+from .explain import Explain
 from .reducer import Reducer
 from .search import Search
 from .transform import Action, Transform
@@ -339,6 +340,36 @@ class Embeddings:
         # Add index and sort desc based on score
         return [sorted(enumerate(score), key=lambda x: x[1], reverse=True) for score in scores]
 
+    def explain(self, query, texts=None, limit=None):
+        """
+        Explains the importance of each input token in text for a query.
+
+        Args:
+            query: input query
+            texts: optional list of (text|list of tokens), otherwise runs search query
+            limit: optional limit if texts is None
+
+        Returns:
+            list of dict per input text where a higher token scores represents higher importance relative to the query
+        """
+
+        return Explain(self)([query], texts, limit)[0]
+
+    def batchexplain(self, queries, texts=None, limit=None):
+        """
+        Explains the importance of each input token in text for a list of queries.
+
+        Args:
+            query: input queries
+            texts: optional list of (text|list of tokens), otherwise runs search queries
+            limit: optional limit if texts is None
+
+        Returns:
+            list of dict per input text per query where a higher token scores represents higher importance relative to the query
+        """
+
+        return Explain(self)(queries, texts, limit)
+
     def exists(self, path, cloud=None):
         """
         Checks if an index exists at path.

src/python/txtai/embeddings/explain.py

@@ -0,0 +1,104 @@
+"""
+Explain module
+"""
+
+import numpy as np
+
+
+class Explain:
+    """
+    Explains the importance of each token in an input text element for a query. This method creates n permutations of the input text, where n
+    is the number of tokens in the input text. This effectively masks each token to determine its importance to the query.
+    """
+
+    def __init__(self, embeddings):
+        """
+        Creates a new explain action.
+
+        Args:
+            embeddings: embeddings instance
+        """
+
+        self.embeddings = embeddings
+        self.content = self.embeddings.config.get("content")
+
+    def __call__(self, queries, texts, limit):
+        """
+        Explains the importance of each input token in text for a list of queries.
+
+        Args:
+            query: input queries
+            texts: optional list of (text|list of tokens), otherwise runs search queries
+            limit: optional limit if texts is None
+
+        Returns:
+            list of dict per input text per query where a higher token scores represents higher importance relative to the query
+        """
+
+        # Construct texts elements per query
+        texts = self.texts(queries, texts, limit)
+
+        # Explain each query-texts combination
+        return [self.explain(query, texts[x]) for x, query in enumerate(queries)]
+
+    def texts(self, queries, texts, limit):
+        """
+        Constructs lists of dict for each input query.
+
+        Args:
+            queries: input queries
+            texts: optional list of texts
+            limit: optional limit if texts is None
+
+        Returns:
+            lists of dict for each input query
+        """
+
+        # Calculate similarity scores per query if texts present
+        if texts:
+            results = []
+            for scores in self.embeddings.batchsimilarity(queries, texts):
+                results.append([{"id": uid, "text": texts[uid], "score": score} for uid, score in scores])
+
+            return results
+
+        # Query for results if texts is None and content is enabled
+        return self.embeddings.batchsearch(queries, limit) if self.content else [[]] * len(queries)
+
+    def explain(self, query, texts):
+        """
+        Explains the importance of each input token in text for a list of queries.
+
+        Args:
+            query: input query
+            texts: list of text
+
+        Returns:
+            list of {"id": value, "text": value, "score": value, "tokens": value} covering each input text element
+        """
+
+        # Explain results
+        results = []
+
+        # Calculate result per input text element
+        for x in texts:
+            text = x["text"]
+            tokens = text if isinstance(text, list) else text.split()
+
+            # Create permutations of input text, masking each token to determine importance
+            permutations = []
+            for i in range(len(tokens)):
+                data = tokens.copy()
+                data.pop(i)
+                permutations.append([" ".join(data)])
+
+            # Calculate similarity for each input text permutation and get score delta as importance
+            scores = [(i, x["score"] - np.abs(s)) for i, s in self.embeddings.similarity(query, permutations)]
+
+            # Add data sorted in index order
+            results.append(
+                {"id": x["id"], "text": text, "score": x["score"], "tokens": [(tokens[i], score) for i, score in sorted(scores, key=lambda x: x[0])]}
+            )
+
+        # Sort score descending and return
+        return sorted(results, key=lambda x: x["score"], reverse=True)

test/python/testapi/testembeddings.py

@@ -145,6 +145,8 @@ class TestEmbeddings(unittest.TestCase):
         self.assertIsNone(api.count())
         self.assertIsNone(api.similarity("test", ["test"]))
         self.assertIsNone(api.batchsimilarity(["test"], ["test"]))
+        self.assertIsNone(api.explain("test"))
+        self.assertIsNone(api.batchexplain(["test"]))
         self.assertIsNone(api.transform("test"))
         self.assertIsNone(api.batchtransform(["test"]))
         self.assertIsNone(api.extract(["test"], ["test"]))
@@ -303,6 +305,27 @@ class TestEmbeddings(unittest.TestCase):
         query = urllib.parse.quote("select length('text') length from txtai limit 1")
         self.assertEqual(self.client.get(f"search?query={query}").json()[0]["length"], 4)
 
+    def testXPlain(self):
+        """
+        Test API instance with explain methods
+        """
+
+        results = self.client.post("explain", json={"query": "feel good story", "limit": 1}).json()
+
+        self.assertEqual(results[0]["text"], self.data[4])
+        self.assertIsNotNone(results[0].get("tokens"))
+
+    def testXPlainBatch(self):
+        """
+        Test batch query explain via API
+        """
+
+        results = self.client.post("batchexplain", json={"queries": ["feel good story", "climate change"], "limit": 1}).json()
+
+        text = [result[0]["text"] for result in results]
+        self.assertEqual(text, [self.data[4], self.data[1]])
+        self.assertIsNotNone(results[0][0].get("tokens"))
+
 
 class Elements:
     """

test/python/testdatabase/testembeddings.py

@@ -15,6 +15,7 @@ from txtai.embeddings import Embeddings
 from txtai.database import Database, SQLError
 
 
+# pylint: disable=R0904
 class TestEmbeddings(unittest.TestCase):
     """
     Embeddings with a database tests.
@@ -182,6 +183,28 @@ class TestEmbeddings(unittest.TestCase):
         embeddings.upsert([])
         self.assertIsNotNone(embeddings.ann)
 
+    def testExplain(self):
+        """
+        Test query explain
+        """
+
+        # Test explain with similarity
+        result = self.embeddings.explain("feel good story", self.data)[0]
+        self.assertEqual(result["text"], self.data[4])
+        self.assertEqual(len(result.get("tokens")), 8)
+
+    def testExplainBatch(self):
+        """
+        Test query explain batch
+        """
+
+        # Test explain with query
+        self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
+
+        result = self.embeddings.batchexplain(["feel good story"], limit=1)[0][0]
+        self.assertEqual(result["text"], self.data[4])
+        self.assertEqual(len(result.get("tokens")), 8)
+
     def testGenerator(self):
         """
         Test index with a generator