Update app.py

app.py

@@ -27,24 +27,47 @@ def load_model(model_type):
     """Load appropriate model based on type"""
     try:
         if model_type == "summarize":
-            base_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
-            # Try loading with ignore_mismatched_sizes=True and other safety parameters
+            base_model = AutoModelForSeq2SeqLM.from_pretrained(
+                "facebook/bart-large-cnn",
+                cache_dir="./models"
+            )
+            # Load scientific lay summarizer model
             model = PeftModel.from_pretrained(
                 base_model, 
                 "pendar02/results",
-                is_trainable=False,
-                config_kwargs={"inference_mode": True}
+                load_in_8bit=False,
+                device_map="auto",
+                torch_dtype=torch.float32
+            )
+            tokenizer = AutoTokenizer.from_pretrained(
+                "facebook/bart-large-cnn",
+                cache_dir="./models"
             )
-            tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
         else:  # question_focused
-            base_model = AutoModelForSeq2SeqLM.from_pretrained("GanjinZero/biobart-base")
+            base_model = AutoModelForSeq2SeqLM.from_pretrained(
+                "GanjinZero/biobart-base",
+                cache_dir="./models"
+            )
+            # Load biobart fine-tuned model
             model = PeftModel.from_pretrained(
                 base_model, 
                 "pendar02/biobart-finetune",
-                is_trainable=False,
-                config_kwargs={"inference_mode": True}
+                load_in_8bit=False,
+                device_map="auto",
+                torch_dtype=torch.float32
             )
-            tokenizer = AutoTokenizer.from_pretrained("GanjinZero/biobart-base")
+            tokenizer = AutoTokenizer.from_pretrained(
+                "GanjinZero/biobart-base",
+                cache_dir="./models"
+            )
+        
+        # Ensure model is in evaluation mode
+        model.eval()
+        return model, tokenizer
+        
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        raise
         
         # Ensure model is in evaluation mode
         model.eval()