214 lines
7.3 KiB
Python
214 lines
7.3 KiB
Python
"""本地嵌入(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",
|
|
}],
|
|
} |