Update app.py

app.py

@@ -38,15 +38,15 @@ def load_model(model_type):
             base_model = AutoModelForSeq2SeqLM.from_pretrained(
                 "facebook/bart-large-cnn",
                 cache_dir="./models",
-                low_cpu_mem_usage=True,
-                device_map={"": device}
-            )
+                torch_dtype=torch.float32
+            ).to(device)
+            
             model = PeftModel.from_pretrained(
                 base_model, 
                 "pendar02/results",
-                device_map={"": device},
-                torch_dtype=torch.float32
+                is_trainable=False
             ).to(device)
+            
             tokenizer = AutoTokenizer.from_pretrained(
                 "facebook/bart-large-cnn",
                 cache_dir="./models"
@@ -55,15 +55,15 @@ def load_model(model_type):
             base_model = AutoModelForSeq2SeqLM.from_pretrained(
                 "GanjinZero/biobart-base",
                 cache_dir="./models",
-                low_cpu_mem_usage=True,
-                device_map={"": device}
-            )
+                torch_dtype=torch.float32
+            ).to(device)
+            
             model = PeftModel.from_pretrained(
                 base_model, 
                 "pendar02/biobart-finetune",
-                device_map={"": device},
-                torch_dtype=torch.float32
+                is_trainable=False
             ).to(device)
+            
             tokenizer = AutoTokenizer.from_pretrained(
                 "GanjinZero/biobart-base",
                 cache_dir="./models"
@@ -74,6 +74,12 @@ def load_model(model_type):
     except Exception as e:
         st.error(f"Error loading model: {str(e)}")
         raise
+        
+        model.eval()
+        return model, tokenizer
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        raise
 
 def cleanup_model(model, tokenizer):
     """Properly cleanup model resources"""