85 lines
2.4 KiB
Python
85 lines
2.4 KiB
Python
#!/usr/bin/env python3
|
|
"""轻量 bge-m3 embedding server — OpenAI 兼容 /v1/embeddings 接口"""
|
|
import os, sys, json, time
|
|
from typing import List
|
|
|
|
# 环境变量
|
|
HOST = os.environ.get("EMBED_HOST", "0.0.0.0")
|
|
PORT = int(os.environ.get("EMBED_PORT", "8000"))
|
|
MODEL = os.environ.get("EMBED_MODEL", "BAAI/bge-m3")
|
|
DEVICE = os.environ.get("EMBED_DEVICE", "cpu")
|
|
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.responses import JSONResponse
|
|
from pydantic import BaseModel, Field
|
|
import uvicorn
|
|
|
|
app = FastAPI(title="bge-m3 Embedding Server", version="1.0")
|
|
|
|
class EmbedRequest(BaseModel):
|
|
input: str | List[str]
|
|
model: str = MODEL
|
|
|
|
class EmbeddingObject(BaseModel):
|
|
object: str = "embedding"
|
|
index: int
|
|
embedding: List[float]
|
|
|
|
class UsageInfo(BaseModel):
|
|
prompt_tokens: int
|
|
total_tokens: int
|
|
|
|
class EmbedResponse(BaseModel):
|
|
object: str = "list"
|
|
data: List[EmbeddingObject]
|
|
model: str = MODEL
|
|
usage: UsageInfo
|
|
|
|
# 延迟加载
|
|
model_pipe = None
|
|
|
|
def load_model():
|
|
global model_pipe
|
|
if model_pipe is not None:
|
|
return
|
|
print(f"[embed] Loading {MODEL} on {DEVICE}...", flush=True)
|
|
t0 = time.time()
|
|
from sentence_transformers import SentenceTransformer
|
|
model_pipe = SentenceTransformer(MODEL, device=DEVICE)
|
|
elapsed = time.time() - t0
|
|
print(f"[embed] Loaded in {elapsed:.1f}s. Dimension: {model_pipe.get_sentence_embedding_dimension()}", flush=True)
|
|
|
|
@app.on_event("startup")
|
|
async def startup():
|
|
load_model()
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {"status": "ok", "model": MODEL, "device": DEVICE,
|
|
"loaded": model_pipe is not None}
|
|
|
|
@app.post("/v1/embeddings")
|
|
async def embed(req: EmbedRequest):
|
|
t0 = time.time()
|
|
texts = req.input if isinstance(req.input, list) else [req.input]
|
|
|
|
# bge-m3 需要加 prefix
|
|
prefixed = [f"为这个句子生成表示以用于检索相关文章：{t}" for t in texts]
|
|
|
|
vecs = model_pipe.encode(prefixed, normalize_embeddings=True,
|
|
show_progress_bar=False)
|
|
vecs = vecs.tolist()
|
|
|
|
total_tokens = sum(len(t) for t in texts)
|
|
elapsed = time.time() - t0
|
|
|
|
data = [
|
|
EmbeddingObject(index=i, embedding=vecs[i])
|
|
for i in range(len(vecs))
|
|
]
|
|
return EmbedResponse(data=data, usage=UsageInfo(prompt_tokens=total_tokens, total_tokens=total_tokens))
|
|
|
|
if __name__ == "__main__":
|
|
load_model()
|
|
uvicorn.run(app, host=HOST, port=PORT, log_level="info")
|