Update app.py

app.py

@@ -4,17 +4,34 @@ import torch
 import time
 import spaces
 
-# ZeroGPU decorator for GPU-intensive functions
+# Model configurations
+MODELS = {
+    "Athena-R3X 8B": "Spestly/Athena-R3X-8B",
+    "Athena-R3X 4B": "Spestly/Athena-R3X-4B",
+    "Athena-R3 7B": "Spestly/Athena-R3-7B",
+    "Athena-3 3B": "Spestly/Athena-3-3B",
+    "Athena-3 7B": "Spestly/Athena-3-7B",
+    "Athena-3 14B": "Spestly/Athena-3-14B",
+    "Athena-2 1.5B": "Spestly/Athena-2-1.5B",
+    "Athena-1 3B": "Spestly/Athena-1-3B",
+    "Athena-1 7B": "Spestly/Athena-1-7B"
+}
+
 @spaces.GPU
-def load_model_gpu(model_id):
-    """Load model on ZeroGPU"""
-    print(f"🚀 Loading {model_id} on ZeroGPU...")
+def generate_response(model_id, conversation, user_message, max_length=512, temperature=0.7):
+    """Generate response using ZeroGPU - all CUDA operations happen here"""
+    
+    # Load model and tokenizer inside the GPU function
+    print(f"🚀 Loading {model_id}...")
     start_time = time.time()
     
     tokenizer = AutoTokenizer.from_pretrained(model_id)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
-        torch_dtype=torch.float16,  # Use float16 for better memory efficiency
+        torch_dtype=torch.float16,
         device_map="auto",
         trust_remote_code=True
     )
@@ -22,15 +39,25 @@ def load_model_gpu(model_id):
     load_time = time.time() - start_time
     print(f"✅ Model loaded in {load_time:.2f}s")
     
-    return model, tokenizer
-
-@spaces.GPU
-def generate_response(model, tokenizer, prompt, max_length=512, temperature=0.7):
-    """Generate response using ZeroGPU"""
-    device = next(model.parameters()).device
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    # Build conversation history
+    conversation_history = []
+    for user_msg, assistant_msg in conversation:
+        if user_msg:
+            conversation_history.append(f"User: {user_msg}")
+        if assistant_msg:
+            conversation_history.append(f"Athena: {assistant_msg}")
     
-    start_time = time.time()
+    # Add current user message
+    conversation_history.append(f"User: {user_message}")
+    conversation_history.append("Athena:")
+    
+    # Create prompt
+    prompt = "\n".join(conversation_history)
+    
+    # Tokenize and generate
+    inputs = tokenizer(prompt, return_tensors="pt")
+    
+    generation_start = time.time()
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -42,26 +69,15 @@ def generate_response(model, tokenizer, prompt, max_length=512, temperature=0.7)
             eos_token_id=tokenizer.eos_token_id
         )
     
-    generation_time = time.time() - start_time
-    output_text = tokenizer.decode(
+    generation_time = time.time() - generation_start
+    
+    # Decode response
+    response = tokenizer.decode(
         outputs[0][inputs['input_ids'].shape[-1]:], 
         skip_special_tokens=True
     ).strip()
     
-    return output_text, generation_time
-
-# Model configurations
-MODELS = {
-    "Athena-R3X 8B": "Spestly/Athena-R3X-8B",
-    "Athena-R3X 4B": "Spestly/Athena-R3X-4B",
-    "Athena-R3 7B": "Spestly/Athena-R3-7B",
-    "Athena-3 3B": "Spestly/Athena-3-3B",
-    "Athena-3 7B": "Spestly/Athena-3-7B",
-    "Athena-3 14B": "Spestly/Athena-3-14B",
-    "Athena-2 1.5B": "Spestly/Athena-2-1.5B",
-    "Athena-1 3B": "Spestly/Athena-1-3B",
-    "Athena-1 7B": "Spestly/Athena-1-7B"
-}
+    return response, load_time, generation_time
 
 def chatbot(conversation, user_message, model_name, max_length=512, temperature=0.7):
     if not user_message.strip():
@@ -74,27 +90,18 @@ def chatbot(conversation, user_message, model_name, max_length=512, temperature=
     model_id = MODELS.get(model_name, MODELS["Athena-R3X 8B"])
     
     try:
-        # Load model and tokenizer using ZeroGPU
-        model, tokenizer = load_model_gpu(model_id)
-        
-        # Append user message to conversation
+        # Add user message to conversation
         conversation.append([user_message, ""])
         
-        # Build prompt from conversation history
-        prompt = ""
-        for user_msg, assistant_msg in conversation[:-1]:  # Exclude the current message
-            prompt += f"User: {user_msg}\nAthena: {assistant_msg}\n"
-        prompt += f"User: {user_message}\nAthena:"
-        
         # Generate response using ZeroGPU
-        output_text, generation_time = generate_response(
-            model, tokenizer, prompt, max_length, temperature
+        response, load_time, generation_time = generate_response(
+            model_id, conversation[:-1], user_message, max_length, temperature
         )
         
-        # Update the last conversation entry with the response
-        conversation[-1][1] = output_text
+        # Update the conversation with the response
+        conversation[-1][1] = response
         
-        stats = f"⚡ Generated in {generation_time:.2f}s | Model: {model_name} | Temp: {temperature}"
+        stats = f"��� Load: {load_time:.1f}s | Gen: {generation_time:.1f}s | Model: {model_name}"
         
         return conversation, "", stats
         
@@ -104,7 +111,7 @@ def chatbot(conversation, user_message, model_name, max_length=512, temperature=
             conversation[-1][1] = error_msg
         else:
             conversation = [[user_message, error_msg]]
-        return conversation, "", f"❌ Error occurred: {str(e)}"
+        return conversation, "", f"❌ Error: {str(e)}"
 
 def clear_chat():
     return [], "", ""