Use num_to_noise with confidence guided noising

app.py

@@ -74,20 +74,36 @@ def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0):
     return noised
 
 # Add new noising function
-def confidence_guided_noising(input_ids, answer_start, confidences, eot_weight):
+def confidence_guided_noising(input_ids, answer_start, confidences, threshold, eot_weight):
     noised = input_ids.copy()
+    answer_len = len(input_ids) - answer_start
+    num_to_noise = int(threshold * answer_len)
+
+    if num_to_noise == 0:
+        return noised
+
+    # Use 1 - confidence as sampling weights
+    weights = 1.0 - np.array(confidences[answer_start:])
+    weights /= weights.sum()
+
+    indices = rng.choice(
+        np.arange(answer_start, len(input_ids)),
+        size=num_to_noise,
+        replace=False,
+        p=weights
+    )
+
     mixed_probs = token_probabilities.copy()
     mixed_probs[eot_token_id] *= eot_weight
     mixed_probs /= mixed_probs.sum()
 
-    for i, conf in enumerate(confidences[answer_start:]):
-        p_noise = 1.0 - conf
-        if rng.random() < p_noise:
-            idx = answer_start + i
-            noised[idx] = rng.choice(np.arange(vocab_size), p=mixed_probs)
+    noise = rng.choice(np.arange(vocab_size), size=num_to_noise, p=mixed_probs)
+    for idx, val in zip(indices, noise):
+        noised[idx] = val
 
     return noised
 
+
 @spaces.GPU
 def generate_diffusion_text(input_ids, answer_start):
     with torch.no_grad():
@@ -160,7 +176,7 @@ def diffusion_chat(question, eot_weight, max_it, sharpness, use_confidence_noisi
 
         threshold = get_noising_schedule(i, max_it, sharpness=sharpness)
         if use_confidence_noising:
-            current_tokens = confidence_guided_noising(generated_tokens, answer_start, confidences, eot_weight)
+            current_tokens = confidence_guided_noising(generated_tokens, answer_start, confidences, threshold, eot_weight)
         else:
             current_tokens = noisify_answer(generated_tokens, answer_start, threshold=threshold, eot_weight=eot_weight)