Update app.py

app.py

@@ -1,7 +1,6 @@
 import torch
-from unsloth import FastLanguageModel
-from peft import PeftConfig
-from transformers import AutoTokenizer, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from peft import PeftConfig, PeftModel
 from threading import Thread
 import gradio as gr
 import spaces
@@ -10,38 +9,38 @@ MODEL_PATH = "Ozaii/zephyr-bae"
 BASE_MODEL = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
 max_seq_length = 2048
 
-print("Attempting to load Zephyr... Cross your fingers! 🤞")
+print("Zephyr is getting ready to charm! 🌟")
+
+model = None
+tokenizer = None
 
 @spaces.GPU
 def load_model():
-    try:
-        peft_config = PeftConfig.from_pretrained(MODEL_PATH)
-        
-        model, tokenizer = FastLanguageModel.from_pretrained(
-            model_name=BASE_MODEL,
-            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,
-        )
-        
-        FastLanguageModel.for_inference(model)
-        print("Zephyr loaded successfully! Time to charm!")
-        return model, tokenizer
-    except Exception as e:
-        print(f"Oops! Zephyr seems to be playing hide and seek. Error: {str(e)}")
-        raise
-
-model, tokenizer = load_model()
+    global model, tokenizer
+    if model is None:
+        try:
+            peft_config = PeftConfig.from_pretrained(MODEL_PATH)
+            
+            base_model = AutoModelForCausalLM.from_pretrained(
+                BASE_MODEL,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                load_in_4bit=True
+            )
+            
+            model = PeftModel.from_pretrained(base_model, MODEL_PATH)
+            tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+            tokenizer.pad_token = tokenizer.eos_token
+            
+            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)}")
+            raise
+    return model, tokenizer
 
 @spaces.GPU
 def generate_response(prompt, max_new_tokens=128):
+    model, tokenizer = load_model()
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     generation_kwargs = dict(