Try to fix chatbot streaming

app.py

@@ -39,34 +39,6 @@ def format_prompt(messages):
     prompt += "Assistant:"
     return prompt
 
-@spaces.GPU
-def chat_with_model(messages):
-    global current_model, current_tokenizer
-    if current_model is None or current_tokenizer is None:
-        yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
-        return
-
-    current_model.to("cuda")
-
-    prompt = format_prompt(messages)
-    inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
-
-    output_ids = []
-    messages = messages.copy()
-    messages.append({"role": "assistant", "content": ""})
-
-    for token_id in current_model.generate(
-        **inputs,
-        max_new_tokens=256,
-        do_sample=False,
-        return_dict_in_generate=True,
-        output_scores=False
-    ).sequences[0][inputs['input_ids'].shape[-1]:]:  # skip input tokens
-        output_ids.append(token_id.item())
-        decoded = current_tokenizer.decode(output_ids, skip_special_tokens=True)
-        messages[-1]["content"] = decoded
-        yield messages
-
 def add_user_message(user_input, history):
     return "", history + [{"role": "user", "content": user_input}]
 
@@ -83,6 +55,34 @@ with gr.Blocks() as demo:
 
     default_model = gr.State("meta-llama/Llama-3.2-3B-Instruct")
 
+    @spaces.GPU
+    def chat_with_model(messages):
+        global current_model, current_tokenizer
+        if current_model is None or current_tokenizer is None:
+            yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
+            return
+
+        current_model.to("cuda")
+
+        prompt = format_prompt(messages)
+        inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
+
+        output_ids = []
+        messages = messages.copy()
+        messages.append({"role": "assistant", "content": ""})
+
+        for token_id in current_model.generate(
+            **inputs,
+            max_new_tokens=256,
+            do_sample=False,
+            return_dict_in_generate=True,
+            output_scores=False
+        ).sequences[0][inputs['input_ids'].shape[-1]:]:  # skip input tokens
+            output_ids.append(token_id.item())
+            decoded = current_tokenizer.decode(output_ids, skip_special_tokens=True)
+            messages[-1]["content"] = decoded
+            yield messages
+
     with gr.Row():
         model_selector = gr.Dropdown(choices=model_choices, label="Select Model")
         model_status = gr.Textbox(label="Model Status", interactive=False)