first commit
This commit is contained in:
@@ -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",
|
||||
}],
|
||||
}
|
||||
Reference in New Issue
Block a user