Update app.py

app.py

@@ -1,4 +1,5 @@
 import torch
+from peft import PeftModel, PeftConfig
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from threading import Thread
 import gradio as gr
@@ -8,8 +9,14 @@ MODEL_PATH = "Ozaii/zephyr-bae"
 print("Attempting to load Zephyr... Cross your fingers! 🤞")
 
 try:
-    model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=torch.float16, device_map="auto")
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+    peft_config = PeftConfig.from_pretrained(MODEL_PATH)
+    base_model = AutoModelForCausalLM.from_pretrained(
+        peft_config.base_model_name_or_path,
+        torch_dtype=torch.float16,
+        device_map="auto"
+    )
+    model = PeftModel.from_pretrained(base_model, MODEL_PATH)
+    tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)
     tokenizer.pad_token = tokenizer.eos_token
     print("Zephyr loaded successfully! Time to charm!")
 except Exception as e: