commit 136163ad888b215693955786fc3b2c11eab49987 Author: Jack Date: Fri Oct 23 18:08:46 2026 +0800 first commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1a5d2b4 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +# 模型权重目录(约 570MB,不进版本库) +models/ +*.onnx +*.part + +__pycache__/ +*.pyc \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..b7a682b --- /dev/null +++ b/Dockerfile @@ -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"] \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..197bb3f --- /dev/null +++ b/README.md @@ -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 里。 \ No newline at end of file diff --git a/app.py b/app.py new file mode 100644 index 0000000..1dd60a6 --- /dev/null +++ b/app.py @@ -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", + }], + } \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..e3cbeb8 --- /dev/null +++ b/docker-compose.yml @@ -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 \ No newline at end of file diff --git a/download_model.py b/download_model.py new file mode 100644 index 0000000..68700d0 --- /dev/null +++ b/download_model.py @@ -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()) \ No newline at end of file diff --git a/entrypoint.sh b/entrypoint.sh new file mode 100755 index 0000000..26e5538 --- /dev/null +++ b/entrypoint.sh @@ -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}" \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..5f7d3d7 --- /dev/null +++ b/requirements.txt @@ -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 \ No newline at end of file