utils/special_token_usage_example.py

"""
Example demonstrating special token usage in chain-of-thought reasoning.

This script shows how to:
1. Generate reasoning chains with special tokens
2. Parse structured reasoning from generated text
3. Use special tokens for PRM evaluation
"""

import torch
from typing import List, Dict, Tuple
from src.reasoning.step_data import ReasoningStep, ReasoningChain, StepType, SPECIAL_TOKENS
from src.reasoning.prm import ProcessRewardModel


class SpecialTokenParser:
    """Parse generated text with special tokens into structured ReasoningChain."""
    
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        
        # Get token IDs
        self.token_ids = {
            name: tokenizer.convert_tokens_to_ids(name)
            for name in SPECIAL_TOKENS
        }
    
    def parse_reasoning_chain(self, generated_text: str, image_path: str, prompt: str) -> ReasoningChain:
        """
        Parse generated text with special tokens into a ReasoningChain.
        
        Args:
            generated_text: Generated text containing special tokens
            image_path: Path to the input image
            prompt: Original prompt
            
        Returns:
            ReasoningChain with parsed steps
        """
        steps = []
        
        # Split by reasoning_start and reasoning_end
        if "<|reasoning_start|>" not in generated_text or "<|reasoning_end|>" not in generated_text:
            # No structured reasoning found
            return ReasoningChain(
                chain_id="parsed_0",
                image_path=image_path,
                prompt=prompt,
                steps=[],
                final_answer=generated_text,
                is_correct=False
            )
        
        # Extract reasoning section
        reasoning_start = generated_text.find("<|reasoning_start|>") + len("<|reasoning_start|>")
        reasoning_end = generated_text.find("<|reasoning_end|>")
        reasoning_text = generated_text[reasoning_start:reasoning_end]
        
        # Extract answer
        answer_start = generated_text.find("<|answer_start|>") + len("<|answer_start|>")
        answer_end = generated_text.find("<|answer_end|>")
        final_answer = generated_text[answer_start:answer_end] if answer_start > 0 and answer_end > 0 else ""
        
        # Parse individual steps
        step_texts = reasoning_text.split("<|step_start|>")
        
        for i, step_text in enumerate(step_texts):
            if not step_text.strip():
                continue
            
            step = self._parse_step(step_text, i)
            if step:
                steps.append(step)
        
        return ReasoningChain(
            chain_id=f"parsed_{hash(generated_text) % 10000}",
            image_path=image_path,
            prompt=prompt,
            steps=steps,
            final_answer=final_answer.strip(),
            is_correct=False  # Will be evaluated later
        )
    
    def _parse_step(self, step_text: str, step_id: int) -> ReasoningStep:
        """Parse a single step from text."""
        try:
            # Extract step type
            step_type = StepType.INFERENCE  # default
            if "<|step_type|>" in step_text:
                type_start = step_text.find("<|step_type|>") + len("<|step_type|>")
                type_end = step_text.find("<|", type_start)
                step_type_str = step_text[type_start:type_end].strip()
                try:
                    step_type = StepType(step_type_str)
                except ValueError:
                    pass
            
            # Extract dependencies
            dependencies = []
            if "<|depends_on|>" in step_text:
                dep_start = step_text.find("<|depends_on|>") + len("<|depends_on|>")
                dep_end = step_text.find("<|", dep_start)
                deps_str = step_text[dep_start:dep_end].strip()
                dependencies = [int(d) for d in deps_str.split(",") if d.strip().isdigit()]
            
            # Extract description
            description = ""
            if "<|description_start|>" in step_text:
                desc_start = step_text.find("<|description_start|>") + len("<|description_start|>")
                desc_end = step_text.find("<|description_end|>")
                description = step_text[desc_start:desc_end].strip()
            
            # Extract confidence
            confidence = 0.5  # default
            if "<|confidence_start|>" in step_text:
                conf_start = step_text.find("<|confidence_start|>") + len("<|confidence_start|>")
                conf_end = step_text.find("<|confidence_end|>")
                try:
                    confidence = float(step_text[conf_start:conf_end].strip())
                except ValueError:
                    pass
            
            return ReasoningStep(
                step_id=step_id,
                step_type=step_type,
                description=description,
                confidence=confidence,
                dependencies=dependencies
            )
        
        except Exception as e:
            print(f"Error parsing step: {e}")
            return None


class SpecialTokenGenerator:
    """Generate text with special tokens for structured reasoning."""
    
    def __init__(self, model, tokenizer, prm: ProcessRewardModel = None):
        self.model = model
        self.tokenizer = tokenizer
        self.prm = prm
        
        # Get special token IDs
        self.reasoning_start_id = tokenizer.convert_tokens_to_ids("<|reasoning_start|>")
        self.reasoning_end_id = tokenizer.convert_tokens_to_ids("<|reasoning_end|>")
        self.step_start_id = tokenizer.convert_tokens_to_ids("<|step_start|>")
        self.step_end_id = tokenizer.convert_tokens_to_ids("<|step_end|>")
        self.ki_id = tokenizer.convert_tokens_to_ids("ки")  # PRM separator
    
    def generate_with_structure(
        self,
        prompt: str,
        image_features: torch.Tensor,
        max_steps: int = 5,
        temperature: float = 0.7,
    ) -> Tuple[str, List[float]]:
        """
        Generate reasoning chain with enforced structure and PRM evaluation.
        
        Args:
            prompt: Input prompt
            image_features: Visual features
            max_steps: Maximum reasoning steps
            temperature: Sampling temperature
            
        Returns:
            (generated_text, step_rewards)
        """
        # Tokenize prompt
        input_ids = self.tokenizer.encode(prompt, return_tensors='pt')
        
        # Start with reasoning token
        reasoning_start = torch.tensor([[self.reasoning_start_id]])
        current_ids = torch.cat([input_ids, reasoning_start], dim=1)
        
        step_rewards = []
        generated_text_parts = [prompt, " <|reasoning_start|>"]
        
        # Generate steps
        for step_num in range(max_steps):
            # Generate step with enforced structure
            step_text, step_reward = self._generate_step(
                current_ids,
                image_features,
                step_num,
                temperature
            )
            
            generated_text_parts.append(step_text)
            step_rewards.append(step_reward)
            
            # Update current_ids
            step_ids = self.tokenizer.encode(step_text, add_special_tokens=False, return_tensors='pt')
            current_ids = torch.cat([current_ids, step_ids], dim=1)
            
            # Check if reasoning should end (model generates low confidence)
            if step_reward < 0.3:
                break
        
        # End reasoning
        generated_text_parts.append(" <|reasoning_end|>")
        
        # Generate final answer
        reasoning_end = torch.tensor([[self.reasoning_end_id]])
        current_ids = torch.cat([current_ids, reasoning_end], dim=1)
        
        answer_text = self._generate_answer(current_ids, temperature)
        generated_text_parts.append(answer_text)
        
        return "".join(generated_text_parts), step_rewards
    
    def _generate_step(
        self,
        current_ids: torch.Tensor,
        image_features: torch.Tensor,
        step_num: int,
        temperature: float,
    ) -> Tuple[str, float]:
        """Generate a single reasoning step with structure."""
        # Add step_start token
        step_start = torch.tensor([[self.step_start_id]])
        step_ids = torch.cat([current_ids, step_start], dim=1)
        
        parts = [" <|step_start|>"]
        
        # Generate step type
        parts.append(" <|step_type|>")
        type_text = self._generate_until_token(step_ids, "<|", max_new_tokens=10, temperature=temperature)
        parts.append(type_text)
        
        # Generate description
        parts.append(" <|description_start|>")
        desc_text = self._generate_until_token(step_ids, "<|description_end|>", max_new_tokens=100, temperature=temperature)
        parts.append(desc_text)
        parts.append("<|description_end|>")
        
        # Generate confidence
        parts.append(" <|confidence_start|>")
        conf_text = self._generate_until_token(step_ids, "<|confidence_end|>", max_new_tokens=5, temperature=temperature)
        parts.append(conf_text)
        parts.append("<|confidence_end|>")
        
        # Add PRM separator
        parts.append(" ки")
        
        # Evaluate with PRM if available
        step_reward = 0.5
        if self.prm:
            step_text = "".join(parts)
            # TODO: Implement PRM evaluation
            # step_reward = self.prm.evaluate_step(...)
        
        parts.append(" <|step_end|>")
        
        return "".join(parts), step_reward
    
    def _generate_until_token(
        self,
        input_ids: torch.Tensor,
        stop_token: str,
        max_new_tokens: int,
        temperature: float,
    ) -> str:
        """Generate text until a specific token appears."""
        # Simplified - in practice, use proper stopping criteria
        outputs = self.model.generate(
            input_ids,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            do_sample=True,
        )
        
        generated_text = self.tokenizer.decode(outputs[0][input_ids.size(1):], skip_special_tokens=False)
        
        # Stop at token
        if stop_token in generated_text:
            generated_text = generated_text[:generated_text.find(stop_token)]
        
        return generated_text
    
    def _generate_answer(self, current_ids: torch.Tensor, temperature: float) -> str:
        """Generate final answer."""
        parts = [" <|answer_start|>"]
        
        # Generate answer text
        outputs = self.model.generate(
            current_ids,
            max_new_tokens=100,
            temperature=temperature,
        )
        
        answer_text = self.tokenizer.decode(outputs[0][current_ids.size(1):], skip_special_tokens=False)
        
        # Stop at answer_end
        if "<|answer_end|>" in answer_text:
            answer_text = answer_text[:answer_text.find("<|answer_end|>")]
        
        parts.append(answer_text)
        parts.append("<|answer_end|>")
        
        return "".join(parts)


def example_usage():
    """Example demonstrating special token usage."""
    from transformers import AutoTokenizer, AutoModel
    
    # Load model with special tokens
    model_path = "path/to/model_with_special_tokens"
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
    
    # Initialize generator
    generator = SpecialTokenGenerator(model, tokenizer)
    
    # Generate structured reasoning
    prompt = "How many red objects are in the image?"
    image_features = torch.randn(1, 768)  # Placeholder
    
    generated_text, step_rewards = generator.generate_with_structure(
        prompt=prompt,
        image_features=image_features,
        max_steps=5,
        temperature=0.7,
    )
    
    print("Generated reasoning:")
    print(generated_text)
    print(f"\nStep rewards: {step_rewards}")
    
    # Parse back into structured format
    parser = SpecialTokenParser(tokenizer)
    chain = parser.parse_reasoning_chain(
        generated_text=generated_text,
        image_path="example.jpg",
        prompt=prompt,
    )
    
    print(f"\nParsed {len(chain.steps)} steps:")
    for step in chain.steps:
        print(f"  Step {step.step_id} ({step.step_type.value}): {step.description}")
    print(f"Final answer: {chain.final_answer}")


if __name__ == "__main__":
    example_usage()