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


class Query(BaseModel):
    sentence: str

class Response(BaseModel):
    embedding: List[float]


model = SentenceTransformer('Alibaba-NLP/gte-large-en-v1.5', trust_remote_code=True)

app = FastAPI()


@app.get("/")
def read_root():
    return {"Hello": "World!"}

@app.post("/embeddings", response_model=Response)
def get_embedding(query: Query) -> Response:
    sentence = query.sentence
    embedding = model.encode(sentence).tolist()
    
    return Response(embedding=embedding)