Blame view

tests/test_rerank_client.py 1.71 KB
a47416ec   tangwang   把融合逻辑改成乘法公式,并把 ES...
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"]