toor/local-embedding
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first commit

Browse Source
This commit is contained in:
toor
2026-10-23 18:08:46 +08:00
commit 136163ad88
8 changed files with 733 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+7
View File
@@ -0,0 +1,7 @@
# 模型权重目录(约 570MB,不进版本库)
models/
*.onnx
*.part
__pycache__/
*.pyc
Dockerfile
+37
View File
@@ -0,0 +1,37 @@
# 本地嵌入服务镜像。
#
# 模型不在构建期打包 —— 体积大(约 570MB)且会频繁升级,改为容器启动时按需下载,
# 已有则复用。这样换模型只需改 compose 里的参数后重启,不用重新 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=900s --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"]
README.md
+189
View File
@@ -0,0 +1,189 @@
# local-embedding
本地嵌入(embedding)服务。**完整替代 Ollama** —— 零 API Key、文档不出内网,
对外提供 OpenAI 兼容的 `POST /v1/embeddings`。
## 为什么不用 Ollama
给 [PandaWiki](https://github.com/chaitin/PandaWiki) 做本地化部署时,嵌入模型
本来是走 Ollama 的,但为了这一件事装了一整套 Ollama,代价有三:
1. **Ollama 不归你管** —— 它由 fnOS 应用中心的 `ai_installer` 托管,
进程和监听地址都不在自己的控制范围内;
2. **官方 registry 不可达** —— `registry.ollama.ai` 的 blob 下载会完全卡死
(实测 30 分钟零进度),`hf-mirror.com` 也连不上,最后只能绕道 ModelScope;
3. **需要一层转发** —— Ollama 只监听 `127.0.0.1`,容器访问不到,得额外加一个
socat 容器把 docker 网关转发过去。
这个服务把这三件事都消掉了:一个容器、一个镜像、启动时自动下模型。
## 与 Ollama 的实测差异
同一批文本,本服务(ONNX int8)与 Ollama(bge-m3 GGUF Q8_0)的向量余弦相似度:
| 文本 | 相似度 |
|---|---|
| 如何部署 PandaWiki 知识库 | 0.9848 |
| PandaWiki 支持 Docker 部署,需要先配置向量模型 | 0.9821 |
| 今天上证指数涨跌如何 | 0.9849 |
内部相对关系也一致(0.7286 / 0.7390、0.3663 / 0.3664)。差异来自量化方式不同,
属于预期范围,不是实现偏差。维度同为 1024。
## 快速开始
```bash
docker compose up -d
```
首次启动会下载约 563MB 模型到 `./models`:
```
[local-embedding] 模型目录: /models
[model_int8.onnx] 0.0B / 542.1MB
[model_int8.onnx] 542.1MB / 542.1MB (100%)
[model_int8.onnx] 校验通过
[local-embedding] 模型准备完成
```
之后重启或重建容器都是零下载:
```
[model_int8.onnx] 已存在且校验通过 (542.1MB)
[local-embedding] 模型已就绪,跳过下载
```
服务监听 **63002**:
```bash
curl -s localhost:63002/health
# {"status":"ok","model":"bge-m3"}
curl -s -X POST localhost:63002/v1/embeddings \
-H 'Content-Type: application/json' \
-d '{"model":"bge-m3","input":["如何部署知识库","今天天气不错"]}'
```
## 设计要点
### 池化:CLS + L2 归一化
ONNX 导出的第一个输出可能是 `[batch, hidden]`(已池化)也可能是
`[batch, seq, hidden]`(token 级),两种都兼容;后者取 `[CLS]`(第 0 个 token)。
bge 系列官方推荐 **CLS 池化 + L2 归一化**,归一化后余弦相似度与内积等价,
调用方随便用哪种度量都对。实测输出 L2 范数为 1.000000。
### 模型不在镜像里,启动时按需下载
`download_model.py` 在 entrypoint 里先跑,做三件事:
1. **检查** —— 每个文件比对 `大小 + sha256`,全通过就立刻返回,零网络请求
2. **续传** —— 下到 `.part`,重试带 `Range` 头从断点继续
3. **原子替换** —— 校验通过才 `os.replace` 成正式文件,
避免中断留下的半个文件被当成有效模型加载
还有**磁盘预检**:空间不够直接报明确错误,而不是刷一片 `no space`。
### 为什么选 int8
带 AVX512-VNNI 的 CPU 上 int8 推理能直接吃到指令集加速。542MB 的 int8
vs 2.1GB 的 fp32,体积差 4 倍,精度损失很小 —— 见上面的对比数据。
注意 `onnx/model.onnx` 只有 0.6MB,它的权重在外置的 `model.onnx_data`(2.1GB)里;
`model_int8.onnx` 是自包含的单文件,少一个踩坑点。
## 接口
### `POST /v1/embeddings`
OpenAI 标准。请求:
```json
{
"model": "bge-m3",
"input": ["文本1", "文本2"],
"encoding_format": "float"
}
```
`input` 可以是单个字符串、字符串数组,或 token id 数组(三种可混排)。
响应:
```json
{
"object": "list",
"data": [
{"object": "embedding", "index": 0, "embedding": [0.0123, ...]}
],
"model": "bge-m3",
"usage": {"prompt_tokens": 7, "total_tokens": 7},
"_meta": {"latency_ms": 141, "dimension": 1024}
}
```
`_meta` 是本服务额外加的非标准字段,方便排查延迟和维度,不影响兼容性。
同时接受 `/embeddings`、`/v1/embed`、`/embed`。
### 其它
- `GET /health` —— 存活探测
- `GET /info` —— 当前模型、输入输出张量名、参数
- `GET /v1/models` —— 兼容部分客户端的启动探测
## 配置
全部写在 `docker-compose.yml` 的 `environment` 里,改完 `docker compose up -d` 生效。
| 变量 | 默认 | 说明 |
|---|---|---|
| 宿主机端口 | `63002` | compose 的 `ports`,映射到容器 8000 |
| `MODEL_DIR` | `/models` | 容器内模型目录 |
| `MODEL_REPO` | `Xenova/bge-m3` | 换模型改这里 |
| `MODEL_BASE` | ModelScope | 换下载源改这里 |
| `ONNX_FILE` | `model_int8.onnx` | 换量化档位 |
| `MAX_LEN` | `512` | 截断长度,bge-m3 原生支持 8192,但 CPU 上长序列很慢 |
| `BATCH_SIZE` | `16` | 每批条数 |
| `WORKERS` | `1` | 调大前先看内存,每个 worker 各加载一份模型 |
同仓库可选量化档位:`model_fp16.onnx`(1081MB)、`model_q4f16.onnx`(668MB)、
`model_q4.onnx`(1190MB)。
## 实测性能
MacBook(Docker Desktop,arm64 模拟)与 i7-1165G7:
| 场景 | 耗时 |
|---|---|
| 40 条批量 | 449 ms(约 90 条/秒) |
| 5 条批量 | 141 ms |
| 镜像大小 | 401 MB |
| 模型大小 | 563 MB |
## 已知限制
- **单 worker**:默认 `--workers 1`,每个 worker 会各加载一份模型(约 2GB 内存峰值)。
高并发应调 `WORKERS`,同时注意内存翻倍。
- **不鉴权**:不校验 `api_key`。放到非受信网络需自行加反向代理。
- **仅 CPU**:未接 ONNXRuntime-GPU。
- **只支持 ONNX**:换 safetensors 需改 `MANIFEST` 并引入 torch。
- **长文本默认截断到 512**:bge-m3 能力上限是 8192,但 CPU 上超过 512 明显变慢。
需要处理长文档时调 `MAX_LEN` 并实测延迟。
## 文件
```
local-embedding/
├── app.py FastAPI 服务(OpenAI 兼容)
├── download_model.py 模型检查/下载/校验(幂等,支持续传)
├── entrypoint.sh 先下模型,再起 uvicorn
├── requirements.txt 运行时依赖(无 torch)
├── Dockerfile
├── docker-compose.yml 参数都在这,端口 63002
├── README.md
└── .gitignore
```
配套的 [local-remarker](../local-remarker) 是同一套思路的重排服务,
接口为 SiliconFlow 格式。两个服务可以独立部署,也可以放在同一个 compose 里。
app.py
+214
View File
@@ -0,0 +1,214 @@
"""本地嵌入(embedding)服务。
对外提供 OpenAI 兼容的 `POST /v1/embeddings`,任何支持自定义 OpenAI 兼容端点
的框架都能直接对接,不需要改 SDK。
与 local-remarker 的区别:
* 端点是 OpenAI 标准(而非 SiliconFlow 私有格式),调用方不会拼错路径;
* 后处理是 CLS 池化 + L2 归一化(bge 系列官方做法),而不是取 logit;
* 完全不依赖 Ollama —— 模型以 ONNX 形式直接跑在 onnxruntime 上。
"""
import os
import time
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-m3")
# bge-m3 原生支持 8192 上下文,但 CPU 上长序列非常慢,默认按 512 截断
MAX_LEN = int(os.environ.get("MAX_LEN", "512"))
BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "16"))
PAD_ID = int(os.environ.get("PAD_ID", "1"))
NUM_THREADS = os.environ.get("NUM_THREADS", "") or None
app = FastAPI(title="local-embedding", 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
options.intra_op_num_threads = int(NUM_THREADS) if NUM_THREADS else (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)
return session, tokenizer
SESSION, TOKENIZER = _build_session()
INPUT_NAMES = {i.name for i in SESSION.get_inputs()}
OUTPUT_NAMES = [o.name for o in SESSION.get_outputs()]
def _pool(outputs: list[np.ndarray]) -> np.ndarray:
"""从 ONNX 输出取句向量并做 L2 归一化。
不同导出方式的第一个输出形状不一致,两种都兼容:
* [batch, hidden] —— 已池化好的句向量
* [batch, seq, hidden] —— token 级输出,取 [CLS](第 0 个 token)
bge 系列官方推荐 CLS 池化 + L2 归一化,归一化后余弦相似度与内积等价。
"""
hidden = outputs[0]
if hidden.ndim == 3:
hidden = hidden[:, 0, :] # [CLS]
norms = np.linalg.norm(hidden, axis=1, keepdims=True)
norms[norms == 0] = 1.0
return hidden / norms
def _forward(rows: list[tuple[list[int], list[int], list[int]]]) -> np.ndarray:
"""把 (ids, attention_mask, type_ids) 组成 batch 跑一次推理,返回归一化向量。"""
width = max(len(ids) for ids, _, _ in rows)
size = len(rows)
input_ids = np.full((size, width), PAD_ID, dtype=np.int64)
attention_mask = np.zeros((size, width), dtype=np.int64)
token_type_ids = np.zeros((size, width), dtype=np.int64)
for row, (ids, mask, types) in enumerate(rows):
length = len(ids)
input_ids[row, :length] = ids
attention_mask[row, :len(mask)] = mask
token_type_ids[row, :len(types)] = types
feeds = {"input_ids": input_ids, "attention_mask": attention_mask}
if "token_type_ids" in INPUT_NAMES:
feeds["token_type_ids"] = token_type_ids
return _pool(SESSION.run(None, feeds)).astype(np.float32)
def _encode(texts: list[str], token_ids: dict[int, list[int]]) -> list[tuple]:
"""统一编码入口。token_ids 里的下标表示该项已是 token id 序列。"""
rows, tokens = [], 0
for index, text in enumerate(texts):
if index in token_ids:
ids = token_ids[index][:MAX_LEN]
mask = [1] * len(ids)
types = [0] * len(ids)
else:
encoding = TOKENIZER.encode(text)
ids, mask, types = encoding.ids, encoding.attention_mask, encoding.type_ids
rows.append((ids, mask, types))
tokens += len(ids)
return rows, tokens
def _invalid(message: str) -> JSONResponse:
return JSONResponse(
{"error": {"message": message, "type": "invalid_request_error"}},
status_code=400,
)
async def _handle(request: Request) -> JSONResponse:
try:
body = await request.json()
except Exception: # noqa: BLE001
return _invalid("invalid json body")
if "input" not in body:
return _invalid("input is required")
raw = body["input"]
items = [raw] if isinstance(raw, str) else raw
if not isinstance(items, list) or not items:
return _invalid("input must be a non-empty string or array")
# 允许混排字符串与 token id 数组(OpenAI 规范允许后者)
texts, token_ids = [], {}
for position, item in enumerate(items):
if isinstance(item, str):
texts.append(item)
elif isinstance(item, list) and item and all(isinstance(x, int) for x in item):
token_ids[position] = item
texts.append("")
else:
return _invalid("input items must be strings or token id arrays")
started = time.time()
try:
vectors, tokens = [], 0
for start in range(0, len(texts), BATCH_SIZE):
chunk = texts[start:start + BATCH_SIZE]
offset = {k - start: v for k, v in token_ids.items() if start <= k < start + len(chunk)}
rows, chunk_tokens = _encode(chunk, offset)
vectors.append(_forward(rows))
tokens += chunk_tokens
matrix = np.concatenate(vectors, axis=0)
except Exception as exc: # noqa: BLE001
return JSONResponse(
{"error": {"message": str(exc), "type": "server_error"}}, status_code=500
)
return JSONResponse({
"object": "list",
"data": [
{"object": "embedding", "index": i, "embedding": matrix[i].tolist()}
for i in range(matrix.shape[0])
],
"model": body.get("model") or MODEL_NAME,
"usage": {"prompt_tokens": tokens, "total_tokens": tokens},
"_meta": {
"latency_ms": int((time.time() - started) * 1000),
"dimension": int(matrix.shape[1]),
},
})
# OpenAI 标准路径 + 常见别名
for _path in ("/v1/embeddings", "/embeddings", "/v1/embed", "/embed"):
app.post(_path)(_handle)
@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),
"output_names": OUTPUT_NAMES,
}
# 部分客户端启动时会先探测 /v1/models
@app.get("/v1/models")
def list_models() -> dict[str, Any]:
return {
"object": "list",
"data": [{
"id": MODEL_NAME,
"object": "model",
"created": int(time.time()),
"owned_by": "local",
}],
}
docker-compose.yml
+63
View File
@@ -0,0 +1,63 @@
# 本地嵌入(embedding)服务 —— 完整替代 Ollama。
#
# 参数全部直接写在本文件里,不再走 .env —— 需要调整时改这里,然后
# docker compose up -d 重建即可。
#
# 首次启动会下载约 570MB 模型权重到 ./models;
# 之后重启/重建都会先校验 sha256,通过就跳过下载,零网络流量。
services:
local-embedding:
build: .
image: local-embedding:latest
container_name: local-embedding
restart: unless-stopped
# 宿主机 63002 -> 容器 8000
ports:
- "63000:8000"
environment:
# ---- 模型存放位置(容器内)----
MODEL_DIR: /models
# ---- 模型来源与选择 ----
MODEL_REPO: Xenova/bge-m3
MODEL_REV: master
MODEL_BASE: https://modelscope.cn/models/Xenova/bge-m3/resolve/master
MODEL_NAME: bge-m3
# 同仓库可选量化档位:
# model_int8.onnx 542MB 默认,带 AVX512-VNNI 的 CPU 上最快
# model_fp16.onnx 1081MB
# model_q4f16.onnx 668MB
# model_q4.onnx 1190MB
# 更大的模型可改用 BAAI/bge-m3(原始 fp32,2.2GB,且需额外的 model.onnx_data)
ONNX_FILE: model_int8.onnx
# ---- 推理参数 ----
MAX_LEN: "512" # bge-m3 原生支持 8192,但 CPU 上长序列很慢
BATCH_SIZE: "16" # 每次推理的条数
PAD_ID: "1" # bge-m3 是 XLM-R 架构,padding 用 [PAD]
# ---- 服务 ----
WORKERS: "1" # 调大前先看内存:每个 worker 会各加载一份模型
PORT: "8000"
LOG_LEVEL: info
TZ: Asia/Shanghai
# 宿主机 ./models -> 容器 /models。指向大盘路径可避免占满系统盘,
# 例如改成 /data/models/embedding:/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: 900s
retries: 3
download_model.py
+201
View File
@@ -0,0 +1,201 @@
"""启动前检查并下载嵌入模型。
设计目标:
* 幂等 —— 模型已存在且 sha256 校验通过就立刻返回,不产生任何网络请求;
* 可续传 —— 弱网下载中断后重启容器能从断点继续;
* 可校验 —— 每个文件都带 sha256,下载后校验,不匹配则丢弃重下。
可配置项(均可用环境变量覆盖):
MODEL_DIR 模型存放目录,默认 /models
MODEL_REPO 仓库名,默认 Xenova/bge-m3
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-m3")
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, 字节数)
#
# 选 model_int8.onnx:568MB 自包含(没有外置的 .onnx_data),且带 AVX512-VNNI
# 的 CPU 上 int8 推理最快。fp32 版需要额外下载 model.onnx_data 共 2.1GB。
MANIFEST = [
(
"model_int8.onnx",
"onnx/model_int8.onnx",
"a206e10e995aa2a833924bcd725ba5dd6c3425cd34bac3cf2b5677cd2a1c51d6",
568456694,
),
(
"tokenizer.json",
"tokenizer.json",
"6710678b12670bc442b99edc952c4d996ae309a7020c1fa0096dd245c2faf790",
17082821,
),
(
"sentencepiece.bpe.model",
"sentencepiece.bpe.model",
"cfc8146abe2a0488e9e2a0c56de7952f7c11ab059eca145a0a727afce0db2865",
5069051,
),
(
"special_tokens_map.json",
"special_tokens_map.json",
"8c785abebea9ae3257b61681b4e6fd8365ceafde980c21970d001e834cf10835",
964,
),
(
"tokenizer_config.json",
"tokenizer_config.json",
"7e4c1cc848840aeccdd763458c18dd525eb0f795c992e00ebe9c28554e7db2d4",
1173,
),
]
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-embedding/1.0")
if have:
request.add_header("Range", f"bytes={have}-")
with urllib.request.urlopen(request, timeout=TIMEOUT) as resp:
# 服务端不支持 Range 时会返回 200 全量,此时从零重写
mode = "ab" if (have and resp.status == 206) 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-embedding] 模型目录: {MODEL_DIR}")
print(f"[local-embedding] 下载源 : {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-embedding] 模型已就绪,跳过下载")
return 0
total_bytes = sum(item[3] for item in pending)
print(f"[local-embedding] 需要下载 {len(pending)} 个文件,共 {human(total_bytes)}")
# 容器内磁盘常较小,给个明确提示而不是让下载器报一堆 no space
free = shutil.disk_usage(MODEL_DIR).free
if free < total_bytes * 1.1:
raise SystemExit(
f"磁盘空间不足:可用 {human(free)},需要约 {human(total_bytes)}。"
"请挂载更大的卷,或改用更小的嵌入模型。"
)
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-embedding] 模型准备完成")
return 0
if __name__ == "__main__":
sys.exit(main())
entrypoint.sh
Executable
+17
View File
@@ -0,0 +1,17 @@
#!/bin/sh
# 容器入口:先确保模型就位,再启动服务。
#
# download_model.py 是幂等的 —— 模型已存在且 sha256 校验通过时会直接返回,
# 因此正常的重启/重建容器不会产生任何下载流量;首次启动才会拉取约 570MB 权重。
set -e
echo "[entrypoint] 检查模型 ..."
python /app/download_model.py
echo "[entrypoint] 启动 embedding 服务 ..."
exec uvicorn app:app \
--host "${HOST:-0.0.0.0}" \
--port "${PORT:-8000}" \
--workers "${WORKERS:-1}" \
--log-level "${LOG_LEVEL:-info}"
requirements.txt
+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