deploy

routes/sms_router.py

@@ -44,4 +44,10 @@ async def calculate_similarity(similarity_request: SimilarityRequest):
 @router.post("/predict_label/", response_model=PredictionResponse)
 async def predict_sms_label(prediction_request: PredictionRequest):
     label, probability = predict_label(prediction_request.message)
+    
+    # Handle prediction errors
+    if label == "Error":
+        raise HTTPException(status_code=500, detail="Prediction failed.")
+    
     return PredictionResponse(label=label, probability=probability)
+

service/prediction_service.py

@@ -23,7 +23,30 @@ else:
 
 # Define predict_label function
 def predict_label(text):
-    # Implement your prediction logic using the model
-    embeddings = model.encode([text])
-    # For demonstration, let's return the embedding shape
-    return embeddings.shape
+    try:
+        # Ensure input is a list for the model
+        if not isinstance(text, list):
+            text = [text]
+        
+        # Generate embeddings
+        embeddings = model.encode(text)
+        
+        # Ensure embeddings are in the correct shape
+        if len(embeddings) == 0:
+            raise ValueError("No embeddings generated.")
+        
+        # Predict using the logistic regression model
+        prediction = clf.predict(embeddings)
+        probability = clf.predict_proba(embeddings).max()
+        
+        # Convert label to string ("0" or "1")
+        label = str(prediction[0])
+        
+        # Return label and probability
+        return label, float(probability)
+
+    except Exception as e:
+        # Log the exception for debugging
+        print(f"Error in predict_label: {e}")
+        return "Error", 0.0
+