Update app.py

app.py

@@ -40,43 +40,48 @@ def load_model(model_type):
     manage_resources()
     
     try:
-        # For CPU-only environment, don't use device_map
         if model_type == "summarize":
             base_model = AutoModelForSeq2SeqLM.from_pretrained(
                 "facebook/bart-large-cnn",
                 cache_dir="./models",
-                torch_dtype=torch.float32,
-                low_cpu_mem_usage=True
-            )
+                device_map=None,  # Explicitly set to None for CPU
+                torch_dtype=torch.float32
+            ).to("cpu")  # Force CPU
+            
             model = PeftModel.from_pretrained(
                 base_model, 
                 "pendar02/results",
-                torch_dtype=torch.float32
-            )
+                device_map=None,  # Explicitly set to None for CPU
+                torch_dtype=torch.float32,
+                is_trainable=False  # Set to inference mode
+            ).to("cpu")  # Force CPU
+            
             tokenizer = AutoTokenizer.from_pretrained(
                 "facebook/bart-large-cnn",
                 cache_dir="./models"
             )
         else:  # question_focused
-            base_model = AutoModelForSeq2SeqLation_model = AutoModelForSeq2SeqLM.from_pretrained(
+            base_model = AutoModelForSeq2SeqLM.from_pretrained(
                 "GanjinZero/biobart-base",
                 cache_dir="./models",
-                torch_dtype=torch.float32,
-                low_cpu_mem_usage=True
-            )
+                device_map=None,  # Explicitly set to None for CPU
+                torch_dtype=torch.float32
+            ).to("cpu")  # Force CPU
+            
             model = PeftModel.from_pretrained(
                 base_model, 
                 "pendar02/biobart-finetune",
-                torch_dtype=torch.float32
-            )
+                device_map=None,  # Explicitly set to None for CPU
+                torch_dtype=torch.float32,
+                is_trainable=False  # Set to inference mode
+            ).to("cpu")  # Force CPU
+            
             tokenizer = AutoTokenizer.from_pretrained(
                 "GanjinZero/biobart-base",
                 cache_dir="./models"
             )
         
-        # Ensure model is on CPU
-        model = model.cpu()
-        model.eval()
+        model.eval()  # Set to evaluation mode
         return model, tokenizer
     except Exception as e:
         st.error(f"Error loading model: {str(e)}")