Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import torch
+import time
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
 from peft import PeftModel
 from fastapi import FastAPI, HTTPException
@@ -7,6 +8,7 @@ from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
 from typing import List
 import gc
+import spaces  # ✅ Import for Zero GPU
 
 # -----------------------------
 # CONFIG
@@ -71,8 +73,8 @@ try:
         token=HF_TOKEN,
         low_cpu_mem_usage=True,
         trust_remote_code=True,
-        offload_folder=offload_dir,  # Enable disk offloading
-        offload_state_dict=True,     # Offload state dict to disk
+        offload_folder=offload_dir,
+        offload_state_dict=True,
     )
     
     print("✅ Base model loaded in 4-bit!")
@@ -83,7 +85,7 @@ try:
         base_model,
         LORA_REPO,
         token=HF_TOKEN,
-        offload_folder=offload_dir,  # Use same offload directory
+        offload_folder=offload_dir,
     )
     
     print("✅ LoRA adapter loaded!")
@@ -170,11 +172,13 @@ def detect_mood(text: str) -> str:
         print(f"Mood detection error: {e}")
         return "neutral"
 
+# ✅ ZERO GPU DECORATOR - This gets you FREE GPU!
+@spaces.GPU(duration=60)  # Max 60 seconds per request
 def generate_shinchan_response(user_input: str, mood: str) -> str:
     """Generate Shinchan's response based on user input and mood"""
     try:
-        # Add conversation context (last 3 exchanges)
-        context = "\n".join(memory[-6:]) if memory else ""
+        # Use last 2 exchanges (4 entries) for context
+        context = "\n".join(memory[-4:]) if memory else ""
         
         prompt = (
             f"<s>[INST] {SYS_PROMPT}\n"
@@ -187,7 +191,7 @@ def generate_shinchan_response(user_input: str, mood: str) -> str:
             prompt, 
             return_tensors="pt", 
             truncation=True, 
-            max_length=512,
+            max_length=384,
             padding=True
         )
         inputs = {k: v.to(model.device) for k, v in inputs.items()}
@@ -197,12 +201,13 @@ def generate_shinchan_response(user_input: str, mood: str) -> str:
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=80,
+                max_new_tokens=50,
                 temperature=temperature,
                 top_p=0.9,
                 top_k=50,
                 repetition_penalty=1.15,
                 do_sample=True,
+                num_beams=1,
                 pad_token_id=tokenizer.eos_token_id,
                 eos_token_id=tokenizer.eos_token_id,
             )
@@ -227,6 +232,8 @@ def generate_shinchan_response(user_input: str, mood: str) -> str:
         
     except Exception as e:
         print(f"Generation error: {e}")
+        import traceback
+        traceback.print_exc()
         return f"Arrey yaar! Something went wrong 🤕 Error: {str(e)[:100]}"
 
 # -----------------------------
@@ -250,34 +257,55 @@ async def health():
         "lora": LORA_REPO,
         "device": str(next(model.parameters()).device),
         "memory_entries": len(memory),
-        "mood_detection": mood_pipe is not None
+        "mood_detection": mood_pipe is not None,
+        "gpu_available": torch.cuda.is_available()
     }
 
 @app.post("/chat", response_model=MessageResponse)
 async def chat_endpoint(req: MessageRequest):
-    """Main chat endpoint"""
+    """Main chat endpoint with Zero GPU acceleration"""
+    start_time = time.time()
+    
+    print(f"\n{'='*60}")
+    print(f"🔵 CHAT REQUEST at {time.strftime('%H:%M:%S')}")
+    print(f"   Input: '{req.user_input[:50]}{'...' if len(req.user_input) > 50 else ''}'")
+    print(f"{'='*60}")
+    
     try:
         if not req.user_input or not req.user_input.strip():
             raise HTTPException(status_code=400, detail="Empty message")
         
         user_text = req.user_input.strip()
         
-        # Detect mood
+        # Step 1: Detect mood (runs on CPU)
+        print("⏱️  [1/2] Detecting mood...")
+        mood_start = time.time()
         mood = detect_mood(user_text)
+        mood_time = time.time() - mood_start
+        print(f"✅ Mood: {mood} ({mood_time:.2f}s)")
         
-        # Generate response
+        # Step 2: Generate response (runs on GPU with @spaces.GPU)
+        print("⏱️  [2/2] Generating response (GPU)...")
+        gen_start = time.time()
         response = generate_shinchan_response(user_text, mood)
+        gen_time = time.time() - gen_start
+        print(f"✅ Generated ({gen_time:.2f}s)")
+        
+        total_time = time.time() - start_time
+        print(f"🎉 TOTAL: {total_time:.2f}s")
+        print(f"{'='*60}\n")
         
         return MessageResponse(
             response=response,
             mood=mood,
-            memory=memory[-10:]  # Return last 5 exchanges
+            memory=memory[-10:]
         )
         
     except HTTPException:
         raise
     except Exception as e:
-        print(f"Chat endpoint error: {e}")
+        elapsed = time.time() - start_time
+        print(f"❌ ERROR after {elapsed:.2f}s: {e}")
         import traceback
         traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"Server error: {str(e)}")
@@ -286,7 +314,9 @@ async def chat_endpoint(req: MessageRequest):
 async def reset_memory():
     """Reset conversation memory"""
     global memory
+    old_size = len(memory)
     memory = []
+    print(f"🔄 Memory reset (cleared {old_size} entries)")
     return {"status": "Memory cleared", "memory_size": 0}
 
 # -----------------------------