Use updated settings for initial and clustered noise

app.py

@@ -57,7 +57,7 @@ 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, noise_start = 0.5):
+def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0, clustering=0.5, noise_start = 1.0):
     noised = input_ids.copy()
     answer_len = len(noised) - answer_start
     num_to_noise = int(threshold * answer_len * noise_start)
@@ -89,10 +89,10 @@ def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0, clust
 
 
 # Add new noising function
-def confidence_guided_noising(input_ids, answer_start, confidences, threshold, eot_weight, noise_clipping):
+def confidence_guided_noising(input_ids, answer_start, confidences, noise_clipping, threshold=1.0, eot_weight = 1.0, noise_start = 1.0):
     noised = input_ids.copy()
     answer_len = len(input_ids) - 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
@@ -164,8 +164,8 @@ def diffusion_chat(question, eot_weight, max_it, pause_length, sharpness, cluste
 
     ori_input_tokens = input_ids
     current_tokens, just_noised_indices = noisify_answer(
-        ori_input_tokens, answer_start, threshold=1.0, eot_weight=eot_weight, clustering=clustering
-    )    
+                current_tokens, answer_start, threshold=1.0, eot_weight=eot_weight, clustering=clustering, noise_start = noise_start,
+            )
     last_tokens = []
     prev_decoded_tokens = []
 
@@ -209,7 +209,7 @@ def diffusion_chat(question, eot_weight, max_it, pause_length, sharpness, cluste
         threshold = get_noising_schedule(i, max_it, sharpness=sharpness)
         if use_confidence_noising:
             noised_answer = confidence_guided_noising(
-                current_tokens, answer_start, confidences, threshold, eot_weight, noise_clipping
+                current_tokens, answer_start, confidences, noise_clipping, threshold=threshold, eot_weight=eot_weight, noise_start=noise_start
             )
             just_noised_indices = []
         else: