test_rerank_client.py
1.71 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
from math import isclose
from search.rerank_client import fuse_scores_and_resort
def test_fuse_scores_and_resort_uses_multiplicative_formula_with_named_query_scores():
hits = [
{
"_id": "1",
"_score": 3.2,
"matched_queries": {
"base_query": 2.4,
"knn_query": 0.8,
},
},
{
"_id": "2",
"_score": 2.8,
"matched_queries": {
"base_query": 1.6,
"knn_query": 0.2,
},
},
]
debug = fuse_scores_and_resort(hits, [0.9, 0.7])
expected_1 = (0.9 + 0.00001) * ((0.8 + 0.6) ** 0.2) * ((2.4 + 0.1) ** 0.75)
expected_2 = (0.7 + 0.00001) * ((0.2 + 0.6) ** 0.2) * ((1.6 + 0.1) ** 0.75)
assert isclose(hits[0]["_fused_score"], expected_1, rel_tol=1e-9)
assert isclose(hits[1]["_fused_score"], expected_2, rel_tol=1e-9)
assert debug[0]["text_score"] == 2.4
assert debug[0]["knn_score"] == 0.8
assert [hit["_id"] for hit in hits] == ["1", "2"]
def test_fuse_scores_and_resort_falls_back_when_matched_queries_missing():
hits = [
{"_id": "1", "_score": 0.5},
{"_id": "2", "_score": 2.0},
]
fuse_scores_and_resort(hits, [0.4, 0.3])
expected_1 = (0.4 + 0.00001) * ((0.0 + 0.6) ** 0.2) * ((0.5 + 0.1) ** 0.75)
expected_2 = (0.3 + 0.00001) * ((0.0 + 0.6) ** 0.2) * ((2.0 + 0.1) ** 0.75)
assert isclose(hits[0]["_text_score"], 2.0, rel_tol=1e-9)
assert isclose(hits[0]["_fused_score"], expected_2, rel_tol=1e-9)
assert isclose(hits[1]["_text_score"], 0.5, rel_tol=1e-9)
assert isclose(hits[1]["_fused_score"], expected_1, rel_tol=1e-9)
assert [hit["_id"] for hit in hits] == ["2", "1"]