| 22 | |
| 23 | class ThinkDeeperProcessor: |
| 24 | def __init__(self, config: Dict[str, Any], tokenizer, model): |
| 25 | self.config = {**DEFAULT_CONFIG, **config} |
| 26 | self.tokenizer = tokenizer |
| 27 | self.model = model |
| 28 | |
| 29 | # Get token IDs for think markers |
| 30 | start_tokens = self.tokenizer.encode(self.config['start_think_token']) |
| 31 | end_tokens = self.tokenizer.encode(self.config['end_think_token']) |
| 32 | self._start_think_token = start_tokens[0] if len(start_tokens) == 1 else start_tokens[1] |
| 33 | self.end_think_token = end_tokens[0] if len(end_tokens) == 1 else end_tokens[1] |
| 34 | |
| 35 | # Store thought switch markers as token sequences |
| 36 | self.thought_switch_sequences = [] |
| 37 | for phrase in self.config["thought_switch_tokens"]: |
| 38 | # Encode without adding special tokens to get exact sequence |
| 39 | token_ids = self.tokenizer.encode(phrase, add_special_tokens=False) |
| 40 | self.thought_switch_sequences.append(token_ids) |
| 41 | logger.debug(f"Encoded '{phrase}' to token sequence: {token_ids}") |
| 42 | logger.debug(f"Decoded back: {self.tokenizer.decode(token_ids)}") |
| 43 | |
| 44 | # Track thought switches |
| 45 | self.thought_count = 0 |
| 46 | self.current_sequence = [] # Track recent tokens for sequence matching |
| 47 | self.max_sequence_length = max(len(seq) for seq in self.thought_switch_sequences) |
| 48 | |
| 49 | for phrase, sequence in zip(self.config["thought_switch_tokens"], self.thought_switch_sequences): |
| 50 | logger.debug(f"Thought switch marker '{phrase}' encoded as: {sequence}") |
| 51 | logger.debug(f"Decoded back as: {self.tokenizer.decode(sequence)}") |
| 52 | |
| 53 | def is_thought_switch(self, token: int) -> bool: |
| 54 | """Check if adding this token creates a thought switch sequence.""" |