Flamehaven commited on
Commit
ef3adba
·
1 Parent(s): aca22b8

fix(tests): Correct logic in test_packer and fix mock in test_rerank

Browse files
Files changed (2) hide show
  1. tests/test_packer.py +1 -1
  2. tests/test_rerank.py +7 -2
tests/test_packer.py CHANGED
@@ -12,4 +12,4 @@ def test_budget_pack_sorting_stable():
12
  {"text": "z", "score": 0.8, "tokens": 10},
13
  ]
14
  sel = budget_pack(chunks, budget=60)
15
- assert [c.text for c in sel] == ["y", "x"]
 
12
  {"text": "z", "score": 0.8, "tokens": 10},
13
  ]
14
  sel = budget_pack(chunks, budget=60)
15
+ assert [c.text for c in sel] == ["y", "x", "z"]
tests/test_rerank.py CHANGED
@@ -1,8 +1,13 @@
1
  from crom_efficientllm.rerank_engine.rerank import hybrid_rerank
2
 
3
  class Dummy:
4
- def encode(self, text, convert_to_numpy=False):
5
- return [ord(c) % 5 for c in str(text)[:8]]
 
 
 
 
 
6
 
7
  def test_hybrid_rerank_returns_scores():
8
  docs = [{"text": "alpha"}, {"text": "beta"}]
 
1
  from crom_efficientllm.rerank_engine.rerank import hybrid_rerank
2
 
3
  class Dummy:
4
+ def encode(self, text_or_list, convert_to_numpy=False):
5
+ if isinstance(text_or_list, list):
6
+ return [self.encode(t) for t in text_or_list]
7
+ vec = [ord(c) % 5 for c in str(text_or_list)[:8]]
8
+ while len(vec) < 8:
9
+ vec.append(0)
10
+ return vec
11
 
12
  def test_hybrid_rerank_returns_scores():
13
  docs = [{"text": "alpha"}, {"text": "beta"}]