local-remarker/app.py

"""本地重排(rerank)服务。

对外暴露 SiliconFlow 兼容的 `POST /v1/services/rerank/text`,方便直接对接
各类 RAG 框架。之所以自己实现而不是用现成推理框架,主要是因为:
很多推理运行时(如 Ollama)并不提供 rerank 端点。

打分逻辑遵循 bge-reranker 系列的惯例:把 (query, document) 拼成句对批量推理,
取 logits[:, 0] 作为相关性分数(原始 logit,**不过 sigmoid**),降序后取 top_n。
"""

import os
import time
import uuid
from typing import Any

import numpy as np
import onnxruntime as ort
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from tokenizers import Tokenizer

# ---------------------------------------------------------------- 配置
MODEL_DIR = os.environ.get("MODEL_DIR", "/models")
ONNX_FILE = os.environ.get("ONNX_FILE", "model_int8.onnx")
TOKENIZER_FILE = os.environ.get("TOKENIZER_FILE", "tokenizer.json")
MODEL_NAME = os.environ.get("MODEL_NAME", "bge-reranker-base")
MAX_LEN = int(os.environ.get("MAX_LEN", "512"))
BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "16"))
NUM_THREADS = os.environ.get("NUM_THREADS", "") or None

app = FastAPI(title="local-remarker", version="1.0")


def _build_session() -> "tuple[ort.InferenceSession, Tokenizer]":
    onnx_path = os.path.join(MODEL_DIR, ONNX_FILE)
    tok_path = os.path.join(MODEL_DIR, TOKENIZER_FILE)
    for path in (onnx_path, tok_path):
        if not os.path.isfile(path):
            raise FileNotFoundError(
                f"缺少模型文件: {path}(entrypoint 应先执行 download_model.py)"
            )

    options = ort.SessionOptions()
    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    if NUM_THREADS:
        options.intra_op_num_threads = int(NUM_THREADS)
    else:
        # 默认吃满可用核;AVX512-VNNI 的机器上 int8 推理收益明显
        options.intra_op_num_threads = os.cpu_count() or 4

    session = ort.InferenceSession(
        onnx_path, sess_options=options, providers=["CPUExecutionProvider"]
    )

    tokenizer = Tokenizer.from_file(tok_path)
    tokenizer.enable_truncation(max_length=MAX_LEN)
    tokenizer.enable_padding(pad_id=1, pad_token="[PAD]")
    return session, tokenizer


SESSION, TOKENIZER = _build_session()
INPUT_NAMES = {i.name for i in SESSION.get_inputs()}


def _score(query: str, documents: list[str]) -> list[float]:
    """对 (query, doc) 句对批量打分,返回与 documents 等长的分数列表。"""
    encodings = TOKENIZER.encode_batch([[query, doc] for doc in documents])
    input_ids = np.array([e.ids for e in encodings], dtype=np.int64)
    attention_mask = np.array([e.attention_mask for e in encodings], dtype=np.int64)

    feeds = {"input_ids": input_ids, "attention_mask": attention_mask}
    if "token_type_ids" in INPUT_NAMES:
        feeds["token_type_ids"] = np.array([e.type_ids for e in encodings], dtype=np.int64)

    logits = SESSION.run(None, feeds)[0]
    # bge-reranker 只有 1 个 logit,取第 0 列即相关性分数(原始 logit)
    return [float(v) for v in np.asarray(logits).reshape(len(documents), -1)[:, 0]]


def _rerank(query: str, documents: list[str], top_n: int,
            return_documents: bool) -> dict[str, Any]:
    started = time.time()

    scores: list[float] = []
    for start in range(0, len(documents), BATCH_SIZE):
        scores.extend(_score(query, documents[start:start + BATCH_SIZE]))

    order = sorted(range(len(documents)), key=lambda i: scores[i],
                   reverse=True)[:top_n]

    results = []
    for index in order:
        item: dict[str, Any] = {
            "index": index,
            "relevance_score": scores[index],
        }
        if return_documents:
            item["document"] = {"text": documents[index]}
        results.append(item)

    total_tokens = sum(len(TOKENIZER.encode(query, doc).ids) for doc in documents)
    return {
        "id": uuid.uuid4().hex,
        "object": "rerank.result",
        "model": MODEL_NAME,
        "results": results,
        "usage": {"total_tokens": total_tokens},
        "meta": {"latency_ms": int((time.time() - started) * 1000)},
    }


async def _handle(request: Request) -> JSONResponse:
    try:
        body = await request.json()
    except Exception:  # noqa: BLE001
        return JSONResponse(
            {"error": {"message": "invalid json body"}}, status_code=400
        )

    query = body.get("query") or ""
    documents = body.get("documents") or []
    top_n = int(body.get("top_n") or len(documents) or 1)
    return_documents = bool(body.get("return_documents", False))

    if not query or not documents:
        return JSONResponse(
            {"error": {"message": "query and documents are required"}},
            status_code=400,
        )

    try:
        result = _rerank(query, documents, top_n, return_documents)
    except Exception as exc:  # noqa: BLE001
        return JSONResponse({"error": {"message": str(exc)}}, status_code=500)

    return JSONResponse(result)


# 显式注册常见路径,避免依赖调用方的拼接方式
for _path in (
    "/v1/services/rerank/text",
    "/services/rerank/text",
    "/v1/rerank",
    "/rerank",
    "/rank",
    "/v1/rank",
):
    app.post(_path)(_handle)


@app.api_route("/{full_path:path}", methods=["POST"])
async def _catch_all(full_path: str, request: Request) -> JSONResponse:
    """兜底路由。

    调用方填写的 base_url 可能是 `http://host:8000` 或 `http://host:8000/v1`,
    框架在其后拼接路径后可能出现 `/v1/v1/services/rerank/text` 之类的变体。
    凡是路径中含 rerank、或以 /text 结尾的 POST 一律按重排请求处理;
    无关路径仍返回 404,不会误吞流量。
    """
    if "rerank" not in full_path and not full_path.endswith("/text"):
        return JSONResponse(
            {"error": {"message": f"not found: /{full_path}"}}, status_code=404
        )
    return await _handle(request)


@app.get("/health")
def health() -> dict[str, Any]:
    return {"status": "ok", "model": MODEL_NAME}


@app.get("/info")
def info() -> dict[str, Any]:
    return {
        "model": MODEL_NAME,
        "onnx_file": ONNX_FILE,
        "model_dir": MODEL_DIR,
        "max_len": MAX_LEN,
        "batch_size": BATCH_SIZE,
        "input_names": sorted(INPUT_NAMES),
    }