Add clustering of noise

app.py

@@ -56,31 +56,40 @@ def get_noising_schedule(i, max_it, sharpness=5.0):
     x = i / max_it
     return (np.exp(-sharpness * x) - np.exp(-sharpness)) / (1 - np.exp(-sharpness))
 
-def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0, num_inserts=0):
+def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0, cluster=0.5):
     noised = input_ids.copy()
     answer_len = len(noised) - answer_start
     num_to_noise = int(threshold * answer_len)
 
-    for _ in range(num_inserts):
-        insert_idx = rng.integers(answer_start + 1, len(noised))
-        insert_token = rng.choice(np.arange(vocab_size), p=token_probabilities)
-        noised = np.concatenate([noised[:insert_idx], [insert_token], noised[insert_idx:]])
-        noised = noised[:len(input_ids)]
+    if num_to_noise == 0:
+        return noised
 
-    if num_to_noise > 0:
-        indices = rng.choice(np.arange(answer_start, len(noised)), size=min(num_to_noise, len(noised) - answer_start), replace=False)
+    mixed_probs = token_probabilities.copy()
+    mixed_probs[eot_token_id] *= eot_weight
+    mixed_probs /= mixed_probs.sum()
 
-        mixed_probs = token_probabilities.copy()
-        mixed_probs[eot_token_id] *= eot_weight
-        mixed_probs /= mixed_probs.sum()
+    # Determine number of clusters and average cluster size
+    num_clusters = max(1, int((1 - cluster) * num_to_noise))  # fewer clusters if more intensity
+    cluster_size = max(1, int(num_to_noise / num_clusters))
 
-        noise = rng.choice(np.arange(vocab_size), size=len(indices), p=mixed_probs)
-        for idx, val in zip(indices, noise):
-            noised[idx] = val
+    noised_indices = set()
+    for _ in range(num_clusters):
+        center = rng.integers(answer_start, len(noised))
+        span_start = max(answer_start, center - cluster_size // 2)
+        span_end = min(len(noised), span_start + cluster_size)
+        noised_indices.update(range(span_start, span_end))
+
+    # Trim in case we overshot due to overlapping clusters
+    noised_indices = sorted(list(noised_indices))[:num_to_noise]
+
+    noise = rng.choice(np.arange(vocab_size), size=len(noised_indices), p=mixed_probs)
+    for idx, val in zip(noised_indices, noise):
+        noised[idx] = val
 
     return noised
 
 
+
 # Add new noising function
 def confidence_guided_noising(input_ids, answer_start, confidences, threshold, eot_weight, noise_clipping):
     noised = input_ids.copy()
@@ -137,7 +146,7 @@ def generate_diffusion_text(input_ids, answer_start):
     return sampled, conf 
 
 # --- Inference Wrapper ---
-def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_confidence_noising, num_inserts):
+def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_confidence_noising, clustering):
     placeholder = "What do you know about the city of New York?"
     if question.strip() == "":
         question = placeholder
@@ -194,7 +203,7 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_
         if use_confidence_noising:
             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, num_inserts=num_inserts)
+            current_tokens = noisify_answer(generated_tokens, answer_start, threshold=threshold, eot_weight=eot_weight, clustering=clustering)
 
         time.sleep(0.01)
 
@@ -218,7 +227,7 @@ demo = gr.Interface(
         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, 100, value=0, step=1, label="Number of tokens to insert randomly (↓ = less structural change)")
+        gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="↑ = more clustered noising (fewer, larger edits)")
     ],
     outputs=[gr.HTML(label="Diffusion Output")],
     title="Diffusion Language Model Chat",