Update logit_processors/logit_.py

logit_processors/logit_.py

@@ -9,6 +9,10 @@ from vllm.sampling_params import SamplingParams
 import os
 from collections import Counter
 
+CHUNK_SIZE=16384
+WINDOW_SIZE=256
+MAX_REPETATION_COUNT=7
+
 class ThinkLogitsProcessor:    
     def __init__(self, think_end_token = 219406, max_len: int = 131072, ratio: float = 0.95):
         self.think_end_token = think_end_token
@@ -23,9 +27,9 @@ class ThinkLogitsProcessor:
         n: n-gram size
         returns dict of {ngram_tuple: count} for repeated n-grams
         """
-        ngrams = [tuple(input_ids[i:i+n]) for i in range(0, len(input_ids) - n + 1, 256)]
+        ngrams = [tuple(input_ids[i:i+n]) for i in range(0, len(input_ids) - n + 1, WINDOW_SIZE)]
         freq = Counter(ngrams)
-        return {ng: c for ng, c in freq.items() if c > 7}
+        return {ng: c for ng, c in freq.items() if c > MAX_REPETATION_COUNT}
 
     def __call__(
         self, 
@@ -38,7 +42,7 @@ class ThinkLogitsProcessor:
             # ngram
             if len(past_token_ids) % self.interval == 0:
                 # If repetation detected, force </think>
-                if self.find_repeated_ngrams(past_token_ids, n=16384):
+                if self.find_repeated_ngrams(past_token_ids, n=CHUNK_SIZE):
                     # Set all other logits to -inf except for </think>
                     logits = torch.full_like(logits, torch.finfo(torch.bfloat16).min)
                     logits[self.think_end_token] = 1.0