loaded model using AutoModel and manual pooling

app.py

@@ -1,6 +1,8 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModel
+import torch
+from torch.nn.functional import F
 import logging
 
 logging.basicConfig(level=logging.INFO)
@@ -9,11 +11,19 @@ logger=logging.getLogger(__name__)
 logger.info("Server Starting")
 try:
     logger.info("Loading model")
-    model=SentenceTransformer("Sid-the-sloth/leetcode_unixcoder_final")
+    MODEL_NAME = "Sid-the-sloth/leetcode_unixcoder_final"
+    device="cpu"
+    tokenizer=AutoTokenizer.from_pretrained(MODEL_NAME)
+    model=AutoModel.from_pretrained(MODEL_NAME)
+
+    model.to(device)
+    model.eval()
+
     logger.info("Model Loaded")
 except:
     logger.error("Failed to load Model")
     model=None
+    tokenizer=None
 
 app=FastAPI()
 
@@ -24,16 +34,33 @@ class EmbedRequest(BaseModel):
 class EmbedResponse(BaseModel):
     embedding: list[float]
 
+def mean_pooling(model_output, attention_mask):
+    """
+    Performs mean pooling on the last hidden state of the model.
+    This turns token-level embeddings into a single sentence-level embedding.
+    """
+    token_embeddings = model_output.last_hidden_state
+    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+
+
 @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")
+    if not model or not tokenizer:
+        HTTPException(status_code=503,detail="/Tokenizer could not be loaded")
     try:
+        encoded_input = tokenizer(request.text, padding=True, truncation=True, return_tensors='pt').to(device)
+        model_output = model(**encoded_input)
         embedding=model.encode(request.text).tolist()
+        sentence_embedding = mean_pooling(model_output, encoded_input['attention_mask'])
+        
+        normalized_embedding = F.normalize(sentence_embedding, p=2, dim=1)
+        
+        embedding_list = normalized_embedding[0].tolist()
         return EmbedResponse(embedding=embedding)
     except Exception as e:
         logger.error("Error during embedding generation %s",e)