Update app.py

app.py

@@ -3,6 +3,8 @@ import gradio as gr
 import spaces
 import torch
 import logging
+import time
+
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from transformers.utils import logging as hf_logging
 
@@ -31,25 +33,39 @@ model = AutoModelForCausalLM.from_pretrained(
 
 @spaces.GPU
 def chat_fn(prompt, max_tokens=512):
+    t0 = time.time()
     max_tokens = min(int(max_tokens), 32_000)
-    messages = [{"role": "user", "content": prompt}]
-    chat_prompt = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-    inputs = tokenizer(chat_prompt, return_tensors="pt").to(model.device)
-    
-    # Generate with proper parameters
-    outputs = model.generate(
-        **inputs, 
-        max_new_tokens=max_tokens,
-        do_sample=True,
-        temperature=0.1,
-        pad_token_id=tokenizer.eos_token_id
-    )
-    
-    # Decode only the new tokens (not the input)
-    generated_text = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
-    return generated_text
+
+    try:
+        messages = [{"role": "user", "content": prompt}]
+        chat_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+        inputs = tokenizer(chat_prompt, return_tensors="pt").to(model.device)
+        t1 = time.time()
+        logging.info(f"🧠 Tokenization complete in {t1 - t0:.2f}s")
+
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            do_sample=True,
+            temperature=0.1,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        t2 = time.time()
+        logging.info(f"⚡️ Generation complete in {t2 - t1:.2f}s (max_tokens={max_tokens})")
+
+        generated_text = tokenizer.decode(
+            outputs[0][inputs['input_ids'].shape[1]:],
+            skip_special_tokens=True
+        )
+        t3 = time.time()
+        logging.info(f"🔓 Decoding complete in {t3 - t2:.2f}s (output length: {len(generated_text)})")
+
+        return generated_text
+
+    except Exception:
+        logging.exception("❌ Exception during generation")
+        return "⚠️ Generation failed"
 
 gr.Interface(
     fn=chat_fn,