change

main.py

@@ -1,47 +1,28 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI
 from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer
-import numpy as np
 
-# Initialize the FastAPI app
-app = FastAPI()
+import os
+os.environ["HF_HOME"] = "/tmp/huggingface"
 
-# Load the pre-trained SentenceTransformer model from Hugging Face
+# Initialize FastAPI app
+app = FastAPI()
 
-model=SentenceTransformer("all-MiniLM-L6-V2")
+# Load pretrained model from Hugging Face (instead of hf_hub_download)
+model = SentenceTransformer("all-MiniLM-L6-v2")  # Updated model
 
-# Define the request body schema
-class TextInput(BaseModel):
+# Define request structure
+class TextRequest(BaseModel):
     text: str
 
-# Home route
-@app.get("/")
-async def home():
-    return {"message": "Welcome to embedding SMS API, use /docs to post SMS text and get dimensions"}
-
-# Define the API endpoint
-@app.post("/embed")
-async def generate_embedding(text_input: TextInput):
-    """
-    Generate a 768-dimensional embedding for the input text.
-    Returns the embedding in a structured format with rounded values.
-    """
-    try:
-        # Generate the embedding
-        embedding = model.encode(text_input.text, convert_to_tensor=True).cpu().numpy()
-
-        # Round embedding values to 2 decimal places
-        rounded_embedding = np.round(embedding, decimals=2).tolist()
-
-        # Get the number of dimensions
-        dimensions = len(rounded_embedding)
+# Define response structure
+class EmbeddingResponse(BaseModel):
+    dimensions: int
+    embedding: list[float]
 
-        # Return structured response
-        return {
-            "dimensions": dimensions,
-            "embeddings": [rounded_embedding]  # Wrap the embedding inside a list
-        }
-    except Exception as e:
-        # Handle any errors
-        raise HTTPException(status_code=500, detail=str(e))
+# Endpoint to get text embedding
+@app.post("/get_embedding", response_model=EmbeddingResponse)
+async def get_embedding(request: TextRequest):
+    embedding = model.encode([request.text])[0]  # Generate embedding
+    return {"dimensions": len(embedding), "embedding": embedding.tolist()}