Spaces:
Running on Zero
Running on Zero
File size: 2,665 Bytes
4164484 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | from __future__ import annotations
import math
from typing import Any
import torch
from transformers.generation.logits_process import LogitsProcessor
class ReasoningBudgetLogitsProcessor(LogitsProcessor):
"""Apply an optional reasoning-token budget for one sequence."""
def __init__(
self,
tokenizer: Any,
prompt_length: int,
reasoning_budget: int,
grace_fraction: float = 0.1,
early_exit: str = ".\n</think>\n\n",
) -> None:
if reasoning_budget < 1:
raise ValueError("reasoning_budget must be positive")
self.tokenizer = tokenizer
self.prompt_length = prompt_length
self.reasoning_budget = reasoning_budget
self.hard_limit = reasoning_budget + max(
1, math.ceil(reasoning_budget * grace_fraction)
)
self.reasoning_end_ids = tokenizer.encode(
"</think>", add_special_tokens=False
)
self.early_exit_ids = tokenizer.encode(
early_exit, add_special_tokens=False
)
self.forced_index: int | None = None
self.done = False
def _force(self, scores: torch.FloatTensor, token_id: int) -> torch.FloatTensor:
scores.fill_(-float("inf"))
scores[:, token_id] = 0
return scores
def __call__(
self,
input_ids: torch.LongTensor,
scores: torch.FloatTensor,
) -> torch.FloatTensor:
if self.done:
return scores
if input_ids.shape[0] != 1:
raise ValueError("Reasoning budget control requires batch size 1")
generated_ids = input_ids[0, self.prompt_length :].tolist()
if self.forced_index is not None:
self.forced_index += 1
if self.forced_index >= len(self.early_exit_ids):
self.done = True
return scores
return self._force(scores, self.early_exit_ids[self.forced_index])
if generated_ids[-len(self.reasoning_end_ids) :] == self.reasoning_end_ids:
self.done = True
return scores
generated_tokens = len(generated_ids)
if generated_tokens < self.reasoning_budget:
return scores
ended_line = bool(
generated_ids
and "\n"
in self.tokenizer.decode(
[generated_ids[-1]],
skip_special_tokens=False,
clean_up_tokenization_spaces=False,
)
)
if not ended_line and generated_tokens < self.hard_limit:
return scores
self.forced_index = 0
return self._force(scores, self.early_exit_ids[0])
|