Update of #54 to return json objects vs tuples to help simplify language binding libraries

src/python/txtai/api.py

@@ -87,8 +87,8 @@ class API(object):
     def search(self, query, request):
         """
         Finds documents in the embeddings model most similar to the input query. Returns
-        a list of (id, score) sorted by highest score, where id is the document id in
-        the embeddings model.
+        a list of {id: value, score: value} sorted by highest score, where id is the
+        document id in the embeddings model.
 
         Downstream applications can override this method to provide enriched search results.
 
@@ -97,30 +97,32 @@ class API(object):
             request: FastAPI request
 
         Returns:
-            list of (id, score)
+            list of {id: value, score: value}
         """
 
         if self.embeddings:
-            return self.embeddings.search(query, self.limit(request.query_params.get("limit")))
+            limit = self.limit(request.query_params.get("limit"))
+            return [{"id": uid, "score": float(score)} for uid, score in self.embeddings.search(query, limit)]
 
         return None
 
     def batchsearch(self, queries, limit):
         """
         Finds documents in the embeddings model most similar to the input queries. Returns
-        a list of (id, score) sorted by highest score per query, where id is the document id
-        in the embeddings model.
+        a list of {id: value, score: value} sorted by highest score per query, where id is
+        the document id in the embeddings model.
 
         Args:
             queries: queries text
             limit: maximum results
 
         Returns:
-            list of (id, score) per query
+            list of {id: value, score: value} per query
         """
 
         if self.embeddings:
-            return self.embeddings.batchsearch(queries, self.limit(limit))
+            results = self.embeddings.batchsearch(queries, self.limit(limit))
+            return [[{"id": uid, "score": float(score)} for uid, score in result] for result in results]
 
         return None
 
@@ -171,7 +173,7 @@ class API(object):
 
             # Build scoring index if scoring method provided
             if self.config.get("scoring"):
-                embeddings.score(self.stream())
+                self.embeddings.score(self.stream())
 
             # Build embeddings index
             self.embeddings.index(self.stream())
@@ -189,46 +191,64 @@ class API(object):
     def similarity(self, query, texts):
         """
         Computes the similarity between query and list of text. Returns a list of
-        (id, score) sorted by highest score, where id is the index in texts.
+        {id: value, score: value} sorted by highest score, where id is the index
+        in texts.
 
         Args:
             query: query text
             texts: list of text
 
         Returns:
-            list of (id, score)
+            list of {id: value, score: value}
         """
 
         # Use similarity instance if available otherwise fall back to embeddings model
         if self.similar:
-            return [(uid, float(text)) for uid, text in self.similar(query, texts)]
+            return [{"id": uid, "score": float(score)} for uid, score in self.similar(query, texts)]
         if self.embeddings:
-            return [(uid, float(text)) for uid, text in self.embeddings.similarity(query, texts)]
+            return [{"id": uid, "score": float(score)} for uid, score in self.embeddings.similarity(query, texts)]
 
         return None
 
     def batchsimilarity(self, queries, texts):
         """
         Computes the similarity between list of queries and list of text. Returns a list
-        of (id, score) sorted by highest score per query, where id is the index in texts.
+        of {id: value, score: value} sorted by highest score per query, where id is the
+        index in texts.
 
         Args:
             queries: queries text
             texts: list of text
 
         Returns:
-            list of (id, score) per query
+            list of {id: value, score: value} per query
         """
 
         # Use similarity instance if available otherwise fall back to embeddings model
         if self.similar:
-            return [[(uid, float(text)) for uid, text in r] for r in self.similar(queries, texts)]
+            return [[{"id": uid, "score": float(score)} for uid, score in r] for r in self.similar(queries, texts)]
         if self.embeddings:
-            return [[(uid, float(text)) for uid, text in r] for r in self.embeddings.batchsimilarity(queries, texts)]
+            return [[{"id": uid, "score": float(score)} for uid, score in r] for r in self.embeddings.batchsimilarity(queries, texts)]
 
         return None
 
-    def transform(self, texts):
+    def transform(self, text):
+        """
+        Transforms input text into embeddings arrays.
+
+        Args:
+            text: input text
+
+        Returns:
+            embeddings array
+        """
+
+        if self.embeddings:
+            return [float(x) for x in self.embeddings.transform((None, text, None))]
+
+        return None
+
+    def batchtransform(self, texts):
         """
         Transforms list of texts into embeddings arrays.
 
@@ -240,7 +260,8 @@ class API(object):
         """
 
         if self.embeddings:
-            return [[float(x) for x in self.embeddings.transform((None, text, None))] for text in texts]
+            documents = [(None, text, None) for text in texts]
+            return [[float(x) for x in result] for result in self.embeddings.batchtransform(documents)]
 
         return None
 
@@ -253,31 +274,36 @@ class API(object):
             texts: list of text
 
         Returns:
-            list of (name, answer)
+            list of {name: value, answer: value}
         """
 
         if self.extractor:
             # Convert queue to tuples
             queue = [(x["name"], x["query"], x.get("question"), x.get("snippet")) for x in queue]
-            return self.extractor(queue, texts)
+            return [{"name": name, "answer": answer} for name, answer in self.extractor(queue, texts)]
 
         return None
 
     def label(self, text, labels):
         """
         Applies a zero shot classifier to text using a list of labels. Returns a list of
-        (id, score) sorted by highest score, where id is the index in labels.
+        {id: value, score: value} sorted by highest score, where id is the index in labels.
 
         Args:
             text: text|list
             labels: list of labels
 
         Returns:
-            list of (id, score) per text element
+            list of {id: value, score: value} per text element
         """
 
         if self.labels:
-            return self.labels(text, labels)
+            # Text is a string
+            if isinstance(text, str):
+                return [{"id": uid, "score": float(score)} for uid, score in self.labels(text, labels)]
+
+            # Text is a list
+            return [[{"id": uid, "score": float(score)} for uid, score in result] for result in self.labels(text, labels)]
 
         return None
 
@@ -410,8 +436,8 @@ def batchsimilarity(queries: List[str] = Body(...), texts: List[str] = Body(...)
 
     return INSTANCE.batchsimilarity(queries, texts)
 
-@app.get("/embeddings")
-def embeddings(text: str):
+@app.get("/transform")
+def transform(text: str):
     """
     Transforms text into an embeddings array.
 
@@ -422,10 +448,10 @@ def embeddings(text: str):
         embeddings array
     """
 
-    return INSTANCE.transform([text])[0]
+    return INSTANCE.transform(text)
 
-@app.post("/batchembeddings")
-def batchembeddings(texts: List[str] = Body(...)):
+@app.post("/batchtransform")
+def batchtransform(texts: List[str] = Body(...)):
     """
     Transforms list of text into embeddings arrays.
 
@@ -436,7 +462,7 @@ def batchembeddings(texts: List[str] = Body(...)):
         embeddings arrays
     """
 
-    return INSTANCE.transform(texts)
+    return INSTANCE.batchtransform(texts)
 
 @app.post("/extract")
 def extract(queue: List[dict] = Body(...), texts: List[str] = Body(...)):

src/python/txtai/embeddings.py

@@ -188,6 +188,19 @@ class Embeddings(object):
 
         return embedding
 
+    def batchtransform(self, documents):
+        """
+        Transforms documents into embeddings vectors. Document text will be tokenized if not pre-tokenized.
+
+        Args:
+            documents: list of (id, text|tokens, tags)
+
+        Returns:
+            embeddings vectors
+        """
+
+        return [self.transform(document) for document in documents]
+
     def search(self, query, limit=3):
         """
         Finds documents in the embeddings model most similar to the input query. Returns

test/python/testapi.py

@@ -91,23 +91,6 @@ class TestAPI(unittest.TestCase):
         cls.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(cls.data)])
         cls.client.get("index")
 
-    def testEmbeddings(self):
-        """
-        Test embeddings transform via API
-        """
-
-        self.assertEqual(len(self.client.get("embeddings?text=testembed").json()), 768)
-
-    def testEmbeddingsBatch(self):
-        """
-        Test batch embeddings transform via API
-        """
-
-        embeddings = self.client.post("batchembeddings", json=self.data).json()
-
-        self.assertEqual(len(embeddings), len(self.data))
-        self.assertEqual(len(embeddings[0]), 768)
-
     def testEmpty(self):
         """
         Test empty API configuration
@@ -119,7 +102,8 @@ class TestAPI(unittest.TestCase):
         self.assertIsNone(api.batchsearch(["test"], None))
         self.assertIsNone(api.similarity("test", ["test"]))
         self.assertIsNone(api.batchsimilarity(["test"], ["test"]))
-        self.assertIsNone(api.transform(["test"]))
+        self.assertIsNone(api.transform("test"))
+        self.assertIsNone(api.batchtransform(["test"]))
         self.assertIsNone(api.extract(["test"], ["test"]))
         self.assertIsNone(api.label("test", ["test"]))
 
@@ -149,8 +133,8 @@ class TestAPI(unittest.TestCase):
         }).json()
 
         answers = execute("Red Sox - Blue Jays")
-        self.assertEqual("Blue Jays", answers[0][1])
-        self.assertEqual("2-1", answers[1][1])
+        self.assertEqual("Blue Jays", answers[0]["answer"])
+        self.assertEqual("2-1", answers[1]["answer"])
 
         # Ad-hoc questions
         question = "What hockey team won?"
@@ -159,7 +143,7 @@ class TestAPI(unittest.TestCase):
             "queue": [{"name": question, "query": question, "question": question, "snippet": False}],
             "texts": data
         }).json()
-        self.assertEqual("Flyers", answers[0][1])
+        self.assertEqual("Flyers", answers[0]["answer"])
 
     def testLabel(self):
         """
@@ -171,7 +155,7 @@ class TestAPI(unittest.TestCase):
             "labels": ["positive", "negative"]
         }).json()
 
-        self.assertEqual(labels[0][0], 0)
+        self.assertEqual(labels[0]["id"], 0)
 
     def testLabelBatch(self):
         """
@@ -183,7 +167,7 @@ class TestAPI(unittest.TestCase):
             "labels": ["positive", "negative"]
         }).json()
 
-        results = [l[0][0] for l in labels]
+        results = [l[0]["id"] for l in labels]
         self.assertEqual(results, [0, 1])
 
     def testSearch(self):
@@ -191,7 +175,7 @@ class TestAPI(unittest.TestCase):
         Test search via API
         """
 
-        uid = self.client.get("search?query=feel%20good%20story&limit=1").json()[0][0]
+        uid = self.client.get("search?query=feel%20good%20story&limit=1").json()[0]["id"]
         self.assertEqual(self.data[uid], self.data[4])
 
     def testSearchBatch(self):
@@ -204,7 +188,7 @@ class TestAPI(unittest.TestCase):
             "limit": 1
         }).json()
 
-        uids = [result[0][0] for result in results]
+        uids = [result[0]["id"] for result in results]
         self.assertEqual(uids, [4, 1])
 
     def testSimilarity(self):
@@ -215,7 +199,7 @@ class TestAPI(unittest.TestCase):
         uid = self.client.post("similarity", json={
             "query": "feel good story",
             "texts": self.data
-        }).json()[0][0]
+        }).json()[0]["id"]
 
         self.assertEqual(self.data[uid], self.data[4])
 
@@ -229,9 +213,26 @@ class TestAPI(unittest.TestCase):
             "texts": self.data
         }).json()
 
-        uids = [result[0][0] for result in results]
+        uids = [result[0]["id"] for result in results]
         self.assertEqual(uids, [4, 1])
 
+    def testTransform(self):
+        """
+        Test embeddings transform via API
+        """
+
+        self.assertEqual(len(self.client.get("transform?text=testembed").json()), 768)
+
+    def testTransformBatch(self):
+        """
+        Test batch embeddings transform via API
+        """
+
+        embeddings = self.client.post("batchtransform", json=self.data).json()
+
+        self.assertEqual(len(embeddings), len(self.data))
+        self.assertEqual(len(embeddings[0]), 768)
+
     def testViewOnly(self):
         """
         Test read-only API instance
@@ -241,13 +242,13 @@ class TestAPI(unittest.TestCase):
         self.client = TestAPI.start(False)
 
         # Test search
-        uid = self.client.get("search?query=feel%20good%20story&limit=1").json()[0][0]
+        uid = self.client.get("search?query=feel%20good%20story&limit=1").json()[0]["id"]
         self.assertEqual(uid, 4)
 
         # Test similarity
         uid = self.client.post("similarity", json={
             "query": "feel good story",
             "texts": self.data
-        }).json()[0][0]
+        }).json()[0]["id"]
 
         self.assertEqual(uid, 4)