Main DeepConf processor implementing online mode with early termination.
| 45 | self.terminated_early = False |
| 46 | |
| 47 | class DeepConfProcessor: |
| 48 | """ |
| 49 | Main DeepConf processor implementing online mode with early termination. |
| 50 | """ |
| 51 | |
| 52 | def __init__(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, |
| 53 | config: Dict[str, Any] = None): |
| 54 | """ |
| 55 | Initialize the DeepConf processor. |
| 56 | |
| 57 | Args: |
| 58 | model: The language model |
| 59 | tokenizer: The tokenizer |
| 60 | config: Configuration dictionary |
| 61 | """ |
| 62 | self.model = model |
| 63 | self.tokenizer = tokenizer |
| 64 | self.config = {**DEFAULT_CONFIG, **(config or {})} |
| 65 | |
| 66 | # Initialize components |
| 67 | self.confidence_calculator = ConfidenceCalculator( |
| 68 | window_size=self.config["window_size"], |
| 69 | top_k=self.config["top_k"] |
| 70 | ) |
| 71 | self.threshold_calibrator = ConfidenceThresholdCalibrator( |
| 72 | variant=self.config["variant"] |
| 73 | ) |
| 74 | |
| 75 | # Track generation state |
| 76 | self.warmup_traces = [] |
| 77 | self.online_traces = [] |
| 78 | self.confidence_threshold = None |
| 79 | self.total_tokens_used = 0 |
| 80 | |
| 81 | logger.info(f"DeepConf processor initialized with variant: {self.config['variant']}") |
| 82 | |
| 83 | def reset(self): |
| 84 | """Reset processor state for new query.""" |
| 85 | self.warmup_traces = [] |
| 86 | self.online_traces = [] |
| 87 | self.confidence_threshold = None |
| 88 | self.total_tokens_used = 0 |
| 89 | self.confidence_calculator.reset() |
| 90 | |
| 91 | def generate_single_trace(self, messages: List[Dict[str, str]], |
| 92 | use_early_termination: bool = False) -> TraceResult: |
| 93 | """ |
| 94 | Generate a single reasoning trace with optional early termination. |
| 95 | |
| 96 | Args: |
| 97 | messages: Input messages |
| 98 | use_early_termination: Whether to apply early termination |
| 99 | |
| 100 | Returns: |
| 101 | TraceResult object containing trace and confidence stats |
| 102 | """ |
| 103 | # Reset confidence calculator for new trace |
| 104 | self.confidence_calculator.reset() |