Update app.py

app.py

@@ -1,67 +1,37 @@
-import spaces
-import threading
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load the model and tokenizer locally
+# Load 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")
 
-# Define the function to handle chat responses
-@spaces.GPU
-def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p):
-    # Prepare the prompt by combining history and system messages
-    if system_message!= "":
-        msg = [
-            {"role": "system", "content": system_message}
-        ]
-    else:
-        msg = []
-    for user_input, assistant_response in history:
-        msg.extend(
-            [
-                {"role": "user", "content": user_input},
-                {"role": "assistant", "content": assistant_response}
-            ]
-        )
-    msg.append({"role": "user", "content": message})
-
-    prompt = tokenizer.apply_chat_template(
-        msg,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-
-    # Tokenize the input prompt
+# Chat function
+def respond(message, history, system_message, max_tokens, temperature, top_p):
+    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})
+    messages.append({"role": "user", "content": message})
+
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
 
-
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-
-    # Use a thread to run the generation in parallel
-    generation_thread = threading.Thread(
-        target=model.generate,
-        kwargs=dict(
-            inputs=inputs.input_ids,
-            max_length=max_tokens,
-            streamer=streamer,
-            do_sample=True,
-            temperature=temperature,
-            top_p=top_p,
-            pad_token_id=tokenizer.eos_token_id,
-        ),
+    outputs = model.generate(
+        input_ids=inputs.input_ids,
+        max_length=max_tokens,
+        do_sample=True,
+        temperature=temperature,
+        top_p=top_p,
+        pad_token_id=tokenizer.eos_token_id
     )
-    generation_thread.start()
-
-    # Stream the tokens as they are generated
-    text_buffer = ""
-    for new_text in streamer:
-        text_buffer+=new_text
-        yield text_buffer
 
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return decoded[len(prompt):]  # return only the generated text after the prompt
 
-# Create the Gradio interface
+# Gradio UI
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
@@ -72,6 +42,6 @@ demo = gr.ChatInterface(
     ]
 )
 
-# Launch the Gradio app
+# Launch
 if __name__ == "__main__":
     demo.launch()