From edf61d1d4a56853ce0a96a0e35a3b6a5cba43170 Mon Sep 17 00:00:00 2001 From: Yuilona <143536279+Yuilona@users.noreply.github.com> Date: Thu, 8 Oct 2026 04:21:36 +0800 Subject: [PATCH] Use each query's hybrid weights when batching similar clauses (#1359) --- src/python/txtai/embeddings/search/base.py | 8 ++++++-- src/python/txtai/embeddings/search/scan.py | 6 +++--- test/python/testdatabase/testrdbms.py | 24 ++++++++++++++++++++++ 3 files changed, 33 insertions(+), 5 deletions(-) diff --git a/src/python/txtai/embeddings/search/base.py b/src/python/txtai/embeddings/search/base.py index c85be551..7d6672c4 100644 --- a/src/python/txtai/embeddings/search/base.py +++ b/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: diff --git a/src/python/txtai/embeddings/search/scan.py b/src/python/txtai/embeddings/search/scan.py index a7803e7b..36369154 100644 --- a/src/python/txtai/embeddings/search/scan.py +++ b/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 diff --git a/test/python/testdatabase/testrdbms.py b/test/python/testdatabase/testrdbms.py index e0466edb..49108894 100644 --- a/test/python/testdatabase/testrdbms.py +++ b/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