gemini-cook / app.py
harish3304's picture
Update app.py
6aa465d verified
raw
history blame contribute delete
No virus
624 Bytes
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)