Update app.py

app.py

@@ -11,39 +11,37 @@ import os
 HF_CACHE_DIR = os.getenv("HF_HOME", "/tmp/hf_cache")
 
 # -----------------------------
-# Load model from Hugging Face
+# Load regression model from Hugging Face
 # -----------------------------
-model_id = "ST-THOMAS-OF-AQUINAS/SCAM"
+model_id = "ST-THOMAS-OF-AQUINAS/impersonation-bart"
 
 tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=HF_CACHE_DIR)
 model = AutoModelForSequenceClassification.from_pretrained(model_id, cache_dir=HF_CACHE_DIR)
 model.eval()
 
-label_map = {0: "MAXWELL KURIA", 1: "Keliv Kuria"}
-
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
 # -----------------------------
 # Helper function
 # -----------------------------
-def predict_author(text: str):
+def predict_score(text: str):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     inputs = {k: v.to(device) for k, v in inputs.items()}
 
     with torch.no_grad():
         outputs = model(**inputs)
-        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-        pred = torch.argmax(probs, dim=1).item()
-        confidence = probs[0][pred].item()
+        # For regression, logits is shape [batch, 1]
+        score = outputs.logits.squeeze().item()
 
-    predicted_author = label_map.get(pred, "unknown")
-    return predicted_author, round(confidence * 100, 2)
+    # Clamp between 0 and 1 (just in case)
+    score = max(0.0, min(1.0, score))
+    return round(score, 3)
 
 # -----------------------------
 # FastAPI app
 # -----------------------------
-app = FastAPI(title="Scam Detector API with Twilio")
+app = FastAPI(title="Impersonation Detector API with Twilio")
 
 # Health-check route
 @app.get("/")
@@ -53,8 +51,8 @@ async def health_check():
 # Simple GET test
 @app.get("/predict")
 async def get_predict(text: str):
-    author, confidence = predict_author(text)
-    return {"prediction": author, "confidence": confidence}
+    score = predict_score(text)
+    return {"impersonation_score": score}
 
 # -----------------------------
 # Twilio WhatsApp POST
@@ -64,8 +62,8 @@ async def whatsapp_reply(Body: str = Form(...)):
     resp = MessagingResponse()
 
     if Body.strip():
-        author, confidence = predict_author(Body)
-        reply = f"Prediction: {author}\nConfidence: {confidence}%"
+        score = predict_score(Body)
+        reply = f"Impersonation Score: {score}\n(0.0 = genuine, 1.0 = impersonation)"
     else:
         reply = "⚠️ No text detected."