d31c7f65
 tangwang
补充云服务reranker
|
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
|
from __future__ import annotations
from typing import Any, Dict, List
from fastapi.testclient import TestClient
class _FakeTopNReranker:
_model_name = "fake-topn-reranker"
def score_with_meta(self, query: str, docs: List[str], normalize: bool = True):
return [0.1 for _ in docs], {"input_docs": len(docs), "path": "base"}
def score_with_meta_topn(
self,
query: str,
docs: List[str],
normalize: bool = True,
top_n: int | None = None,
):
scores = [0.0 for _ in docs]
if docs and top_n:
scores[0] = 1.0
return scores, {"input_docs": len(docs), "path": "topn", "top_n": top_n}
def test_reranker_server_forwards_top_n():
import reranker.server as reranker_server
reranker_server.app.router.on_startup.clear()
reranker_server._reranker = _FakeTopNReranker()
reranker_server._backend_name = "fake_topn"
with TestClient(reranker_server.app) as client:
response = client.post(
"/rerank",
json={
"query": "wireless mouse",
"docs": ["a", "b", "c"],
"top_n": 2,
},
)
assert response.status_code == 200
data: Dict[str, Any] = response.json()
assert data["scores"] == [1.0, 0.0, 0.0]
assert data["meta"]["path"] == "topn"
assert data["meta"]["requested_top_n"] == 2
assert data["meta"]["top_n"] == 2
|