Update app.py

fixing the model loading bug

app.py

@@ -94,20 +94,17 @@ class ModelWrapper:
         self.tokenizer.pad_token_id = self.tokenizer.pad_token_id or self.tokenizer.eos_token_id
 
         print(f"Loading model: {model_name}...")
-        if "8b" in model_name.lower():
-            config = AutoConfig.from_pretrained(model_name, torch_dtype=torch.bfloat16)
-            with init_empty_weights():
-                model_empty = AutoModelForCausalLM.from_config(config)
-
-            self.model = load_checkpoint_and_dispatch(
-                model_empty,
-                model_name,
-                device_map="auto",
-                offload_folder="offload",
-            ).eval()
-        else:
-            self.model = AutoModelForCausalLM.from_pretrained(
-                model_name, device_map="auto", torch_dtype=torch.bfloat16).eval()
+        
+        # We can now use the same, simpler loading logic for all models.
+        # The `from_pretrained` method will handle downloading from the Hub
+        # and applying the device_map.
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map="auto",
+            torch_dtype=torch.bfloat16,
+            offload_folder="offload"  # Keep this for memory management
+        ).eval()
+        
         print(f"Model {model_name} loaded successfully.")
 
     def get_message_template(self, system_content=None, user_content=None, assistant_content=None):