Change interface

app.py

@@ -73,6 +73,10 @@ def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0):
             noised[idx] = val
     return noised
 
+print("Loading model...")
+model = load_model()
+print("✅ Model loaded.")
+
 def generate_diffusion_text(input_ids, answer_start):
     with torch.no_grad():
         input_tensor = torch.tensor([input_ids], dtype=torch.long).to(model.device)
@@ -82,34 +86,20 @@ def generate_diffusion_text(input_ids, answer_start):
         sampled = torch.multinomial(probs, num_samples=1).squeeze().tolist()
     return input_ids[:answer_start] + sampled[answer_start:]
 
-# --- Inference Wrapper ---
-
-
-
-# --- Gradio Interface ---
-
-print("Loading model...")
-model = load_model()
-print("✅ Model loaded.")
-
-# --- Generation logic ---
+# --- Diffusion Chat Function ---
 @spaces.GPU
-def run_diffusion_loop(question, eot_weight, max_it, sharpness):
-    placeholder = "What do you know about the city of New York?"
-    if question.strip() == "":
-        question = placeholder
-
-    prompt = f"User: {question}\nAssistant:"
+def diffusion_chat(message, system_prompt, eot_weight, max_it, sharpness):
+    prompt = f"{system_prompt}\nUser: {message}\nAssistant:"
     input_ids = tokenizer.encode(prompt, add_special_tokens=False)
     answer_start = find_answer_start(input_ids, assistant_marker_ids)
     if answer_start is None:
-        return [], "Error: Could not find Assistant marker in input."
+        yield "<span style='color:red'><b>Error:</b> Could not find Assistant marker in input.</span>"
+        return
 
     input_ids = (input_ids + [pad_token] * (256 - len(input_ids)))[:256]
     current_tokens = noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=eot_weight)
     prev_decoded_tokens = []
     last_tokens = []
-    history = ["**User:** " + question]
 
     for i in range(max_it):
         generated_tokens = generate_diffusion_text(current_tokens, answer_start)
@@ -129,8 +119,12 @@ def run_diffusion_loop(question, eot_weight, max_it, sharpness):
                 highlighted.append(text)
 
         prev_decoded_tokens = decoded_tokens
+        yield ("<div style='padding:0.5em'><b>Iteration {}</b><br>"
+               "<div style='background:#f5f5f5;padding:0.5em;border-radius:0.5em'>{}</div></div>").format(i+1, ''.join(highlighted))
+
         last_tokens.append(generated_tokens)
         if len(last_tokens) == 3 and all(t == last_tokens[0] for t in last_tokens):
+            yield f"<div style='color:gray'><i>Stopped early after {i+1} iterations (converged).</i></div>"
             break
 
         threshold = get_noising_schedule(i, max_it, sharpness=sharpness)
@@ -140,31 +134,19 @@ def run_diffusion_loop(question, eot_weight, max_it, sharpness):
     final_tokens = tokenizer.convert_ids_to_tokens(current_tokens[answer_start:])
     final_tokens = [tok for tok in final_tokens if tokenizer.convert_tokens_to_ids(tok) != eot_token_id]
     final_output = tokenizer.convert_tokens_to_string(final_tokens)
-    history.append("**Assistant:** " + final_output)
-    return history, final_output
-
-# --- UI Layout ---
-css = ".category-legend{display:none}"
-with gr.Blocks(css=css) as demo:
-    gr.Markdown("# Tini Diffusion LLM 🌀")
-    with gr.Row():
-        with gr.Column(scale=3):
-            chatbox = gr.Chatbot(label="Conversation", value=[], height=400)
-            question_input = gr.Textbox(label="Your Question", placeholder="What do you want to ask?", scale=8)
-            send_btn = gr.Button("Generate")
-        with gr.Column(scale=2):
-            eot_weight = gr.Slider(0, 1, value=0.4, step=0.05, label="EOT weight")
-            max_iters = gr.Slider(1, 512, value=64, step=1, label="Iterations")
-            sharpness = gr.Slider(1.0, 20.0, value=5.0, step=0.5, label="Sharpness")
-
-    def handle_submit(question, eot, max_it, sharp):
-        history, _ = run_diffusion_loop(question, eot, max_it, sharp)
-        return history
-
-    send_btn.click(
-        fn=handle_submit,
-        inputs=[question_input, eot_weight, max_iters, sharpness],
-        outputs=[chatbox]
-    )
-
-demo.queue().launch()
+    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."
+)
+
+demo.launch()