ReplaceStrategy: Add docstring and max_prob param

Idea provided by Alex Sebastian: https://x.com/alexbastian_ai/status/1886253431557587419

Author: Mihaiii

strategy/replace_strategy.py

@@ -14,11 +14,21 @@ def __init__(
         max_replacements: int = sys.maxsize,
         max_new_tokens_for_replace: int = sys.maxsize,
         skip_tokens: int = 0,
+        max_prob: float = 1
     ):
+        """
+        find: The list of texts to be replaced.
+        replace: The text to use as a replacement.
+        max_replacements: Maximum number of replacements allowed.
+        max_new_tokens_for_replace: Only allow find-and-replace within the first N generated tokens.
+        skip_tokens: Only allow find-and-replace after N generated tokens.
+        max_prob: Only allow find-and-replace if the probability of the first token in the found text is lower than N.
+        """
         super().__init__(provider, find, with_variants=False, skip_tokens=skip_tokens)
         self.replace_ids = self.provider.encode(replace, add_special_tokens=False)
         self.max_replacements = max_replacements
         self.max_new_tokens_for_replace = max_new_tokens_for_replace
+        self.max_prob = max_prob
         self.reset()
 
     def reset(self) -> None:
@@ -29,16 +39,23 @@ def reset(self) -> None:
     def on_logits(
         self, logits: torch.Tensor, continuation_tokens: List[int]
     ) -> torch.Tensor:
+        return logits
+
+    def on_probs(
+        self, probs: torch.FloatTensor, continuation_tokens: List[int]
+    ) -> torch.FloatTensor:
         if self.slop_start_pos is not None:
-            self.replace_index = 0
-            self.replaced += 1
+            for slop_token in self.found_slop_tokens[self.slop_start_pos]:
+                if probs[:, slop_token] <= self.max_prob:
+                    self.replace_index = 0
+                    self.replaced += 1
 
         if self.replace_index is not None:
-            logits[:, self.replace_ids[self.replace_index]] = 1e9
+            probs[:, self.replace_ids[self.replace_index]] = 1e9
             self.replace_index = self.replace_index + 1
             if self.replace_index >= len(self.replace_ids):
                 self.replace_index = None
-        return logits
+        return probs
 
     def backtrack(self, continuation_tokens: List[int]) -> List[int]:
         self.slop_start_pos = None