Update app.py

app.py

@@ -9,9 +9,6 @@ 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.
@@ -23,71 +20,36 @@ from llm_client import LLMClient
 
 # --------- Config ----------
 REPO_ID = "dungeon29/deberta-lstm-detect-phishing"       # HF repo that holds the checkpoint
-CKPT_NAME = "model.safetensors"               # the .safetensors file name
+CKPT_NAME = "pytorch_model.bin"               # the .pt 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)
 
-# 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 = {}
+ckpt_path = hf_hub_download(repo_id=REPO_ID, filename=CKPT_NAME)
+checkpoint = torch.load(ckpt_path, map_location=device)
 
-# Initialize model
+# If you saved hyperparams in the checkpoint, use them:
+model_args = checkpoint.get("model_args", {})  # e.g., {"lstm_hidden":256, "num_labels":2, ...}
 model = DeBERTaLSTMClassifier(**model_args)
 
-# Load state dict
+# Load weights
 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)
     
-    # Check attention layer
+    # Kiểm tra layer attention
     if hasattr(model, 'attention') and 'attention.weight' not in state_dict:
          print("⚠️ Loaded model without attention layer, using newly initialized attention weights")
     else: