File size: 1,580 Bytes
b5db444 649b044 b5db444 df5822b 649b044 5e730ed df5822b 5e730ed 649b044 df5822b b5db444 ea47ed4 b5db444 649b044 b5db444 df5822b 649b044 b5db444 649b044 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoModel
from sentence_transformers import SentenceTransformer, models
import logging
logging.basicConfig(level=logging.INFO)
logger=logging.getLogger(__name__)
logger.info("Server Starting")
try:
logger.info("Loading model")
MODEL_NAME = "Sid-the-sloth/leetcode_unixcoder_final"
device="cpu"
# tokenizer=AutoTokenizer.from_pretrained(MODEL_NAME)
# embedding_model=AutoModel.from_pretrained(MODEL_NAME)
# pooling_model=models.Pooling(embedding_model.get_word_embedding_dimension())
model=SentenceTransformer(MODEL_NAME,
# modules=[embedding_model,pooling_model],
device=device
)
logger.info("Model Loaded")
except Exception as e:
logger.error("Failed to load Model %s",e)
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:
raise 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)
raise HTTPException(status_code=500,detail="Error generating embeddings") |