adjustable noise clipping

app.py

@@ -74,7 +74,7 @@ def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0):
     return noised
 
 # Add new noising function
-def confidence_guided_noising(input_ids, answer_start, confidences, threshold, eot_weight):
+def confidence_guided_noising(input_ids, answer_start, confidences, threshold, eot_weight, noise_clipping):
     noised = input_ids.copy()
     answer_len = len(input_ids) - answer_start
     num_to_noise = int(threshold * answer_len)
@@ -84,7 +84,7 @@ def confidence_guided_noising(input_ids, answer_start, confidences, threshold, e
 
     # Avoid zero-probability weights for selection
     raw_weights = 1.0 - np.array(confidences[answer_start:])
-    raw_weights = np.clip(raw_weights, 0.01, None)  # avoid 0s
+    raw_weights = np.clip(raw_weights, noise_clipping, None)  # avoid 0s
     weights = raw_weights / raw_weights.sum()
 
     if num_to_noise > len(weights):
@@ -100,7 +100,7 @@ def confidence_guided_noising(input_ids, answer_start, confidences, threshold, e
     # Avoid zero-probability for token sampling
     mixed_probs = token_probabilities.copy()
     mixed_probs[eot_token_id] *= eot_weight
-    mixed_probs = np.clip(mixed_probs, 1e-5, None)  # fix for EOT weight near 0
+    # mixed_probs = np.clip(mixed_probs, 1e-6, None)  # fix for EOT weight near 0
     mixed_probs /= mixed_probs.sum()
 
     noise = rng.choice(np.arange(vocab_size), size=num_to_noise, p=mixed_probs)
@@ -129,7 +129,7 @@ def generate_diffusion_text(input_ids, answer_start):
 
 
 # Modify diffusion_chat to use new noising conditionally
-def diffusion_chat(question, eot_weight, max_it, sharpness, use_confidence_noising):
+def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_confidence_noising):
     placeholder = "What do you know about the city of New York?"
     if question.strip() == "":
         question = placeholder
@@ -184,7 +184,7 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, use_confidence_noisi
 
         threshold = get_noising_schedule(i, max_it, sharpness=sharpness)
         if use_confidence_noising:
-            current_tokens = confidence_guided_noising(generated_tokens, answer_start, confidences, threshold, eot_weight)
+            current_tokens = confidence_guided_noising(generated_tokens, answer_start, confidences, threshold, eot_weight, noise_clipping)
         else:
             current_tokens = noisify_answer(generated_tokens, answer_start, threshold=threshold, eot_weight=eot_weight)
 
@@ -209,6 +209,7 @@ demo = gr.Interface(
         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)"),
+        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")  # ✅ NEW
     ],
     outputs=[gr.HTML(label="Diffusion Output")],