first commit
This commit is contained in:
@@ -0,0 +1,7 @@
|
||||
# 模型权重目录(约 290MB,不进版本库)
|
||||
models/
|
||||
*.onnx
|
||||
*.part
|
||||
|
||||
__pycache__/
|
||||
*.pyc
|
||||
+37
@@ -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"]
|
||||
@@ -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
|
||||
```
|
||||
@@ -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),
|
||||
}
|
||||
@@ -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
|
||||
@@ -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
@@ -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}"
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user