neuml/txtai
1
0
Fork 0
mirror of https://github.com/neuml/txtai.git synced 2026-10-12 02:19:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Code standardization and added tests #1041

Browse Source
This commit is contained in:
davidmezzetti
2026-02-18 14:18:11 -05:00
parent 48f0962a44
commit 00acd243b0
6 changed files with 170 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/python/txtai/embeddings/search/hybrid.py
+151 -98
View File
@@ -5,10 +5,6 @@ Hybrid module
import math
# Numerical clamp for log-odds computation
_EPSILON = 1e-10
class Hybrid:
"""
Hybrid score fusion strategies for combining dense and sparse search results.
@@ -49,110 +45,20 @@ class Hybrid:
return self.method(vectors, weights, limit)
def calibrate(self, dense_raw):
"""
Computes per-query calibration parameters for dense cosine scores.
Uses the same approach as BB25: beta=median, alpha_eff=1/std so the
logit for a dense score is alpha * (score - median), centering the
median candidate at logit 0.
Args:
dense_raw: list of raw dense cosine scores
Returns:
(median, alpha) calibration parameters
"""
d_median, d_alpha = 0.0, 1.0
dense_arr = [s for s in dense_raw if s > 0]
if dense_arr:
d_median = sorted(dense_arr)[len(dense_arr) // 2]
d_std = (sum((x - sum(dense_arr) / len(dense_arr)) ** 2 for x in dense_arr) / len(dense_arr)) ** 0.5
d_alpha = 1.0 / d_std if d_std > 0 else 1.0
return d_median, d_alpha
def logodds(self, vectors, weights, limit):
"""
Log-odds conjunction fusion for Bayesian (BB25) normalized scores.
Implements the framework from "From Bayesian Inference to Neural Computation"
(Jeong, 2026) with asymmetric dynamic calibration:
1. Calibrate dense cosine scores via per-query dynamic sigmoid
(beta=median, alpha_eff=1/std) to produce logits centered at 0.
2. Convert sparse BB25 probabilities to logits.
3. Fuse via weighted mean log-odds with confidence scaling.
Scores are returned as raw logits (not mapped back through sigmoid) to
preserve ranking resolution among top candidates.
Args:
vectors: tuple of (dense_results, sparse_results)
weights: [dense_weight, sparse_weight]
vectors: tuple of (dense results, sparse results)
weights: [dense weight, sparse weight]
limit: maximum results
Returns:
sorted list of (uid, score) where score is a fused logit
sorted list (uid, score)
"""
# Phase 1: Collect raw scores per document
uids = {}
dense_raw = []
for v, scores in enumerate(vectors):
for uid, score in scores if weights[v] > 0 else []:
if uid not in uids:
uids[uid] = [None, None]
if v == 0:
uids[uid][0] = score
dense_raw.append(score)
else:
# Sparse BB25 score: already a calibrated probability
uids[uid][1] = score
# Phase 2: Compute per-query calibration parameters for dense cosine scores.
# Same approach as BB25: beta=median, alpha_eff=1/std. The logit for a dense
# score is alpha * (score - median), centering the median candidate at logit 0.
d_median, d_alpha = self.calibrate(dense_raw)
# Phase 3: Fuse via weighted mean log-odds with confidence scaling.
# Raw logit scores are used for ranking instead of sigmoid(logit) to
# preserve fine-grained ordering among top candidates.
fused = {}
n = 2
alpha = 0.5
scale = n**alpha
for uid, pair in uids.items():
raw_dense = pair[0]
p_sparse = pair[1]
if raw_dense is not None and p_sparse is not None:
# Calibrate dense score via dynamic sigmoid
logit_d = d_alpha * (raw_dense - d_median)
logit_d = max(min(logit_d, 500), -500)
# Sparse BB25 score -> logit
p_s = min(max(p_sparse, _EPSILON), 1.0 - _EPSILON)
logit_s = math.log(p_s / (1.0 - p_s))
# Weighted mean log-odds with confidence scaling (Paper 2, Def 4.2.1)
l_bar = weights[0] * logit_d + weights[1] * logit_s
fused[uid] = l_bar * scale
elif raw_dense is not None:
# Only dense signal: calibrated logit scaled by weight
logit_d = d_alpha * (raw_dense - d_median)
logit_d = max(min(logit_d, 500), -500)
fused[uid] = logit_d * weights[0]
else:
# Only sparse signal: logit scaled by weight
p_s = min(max(p_sparse, _EPSILON), 1.0 - _EPSILON)
fused[uid] = math.log(p_s / (1.0 - p_s)) * weights[1]
return sorted(fused.items(), key=lambda x: x[1], reverse=True)[:limit]
return LogOdds()(vectors, weights, limit)
def convex(self, vectors, weights, limit):
"""
@@ -197,3 +103,150 @@ class Hybrid:
uids[uid] += (1.0 / (r + 1)) * weights[v]
return sorted(uids.items(), key=lambda x: x[1], reverse=True)[:limit]
class LogOdds:
"""
Log-odds conjunction fusion for Bayesian (BB25) normalized scores.
Implements the framework from "From Bayesian Inference to Neural Computation"
(Jeong, 2026) with asymmetric dynamic calibration:
1. Calibrate dense cosine scores via per-query dynamic sigmoid
(beta=median, alpha_eff=1/std) to produce logits centered at 0.
2. Convert sparse BB25 probabilities to logits.
3. Fuse via weighted mean log-odds with confidence scaling.
Scores are returned as raw logits (not mapped back through sigmoid) to
preserve ranking resolution among top candidates.
"""
# Numerical clamp for log-odds computation
EPSILON = 1e-10
def __call__(self, vectors, weights, limit):
"""
Log-odds fusion.
Args:
vectors: tuple of (dense vectors, sparse vectors)
weights: [dense weights, sparse weights]
limit: maximum results
Returns:
sorted list of (uid, score) where score is a fused logit
"""
# Phase 1: Collect raw scores per document
uids, denseraw = self.rawscores(vectors, weights)
# Phase 2: Compute per-query calibration parameters for dense cosine scores.
# Same approach as BB25: beta=median, alpha_eff=1/std. The logit for a dense
# score is alpha * (score - median), centering the median candidate at logit 0.
densemedian, densealpha = self.calibrate(denseraw)
# Phase 3: Fuse via weighted mean log-odds with confidence scaling.
# Raw logit scores are used for ranking instead of sigmoid(logit) to
# preserve fine-grained ordering among top candidates.
fused = self.fuse(uids, weights, densemedian, densealpha)
return sorted(fused.items(), key=lambda x: x[1], reverse=True)[:limit]
def rawscores(self, vectors, weights):
"""
Collects raw scores.
Args:
vectors: tuple of (dense vectors, sparse vectors)
weights: [dense weights, sparse weights]
Returns:
(uids, raw dense scores)
"""
uids, denseraw = {}, []
for v, scores in enumerate(vectors):
for uid, score in scores if weights[v] > 0 else []:
if uid not in uids:
uids[uid] = [None, None]
if v == 0:
uids[uid][0] = score
denseraw.append(score)
else:
# Sparse BB25 score: already a calibrated probability
uids[uid][1] = score
return uids, denseraw
def calibrate(self, raw):
"""
Computes per-query calibration parameters for dense cosine scores.
Uses the same approach as BB25: beta=median, alpha_eff=1/std so the
logit for a dense score is alpha * (score - median), centering the
median candidate at logit 0.
Args:
raw: list of raw dense cosine scores
Returns:
(median, alpha) calibration parameters
"""
median, alpha = 0.0, 1.0
array = [s for s in raw if s > 0]
if array:
median = sorted(array)[len(array) // 2]
std = (sum((x - sum(array) / len(array)) ** 2 for x in array) / len(array)) ** 0.5
alpha = 1.0 / std if std > 0 else 1.0
return median, alpha
def fuse(self, uids, weights, densemedian, densealpha):
"""
Fuses scores together.
Args:
uids: result ids
weights: [dense weights, sparse weights]
densemedian: dense median
densealpha: dense alpha
Returns:
fused scores
"""
fused, n, alpha = {}, 2, 0.5
scale = n**alpha
for uid, pair in uids.items():
# Unpair scores
rawdense, psparse = pair
if rawdense is not None and psparse is not None:
# Calibrate dense score via dynamic sigmoid
logitdense = densealpha * (rawdense - densemedian)
logitdense = max(min(logitdense, 500), -500)
# Sparse BB25 score -> logit
psparse = min(max(psparse, LogOdds.EPSILON), 1.0 - LogOdds.EPSILON)
logitsparse = math.log(psparse / (1.0 - psparse))
# Weighted mean log-odds with confidence scaling (Paper 2, Def 4.2.1)
lbar = weights[0] * logitdense + weights[1] * logitsparse
fused[uid] = lbar * scale
elif rawdense is not None:
# Only dense signal: calibrated logit scaled by weight
logitdense = densealpha * (rawdense - densemedian)
logitdense = max(min(logitdense, 500), -500)
fused[uid] = logitdense * weights[0]
else:
# Only sparse signal: logit scaled by weight
psparse = min(max(psparse, LogOdds.EPSILON), 1.0 - LogOdds.EPSILON)
fused[uid] = math.log(psparse / (1.0 - psparse)) * weights[1]
return fused
src/python/txtai/scoring/base.py
+1 -1
View File
@@ -195,4 +195,4 @@ class Scoring:
True if BB25/Bayesian normalization is active, False otherwise
"""
return False
raise NotImplementedError
src/python/txtai/scoring/normalize.py
+1 -4
View File
@@ -33,9 +33,6 @@ class Normalize:
self.beta = self.config.get("beta")
self.beta = float(self.beta) if self.beta is not None else self.beta
# BB25-compatible aliases for Bayesian normalization mode.
BAYESIAN_METHODS = ("bayes", "bayesian", "bayesian-bm25", "bb25")
def isbayes(self):
"""
Checks if Bayesian normalization mode is active.
@@ -44,7 +41,7 @@ class Normalize:
True if using BB25/Bayesian normalization
"""
return self.method in self.BAYESIAN_METHODS
return self.method in ("bayes", "bayesian", "bayesian-bm25", "bb25")
def __call__(self, scores, avgscore):
"""
test/python/testdatabase/testrdbms.py
+8
View File
@@ -299,6 +299,14 @@ class Common:
result = embeddings.search("feel good story", 1)[0]
self.assertEqual(result["text"], data[4][1])
# Index data with sparse + dense vectors and bb25 normalized scores
embeddings.config["scoring"]["normalize"] = "bb25"
embeddings.index(data)
# Run search
result = embeddings.search("canada intact iceberg a", 1)[0]
self.assertEqual(result["text"], data[1][1])
# Test upsert
data[0] = (0, "Feel good story: baby panda born", None)
embeddings.upsert([data[0]])
test/python/testembeddings.py
+8
View File
@@ -238,6 +238,14 @@ class TestEmbeddings(unittest.TestCase):
uid = embeddings.search("feel good story", 1)[0][0]
self.assertEqual(uid, 4)
# Index data with sparse + dense vectors and bb25 normalization
embeddings.config["scoring"]["normalize"] = "bb25"
embeddings.index(data)
# Run search
uid = embeddings.search("canada intact iceberg a", 1)[0][0]
self.assertEqual(uid, 1)
# Test upsert
data[0] = (0, "Feel good story: baby panda born", None)
embeddings.upsert([data[0]])
test/python/testscoring/testkeyword.py
+1
View File
@@ -75,6 +75,7 @@ class TestKeyword(unittest.TestCase):
self.assertRaises(NotImplementedError, scoring.close)
self.assertRaises(NotImplementedError, scoring.issparse)
self.assertRaises(NotImplementedError, scoring.isnormalized)
self.assertRaises(NotImplementedError, scoring.isbayes)
@patch("sqlalchemy.orm.Query.params")
def testPGText(self, query):