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
-CKPT_NAME = "model.safetensors"               # the .safetensors file name
+FILE_NAME = "model.safetensors"               # the .safetensors file name
 MODEL_NAME = "microsoft/deberta-base"         # base tokenizer/backbone
 LABELS = ["benign", "phishing"]               # adjust to your classes
 
@@ -31,73 +31,26 @@ 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_args)
+model = DeBERTaLSTMClassifier()
+model.to(device)
 
 # 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)
     
-    # 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!")
+    print("✅ Load weights successfully!")
          
 except Exception as e:
     print(f"❌ Error when loading weights: {e}")
     raise e
 
-model.to(device).eval()
+model.eval()
 
 # --------- Initialize RAG & LLM ----------
 print("Initializing RAG Engine (LangChain)...")