from fastapi import FastAPI
from pydantic import BaseModel
from typing import List
from sentence_transformers import SentenceTransformer

app = FastAPI()

model = SentenceTransformer('sentence-transformers/gtr-t5-base')

class TextBatch(BaseModel):
    texts: List[str]

@app.post("/embed")
async def get_embeddings(batch: TextBatch):
    embeddings = model.encode(batch.texts)
    return {"embeddings": embeddings.tolist()}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=2000)