Add top_p and top_k sliders

app.py

@@ -179,11 +179,11 @@ def filter_logits(logits, top_k=0, top_p=0.0):
     return logits
 
 @spaces.GPU
-def generate_diffusion_text(input_ids):
+def generate_diffusion_text(input_ids, top_p, top_k):
     with torch.no_grad():
         input_tensor = torch.tensor([input_ids], dtype=torch.long).to(model.device)
         logits = model(input_ids=input_tensor)["logits"]
-        logits = filter_logits(logits, top_k=100, top_p=0.9) 
+        logits = filter_logits(logits, top_k=top_p, top_p=top_k) 
         logits = logits.clamp(min=-1e8, max=1e4)
         probs = torch.nn.functional.softmax(logits, dim=-1)[0]
         probs = torch.clamp(probs, min=1e-8, max=1.0)
@@ -196,7 +196,9 @@ def generate_diffusion_text(input_ids):
     return sampled, conf 
 
 # --- Inference Wrapper ---
-def diffusion_chat(question, max_it, pause_length, sharpness, clustering, noise_start, use_confidence_noising, noise_clipping):
+def diffusion_chat(question, max_it, pause_length, sharpness, 
+                   clustering, noise_start, use_confidence_noising, 
+                   noise_clipping, top_p, top_k):
     placeholder = "What do you know about the city of Amsterdam?"
     if question.strip() == "":
         question = placeholder
@@ -229,7 +231,7 @@ def diffusion_chat(question, max_it, pause_length, sharpness, clustering, noise_
         print('Generating output')
 
         # Model step
-        generated_tokens, confidences = generate_diffusion_text(current_tokens)
+        generated_tokens, confidences = generate_diffusion_text(current_tokens, top_p, top_k)
 
         elapsed = time.time() - generation_start
         remaining = pause_length - elapsed
@@ -322,14 +324,15 @@ demo = gr.Interface(
     fn=diffusion_chat,
     inputs=[
         gr.Textbox(label="User Question", lines=2, placeholder="What do you know about the city of Amsterdam?"),
-        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=1.0, step=0.5, label="↓ = more noising (sharpness)"),
-        gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="↑ = more clustered noising (fewer, larger edits)"),
-        gr.Slider(0.0, 1.0, value=0.2, step=0.05, label="↑ = more noise (noise start)"),
+        gr.Slider(1, 512, value=64, step=1, label="Number of iterarions: ↑ = more iterations"),
+        gr.Slider(0.01, 5, value=0.01, step=0.01, label="Pause between iteration ↑ = longer pause"),
+        gr.Slider(1.0, 20.0, value=1.0, step=0.5, label="Noise decay sharpness: ↓ = more noise in later iterations"),
+        gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Clustering: ↑ = more clustered noising"),
+        gr.Slider(0.0, 1.0, value=0.2, step=0.05, label="Noise start fraction: ↑ = more noise"),
         gr.Checkbox(value=False, label="Use confidence-guided noising"),
-        gr.Slider(0.01, 1.0, value=0.01, step=0.01, label="↓ = more confidence guidance (noise clipping)"),
-        
+        gr.Slider(0.01, 1.0, value=0.01, step=0.01, label="Noise clipping: ↓ = more confidence guidance"),
+        gr.Slider(1, 1000, value = 100, step = 1, label = "Top-p: ↑ = more random answers"]),
+        gr.Slider(0.0, 1.0, value = 0.9, step = 0.01, label = "Top-k: ↑ = more random answers"])
     ],
     outputs=[gr.HTML(label="Diffusion Output")],
     title="Diffusion Language Model Chat",