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)