Update app.py

app.py

@@ -8,23 +8,29 @@ from sklearn.model_selection import train_test_split
 
 MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    device_map="auto",
-    trust_remote_code=True
-)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+try:
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map="auto",
+        trust_remote_code=True
+    ).to(device)
+except Exception as e:
+    print(f"Error loading base model: {e}")
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 tokenizer.pad_token = tokenizer.eos_token
 model.gradient_checkpointing_enable()
 
-# Load the pre-trained model with PEFT
-peft_config = PeftConfig.from_pretrained("trungtienluong/experiments500czephymodelngay11t6l1")
-model = PeftModel.from_pretrained(model, "trungtienluong/experiments500czephymodelngay11t6l1")
-
-# Move the model to the appropriate device
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+try:
+    # Load the pre-trained model with PEFT
+    peft_config = PeftConfig.from_pretrained("trungtienluong/experiments500czephymodelngay11t6l1")
+    model = PeftModel.from_pretrained(model, "trungtienluong/experiments500czephymodelngay11t6l1").to(device)
+except KeyError as e:
+    print(f"KeyError during PEFT model loading: {e}")
+except Exception as e:
+    print(f"Error loading PEFT model: {e}")
 
 # Load the dataset
 dataset = load_dataset("trungtienluong/500cau")