Update app.py

app.py

@@ -1,34 +1,20 @@
-import spaces
 import torch
-from peft import PeftModel, PeftConfig
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from threading import Thread
 import gradio as gr
 
 MODEL_PATH = "Ozaii/zephyr-bae"
 
-@spaces.GPU
-def load_model():
-    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",
-            load_in_8bit=True
-        )
-        model = PeftModel.from_pretrained(base_model, MODEL_PATH, torch_dtype=torch.float16)
-        tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)
-        tokenizer.pad_token = tokenizer.eos_token
-        tokenizer.padding_side = "right"
-        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
+print("Attempting to load Zephyr... Cross your fingers! 🤞")
 
-model, tokenizer = load_model()
+try:
+    model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=torch.float16, device_map="auto")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+    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
 
 def generate_response(prompt, max_new_tokens=128):
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
@@ -87,4 +73,5 @@ with gr.Blocks(css=css) as iface:
     Remember, Zephyr is an AI and this is for fun and entertainment. Enjoy your chat! 💖
     """)
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()