Embedding / app.py
Bui Vu Duc Nghia
fix: add necessary to readme, and trust remote code to app
f6c8f41
raw
history blame contribute delete
840 Bytes
from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
import uvicorn
MODEL_NAME = "Alibaba-NLP/gte-multilingual-base"
app = FastAPI(title="Text Embedding API")
model = SentenceTransformer(MODEL_NAME, trust_remote_code=True)
class EmbedRequest(BaseModel):
text: str
class EmbedResponse(BaseModel):
embedding: list[float]
dim: int
model: str
@app.post("/embed", response_model=EmbedResponse)
def embed(req: EmbedRequest):
embedding = model.encode(req.text, normalize_embeddings=True)
return {
"embedding": embedding.tolist(),
"dim": len(embedding),
"model": MODEL_NAME
}
@app.get("/")
def health():
return {"status": "ok", "model": MODEL_NAME}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)