Update app.py

app.py

@@ -2,16 +2,16 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 import gradio as gr
+import spaces  # Important for ZeroGPU
 
-# Load base model
+# Load models (will be moved to GPU when needed)
 base_model = AutoModelForCausalLM.from_pretrained(
     "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
     torch_dtype=torch.float16,
-    device_map="auto",
-    load_in_4bit=True
+    device_map="auto",  # ZeroGPU handles this
+    trust_remote_code=True
 )
 
-# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained("unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit")
 
 # Add padding token if missing
@@ -21,6 +21,7 @@ if tokenizer.pad_token is None:
 # Load LoRA adapter
 model = PeftModel.from_pretrained(base_model, "rezaenayati/RezAi-Model")
 
+@spaces.GPU  # This decorator is CRITICAL for ZeroGPU
 def chat_with_rezAi(messages, history):
     conversation = "<|start_header_id|>system<|end_header_id|>\nYou are Reza Enayati, a Computer Science student and entrepreneur from Los Angeles, who is eager to work as a software engineer or machine learning engineer. Answer these questions as if you are in an interview.<|eot_id|>"
     
@@ -32,13 +33,13 @@ def chat_with_rezAi(messages, history):
     # Add current message
     conversation += f"<|start_header_id|>user<|end_header_id|>\n{messages}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n"
     
-    # Tokenize with proper handling
+    # Tokenize
     inputs = tokenizer(conversation, return_tensors="pt", truncate=True, max_length=2048)
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}  # Move to GPU
     
+    # Generate response
     with torch.no_grad():
         outputs = model.generate(
-            **inputs,  # Unpack inputs properly
+            **inputs,
             max_new_tokens=128,
             temperature=0.5,
             do_sample=True,