Smooth confidence guided noising

app.py

@@ -82,9 +82,13 @@ def confidence_guided_noising(input_ids, answer_start, confidences, threshold, e
     if num_to_noise == 0:
         return noised
 
-    # Use 1 - confidence as sampling weights
-    weights = 1.0 - np.array(confidences[answer_start:])
-    weights /= weights.sum()
+    # Avoid zero-probability weights
+    raw_weights = 1.0 - np.array(confidences[answer_start:])
+    raw_weights = np.clip(raw_weights, 1e-6, None)  # prevent exact 0s
+    weights = raw_weights / raw_weights.sum()
+
+    if num_to_noise > len(weights):
+        num_to_noise = len(weights)  # safety: can’t sample more than available
 
     indices = rng.choice(
         np.arange(answer_start, len(input_ids)),
@@ -104,6 +108,7 @@ def confidence_guided_noising(input_ids, answer_start, confidences, threshold, e
     return noised
 
 
+
 @spaces.GPU
 def generate_diffusion_text(input_ids, answer_start):
     with torch.no_grad():