Last try interface

app.py

@@ -73,10 +73,6 @@ 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)
@@ -86,22 +82,33 @@ 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:]
 
-# --- Diffusion Chat Function ---
+# --- Inference Wrapper ---
+
 @spaces.GPU
-def diffusion_chat(message, system_prompt, eot_weight, max_it, sharpness):
-    prompt = f"{system_prompt}\nUser: {message}\nAssistant:"
+def diffusion_chat(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:"
     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 "<span style='color:red'><b>Error:</b> Could not find Assistant marker in input.</span>"
+        yield "Error: Could not find Assistant marker in input."
         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)
+    if len(input_ids) < 256:
+        input_ids += [pad_token] * (256 - len(input_ids))
+    else:
+        input_ids = input_ids[:256]
+
+    ori_input_tokens = input_ids
+    current_tokens = noisify_answer(ori_input_tokens, answer_start, threshold=1.0, eot_weight=eot_weight)
     prev_decoded_tokens = []
     last_tokens = []
 
     for i in range(max_it):
+        print('Generating output')
         generated_tokens = generate_diffusion_text(current_tokens, answer_start)
         current_tokens = generated_tokens
 
@@ -110,21 +117,24 @@ def diffusion_chat(message, system_prompt, 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 []
 
-        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)
+        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]
 
         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))
+        yield f"<b>Iteration {i+1}/{max_it} (running):</b><br>" + "".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>"
+        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>"
             break
 
         threshold = get_noising_schedule(i, max_it, sharpness=sharpness)
@@ -134,33 +144,27 @@ def diffusion_chat(message, system_prompt, 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)
-    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>"
-
-
-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,
-    )
-    
+    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.")
 
-if __name__ == "__main__":
-    demo.launch()
+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(share=True)