Sid-the-sloth's picture
Initial model sever setup
b5db444
raw
history blame
1.17 kB
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
import logging
logging.basicConfig(level=logging.INFO)
logger=logging.getLogger(__name__)
logger.info("Server Starting")
try:
logger.info("Loading model")
model=SentenceTransformer("Sid-the-sloth/leetcode_unixcoder_final")
logger.info("Model Loaded")
except:
logger.error("Failed to load Model")
model=None
app=FastAPI()
#Req and Response Pydantic models
class EmbedRequest(BaseModel):
text : str
class EmbedResponse(BaseModel):
embedding: list[float]
@app.get("/")
def root_status():
return {"status":"ok","model":model is not None}
@app.post("/embed",response_model=EmbedResponse)
def get_embedding(request: EmbedRequest):
if model is None:
HTTPException(status_code=503,detail="Model could not be loaded")
try:
embedding=model.encode(request.text).tolist()
return EmbedResponse(embedding=embedding)
except Exception as e:
logger.error("Error during embedding generation %s",e)
return HTTPException(status_code=500,detail="Error generating embeddings")