Change interface

app.py

@@ -84,8 +84,17 @@ def generate_diffusion_text(input_ids, answer_start):
 
 # --- Inference Wrapper ---
 
+
+
+# --- Gradio Interface ---
+
+print("Loading model...")
+model = load_model()
+print("✅ Model loaded.")
+
+# --- Generation logic ---
 @spaces.GPU
-def diffusion_chat(question, eot_weight, max_it, sharpness):
+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
@@ -94,21 +103,15 @@ def diffusion_chat(question, eot_weight, max_it, sharpness):
     input_ids = tokenizer.encode(prompt, add_special_tokens=False)
     answer_start = find_answer_start(input_ids, assistant_marker_ids)
     if answer_start is None:
-        yield "Error: Could not find Assistant marker in input."
-        return
-
-    if len(input_ids) < 256:
-        input_ids += [pad_token] * (256 - len(input_ids))
-    else:
-        input_ids = input_ids[:256]
+        return [], "Error: Could not find Assistant marker in input."
 
-    ori_input_tokens = input_ids
-    current_tokens = noisify_answer(ori_input_tokens, answer_start, threshold=1.0, eot_weight=eot_weight)
+    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):
-        print('Generating output')
         generated_tokens = generate_diffusion_text(current_tokens, answer_start)
         current_tokens = generated_tokens
 
@@ -117,24 +120,17 @@ def diffusion_chat(question, eot_weight, max_it, sharpness):
         filtered_tokens = [tok for tok in decoded_tokens if tokenizer.convert_tokens_to_ids(tok) != eot_token_id]
         filtered_prev_tokens = [tok for tok in prev_decoded_tokens if tokenizer.convert_tokens_to_ids(tok) != eot_token_id] if prev_decoded_tokens else []
 
-        if filtered_prev_tokens:
-            highlighted = []
-            for tok_new, tok_old in zip(filtered_tokens, filtered_prev_tokens):
-                if tok_new != tok_old:
-                    highlighted.append(f'<span style="color:green">{tokenizer.convert_tokens_to_string([tok_new])}</span>')
-                else:
-                    highlighted.append(tokenizer.convert_tokens_to_string([tok_new]))
-        else:
-            highlighted = [tokenizer.convert_tokens_to_string([tok]) for tok in filtered_tokens]
+        highlighted = []
+        for tok_new, tok_old in zip(filtered_tokens, filtered_prev_tokens):
+            text = tokenizer.convert_tokens_to_string([tok_new])
+            if tok_new != tok_old:
+                highlighted.append(f"<span style='color:green'>{text}</span>")
+            else:
+                highlighted.append(text)
 
         prev_decoded_tokens = decoded_tokens
-        yield f"<b>Iteration {i+1}/{max_it} (running):</b><br>" + "".join(highlighted)
-
         last_tokens.append(generated_tokens)
-        if len(last_tokens) > 3:
-            last_tokens.pop(0)
-        if len(last_tokens) == 3 and last_tokens[0] == last_tokens[1] == last_tokens[2]:
-            yield f"<b>Stopped early after {i+1} iterations.</b>"
+        if len(last_tokens) == 3 and all(t == last_tokens[0] for t in last_tokens):
             break
 
         threshold = get_noising_schedule(i, max_it, sharpness=sharpness)
@@ -144,27 +140,31 @@ def diffusion_chat(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)
-    print(final_output)
-    yield f"<b>Final Output (after {i+1} iterations):</b><br>" + final_output
-
-# --- Gradio Interface ---
-
-print("Loading model...")
-model = load_model()
-print("✅ Model loaded.")
+    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 = gr.Interface(
-    fn=diffusion_chat,
-    inputs=[
-        gr.Textbox(label="User Question", lines=2, placeholder="What do you know about the city of New York?"),
-        gr.Slider(0, 1, value=0.4, step=0.05, label="↓ = longer answers (EOT weight)"),
-        gr.Slider(1, 512, value=64, step=1, label="↑ = more iterations"),
-        gr.Slider(1.0, 20.0, value=5.0, step=0.5, label="↓ = more noising (sharpness)")
-    ],
-    outputs=gr.HTML(label="Diffusion Output"),
-    title="Diffusion Language Model Chat",
-    theme="default",
-    description="This interface runs a diffusion-based language model to generate answers progressively."
-)
-
-demo.launch()
+demo.queue().launch()