Update app.py

app.py

@@ -324,44 +324,34 @@ if not check_password():
 
 # ========== Load Model (Cached) ==========
 @st.cache_resource
-
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModel.from_pretrained(model_name).to(device)
-
-    if os.path.exists(MODEL_PATH):
-        try:
-            ckpt = torch.load(MODEL_PATH, map_location=device)
-            
-            # Case 1: checkpoint is a dict with model_state
-            if isinstance(ckpt, dict) and "model_state" in ckpt:
-                model.load_state_dict(ckpt["model_state"], strict=False)
-                print("Loaded checkpoint with model_state ✅")
-            
-            # Case 2: checkpoint is a state_dict directly
-            elif isinstance(ckpt, dict):
-                model.load_state_dict(ckpt, strict=False)
-                print("Loaded raw state_dict ✅")
-            
-            # Case 3: checkpoint is a full model object
-            else:
-                print("Checkpoint is a full model object — using directly ✅")
-                model = ckpt.to(device)
-
-        except Exception as e:
-            print(f"⚠️ Could not load checkpoint ({e}). Using pretrained TinyBERT model.")
-    
-    model.eval()
-
-    chatbot = DefenseQAChatbot(
-        model=model,
-        tokenizer=tokenizer,
-        device=device,
-        embeddings_path=EMBEDDINGS_PATH
-    )
-
-    return chatbot
-
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModel.from_pretrained(model_name).to(device)
+        
+        if os.path.exists(MODEL_PATH):
+            # Try loading with weights_only=False
+            try:
+                ckpt = torch.load(MODEL_PATH, map_location=device, weights_only=False)
+            except:
+                # Fallback to old method
+                ckpt = torch.load(MODEL_PATH, map_location=device)
+                
+            model.load_state_dict(ckpt["model_state"])
+            model.eval()
+        
+        chatbot = DefenseQAChatbot(
+            model=model,
+            tokenizer=tokenizer,
+            device=device,
+            embeddings_path=EMBEDDINGS_PATH
+        )
+        
+        return chatbot
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        st.stop()
+        
 chatbot = load_model()
 
 # ========== Initialize Session State ==========