Only increment index offset by number of items passed to ANN index, closes #294

src/python/txtai/embeddings/transform.py

@@ -89,35 +89,39 @@ class Transform:
             documents: iterable collection (id, data, tags)
         """
 
-        batch = []
+        batch, offset = [], 0
 
         # Iterate and process documents stream
         for document in documents:
             if isinstance(document[1], dict):
                 if "text" in document[1]:
                     yield (document[0], document[1]["text"], document[2])
+                    offset += 1
                 elif "object" in document[1]:
                     yield (document[0], document[1]["object"], document[2])
+                    offset += 1
             else:
                 yield document
+                offset += 1
 
             # Batch document
             batch.append(document)
             if len(batch) == self.batch:
-                self.load(batch)
-                batch = []
+                self.load(batch, offset)
+                batch, offset = [], 0
 
         # Final batch
         if batch:
-            self.load(batch)
+            self.load(batch, offset)
 
-    def load(self, batch):
+    def load(self, batch, offset):
         """
         Loads a document batch. This method deletes existing ids from an embeddings index and
         load into a database, if applicable.
 
         Args:
             batch: list of (id, data, tags)
+            offset: index offset for batch
         """
 
         # Delete from embeddings index first (which deletes from underlying ANN index and database) if this is an upsert
@@ -134,7 +138,7 @@ class Transform:
         # Load batch into database except if this is a reindex
         if self.database and self.action != Action.REINDEX:
             self.database.insert(batch, self.offset)
-            self.offset += len(batch)
+            self.offset += offset
 
 
 class Action(Enum):

test/python/testdatabase/testembeddings.py

@@ -330,6 +330,26 @@ class TestEmbeddings(unittest.TestCase):
         self.assertRaises(NotImplementedError, database.embed, None, None)
         self.assertRaises(NotImplementedError, database.query, None, None)
 
+    def testMultiData(self):
+        """
+        Test indexing with multiple data types (text, documents)
+        """
+
+        embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": True, "batch": len(self.data)})
+
+        # Create an index using mixed data (text and documents)
+        data = []
+        for uid, text in enumerate(self.data):
+            data.append((uid, text, None))
+            data.append((uid, {"content": text}, None))
+
+        embeddings.index(data)
+
+        # Search for best match
+        result = embeddings.search("feel good story", 1)[0]
+
+        self.assertEqual(result["text"], self.data[4])
+
     def testQueryModel(self):
         """
         Test index