Add noise start

app.py

@@ -57,10 +57,10 @@ 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, clustering=0.5):
+def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0, clustering=0.5, noise_start = 0.5):
     noised = input_ids.copy()
     answer_len = len(noised) - answer_start
-    num_to_noise = int(threshold * answer_len)
+    num_to_noise = int(threshold * answer_len * noise_start)
 
     if num_to_noise == 0:
         return noised, []
@@ -144,7 +144,7 @@ def generate_diffusion_text(input_ids):
     return sampled, conf 
 
 # --- Inference Wrapper ---
-def diffusion_chat(question, eot_weight, max_it, pause_length, sharpness, clustering, use_confidence_noising, noise_clipping):
+def diffusion_chat(question, eot_weight, max_it, pause_length, sharpness, clustering, noise_start, use_confidence_noising, noise_clipping):
     placeholder = "What do you know about the city of New York?"
     if question.strip() == "":
         question = placeholder
@@ -214,7 +214,7 @@ def diffusion_chat(question, eot_weight, max_it, pause_length, sharpness, cluste
             just_noised_indices = []
         else:
             noised_answer, just_noised_indices = noisify_answer(
-                current_tokens, answer_start, threshold=threshold, eot_weight=eot_weight, clustering=clustering
+                current_tokens, answer_start, threshold=threshold, eot_weight=eot_weight, clustering=clustering, noise_start = noise_start,
             )
 
         # Compose full input again: prompt + noised answer
@@ -258,6 +258,7 @@ demo = gr.Interface(
         gr.Slider(0.01, 5, value=0.01, step=0.01, label="↑ = longer pause (for visualization)"),
         gr.Slider(1.0, 20.0, value=5.0, step=0.5, label="↓ = more noising (sharpness)"),
         gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="↑ = more clustered noising (fewer, larger edits)"),
+        gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="↑ = more noise (noise start)"),
         gr.Checkbox(value=False, label="Use confidence-guided noising"),
         gr.Slider(0.01, 1.0, value=0.05, step=0.01, label="↓ = more confidence guidance (noise clipping)"),