feat: Add capabilities/reasoning.py

capabilities/reasoning.py

@@ -0,0 +1,432 @@
+"""
+Complex Reasoning Module for MiniMind Max2
+Chain-of-Thought distillation from larger models (DeepSeek-R1, OpenAI o1).
+"""
+
+from dataclasses import dataclass, field
+from typing import List, Optional, Dict, Any, Tuple
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+import json
+import re
+
+
+@dataclass
+class ReasoningConfig:
+    """Configuration for reasoning capabilities."""
+    # Special tokens for reasoning
+    think_start_token: str = "<think>"
+    think_end_token: str = "</think>"
+    step_token: str = "<step>"
+
+    # Training settings
+    max_reasoning_steps: int = 10
+    reasoning_temperature: float = 0.7
+    distillation_temperature: float = 2.0
+    alpha_reasoning: float = 0.5  # Weight for reasoning loss vs answer loss
+
+    # Reasoning patterns
+    enable_self_reflection: bool = True
+    enable_step_verification: bool = True
+    min_reasoning_tokens: int = 50
+    max_reasoning_tokens: int = 512
+
+
+class ReasoningTokenizer:
+    """Handles special reasoning tokens."""
+
+    SPECIAL_TOKENS = {
+        "think_start": "<think>",
+        "think_end": "</think>",
+        "step": "<step>",
+        "verify": "<verify>",
+        "reflect": "<reflect>",
+        "conclude": "<conclude>",
+    }
+
+    @classmethod
+    def wrap_reasoning(cls, reasoning_text: str) -> str:
+        """Wrap reasoning in think tokens."""
+        return f"{cls.SPECIAL_TOKENS['think_start']}{reasoning_text}{cls.SPECIAL_TOKENS['think_end']}"
+
+    @classmethod
+    def extract_reasoning(cls, text: str) -> Tuple[str, str]:
+        """Extract reasoning and answer from model output."""
+        pattern = rf"{re.escape(cls.SPECIAL_TOKENS['think_start'])}(.*?){re.escape(cls.SPECIAL_TOKENS['think_end'])}"
+        match = re.search(pattern, text, re.DOTALL)
+
+        if match:
+            reasoning = match.group(1).strip()
+            answer = text[match.end():].strip()
+            return reasoning, answer
+        return "", text
+
+    @classmethod
+    def format_cot_prompt(cls, question: str, reasoning_steps: List[str], answer: str) -> str:
+        """Format a Chain-of-Thought training example."""
+        steps_text = f"\n{cls.SPECIAL_TOKENS['step']}".join(reasoning_steps)
+        reasoning = f"{cls.SPECIAL_TOKENS['step']}{steps_text}"
+        return f"{question}\n{cls.wrap_reasoning(reasoning)}\n{answer}"
+
+
+class ReasoningModule(nn.Module):
+    """
+    Reasoning enhancement module for MiniMind Max2.
+    Adds internal monologue capability for complex reasoning tasks.
+    """
+
+    def __init__(self, config: ReasoningConfig, hidden_size: int):
+        super().__init__()
+        self.config = config
+        self.hidden_size = hidden_size
+
+        # Reasoning state classifier
+        self.reasoning_gate = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size // 2),
+            nn.GELU(),
+            nn.Linear(hidden_size // 2, 3),  # [continue_reasoning, stop_reasoning, output_answer]
+        )
+
+        # Step quality predictor (for self-verification)
+        self.step_verifier = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size // 4),
+            nn.GELU(),
+            nn.Linear(hidden_size // 4, 1),
+            nn.Sigmoid(),
+        )
+
+        # Reasoning depth adapter
+        self.depth_adapter = nn.Linear(hidden_size, hidden_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        reasoning_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        """
+        Process hidden states with reasoning awareness.
+
+        Args:
+            hidden_states: [batch, seq_len, hidden_size]
+            reasoning_mask: Binary mask indicating reasoning tokens
+
+        Returns:
+            Enhanced hidden states and reasoning metrics
+        """
+        batch_size, seq_len, _ = hidden_states.shape
+
+        # Compute reasoning gate decisions
+        gate_logits = self.reasoning_gate(hidden_states)
+        gate_probs = F.softmax(gate_logits, dim=-1)
+
+        # Verify step quality
+        step_quality = self.step_verifier(hidden_states).squeeze(-1)
+
+        # Apply depth adaptation for reasoning tokens
+        if reasoning_mask is not None:
+            adapted = self.depth_adapter(hidden_states)
+            reasoning_mask_expanded = reasoning_mask.unsqueeze(-1).float()
+            hidden_states = hidden_states + adapted * reasoning_mask_expanded
+
+        metrics = {
+            "gate_probs": gate_probs,
+            "step_quality": step_quality,
+            "reasoning_ratio": reasoning_mask.float().mean() if reasoning_mask is not None else torch.tensor(0.0),
+        }
+
+        return hidden_states, metrics
+
+    def compute_reasoning_loss(
+        self,
+        hidden_states: torch.Tensor,
+        reasoning_labels: torch.Tensor,
+        step_boundaries: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Compute auxiliary loss for reasoning quality."""
+        # Gate prediction loss
+        gate_logits = self.reasoning_gate(hidden_states)
+        gate_loss = F.cross_entropy(
+            gate_logits.view(-1, 3),
+            reasoning_labels.view(-1),
+            ignore_index=-100,
+        )
+
+        # Step verification loss (if boundaries provided)
+        if step_boundaries is not None:
+            step_quality = self.step_verifier(hidden_states).squeeze(-1)
+            verification_loss = F.binary_cross_entropy(
+                step_quality,
+                step_boundaries.float(),
+            )
+            gate_loss = gate_loss + 0.1 * verification_loss
+
+        return gate_loss
+
+
+class ChainOfThoughtDataset(Dataset):
+    """Dataset for Chain-of-Thought training."""
+
+    def __init__(
+        self,
+        data_path: str,
+        tokenizer,
+        max_length: int = 2048,
+        config: Optional[ReasoningConfig] = None,
+    ):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.config = config or ReasoningConfig()
+        self.examples = []
+
+        # Load data
+        with open(data_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                if line.strip():
+                    example = json.loads(line)
+                    self.examples.append(example)
+
+    def __len__(self) -> int:
+        return len(self.examples)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        example = self.examples[idx]
+
+        # Format: question, reasoning trace, answer
+        question = example.get("question", example.get("prompt", ""))
+        reasoning = example.get("reasoning", example.get("thinking", ""))
+        answer = example.get("answer", example.get("response", ""))
+
+        # Build full text with reasoning tokens
+        full_text = ReasoningTokenizer.format_cot_prompt(
+            question,
+            reasoning.split("\n") if isinstance(reasoning, str) else reasoning,
+            answer,
+        )
+
+        # Tokenize
+        encodings = self.tokenizer(
+            full_text,
+            max_length=self.max_length,
+            truncation=True,
+            padding="max_length",
+            return_tensors="pt",
+        )
+
+        input_ids = encodings["input_ids"].squeeze(0)
+        attention_mask = encodings["attention_mask"].squeeze(0)
+
+        # Create reasoning mask (tokens between <think> and </think>)
+        reasoning_mask = self._create_reasoning_mask(input_ids)
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": input_ids.clone(),
+            "reasoning_mask": reasoning_mask,
+        }
+
+    def _create_reasoning_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """Create binary mask for reasoning tokens."""
+        # This is a simplified version - actual implementation would use token IDs
+        mask = torch.zeros_like(input_ids)
+        # In practice, find think_start and think_end token positions
+        return mask
+
+
+class ChainOfThoughtTrainer:
+    """
+    Trainer for Chain-of-Thought distillation.
+    Distills reasoning capabilities from larger models.
+    """
+
+    def __init__(
+        self,
+        student_model: nn.Module,
+        teacher_model: Optional[nn.Module] = None,
+        config: Optional[ReasoningConfig] = None,
+        learning_rate: float = 1e-5,
+        device: str = "cuda",
+    ):
+        self.student = student_model
+        self.teacher = teacher_model
+        self.config = config or ReasoningConfig()
+        self.device = device
+
+        # Add reasoning module to student
+        if hasattr(student_model, 'config'):
+            hidden_size = student_model.config.hidden_size
+        else:
+            hidden_size = 1024  # Default
+
+        self.reasoning_module = ReasoningModule(self.config, hidden_size).to(device)
+
+        # Optimizer
+        params = list(student_model.parameters()) + list(self.reasoning_module.parameters())
+        self.optimizer = torch.optim.AdamW(params, lr=learning_rate)
+
+        # Freeze teacher if provided
+        if self.teacher is not None:
+            self.teacher.eval()
+            for param in self.teacher.parameters():
+                param.requires_grad = False
+
+    def distillation_loss(
+        self,
+        student_logits: torch.Tensor,
+        teacher_logits: torch.Tensor,
+        temperature: float = 2.0,
+    ) -> torch.Tensor:
+        """Compute KL divergence distillation loss."""
+        student_probs = F.log_softmax(student_logits / temperature, dim=-1)
+        teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
+
+        loss = F.kl_div(student_probs, teacher_probs, reduction="batchmean")
+        return loss * (temperature ** 2)
+
+    def train_step(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]:
+        """Single training step."""
+        self.student.train()
+        self.reasoning_module.train()
+
+        input_ids = batch["input_ids"].to(self.device)
+        attention_mask = batch["attention_mask"].to(self.device)
+        labels = batch["labels"].to(self.device)
+        reasoning_mask = batch.get("reasoning_mask", None)
+        if reasoning_mask is not None:
+            reasoning_mask = reasoning_mask.to(self.device)
+
+        # Student forward
+        loss, student_logits, _, aux_loss = self.student(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+        )
+
+        total_loss = loss
+        metrics = {"ce_loss": loss.item(), "aux_loss": aux_loss.item()}
+
+        # Distillation from teacher
+        if self.teacher is not None:
+            with torch.no_grad():
+                _, teacher_logits, _, _ = self.teacher(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                )
+
+            distill_loss = self.distillation_loss(
+                student_logits,
+                teacher_logits,
+                self.config.distillation_temperature,
+            )
+            total_loss = (1 - self.config.alpha_reasoning) * loss + self.config.alpha_reasoning * distill_loss
+            metrics["distill_loss"] = distill_loss.item()
+
+        # Backward
+        self.optimizer.zero_grad()
+        total_loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.student.parameters(), 1.0)
+        self.optimizer.step()
+
+        metrics["total_loss"] = total_loss.item()
+        return metrics
+
+    def train(
+        self,
+        train_dataloader: DataLoader,
+        num_epochs: int = 3,
+        eval_dataloader: Optional[DataLoader] = None,
+    ) -> Dict[str, List[float]]:
+        """Full training loop."""
+        history = {"train_loss": [], "eval_loss": []}
+
+        for epoch in range(num_epochs):
+            epoch_losses = []
+
+            for batch in train_dataloader:
+                metrics = self.train_step(batch)
+                epoch_losses.append(metrics["total_loss"])
+
+            avg_loss = sum(epoch_losses) / len(epoch_losses)
+            history["train_loss"].append(avg_loss)
+            print(f"Epoch {epoch+1}/{num_epochs} - Train Loss: {avg_loss:.4f}")
+
+            # Evaluation
+            if eval_dataloader is not None:
+                eval_loss = self.evaluate(eval_dataloader)
+                history["eval_loss"].append(eval_loss)
+                print(f"  Eval Loss: {eval_loss:.4f}")
+
+        return history
+
+    def evaluate(self, dataloader: DataLoader) -> float:
+        """Evaluate on validation set."""
+        self.student.eval()
+        total_loss = 0.0
+        num_batches = 0
+
+        with torch.no_grad():
+            for batch in dataloader:
+                input_ids = batch["input_ids"].to(self.device)
+                attention_mask = batch["attention_mask"].to(self.device)
+                labels = batch["labels"].to(self.device)
+
+                loss, _, _, _ = self.student(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    labels=labels,
+                )
+                total_loss += loss.item()
+                num_batches += 1
+
+        return total_loss / num_batches if num_batches > 0 else 0.0
+
+
+def prepare_openr1_dataset(
+    raw_data_path: str,
+    output_path: str,
+    config: Optional[ReasoningConfig] = None,
+) -> int:
+    """
+    Prepare OpenR1 or DeepSeek-R1 distillation data.
+    Converts raw reasoning traces to training format.
+    """
+    config = config or ReasoningConfig()
+    processed = 0
+
+    with open(raw_data_path, 'r', encoding='utf-8') as fin, \
+         open(output_path, 'w', encoding='utf-8') as fout:
+
+        for line in fin:
+            if not line.strip():
+                continue
+
+            data = json.loads(line)
+
+            # Extract components (format varies by source)
+            question = data.get("question", data.get("prompt", data.get("input", "")))
+
+            # Handle different reasoning formats
+            if "thinking" in data:
+                reasoning = data["thinking"]
+            elif "reasoning" in data:
+                reasoning = data["reasoning"]
+            elif "chain_of_thought" in data:
+                reasoning = data["chain_of_thought"]
+            else:
+                continue  # Skip if no reasoning trace
+
+            answer = data.get("answer", data.get("response", data.get("output", "")))
+
+            # Format for training
+            processed_example = {
+                "question": question,
+                "reasoning": reasoning,
+                "answer": answer,
+            }
+
+            fout.write(json.dumps(processed_example, ensure_ascii=False) + "\n")
+            processed += 1
+
+    return processed