Update app.py

app.py

@@ -23,7 +23,7 @@ from llm_client import LLMClient
 
 # --------- Config ----------
 REPO_ID = "dungeon29/deberta-lstm-detect-phishing"       # HF repo that holds the checkpoint
-FILE_NAME = "model.safetensors"               # the .safetensors 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
 
@@ -31,26 +31,73 @@ LABELS = ["benign", "phishing"]               # adjust to your classes
 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 = {}
 
 # Initialize model
-model = DeBERTaLSTMClassifier()
-model.to(device)
+model = DeBERTaLSTMClassifier(**model_args)
 
 # Load state dict
 try:
-    model_path = hf_hub_download(repo_id=REPO_ID, filename=FILE_NAME)
-
-    state_dict = load_file(model_path)
-
     model.load_state_dict(state_dict, strict=False)
     
-    print("✅ Load weights successfully!")
+    # 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:
+         print("✅ Load weights successfully!")
          
 except Exception as e:
     print(f"❌ Error when loading weights: {e}")
     raise e
 
-model.eval()
+model.to(device).eval()
 
 # --------- Initialize RAG & LLM ----------
 print("Initializing RAG Engine (LangChain)...")