Use each query's hybrid weights when batching similar clauses (#1359)

src/python/txtai/embeddings/search/base.py

@@ -92,7 +92,7 @@ class Search:
         Args:
             queries: list of queries
             limit: maximum results
-            weights: hybrid score weights
+            weights: hybrid score weights, either one set of weights for all queries or a list of [dense, sparse] weights per query
             index: index name
 
         Returns:
@@ -114,9 +114,13 @@ class Search:
             if isinstance(weights, (int, float)):
                 weights = [weights, 1 - weights]
 
+            # Use the same weights for each query, unless a list of weights per query is passed
+            if not isinstance(weights[0], (list, tuple)):
+                weights = [weights] * len(queries)
+
             # Create weighted scores via hybrid fusion strategy
             fusion = Hybrid(self.scoring)
-            return [fusion(vectors, weights, limit) for vectors in zip(dense, sparse)]
+            return [fusion(vectors, weights[x], limit) for x, vectors in enumerate(zip(dense, sparse))]
 
         # Raise an error if when no indexes are available
         if not sparse and not dense:

src/python/txtai/embeddings/search/scan.py

@@ -61,9 +61,9 @@ class Scan:
 
             candidates = max(candidates) if candidates else default
 
-            # Query weights to pass to batch search
-            weights = [query.weights for query in iqueries if query.weights is not None]
-            weights = max(weights) if weights else self.weights
+            # Query weights to pass to batch search. Each query keeps its own weights, queries without weights use the default.
+            weights = [query.weights if query.weights is not None else self.weights for query in iqueries]
+            weights = [[weight, 1 - weight] if isinstance(weight, (int, float)) else weight for weight in weights]
 
             # Index to run query against
             index = index if index else self.index

test/python/testdatabase/testrdbms.py

@@ -329,6 +329,30 @@ class Common:
             result = embeddings.search("feel good story", 1)[0]
             self.assertEqual(result["text"], data[0][1])
 
+        def testHybridWeights(self):
+            """
+            Test hybrid search with different similar clause weights in the same batch
+            """
+
+            queries = [
+                "select id, score from txtai where similar('a day at work', 0.1)",
+                "select id, score from txtai where similar('feel good story', 0.9)",
+            ]
+
+            # Test with sparse + dense vectors in the top-level index and in a subindex
+            hybrid = {"path": "sentence-transformers/nli-mpnet-base-v2", "hybrid": True}
+            for config in [hybrid, {"defaults": False, "indexes": {"index1": hybrid}}]:
+                embeddings = Embeddings({**config, "content": self.backend})
+                embeddings.index(self.data)
+
+                # Each query in a batch must use its own weights
+                batch = embeddings.batchsearch(queries, 6)
+                for query, results in zip(queries, batch):
+                    expected = embeddings.search(query, 6)
+                    self.assertEqual([r["id"] for r in results], [r["id"] for r in expected])
+                    for result, value in zip(results, expected):
+                        self.assertAlmostEqual(result["score"], value["score"], places=5)
+
         def testIndex(self):
             """
             Test index