拒绝“想太多”：大模型Thinking Budget控制方案解析

import torch
from typing import List, Dict, Tuple, Set, Optional
from collections import defaultdict

class UncertaintyThinkLogitsProcessor:
    """
    A logits processor that limits thinking based on detected uncertainty patterns.
    Optimized for efficiency and memory usage.
    """

    def __init__(
        self,
        tokenizer,
        think_start_token: int,
        think_end_token: int,
        uncertainty_phrases: Dict[str, float],
        uncertainty_threshold: float = 5.0,
        growth_strategy: str = "linear",
        growth_factor: float = 1.0,
        max_think_tokens: int = 1000,
        debug: bool = False
    ):
        """Initialize the uncertainty-based think logits processor."""
        self.tokenizer = tokenizer
        self.think_start_token = think_start_token
        self.think_end_token = think_end_token
        self.uncertainty_threshold = uncertainty_threshold
        self.growth_strategy = growth_strategy
        self.growth_factor = growth_factor
        self.max_think_tokens = max_think_tokens
        self.debug = debug

        # Process uncertainty phrases and build efficient lookup structures
        self.uncertainty_phrases = {}
        self.phrase_by_length = defaultdict(list)  # Group phrases by length for efficient lookup
        
        for phrase, score in uncertainty_phrases.items():
            token_ids = tuple(tokenizer.encode(phrase, add_special_tokens=False))
            self.uncertainty_phrases[token_ids] = score
            self.phrase_by_length[len(token_ids)].append(token_ids)

        # Sort by length for more efficient searching (shorter phrases first)
        self.sorted_lengths = sorted(self.phrase_by_length.keys())
        self.max_phrase_length = max(self.sorted_lengths) if self.sorted_lengths else 0

        # State tracking - optimized data structures
        self.accumulated_uncertainty = 0.0
        self.last_processed_position = 0  # Track last processed position instead of full input length
        self.detected_phrase_positions = set()  # Use set for O(1) lookup, store positions only
        
        # Pre-allocate tensor for logits modification to avoid repeated allocation
        self._force_think_end_logits = None

        if self.debug:
            print(f"Initialized with {len(uncertainty_phrases)} uncertainty phrases")
            print(f"Max phrase length: {self.max_phrase_length}")
            print(f"Uncertainty threshold: {uncertainty_threshold}")

    def _find_uncertainty_phrases(self, tokens: List[int], start_scan_pos: int = 0) -> List[Tuple[int, int, float]]:
        """
        Efficiently find uncertainty phrases in tokens using optimized algorithm.
        Returns list of (start_pos, end_pos, uncertainty_score) tuples.
        """
        found_phrases = []
        tokens_len = len(tokens)
        
        # Only scan from the specified start position
        for pos in range(start_scan_pos, tokens_len):
            # Check phrases of each length, starting from shortest
            for phrase_len in self.sorted_lengths:
                if pos + phrase_len > tokens_len:
                    break  # Can't fit this phrase length
                
                # Extract candidate phrase
                candidate = tuple(tokens[pos:pos + phrase_len])
                
                # Check if this phrase exists in our lookup for this length
                if candidate in self.phrase_by_length[phrase_len]:
                    # Verify it's actually in our main dictionary (should always be true)
                    if candidate in self.uncertainty_phrases:
                        found_phrases.append((pos, pos + phrase_len - 1, self.uncertainty_phrases[candidate]))
        
        return found_phrases

    def __call__(
        self,
        input_ids: List[int],
        logits: torch.Tensor
    ) -> torch.Tensor:
        """Process logits based on uncertainty in thinking."""
        # Convert input_ids to list if it's a tensor (avoid repeated conversion)
        if isinstance(input_ids, torch.Tensor):
            input_ids = input_ids.tolist()

        # Early exit if no uncertainty phrases configured
        if not self.uncertainty_phrases:
            return logits

        # Check if we're in a thinking section
        try:
            think_start_pos = len(input_ids) - 1 - input_ids[::-1].index(self.think_start_token)
            in_thinking = self.think_end_token not in input_ids[think_start_pos:]
        except ValueError:
            # No think_start_token found
            return logits

        if in_thinking:
            thinking_tokens = input_ids[think_start_pos + 1:]
            tokens_since_think = len(thinking_tokens)
            
            # Only process new tokens since last call
            new_tokens_start = max(0, self.last_processed_position)
            
            if tokens_since_think > new_tokens_start:
                # Calculate scan start position to catch phrases spanning the boundary
                scan_start = max(0, new_tokens_start - self.max_phrase_length + 1)
                
                # Find new uncertainty phrases
                new_phrases = self._find_uncertainty_phrases(thinking_tokens, scan_start)
                
                # Process only truly new phrases (not already detected)
                for start_pos, end_pos, base_score in new_phrases:
                    phrase_id = (start_pos, end_pos)
                    
                    if phrase_id not in self.detected_phrase_positions:
                        # Apply growth strategy
                        if self.growth_strategy == "poly":
                            uncertainty_increase = base_score * (tokens_since_think ** self.growth_factor)
                        else:  # linear or default
                            uncertainty_increase = base_score
                        
                        self.accumulated_uncertainty += uncertainty_increase
                        self.detected_phrase_positions.add(phrase_id)
                        
                        if self.debug:
                            phrase_tokens = thinking_tokens[start_pos:end_pos + 1]
                            phrase_text = self.tokenizer.decode(phrase_tokens)
                            print(f"Detected: '{phrase_text}' pos:{start_pos}-{end_pos} "
                                  f"score:+{uncertainty_increase:.2f} total:{self.accumulated_uncertainty:.2f}")
                
                # Update last processed position
                self.last_processed_position = tokens_since_think

            # Check if we should force </think>
            should_force_end = (
                self.accumulated_uncertainty >= self.uncertainty_threshold or
                tokens_since_think >= self.max_think_tokens
            )
            
            if should_force_end:
                if self.debug:
                    reason = ("uncertainty threshold" if self.accumulated_uncertainty >= self.uncertainty_threshold 
                             else "max tokens")
                    print(f"Forcing </think> due to {reason}")
                
                # Efficiently create force-end logits (reuse tensor if possible)
                if self._force_think_end_logits is None or self._force_think_end_logits.shape != logits.shape:
                    self._force_think_end_logits = torch.full_like(logits, float('-inf'))
                else:
                    self._force_think_end_logits.fill_(float('-inf'))
                
                self._force_think_end_logits[self.think_end_token] = 1.0
                return self._force_think_end_logits

        else:
            # Reset state when not in thinking mode or when thinking has ended
            self._reset_state()

        return logits

    def _reset_state(self):
        """Reset the processor state efficiently."""
        self.accumulated_uncertainty = 0.0
        self.last_processed_position = 0
        self.detected_phrase_positions.clear()  # More efficient than creating new set

    def reset(self):
        """Reset the processor state."""
        self._reset_state()
        # Also clear the cached tensor
        self._force_think_end_logits = None

from vllm import LLM, SamplingParams
from transformers import AutoTokenizer

# Define uncertainty phrases
uncertainty_phrases = {
    # Chinese uncertainty phrases
    "可能我哪里算错了": 2.0,
    "我不确定": 1.5,
    "让我想想": 1.0,
    "可能": 0.5,
    "也许": 0.5,
    "或许": 0.5,
    "等等": 0.8,
    "我需要重新思考": 1.2,
    "有点困惑": 1.0,
    "我搞错了": 1.5,
    "我可能理解错了": 1.8,
    "这个有点棘手": 1.0,
    "需要再想想": 1.0,
    "可能我的分析有误？": 2.0,
    "可能我的理解有误？":2.0,
    "？":1.0,
    "正确？":1.0,
    "或者": 1.0,
    "这似乎不符":0.8,

    # English uncertainty phrases
    "I'm not sure": 1.5,
    "I need to reconsider": 1.2,
    "Let me think": 1.0,
    "maybe": 0.5,
    "perhaps": 0.5,
    "wait": 0.8,
    "hmm": 0.7,
    "let me check again": 1.2,
    "I might be wrong": 1.5,
    "I'm confused": 1.2,
    "I'm not certain": 1.5,
    "?":1.0,
    "That can't be right":1.0,
    "Wait":0.7,
    "right?":0.5
}

# Setup

llm = LLM(model="Qwen/Qwen3-4B")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")

# Get token IDs
think_start_token = tokenizer.convert_tokens_to_ids("<think>")
think_end_token = tokenizer.convert_tokens_to_ids("</think>")

uncertainty_processor = UncertaintyThinkLogitsProcessor(
    tokenizer=tokenizer,
    think_start_token=think_start_token,
    think_end_token=think_end_token,
    uncertainty_phrases=uncertainty_phrases,
    uncertainty_threshold=5.0,
    growth_strategy="linear",
    max_think_tokens=1000,
    debug=False
)

sampling_params = SamplingParams(
    n=5,
    temperature=0.7,
    max_tokens=5000,
    # top_p = 0.95,
    # top_k = 20,
    logits_processors=[uncertainty_processor]
)