Update app.py

app.py

@@ -5,32 +5,31 @@ from peft import PeftConfig, PeftModel
 import pandas as pd
 from datasets import Dataset, load_dataset
 from sklearn.model_selection import train_test_split
+from accelerate import Accelerator
 
-MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Initialize the accelerator
+accelerator = Accelerator()
 
-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}")
+MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
 
+# Load the tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 tokenizer.pad_token = tokenizer.eos_token
-model.gradient_checkpointing_enable()
 
-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 base model with accelerate
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    trust_remote_code=True
+)
+model = accelerator.prepare(model)
+
+# Load the pre-trained model with PEFT
+peft_config = PeftConfig.from_pretrained("trungtienluong/experiments500czephymodelngay11t6l1")
+model = PeftModel.from_pretrained(model, "trungtienluong/experiments500czephymodelngay11t6l1")
+model = accelerator.prepare(model)
+
+# Enable gradient checkpointing
+model.gradient_checkpointing_enable()
 
 # Load the dataset
 dataset = load_dataset("trungtienluong/500cau")
@@ -61,7 +60,8 @@ def post_process_answer(answer):
 def generate_answer(question):
     try:
         prompt = create_prompt(question)
-        encoding = tokenizer(prompt, return_tensors="pt").to(device)
+        encoding = tokenizer(prompt, return_tensors="pt")
+        encoding = accelerator.prepare(encoding)
         with torch.inference_mode():
             outputs = model.generate(
                 input_ids=encoding.input_ids,