Update app.py

app.py

@@ -1,23 +1,35 @@
 import torch
-from peft import PeftModel, PeftConfig
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from unsloth import FastLanguageModel
+from peft import PeftConfig
+from transformers import AutoTokenizer, TextIteratorStreamer
 from threading import Thread
 import gradio as gr
 
 MODEL_PATH = "Ozaii/zephyr-bae"
+max_seq_length = 2048
 
 print("Attempting to load Zephyr... Cross your fingers! 🤞")
 
 try:
     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"
+    base_model_name = peft_config.base_model_name_or_path
+    print(f"Loading base model: {base_model_name}")
+    
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=base_model_name,
+        max_seq_length=max_seq_length,
+        dtype=None,  # Auto-detect
+        load_in_4bit=True,
+    )
+    
+    model = FastLanguageModel.get_peft_model(
+        model,
+        peft_config=peft_config,
+        adapter_name="default",
+        use_gradient_checkpointing=True,
     )
-    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
+    
+    FastLanguageModel.for_inference(model)
     print("Zephyr loaded successfully! Time to charm!")
 except Exception as e:
     print(f"Oops! Zephyr seems to be playing hide and seek. Error: {str(e)}")