Update app.py

app.py

@@ -1,38 +1,61 @@
 from fastapi import FastAPI, Form
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 from twilio.twiml.messaging_response import MessagingResponse
-from transformers import pipeline
+import os
 
-# --- Hugging Face Inference API setup ---
-HF_MODEL_ID = "ST-THOMAS-OF-AQUINAS/SCAM"
+# --- Fix cache permission issue by setting cache directory ---
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/transformers_cache"
 
+# --- Load model ---
+model_id = "ST-THOMAS-OF-AQUINAS/SCAM"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForSequenceClassification.from_pretrained(model_id)
+model.eval()
 
-classifier = pipeline(
-    "text-classification",
-    model=HF_MODEL_ID,
-  
-)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# --- Label map (update according to your trained model) ---
+label_map = {0: "author1", 1: "author2"}
 
 # --- FastAPI app ---
 app = FastAPI(title="Scam Detector API with Twilio")
 
-# --- Prediction function ---
+# --- Helper prediction function ---
 def predict_author(text: str):
-    if not text.strip():
-        return "No text", 0.0
-    result = classifier(text)[0]  # {'label': 'author1', 'score': 0.95}
-    return result['label'], round(result['score'] * 100, 2)
+    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()
+
+    predicted_author = label_map[pred]
+    return predicted_author, round(confidence * 100, 2)
 
 # --- Twilio WhatsApp webhook ---
 @app.post("/whatsapp")
 async def whatsapp_reply(Body: str = Form(...)):
     resp = MessagingResponse()
-    author, confidence = predict_author(Body)
-    reply = f"Prediction: {author}\nConfidence: {confidence}%"
+    if Body.strip():
+        author, confidence = predict_author(Body)
+        reply = f"Prediction: {author}\nConfidence: {confidence}%"
+    else:
+        reply = "⚠️ No text detected."
     resp.message(reply)
     return str(resp)
 
-# --- Simple GET test endpoint ---
+# --- Simple test endpoint ---
 @app.get("/predict")
 async def predict(text: str):
     author, confidence = predict_author(text)
     return {"prediction": author, "confidence": confidence}
+
+# --- Run locally (optional) ---
+if __name__ == "__main__":
+    import uvicorn
+    port = int(os.environ.get("PORT", 8000))
+    uvicorn.run(app, host="0.0.0.0", port=port)