Update app.py

app.py

@@ -1,26 +1,35 @@
+import spaces
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load model and tokenizer
+# Load the model and tokenizer
 model_name = "krish10/Qwen3_0.6B_16bit_TA_screen"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")
 
-# Chat function
-def respond(message, history, system_message, max_tokens, temperature, top_p):
+# Non-streaming chat function
+@spaces.GPU
+def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p):
+    # Construct messages from history + system message
     messages = []
     if system_message:
         messages.append({"role": "system", "content": system_message})
-    for user, assistant in history:
-        messages.append({"role": "user", "content": user})
-        messages.append({"role": "assistant", "content": assistant})
+    for user_msg, bot_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": bot_msg})
     messages.append({"role": "user", "content": message})
 
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    # Build prompt
+    prompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    # Tokenize and move to GPU
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
 
+    # Generate response
     outputs = model.generate(
-        input_ids=inputs.input_ids,
+        input_ids=inputs["input_ids"],
         max_length=max_tokens,
         do_sample=True,
         temperature=temperature,
@@ -28,8 +37,9 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         pad_token_id=tokenizer.eos_token_id
     )
 
+    # Decode and return only new content
     decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return decoded[len(prompt):]  # return only the generated text after the prompt
+    return decoded[len(prompt):]  # strip prompt prefix
 
 # Gradio UI
 demo = gr.ChatInterface(
@@ -42,6 +52,5 @@ demo = gr.ChatInterface(
     ]
 )
 
-# Launch
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()