Update app.py

app.py

@@ -1,41 +1,66 @@
-import os
-
-# --- FIX CACHE PERMISSIONS ---
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/transformers_cache"
-os.environ["HF_HOME"] = "/tmp/hf_home"  # optional for Hugging Face hub files
-
 from fastapi import FastAPI, Form
+from fastapi.responses import JSONResponse
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 from twilio.twiml.messaging_response import MessagingResponse
+import os
+
+# -----------------------------
+# Environment-safe cache path
+# -----------------------------
+HF_CACHE_DIR = os.getenv("HF_HOME", "/tmp/hf_cache")
 
-# --- Load model ---
+# -----------------------------
+# Load model from Hugging Face
+# -----------------------------
 model_id = "ST-THOMAS-OF-AQUINAS/SCAM"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForSequenceClassification.from_pretrained(model_id)
+
+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: "author1", 1: "author2"}
 
-# --- FastAPI app ---
-app = FastAPI(title="Scam Detector API with Twilio")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
-# --- Helper prediction function ---
+# -----------------------------
+# Helper function
+# -----------------------------
 def predict_author(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()
-    predicted_author = label_map[pred]
+
+    predicted_author = label_map.get(pred, "unknown")
     return predicted_author, round(confidence * 100, 2)
 
-# --- Twilio WhatsApp webhook ---
+# -----------------------------
+# FastAPI app
+# -----------------------------
+app = FastAPI(title="Scam Detector API with Twilio")
+
+# Health-check route
+@app.get("/")
+async def health_check():
+    return {"status": "✅ API is running"}
+
+# Simple GET test
+@app.get("/predict")
+async def get_predict(text: str):
+    author, confidence = predict_author(text)
+    return {"prediction": author, "confidence": confidence}
+
+# Twilio WhatsApp POST
 @app.post("/whatsapp")
 async def whatsapp_reply(Body: str = Form(...)):
     resp = MessagingResponse()
-    
+
     if Body.strip():
         author, confidence = predict_author(Body)
         reply = f"Prediction: {author}\nConfidence: {confidence}%"
@@ -43,10 +68,4 @@ async def whatsapp_reply(Body: str = Form(...)):
         reply = "⚠️ No text detected."
 
     resp.message(reply)
-    return str(resp)
-
-# --- Simple test endpoint ---
-@app.get("/predict")
-async def predict(text: str):
-    author, confidence = predict_author(text)
-    return {"prediction": author, "confidence": confidence}
+    return JSONResponse(content=str(resp))