51 lines
1.2 KiB
Python
51 lines
1.2 KiB
Python
"""
|
|
E5-small-v2 embedding server — OpenAI-compatible /v1/embeddings endpoint.
|
|
Minimal, CPU-only, ~100MB RAM.
|
|
"""
|
|
from fastapi import FastAPI
|
|
from pydantic import BaseModel
|
|
from sentence_transformers import SentenceTransformer
|
|
import numpy as np
|
|
|
|
app = FastAPI()
|
|
model = SentenceTransformer("intfloat/e5-small-v2", device="cpu")
|
|
|
|
|
|
class EmbedRequest(BaseModel):
|
|
model: str = "intfloat/e5-small-v2"
|
|
input: list[str]
|
|
|
|
|
|
class EmbedData(BaseModel):
|
|
object: str = "embedding"
|
|
index: int
|
|
embedding: list[float]
|
|
|
|
|
|
class EmbedResponse(BaseModel):
|
|
object: str = "list"
|
|
data: list[EmbedData]
|
|
model: str
|
|
|
|
|
|
@app.post("/v1/embeddings")
|
|
def embed(req: EmbedRequest):
|
|
# E5 requires "query: " or "passage: " prefix
|
|
prefixed = [f"passage: {t}" if "\n" not in t else t for t in req.input]
|
|
embs = model.encode(prefixed, normalize_embeddings=True, show_progress_bar=False)
|
|
data = [
|
|
EmbedData(index=i, embedding=emb.tolist())
|
|
for i, emb in enumerate(embs)
|
|
]
|
|
return EmbedResponse(data=data, model=req.model)
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok"}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=8081)
|