Add pause length

app.py

@@ -144,7 +144,7 @@ def generate_diffusion_text(input_ids):
     return sampled, conf 
 
 # --- Inference Wrapper ---
-def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_confidence_noising, clustering):
+def diffusion_chat(question, eot_weight, max_it, pause_length, sharpness, clustering, use_confidence_noising, noise_clipping):
     placeholder = "What do you know about the city of New York?"
     if question.strip() == "":
         question = placeholder
@@ -193,7 +193,7 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_
 
         prev_decoded_tokens = decoded_tokens
         yield f"<b>Iteration {i+1}/{max_it} (after generation):</b><br>" + "".join(highlighted).replace('\n', '<br>')
-        time.sleep(0.1)
+        time.sleep(pause_length)
 
         # --- Early stopping ---
         last_tokens.append(current_tokens)
@@ -235,7 +235,7 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_
                 highlighted.append(token_str)
 
         yield f"<b>Iteration {i+1}/{max_it} (before noising):</b><br>" + "".join(highlighted).replace('\n', '<br>')
-        time.sleep(0.1)
+        time.sleep(pause_length)
 
 
     final_tokens = tokenizer.convert_ids_to_tokens(current_tokens[answer_start:])
@@ -255,10 +255,12 @@ demo = gr.Interface(
         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(0.01, 5, value=0.01, step=0.01, label="↑ = longer pause (for visualization)"),
         gr.Slider(1.0, 20.0, value=5.0, step=0.5, label="↓ = more noising (sharpness)"),
-        gr.Slider(0.01, 1.0, value=0.05, step=0.01, label="↓ = more confidence guidance (noise clipping)"),
-        gr.Checkbox(value=False, label="Use confidence-guided noising"),
         gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="↑ = more clustered noising (fewer, larger edits)")
+        gr.Checkbox(value=False, label="Use confidence-guided noising"),
+        gr.Slider(0.01, 1.0, value=0.05, step=0.01, label="↓ = more confidence guidance (noise clipping)"),
+        
     ],
     outputs=[gr.HTML(label="Diffusion Output")],
     title="Diffusion Language Model Chat",