Change interface

app.py

@@ -136,19 +136,31 @@ def diffusion_chat(message, system_prompt, eot_weight, max_it, sharpness):
     final_output = tokenizer.convert_tokens_to_string(final_tokens)
     yield f"<div style='padding:0.5em'><b>Final Output:</b><br><div style='background:#e0ffe0;padding:0.5em;border-radius:0.5em'>{final_output}</div></div>"
 
-# --- Chat Interface ---
-demo = gr.ChatInterface(
-    diffusion_chat,
-    additional_inputs=[
-        gr.Textbox(value="You are a helpful assistant.", label="System message"),
-        gr.Slider(0, 1, value=0.4, step=0.05, label="EOT token weight (lower = longer output)"),
-        gr.Slider(1, 512, value=64, step=1, label="Max Iterations"),
-        gr.Slider(1.0, 20.0, value=5.0, step=0.5, label="Noising sharpness (lower = more noise)")
-    ],
-    title="Diffusion Language Model Chat",
-    description="Iterative denoising chat interface using a fine-tuned LLaMA model."
-)
 
+with gr.Blocks() as demo:
+    gr.Markdown("## Diffusion Language Model Chat")
+    with gr.Row():
+        with gr.Column(scale=3):
+            chatbot = gr.Chatbot()
+            message = gr.Textbox(label="User Message")
+            submit = gr.Button("Send")
+        with gr.Column(scale=1):
+            system_prompt = gr.Textbox(value="You are a helpful assistant.", label="System Message")
+            eot_weight = gr.Slider(0, 1, value=0.4, step=0.05, label="EOT token weight")
+            max_it = gr.Slider(1, 512, value=64, step=1, label="Max Iterations")
+            sharpness = gr.Slider(1.0, 20.0, value=5.0, step=0.5, label="Noising Sharpness")
+
+    def wrapped_chat(message, history, system_prompt, eot_weight, max_it, sharpness):
+        history = history or []
+        for update in diffusion_chat(message, system_prompt, eot_weight, max_it, sharpness):
+            yield history + [(message, update)]
+
+    submit.click(
+        fn=wrapped_chat,
+        inputs=[message, chatbot, system_prompt, eot_weight, max_it, sharpness],
+        outputs=chatbot,
+    )
+    
 
 if __name__ == "__main__":
     demo.launch()