transformer_api / app.py
JohnKouf's picture
Update app.py
cbd145d verified
from fastapi import FastAPI
from pydantic import BaseModel
from Similarity import Similarity
app = FastAPI()
similarity_model = Similarity()
class TextRequest(BaseModel):
text: str
@app.post("/embed")
def embed_text_endpoint(request: TextRequest):
try:
embedding = similarity_model.embed_text(request.text)
# Convert tensor to list so it can be JSON serialized
embedding_list = embedding.cpu().tolist() if hasattr(embedding, 'cpu') else embedding.tolist()
return {"embedding": embedding_list}
except Exception as e:
return {"error": str(e)}