first commit

This commit is contained in:
2026-10-23 18:09:59 +08:00
commit e0248e810f
8 changed files with 685 additions and 0 deletions
+7
View File
@@ -0,0 +1,7 @@
# 模型权重目录(约 290MB,不进版本库)
models/
*.onnx
*.part
__pycache__/
*.pyc
+37
View File
@@ -0,0 +1,37 @@
# 本地重排服务镜像。
#
# 模型不在构建期打包 —— 体积大(约 290MB)且会频繁升级,改为容器启动时按需下载,
# 已有则复用。这样换模型只需改环境变量后重启,不用重新 build。
FROM python:3.11-slim
# pip 源做成构建参数:国内机器可用清华源显著提速
ARG PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple
ARG PIP_TRUSTED_HOST=pypi.tuna.tsinghua.edu.cn
ENV PYTHONUNBUFFERED=1 \
PYTHONDONTWRITEBYTECODE=1 \
PIP_NO_CACHE_DIR=1 \
PIP_DISABLE_PIP_VERSION_CHECK=1
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir \
-i "${PIP_INDEX_URL}" \
--trusted-host "${PIP_TRUSTED_HOST}" \
-r requirements.txt
COPY app.py download_model.py entrypoint.sh ./
RUN chmod +x entrypoint.sh
# 模型目录。挂卷到这里,重启容器时模型不会丢。
ENV MODEL_DIR=/models
RUN mkdir -p /models
EXPOSE 8000
# slim 镜像里没有 curl,直接用 python 做健康检查
HEALTHCHECK --interval=30s --timeout=5s --start-period=600s --retries=3 \
CMD python -c "import urllib.request,sys; sys.exit(0 if urllib.request.urlopen('http://127.0.0.1:8000/health',timeout=4).status==200 else 1)"
CMD ["./entrypoint.sh"]
+177
View File
@@ -0,0 +1,177 @@
# local-remarker
本地重排(rerank)服务。零 API Key、文档不出内网,对外提供 SiliconFlow 兼容接口。
起因是在给 [PandaWiki](https://github.com/chaitin/PandaWiki) 做本地化部署时,
「重排序模型」被卡住了,一共有四个独立的阻塞点:
| # | 阻塞点 | 实测结论 |
|---|---|---|
| 1 | Ollama 不提供重排 | Ollama 0.40.8 的 `/api/rerank`、`/v1/rerank` 都是纯 404 |
| 2 | 调用方把路径写死 | PandaWiki 的 raglite 只认 `<base_url>/v1/services/rerank/text`(SilhouFlow 格式) |
| 3 | 装不上 torch | `download.pytorch.org/whl/cpu` 与阿里云 pytorch-wheels 下载均 20 秒零进度 |
| 4 | 官方模型站不可达 | `registry.ollama.ai` 的 blob 下载完全卡死,`hf-mirror.com` 连不上 |
最终方案是**用 ONNX Runtime,绕开 torch**。ModelScope 上 `Xenova/bge-reranker-base`
自带量化 ONNX,于是第 3、4 两个阻塞点同时消失 —— 整条链路只依赖
**ModelScope 一个国内源**,实测 15~25 MB/s。
## 快速开始
```bash
docker compose up -d
```
首次启动会下载约 290MB 模型到 `./models`,日志形如:
```
[local-remarker] 模型目录: /models
[model_int8.onnx] 需要下载 ...
[model_int8.onnx] 266.0MB / 266.0MB (100%)
[model_int8.onnx] 校验通过
[local-remarker] 模型准备完成
```
之后无论重启还是重建容器,都不会再产生下载流量。
服务监听 **63001**:
```bash
curl -s localhost:63001/health
# {"status":"ok","model":"bge-reranker-base"}
curl -s -X POST localhost:63001/v1/services/rerank/text \
-H 'Content-Type: application/json' \
-d '{
"model": "bge-reranker-base",
"query": "如何部署 PandaWiki 知识库",
"documents": [
"PandaWiki 支持通过 Docker 部署,数据保存在 postgres 与 minio",
"今天股市行情如何,上证指数涨跌",
"部署时需要先配置 embedding 模型,否则无法创建知识库"
],
"top_n": 3,
"return_documents": true
}'
```
## 设计要点
### 模型不在镜像里,启动时按需下载
`download_model.py` 在 entrypoint 里跑,做三件事:
1. **检查** —— 每个文件比对 `大小 + sha256`,全部通过就立刻返回,零网络请求
2. **续传** —— 下载落到 `.part` 文件,重试时带 `Range` 头从断点继续
3. **原子替换** —— 下载完先校验 sha256,通过才 `os.replace` 成正式文件,
避免中断留下的半个文件被当成有效模型加载
模型仓库、版本、下载源都是环境变量,换模型不用重新 build。
### 路径兜底
调用方填的 base_url 可能是 `http://host:8000`,也可能是 `http://host:8000/v1`,
框架在其后拼接路径后可能出现 `/v1/v1/services/rerank/text` 之类的变体。所以除了
显式注册常见路径,还加了一个兜底路由:路径中含 `rerank` 或以 `/text` 结尾的
POST 一律按重排处理;无关路径照常返回 404,不会误吞流量。
### 打分逻辑
遵循 bge-reranker 系列的惯例:把 `(query, document)` 拼成句对批量推理,
取 `logits[:, 0]` 作为相关性分数 —— 用的是**原始 logit,不过 sigmoid**,
然后降序取 top_n。
### 为什么选 int8
宿主机是 i7-1165G7,带 **AVX512-VNNI**,int8 推理能直接吃到该指令集加速。
266MB 的 int8 vs 1061MB 的 fp32,精度损失很小但速度和体积都明显更优。
同仓库可选:`model_quantized.onnx` (266MB)、`model_fp16.onnx` (531MB)、
`model.onnx` (1061MB)。换法见 `docker-compose.yml` 里的注释。
## 接口
### `POST /v1/services/rerank/text`
兼容 SiliconFlow,请求:
```json
{
"model": "bge-reranker-base",
"query": "查询词",
"documents": ["片段1", "片段2"],
"top_n": 5,
"return_documents": false
}
```
响应:
```json
{
"id": "…",
"object": "rerank.result",
"model": "bge-reranker-base",
"results": [
{"index": 0, "relevance_score": 6.8775}
],
"usage": {"total_tokens": 44},
"meta": {"latency_ms": 18}
}
```
`relevance_score` 是原始 logit,**无固定取值范围**,只保证相对大小关系正确。
需要 0~1 区间的话调用方自行过 sigmoid。
### 其它
- `GET /health` —— 存活探测
- `GET /info` —— 当前模型、路径、输入张量名
- 同样接受 `POST /services/rerank/text`、`/v1/rerank`、`/rerank`、`/rank`
## 配置
全部写在 `docker-compose.yml` 的 `environment` 里,改完 `docker compose up -d` 生效。常用的几个:
| 变量 | 默认 | 说明 |
|---|---|---|
| `PORT`(宿主机侧) | `63001` | compose 的 `ports`,映射到容器 8000 |
| `MODEL_DIR` | `/models` | 容器内模型目录 |
| `MODEL_REPO` | `Xenova/bge-reranker-base` | 换模型改这里 |
| `MODEL_BASE` | ModelScope | 换下载源改这里 |
| `ONNX_FILE` | `model_int8.onnx` | 换量化档位改这里 |
| `MAX_LEN` | `512` | 截断长度,长文档可调大(吃内存) |
| `BATCH_SIZE` | `16` | 每次推理的句对数 |
| `WORKERS` | `1` | 调大前先看内存,每个 worker 各加载一份模型 |
## 实测性能
i7-1165G7 / 8 逻辑核 / 15GB 内存:
| 场景 | 耗时 |
|---|---|
| 3 个文档重排 | 14~18 ms |
| 镜像大小 | 401 MB |
| 模型大小 | 290 MB |
## 已知限制
- **单 worker**:默认 `--workers 1`,每个 worker 会各自加载一份模型(约 400MB 内存)。
高并发场景应调 `WORKERS`,同时注意内存翻倍。
- **不鉴权**:不校验 `api_key`。放到非受信网络需要自行加反向代理。
- **仅 CPU**:未做 GPU 加速,也没接 TensorRT/ONNXRuntime-GPU。
- **只支持 ONNX**:换 safetensors 格式的模型需要调整 `download_model.py` 的
`MANIFEST` 并改用 torch 运行时。
## 文件
```
local-remarker/
├── app.py FastAPI 服务本体
├── download_model.py 模型检查/下载/校验(幂等,支持续传)
├── entrypoint.sh 先下模型,再起 uvicorn
├── requirements.txt 运行时依赖(无 torch)
├── Dockerfile
├── docker-compose.yml 参数全部写在这里
└── .gitignore
```
+181
View File
@@ -0,0 +1,181 @@
"""本地重排(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),
}
+62
View File
@@ -0,0 +1,62 @@
# 本地重排服务。
#
# 参数全部直接写在本文件里,不再走 .env —— 需要调整时改这里,然后
# docker compose up -d 重建即可。
#
# 首次启动会下载约 290MB 模型权重到 ./models;
# 之后重启/重建都会先校验 sha256,通过就跳过下载,零网络流量。
services:
local-remarker:
build: .
image: local-remarker:latest
container_name: local-remarker
restart: unless-stopped
# 宿主机 63001 -> 容器 8000
ports:
- "63001:8000"
environment:
# ---- 模型存放位置(容器内)----
MODEL_DIR: /models
# ---- 模型来源与选择 ----
MODEL_REPO: Xenova/bge-reranker-base
MODEL_REV: master
MODEL_BASE: https://modelscope.cn/models/Xenova/bge-reranker-base/resolve/master
MODEL_NAME: bge-reranker-base
# 同仓库可选量化档位:
# model_int8.onnx 266MB 默认,带 AVX512-VNNI 的 CPU 上最快
# model_quantized.onnx 266MB
# model_fp16.onnx 531MB
# model.onnx 1061MB 原始精度,CPU 上明显变慢
# 更大的模型可改用 BAAI/bge-reranker-v2-m3(2.2GB,质量最好但重排较慢)
ONNX_FILE: model_int8.onnx
# ---- 推理参数 ----
MAX_LEN: "512" # 截断长度,长文档可调大(吃内存)
BATCH_SIZE: "16" # 每次推理的句对数
# ---- 服务 ----
WORKERS: "1" # 调大前先看内存:每个 worker 会各加载一份模型
PORT: "8000"
LOG_LEVEL: info
TZ: Asia/Shanghai
# 宿主机 ./models -> 容器 /models。指向大盘路径可避免占满系统盘,
# 例如改成 /data/models/remarker:/models
volumes:
- ./models:/models
healthcheck:
test:
- CMD
- python
- -c
- "import urllib.request,sys; sys.exit(0 if urllib.request.urlopen('http://127.0.0.1:8000/health',timeout=4).status==200 else 1)"
interval: 30s
timeout: 5s
# 首次启动要下模型,给足预热时间
start_period: 600s
retries: 3
+199
View File
@@ -0,0 +1,199 @@
"""启动前检查并下载 rerank 模型。
设计目标:
* 幂等 —— 模型已存在且校验通过就立刻返回,不产生任何网络请求;
* 可续传 —— 弱网下载中断后重启容器能从断点继续;
* 可校验 —— 每个文件都带 sha256,下载后校验,不匹配则丢弃重下。
可配置项(均可用环境变量覆盖):
MODEL_DIR 模型存放目录,默认 /models
MODEL_REPO 仓库名,默认 Xenova/bge-reranker-base
MODEL_BASE 下载源前缀
FORCE 设为 1 时忽略本地文件强制重下
MAX_RETRIES 单个文件最大重试次数
"""
import hashlib
import os
import shutil
import sys
import time
import urllib.error
import urllib.request
MODEL_DIR = os.environ.get("MODEL_DIR", "/models")
MODEL_REPO = os.environ.get("MODEL_REPO", "Xenova/bge-reranker-base")
MODEL_REV = os.environ.get("MODEL_REV", "master")
MODEL_BASE = os.environ.get(
"MODEL_BASE", f"https://modelscope.cn/models/{MODEL_REPO}/resolve/{MODEL_REV}"
)
FORCE = os.environ.get("FORCE", "0") == "1"
MAX_RETRIES = int(os.environ.get("MAX_RETRIES", "5"))
TIMEOUT = int(os.environ.get("DOWNLOAD_TIMEOUT", "60"))
# (目标文件名, 仓库内路径, sha256, 字节数)
MANIFEST = [
(
"model_int8.onnx",
"onnx/model_int8.onnx",
"2059d8ef0b6e935b4845e11b38c9af9e9e2e7b91f69fc99efe03254e0a7da8d3",
278825308,
),
(
"tokenizer.json",
"tokenizer.json",
"48564c5c7d3fa64d85d95e65414a542385f88b0f128fd8d4163fd7a57f2be05c",
17098079,
),
(
"sentencepiece.bpe.model",
"sentencepiece.bpe.model",
"cfc8146abe2a0488e9e2a0c56de7952f7c11ab059eca145a0a727afce0db2865",
5069051,
),
(
"special_tokens_map.json",
"special_tokens_map.json",
"d5469a60db23249c7f8945013d78df30b44b6bf686c6bb4740f4223f77b1b535",
279,
),
(
"tokenizer_config.json",
"tokenizer_config.json",
"a1d6bc8734a6f635dc158508bef000f8e2e5a759c7d92f984b2c86e5ff53425b",
443,
),
]
def human(size: float) -> str:
for unit in ("B", "KB", "MB", "GB"):
if size < 1024 or unit == "GB":
return f"{size:.1f}{unit}"
size /= 1024
return f"{size:.1f}GB"
def sha256_of(path: str, chunk: int = 1 << 20) -> str:
digest = hashlib.sha256()
with open(path, "rb") as fh:
while True:
block = fh.read(chunk)
if not block:
break
digest.update(block)
return digest.hexdigest()
def is_valid(path: str, expect_sha: str, expect_size: int) -> bool:
if not os.path.isfile(path):
return False
if os.path.getsize(path) != expect_size:
return False
return sha256_of(path) == expect_sha
def download(name: str, remote_path: str) -> None:
"""带断点续传的下载,落到 .part 临时文件,成功后原子改名。"""
target = os.path.join(MODEL_DIR, name)
part = target + ".part"
url = f"{MODEL_BASE}/{remote_path}"
last_error: Exception | None = None
for attempt in range(1, MAX_RETRIES + 1):
try:
have = os.path.getsize(part) if os.path.isfile(part) else 0
total = next((m[3] for m in MANIFEST if m[0] == name), 0)
request = urllib.request.Request(url)
request.add_header("User-Agent", "local-remarker/1.0")
if have:
request.add_header("Range", f"bytes={have}-")
with urllib.request.urlopen(request, timeout=TIMEOUT) as resp:
# 服务端不支持 Range 时会返回 200 全量,此时从零重写
resuming = resp.status == 206
mode = "ab" if (have and resuming) else "wb"
if mode == "wb":
have = 0
downloaded = have
print(f" [{name}] {human(downloaded)} / {human(total)}", flush=True)
with open(part, mode) as out:
while True:
block = resp.read(1 << 20)
if not block:
break
out.write(block)
downloaded += len(block)
pct = downloaded * 100 // total if total else 0
print(
f"\r [{name}] {human(downloaded)} / {human(total)} ({pct}%)",
end="",
flush=True,
)
print(flush=True)
os.replace(part, target)
return
except (urllib.error.URLError, OSError, TimeoutError) as exc:
last_error = exc
wait = min(2 ** attempt, 30)
print(
f"\n [{name}] 第 {attempt}/{MAX_RETRIES} 次失败: {exc},{wait}s 后重试",
flush=True,
)
if attempt < MAX_RETRIES:
time.sleep(wait)
raise SystemExit(f"下载 {name} 失败: {last_error}")
def main() -> int:
print(f"[local-remarker] 模型目录: {MODEL_DIR}")
print(f"[local-remarker] 下载源 : {MODEL_BASE}")
os.makedirs(MODEL_DIR, exist_ok=True)
pending = []
for name, remote, sha, size in MANIFEST:
target = os.path.join(MODEL_DIR, name)
if FORCE or not is_valid(target, sha, size):
if not FORCE and os.path.isfile(target):
print(f" [{name}] 校验不通过,将重新下载")
pending.append((name, remote, sha, size))
else:
print(f" [{name}] 已存在且校验通过 ({human(size)})")
if not pending:
print("[local-remarker] 模型已就绪,跳过下载")
return 0
total_bytes = sum(item[3] for item in pending)
print(f"[local-remarker] 需要下载 {len(pending)} 个文件,共 {human(total_bytes)}")
# 容器内磁盘常较小,给个明确提示而不是让 curl 报一堆 no space
free = shutil.disk_usage(MODEL_DIR).free
if free < total_bytes * 1.1:
raise SystemExit(
f"磁盘空间不足:可用 {human(free)},需要约 {human(total_bytes)}。"
"请挂载更大的卷或改用 fp16/其它量化版本。"
)
for name, remote, sha, size in pending:
download(name, remote)
# 最终校验,任何文件对不上都不允许带着半个模型启动
for name, _remote, sha, size in MANIFEST:
target = os.path.join(MODEL_DIR, name)
if not is_valid(target, sha, size):
raise SystemExit(f"{name} 校验失败,拒绝启动")
print(f" [{name}] 校验通过")
print("[local-remarker] 模型准备完成")
return 0
if __name__ == "__main__":
sys.exit(main())
Executable
+17
View File
@@ -0,0 +1,17 @@
#!/bin/sh
# 容器入口:先确保模型就位,再启动服务。
#
# download_model.py 是幂等的 —— 模型已存在且 sha256 校验通过时会直接返回,
# 因此正常的重启/重建容器不会产生任何下载流量;首次启动才会拉取约 290MB 权重。
set -e
echo "[entrypoint] 检查模型 ..."
python /app/download_model.py
echo "[entrypoint] 启动 rerank 服务 ..."
exec uvicorn app:app \
--host "${HOST:-0.0.0.0}" \
--port "${PORT:-8000}" \
--workers "${WORKERS:-1}" \
--log-level "${LOG_LEVEL:-info}"
+5
View File
@@ -0,0 +1,5 @@
fastapi==0.115.6
uvicorn[standard]==0.34.0
onnxruntime==1.20.1
tokenizers==0.21.0
numpy==2.2.1