Add insert noising

app.py

@@ -4,7 +4,6 @@ import numpy as np
 import json
 import time
 from transformers import AutoTokenizer
-from llama_diffusion_model import disable_dropout
 import os
 import importlib
 from huggingface_hub import hf_hub_download
@@ -57,22 +56,32 @@ 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):
+def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0, insert_prob=0.05):
     noised = input_ids.copy()
-    answer_len = len(input_ids) - answer_start
+    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
+        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:
-        indices = rng.choice(np.arange(answer_start, len(input_ids)), size=num_to_noise, replace=False)
+        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()
 
-        noise = rng.choice(np.arange(vocab_size), size=num_to_noise, p=mixed_probs)
+        noise = rng.choice(np.arange(vocab_size), size=len(indices), p=mixed_probs)
         for idx, val in zip(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()
@@ -82,9 +91,14 @@ def confidence_guided_noising(input_ids, answer_start, confidences, threshold, e
     if num_to_noise == 0:
         return noised
 
-    # Avoid zero-probability weights for selection
+    
     raw_weights = 1.0 - np.array(confidences[answer_start:])
-    raw_weights = np.clip(raw_weights, noise_clipping, None)  # avoid 0s
+
+    # Avoid zero-probability weights for selection
+    # If noise clipping == 1, all tokens have equal chance to be noised. 
+    # If noise_clipping == 0.00001, all tokens are noised according to the confidence of the past prediction
+    raw_weights = np.clip(raw_weights, a_min = noise_clipping, a_max = None)
+
     weights = raw_weights / raw_weights.sum()
 
     if num_to_noise > len(weights):
@@ -97,10 +111,8 @@ def confidence_guided_noising(input_ids, answer_start, confidences, threshold, e
         p=weights
     )
 
-    # 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-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)
@@ -123,13 +135,10 @@ def generate_diffusion_text(input_ids, answer_start):
 
         # Extract confidence of selected tokens
         conf = probs[range(len(sampled)), sampled].cpu().numpy()
-    return sampled, conf  # ✅ NEW: Return confidence
+    return sampled, conf 
 
 # --- Inference Wrapper ---
-
-
-# Modify diffusion_chat to use new noising conditionally
-def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_confidence_noising):
+def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_confidence_noising, insert_prob):
     placeholder = "What do you know about the city of New York?"
     if question.strip() == "":
         question = placeholder
@@ -186,7 +195,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)
+            current_tokens = noisify_answer(generated_tokens, answer_start, threshold=threshold, eot_weight=eot_weight, insert_prob=insert_prob)
 
         time.sleep(0.01)
 
@@ -197,7 +206,6 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, noise_clipping, use_
     yield f"<b>Final Output (after {i+1} iterations):</b><br>" + final_output.replace('\n', '<br>')
 
 # --- Gradio Interface ---
-
 print("Loading model...")
 model = load_model()
 print("✅ Model loaded.")
@@ -209,8 +217,9 @@ 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
+        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)")
     ],
     outputs=[gr.HTML(label="Diffusion Output")],
     title="Diffusion Language Model Chat",