Update app.py

app.py

@@ -9,6 +9,9 @@ from urllib.parse import urlparse
 import time
 import os
 
+from safetensors.torch import load_file
+import json
+
 # --- import your architecture ---
 # Make sure this file is in the repo (e.g., models/deberta_lstm_classifier.py)
 # and update the import path accordingly.
@@ -20,36 +23,71 @@ from llm_client import LLMClient
 
 # --------- Config ----------
 REPO_ID = "dungeon29/deberta-lstm-detect-phishing"       # HF repo that holds the checkpoint
-CKPT_NAME = "pytorch_model.bin"               # the .pt file name
+CKPT_NAME = "model.safetensors"               # the .safetensors file name
 MODEL_NAME = "microsoft/deberta-base"         # base tokenizer/backbone
 LABELS = ["benign", "phishing"]               # adjust to your classes
 
-# If your checkpoint contains hyperparams, you can fetch them like:
-# checkpoint.get("config") or checkpoint.get("model_args")
-# and pass into DeBERTaLSTMClassifier(**model_args)
-
 # --------- Load model/tokenizer once (global) ----------
 device = "cuda" if torch.cuda.is_available() else "cpu"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-ckpt_path = hf_hub_download(repo_id=REPO_ID, filename=CKPT_NAME)
-checkpoint = torch.load(ckpt_path, map_location=device)
+# Check if checkpoint exists locally, otherwise download from HF
+if os.path.exists(CKPT_NAME):
+    print(f"📂 Found local checkpoint: {CKPT_NAME}")
+    ckpt_path = CKPT_NAME
+else:
+    print(f"⬇️ Downloading checkpoint {CKPT_NAME} from HF Hub...")
+    try:
+        ckpt_path = hf_hub_download(repo_id=REPO_ID, filename=CKPT_NAME)
+    except Exception as e:
+        print(f"⚠️ Could not download from HF: {e}")
+        # Fallback to pytorch_model.bin if the new name fails
+        print("🔄 Trying fallback to pytorch_model.bin...")
+        ckpt_path = hf_hub_download(repo_id=REPO_ID, filename="pytorch_model.bin")
+
+# Load weights based on file extension
+if ckpt_path.endswith(".safetensors"):
+    print("📦 Loading weights from safetensors...")
+    state_dict = load_file(ckpt_path, device=device)
+    
+    # Try to load config.json for model_args
+    try:
+        if os.path.exists("config.json"):
+            config_path = "config.json"
+        else:
+            config_path = hf_hub_download(repo_id=REPO_ID, filename="config.json")
+            
+        with open(config_path, "r") as f:
+            config = json.load(f)
+            # Extract model_args from config if they exist, otherwise use defaults
+            # Assuming config might have custom keys or we just use defaults
+            model_args = config.get("model_args", {})
+    except Exception as e:
+        print(f"⚠️ Could not load config.json: {e}. Using default model args.")
+        model_args = {}
+        
+else:
+    # Legacy loading for .bin/.pt
+    print("📦 Loading weights from torch checkpoint...")
+    checkpoint = torch.load(ckpt_path, map_location=device)
+    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
+        state_dict = checkpoint["model_state_dict"]
+        model_args = checkpoint.get("model_args", {})
+    elif isinstance(checkpoint, dict):
+        state_dict = checkpoint
+        model_args = checkpoint.get("model_args", {})
+    else:
+        state_dict = checkpoint
+        model_args = {}
 
-# If you saved hyperparams in the checkpoint, use them:
-model_args = checkpoint.get("model_args", {})  # e.g., {"lstm_hidden":256, "num_labels":2, ...}
+# Initialize model
 model = DeBERTaLSTMClassifier(**model_args)
 
-# Load weights
+# Load state dict
 try:
-    state_dict = torch.load(ckpt_path, map_location=device)
-    
-    # Xử lý nếu file lưu dạng checkpoint đầy đủ (có key "model_state_dict")
-    if "model_state_dict" in state_dict:
-        state_dict = state_dict["model_state_dict"]
-
     model.load_state_dict(state_dict, strict=False)
     
-    # Kiểm tra layer attention
+    # Check attention layer
     if hasattr(model, 'attention') and 'attention.weight' not in state_dict:
          print("⚠️ Loaded model without attention layer, using newly initialized attention weights")
     else: