Add insert noise: number of inserts

app.py

@@ -56,14 +56,13 @@ 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, insert_prob=0.05):
+def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0, num_inserts=0):
     noised = input_ids.copy()
     answer_len = len(noised) - answer_start
     num_to_noise = int(threshold * answer_len)
 
-    # Randomly insert one token with probability `insert_prob`
-    if rng.random() < insert_prob:
-        insert_idx = rng.integers(answer_start + 1, len(noised))  # Avoid inserting at the very start
+    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)]
@@ -138,7 +137,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, insert_prob):
+def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_confidence_noising, num_inserts):
     placeholder = "What do you know about the city of New York?"
     if question.strip() == "":
         question = placeholder
@@ -195,7 +194,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, insert_prob=insert_prob)
+            current_tokens = noisify_answer(generated_tokens, answer_start, threshold=threshold, eot_weight=eot_weight, num_inserts=num_inserts)
 
         time.sleep(0.01)
 
@@ -219,7 +218,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.0, 1.0, value=0.05, step=0.01, label="Chance to insert token (↓ = less structural change)")
+        gr.Slider(0, 100, value=0, step=1, label="Number of tokens to insert randomly (↓ = less structural change)")
     ],
     outputs=[gr.HTML(label="Diffusion Output")],
     title="Diffusion Language Model Chat",