change n-gram and ratio check order

logit_processors/logit_.py

@@ -38,24 +38,24 @@ class ThinkLogitsProcessor:
         logits: torch.Tensor
     ) -> torch.Tensor:
         if self.think_end_token not in past_token_ids:
+            # ratio
+            tokens_since_think = len(past_token_ids)
+
+            response_budget = max(self.min_answer_budget, int((self.max_len - len(prompt_token_ids)) * (1-self.ratio)))
+            remaining_budget = self.max_len - len(prompt_token_ids) - response_budget - tokens_since_think
+        
+            if 0 >= remaining_budget:
+                logits = torch.full_like(logits, torch.finfo(logits.dtype).min)
+                logits[self.think_end_token] = 1.0
 
             # ngram
-            if len(past_token_ids) % self.interval == 0:
+            elif len(past_token_ids) % self.interval == 0:
                 # If repetation detected, force </think>
                 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(logits.dtype).min)
                     logits[self.think_end_token] = 1.0
-            else:
-                # ratio
-                tokens_since_think = len(past_token_ids)
-
-                response_budget = max(self.min_answer_budget, int((self.max_len - len(prompt_token_ids)) * (1-self.ratio)))
-                remaining_budget = self.max_len - len(prompt_token_ids) - response_budget - tokens_since_think
-            
-                if 0 >= remaining_budget:
-                    logits = torch.full_like(logits, torch.finfo(logits.dtype).min)
-                    logits[self.think_end_token] = 1.0
+                
         return logits