Update app.py

app.py

@@ -4,38 +4,41 @@ from peft import PeftConfig
 from transformers import AutoTokenizer, TextIteratorStreamer
 from threading import Thread
 import gradio as gr
-from huggingface import spaces
+import spaces
 
 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! 🤞")
-try:
-    @spaces.GPU
-    peft_config = PeftConfig.from_pretrained(MODEL_PATH)
-    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,
-    )
-    
-    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)}")
-    raise
+
+@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()
 
 @spaces.GPU
 def generate_response(prompt, max_new_tokens=128):