Reimplement EOT weighting

app.py

@@ -100,6 +100,9 @@ def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0, mask_
 
     mixed_probs = token_probabilities.copy()
 
+    # Apply EOT weighting 
+    mixed_probs[eot_token_id] *= eot_weight
+
     # Scale all other probabilities so they sum to 1 - mask_weight
     total_other = mixed_probs.sum() - mixed_probs[mask_token_id]
     scale = (1.0 - mask_weight) / total_other
@@ -159,6 +162,9 @@ def confidence_guided_noising(input_ids, answer_start, confidences, noise_clippi
 
     mixed_probs = token_probabilities.copy()
 
+    # Apply EOT weighting 
+    mixed_probs[eot_token_id] *= eot_weight
+
     # Scale all other probabilities so they sum to 1 - mask_weight
     total_other = mixed_probs.sum() - mixed_probs[mask_token_id]
     scale = (1.0 - mask_weight) / total_other