d90e7428
 tangwang
补充重排
|
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
|
"""
FastAPI service for BGE reranking.
POST /rerank
Request:
{
"query": "...",
"docs": ["doc1", "doc2", ...]
}
Response:
{
"scores": [0.98, 0.12, ...],
"meta": {...}
}
"""
import logging
import time
from typing import Any, Dict, List, Optional
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from reranker.bge_reranker import BGEReranker
from reranker.config import CONFIG
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s %(levelname)s %(name)s | %(message)s",
)
logger = logging.getLogger("reranker.service")
|
|
d90e7428
tangwang
补充重排
|
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
|
_reranker: Optional[BGEReranker] = None
class RerankRequest(BaseModel):
query: str = Field(..., description="Search query")
docs: List[str] = Field(..., description="Documents/passages to rerank")
normalize: Optional[bool] = Field(
default=CONFIG.NORMALIZE, description="Apply sigmoid normalization"
)
class RerankResponse(BaseModel):
scores: List[float] = Field(..., description="Scores aligned to input docs order")
meta: Dict[str, Any] = Field(default_factory=dict)
@app.on_event("startup")
def load_model() -> None:
global _reranker
logger.info("Starting reranker service on port %s", CONFIG.PORT)
try:
_reranker = BGEReranker(
model_name=CONFIG.MODEL_NAME,
device=CONFIG.DEVICE,
batch_size=CONFIG.BATCH_SIZE,
use_fp16=CONFIG.USE_FP16,
max_length=CONFIG.MAX_LENGTH,
cache_dir=CONFIG.CACHE_DIR,
enable_warmup=CONFIG.ENABLE_WARMUP,
)
logger.info(
"Reranker ready | model=%s device=%s fp16=%s batch=%s max_len=%s",
CONFIG.MODEL_NAME,
_reranker.device,
_reranker.use_fp16,
_reranker.batch_size,
_reranker.max_length,
)
except Exception as exc:
logger.error("Failed to initialize reranker: %s", exc, exc_info=True)
raise
@app.get("/health")
def health() -> Dict[str, Any]:
return {
"status": "ok" if _reranker is not None else "unavailable",
"model_loaded": _reranker is not None,
"model": CONFIG.MODEL_NAME,
"device": CONFIG.DEVICE,
}
@app.post("/rerank", response_model=RerankResponse)
def rerank(request: RerankRequest) -> RerankResponse:
if _reranker is None:
raise HTTPException(status_code=503, detail="Reranker model not loaded")
query = (request.query or "").strip()
if not query:
raise HTTPException(status_code=400, detail="query cannot be empty")
if request.docs is None or len(request.docs) == 0:
raise HTTPException(status_code=400, detail="docs cannot be empty")
if len(request.docs) > CONFIG.MAX_DOCS:
raise HTTPException(
status_code=400,
detail=f"Too many docs: {len(request.docs)} > {CONFIG.MAX_DOCS}",
)
normalize = CONFIG.NORMALIZE if request.normalize is None else bool(request.normalize)
start_ts = time.time()
logger.info(
"Rerank request | docs=%d normalize=%s",
len(request.docs),
normalize,
)
scores, meta = _reranker.score_with_meta(query, request.docs, normalize=normalize)
meta = dict(meta)
meta.update({"service_elapsed_ms": round((time.time() - start_ts) * 1000.0, 3)})
logger.info(
"Rerank done | docs=%d unique=%s dedup=%s elapsed_ms=%s",
meta.get("input_docs"),
meta.get("unique_docs"),
meta.get("dedup_ratio"),
meta.get("service_elapsed_ms"),
)
return RerankResponse(scores=scores, meta=meta)
if __name__ == "__main__":
import uvicorn
uvicorn.run(
"reranker.server:app",
host=CONFIG.HOST,
port=CONFIG.PORT,
reload=False,
log_level="info",
)
|