From 36f55986ef2851c6b1898c3475dcd01c0b709ece Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 26 Nov 2025 20:03:34 +0000 Subject: [PATCH 1/5] Add reasoning model architecture for efficient chain-of-thought reasoning New features: - ReasoningController: decides when to continue reasoning vs answer - ProcessRewardModel: evaluates intermediate reasoning steps - ChainOfThoughtLoss: multi-component loss for CoT training - ReasoningLandmarkSelector: specialized landmarks for reasoning - SelfConsistencyDecoder: multi-path generation with voting - Special tokens: , , , Architecture enables smaller models to reason efficiently by: - Using sparse local-global attention (O(L*(W+G)) vs O(L^2)) - Learning to select important reasoning anchors as landmarks - Progressive curriculum on reasoning depth - Process supervision for intermediate steps Files added: - src/reasoning.py: Core reasoning components - src/reasoning_model.py: Integrated SLGA-Reasoning model - config/config_reasoning.yaml: Training configuration - scripts/train_reasoning.py: Training script with GSM8K/MATH support --- config/config_reasoning.yaml | 250 ++++++++++++++ scripts/train_reasoning.py | 533 +++++++++++++++++++++++++++++ src/__init__.py | 40 ++- src/reasoning.py | 640 +++++++++++++++++++++++++++++++++++ src/reasoning_model.py | 506 +++++++++++++++++++++++++++ 5 files changed, 1967 insertions(+), 2 deletions(-) create mode 100644 config/config_reasoning.yaml create mode 100644 scripts/train_reasoning.py create mode 100644 src/reasoning.py create mode 100644 src/reasoning_model.py diff --git a/config/config_reasoning.yaml b/config/config_reasoning.yaml new file mode 100644 index 0000000..466e522 --- /dev/null +++ b/config/config_reasoning.yaml @@ -0,0 +1,250 @@ +# ============================================================================= +# SLGA-Reasoning Model Configuration +# Configuration optimisée pour l'entraînement d'un modèle de raisonnement +# ============================================================================= + +# ----------------------------------------------------------------------------- +# Architecture du modèle +# ----------------------------------------------------------------------------- +model: + vocab_size: 50261 # 50257 (GPT-2) + 4 tokens spéciaux reasoning + max_seq_len: 4096 # Plus long pour les chaînes de raisonnement + embed_dim: 768 # Dimension plus grande pour capturer le raisonnement + num_heads: 12 + ff_hidden_multiplier: 4 + n_layers: 16 # Plus de couches pour un raisonnement profond + dropout_rate: 0.1 + +# ----------------------------------------------------------------------------- +# SLGA - Sparse Local-Global Attention +# Optimisé pour le raisonnement +# ----------------------------------------------------------------------------- +slga: + local_window: 256 # Fenêtre plus grande pour le contexte local + global_k: 48 # Plus de landmarks pour les références + gated_fusion: true + learned_landmarks: true + dilated_windows: true + diverse_topk: true + + # Nouveau: Landmarks spécialisés raisonnement + reasoning_landmarks: true # Active ReasoningLandmarkSelector + num_reasoning_types: 4 # prémisse, définition, étape, conclusion + +# ----------------------------------------------------------------------------- +# Module de Raisonnement +# ----------------------------------------------------------------------------- +reasoning: + enabled: true + + # Tokens spéciaux (ajoutés au vocabulaire) + think_token_id: 50257 # + end_think_token_id: 50258 # + step_token_id: 50259 # + answer_token_id: 50260 # + + # Contrôle du raisonnement + max_reasoning_steps: 16 + min_reasoning_steps: 2 + reasoning_depth_penalty: 0.01 + + # Process Reward Model + use_prm: true + prm_hidden_dim: 256 + prm_weight: 0.5 # Poids de la loss PRM + + # Self-Consistency (inférence) + self_consistency: + enabled: true + num_samples: 5 + temperature: 0.7 + vote_threshold: 0.6 + +# ----------------------------------------------------------------------------- +# Données d'entraînement +# ----------------------------------------------------------------------------- +data: + # Datasets de raisonnement (à télécharger) + datasets: + # Mathématiques + - name: "gsm8k" + path: "openai/gsm8k" + split: "train" + weight: 1.0 + + - name: "math" + path: "hendrycks/competition_math" + split: "train" + weight: 0.8 + + - name: "aqua_rat" + path: "aqua_rat" + split: "train" + weight: 0.5 + + # Raisonnement logique + - name: "logiqa" + path: "lucasmccabe/logiqa" + split: "train" + weight: 0.7 + + - name: "reclor" + path: "reclor" + split: "train" + weight: 0.6 + + # Science + - name: "sciq" + path: "allenai/sciq" + split: "train" + weight: 0.5 + + - name: "arc_challenge" + path: "ai2_arc" + config: "ARC-Challenge" + split: "train" + weight: 0.6 + + # Chain-of-Thought synthétiques + - name: "cot_collection" + path: "kaist-ai/CoT-Collection" + split: "train" + weight: 1.5 # Plus important + + # Prétraitement + preprocessing: + add_reasoning_tokens: true + format_as_cot: true + max_steps_per_example: 10 + + # Augmentation + augmentation: + enabled: true + step_dropout: 0.1 # Supprimer aléatoirement des étapes + step_shuffle: 0.05 # Mélanger l'ordre (pour robustesse) + paraphrase_steps: 0.2 # Paraphraser les étapes + +# ----------------------------------------------------------------------------- +# Entraînement +# ----------------------------------------------------------------------------- +training: + # Batch et accumulation + batch_size: 4 # Petits batches (séquences longues) + accum_steps: 8 # Accumulation pour batch effectif 32 + + # Optimiseur + optimizer: "adamw" + lr: 1.0e-4 + weight_decay: 0.1 + betas: [0.9, 0.95] + + # Scheduler + scheduler: "cosine" + warmup_steps: 1000 + max_steps: 100000 + + # Pertes + loss: + lm_weight: 1.0 + process_reward_weight: 0.5 + depth_penalty_weight: 0.1 + consistency_weight: 0.1 + landmark_spacing_weight: 0.01 + + # Curriculum d'apprentissage du raisonnement + curriculum: + enabled: true + # Phase 1: Étapes courtes (2-4 steps) + phase1_steps: 0 + phase1_max_reasoning: 4 + # Phase 2: Étapes moyennes (4-8 steps) + phase2_steps: 20000 + phase2_max_reasoning: 8 + # Phase 3: Étapes longues (8-16 steps) + phase3_steps: 50000 + phase3_max_reasoning: 16 + + # Régularisation + gradient_clip: 1.0 + label_smoothing: 0.1 + + # Checkpoints + save_every_steps: 5000 + eval_every_steps: 1000 + + # Mixed precision + amp: true + amp_dtype: "bfloat16" + +# ----------------------------------------------------------------------------- +# Évaluation +# ----------------------------------------------------------------------------- +evaluation: + # Benchmarks de raisonnement + benchmarks: + - name: "gsm8k_test" + dataset: "openai/gsm8k" + split: "test" + metric: "exact_match" + + - name: "math_test" + dataset: "hendrycks/competition_math" + split: "test" + metric: "exact_match" + + - name: "arc_challenge_test" + dataset: "ai2_arc" + config: "ARC-Challenge" + split: "test" + metric: "accuracy" + + # Métriques spécifiques + metrics: + - reasoning_accuracy # Réponse finale correcte + - step_accuracy # Étapes individuelles correctes + - reasoning_length # Nombre moyen d'étapes + - reasoning_consistency # Cohérence entre étapes + + # Self-consistency évaluation + self_consistency_eval: + enabled: true + num_samples: 10 + report_variance: true + +# ----------------------------------------------------------------------------- +# Inférence +# ----------------------------------------------------------------------------- +inference: + # Génération + temperature: 0.0 # Déterministe par défaut + top_k: 50 + top_p: 0.9 + max_new_tokens: 1024 + + # Arrêt + stop_tokens: + - "" + - "<|endoftext|>" + + # Self-consistency + use_self_consistency: true + num_reasoning_paths: 5 + + # Beam search (optionnel) + beam_search: + enabled: false + num_beams: 4 + length_penalty: 1.0 + +# ----------------------------------------------------------------------------- +# Hardware +# ----------------------------------------------------------------------------- +hardware: + # Pour RTX 3090 (24GB) + device: "cuda" + precision: "bf16" + grad_checkpointing: true # Nécessaire pour seq_len=4096 + + # Optimisations + torch_compile: false # Désactivé pour debugging + flash_attention: true # Si disponible diff --git a/scripts/train_reasoning.py b/scripts/train_reasoning.py new file mode 100644 index 0000000..23e8081 --- /dev/null +++ b/scripts/train_reasoning.py @@ -0,0 +1,533 @@ +#!/usr/bin/env python3 +""" +Script d'entraînement pour SLGA-Reasoning Model + +Entraîne un modèle de raisonnement avec: +- Chain-of-Thought (CoT) sur datasets de raisonnement +- Process Reward Model (PRM) pour superviser les étapes +- Curriculum learning sur la profondeur de raisonnement +- Self-consistency pour l'évaluation + +Usage: + python scripts/train_reasoning.py --config config/config_reasoning.yaml +""" + +import os +import sys +import argparse +import yaml +import math +import time +from pathlib import Path +from typing import Dict, Any, Optional, List +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import DataLoader, Dataset +from torch.cuda.amp import autocast, GradScaler + +# Add src to path +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from src.reasoning_model import SLGAReasoningModel, ReasoningModelConfig, create_reasoning_model +from src.reasoning import format_cot_example, extract_reasoning_steps + + +# ============================================================================= +# Datasets de Raisonnement +# ============================================================================= + +class ReasoningDataset(Dataset): + """ + Dataset pour l'entraînement au raisonnement. + + Supporte les formats: + - GSM8K (mathématiques) + - MATH (mathématiques avancées) + - LogiQA (raisonnement logique) + - ARC (science) + """ + + def __init__( + self, + data: List[Dict[str, Any]], + tokenizer, + max_seq_len: int = 4096, + special_tokens: Dict[str, int] = None, + ): + self.data = data + self.tokenizer = tokenizer + self.max_seq_len = max_seq_len + self.special_tokens = special_tokens or { + "": 50257, + "": 50258, + "": 50259, + "": 50260, + } + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + item = self.data[idx] + + # Formater en Chain-of-Thought + if "reasoning_steps" in item: + # Format explicite avec étapes + text = format_cot_example( + question=item["question"], + reasoning_steps=item["reasoning_steps"], + answer=item["answer"], + ) + else: + # Format simple question -> réponse + text = f"{item['question']} {item.get('solution', '')} {item['answer']}" + + # Tokenizer + tokens = self.tokenizer.encode(text) + + # Remplacer les tokens spéciaux par leurs IDs + for token_str, token_id in self.special_tokens.items(): + token_encoded = self.tokenizer.encode(token_str) + if len(token_encoded) == 1: + # Token déjà dans le vocabulaire + continue + # Sinon, on doit gérer manuellement (simplifié ici) + + # Truncate / pad + if len(tokens) > self.max_seq_len: + tokens = tokens[:self.max_seq_len] + + input_ids = torch.tensor(tokens[:-1], dtype=torch.long) + targets = torch.tensor(tokens[1:], dtype=torch.long) + + return { + "input_ids": input_ids, + "targets": targets, + } + + +def collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]: + """Collate function avec padding.""" + max_len = max(item["input_ids"].size(0) for item in batch) + + input_ids = [] + targets = [] + + for item in batch: + L = item["input_ids"].size(0) + pad_len = max_len - L + + # Pad à droite + input_ids.append(F.pad(item["input_ids"], (0, pad_len), value=50256)) + targets.append(F.pad(item["targets"], (0, pad_len), value=-100)) + + return { + "input_ids": torch.stack(input_ids), + "targets": torch.stack(targets), + } + + +# ============================================================================= +# Chargement des données +# ============================================================================= + +def load_gsm8k(split: str = "train") -> List[Dict[str, Any]]: + """Charge le dataset GSM8K.""" + try: + from datasets import load_dataset + ds = load_dataset("openai/gsm8k", "main", split=split) + + data = [] + for item in ds: + # Extraire les étapes du "answer" (format: step1 #### final_answer) + answer_text = item["answer"] + parts = answer_text.split("####") + + if len(parts) == 2: + steps_text = parts[0].strip() + final_answer = parts[1].strip() + + # Diviser en étapes (par lignes ou phrases) + steps = [s.strip() for s in steps_text.split("\n") if s.strip()] + + data.append({ + "question": item["question"], + "reasoning_steps": steps, + "answer": final_answer, + }) + + return data + except ImportError: + print("Warning: datasets library not installed. Using dummy data.") + return [ + { + "question": "What is 2 + 2?", + "reasoning_steps": ["2 + 2 = 4"], + "answer": "4", + } + ] + + +def load_math_dataset(split: str = "train") -> List[Dict[str, Any]]: + """Charge le dataset MATH.""" + try: + from datasets import load_dataset + ds = load_dataset("hendrycks/competition_math", split=split) + + data = [] + for item in ds: + data.append({ + "question": item["problem"], + "solution": item["solution"], + "answer": item["solution"].split("\\boxed{")[-1].split("}")[0] if "\\boxed{" in item["solution"] else "", + }) + + return data + except ImportError: + return [] + + +def load_cot_collection(split: str = "train") -> List[Dict[str, Any]]: + """Charge CoT-Collection.""" + try: + from datasets import load_dataset + ds = load_dataset("kaist-ai/CoT-Collection", split=split) + + data = [] + for item in ds: + data.append({ + "question": item["source"], + "reasoning_steps": item["rationale"].split("\n"), + "answer": item["target"], + }) + + return data[:10000] # Limiter pour mémoire + except ImportError: + return [] + + +# ============================================================================= +# Entraînement +# ============================================================================= + +@dataclass +class TrainingArgs: + """Arguments d'entraînement.""" + batch_size: int = 4 + accum_steps: int = 8 + lr: float = 1e-4 + weight_decay: float = 0.1 + warmup_steps: int = 1000 + max_steps: int = 100000 + eval_every: int = 1000 + save_every: int = 5000 + gradient_clip: float = 1.0 + amp: bool = True + output_dir: str = "checkpoints/reasoning" + + +class ReasoningTrainer: + """Trainer pour le modèle de raisonnement.""" + + def __init__( + self, + model: SLGAReasoningModel, + train_dataloader: DataLoader, + eval_dataloader: Optional[DataLoader], + args: TrainingArgs, + device: str = "cuda", + ): + self.model = model.to(device) + self.train_dataloader = train_dataloader + self.eval_dataloader = eval_dataloader + self.args = args + self.device = device + + # Optimizer + self.optimizer = torch.optim.AdamW( + model.parameters(), + lr=args.lr, + weight_decay=args.weight_decay, + betas=(0.9, 0.95), + ) + + # Scheduler + self.scheduler = self._create_scheduler() + + # AMP + self.scaler = GradScaler() if args.amp else None + + # Tracking + self.global_step = 0 + self.best_eval_loss = float('inf') + + # Create output dir + os.makedirs(args.output_dir, exist_ok=True) + + def _create_scheduler(self): + """Crée le scheduler avec warmup + cosine decay.""" + def lr_lambda(step): + if step < self.args.warmup_steps: + return step / self.args.warmup_steps + else: + progress = (step - self.args.warmup_steps) / (self.args.max_steps - self.args.warmup_steps) + return 0.5 * (1 + math.cos(math.pi * progress)) + + return torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda) + + def train(self): + """Boucle d'entraînement principale.""" + self.model.train() + accum_loss = 0.0 + accum_steps = 0 + + train_iter = iter(self.train_dataloader) + + print(f"Starting training for {self.args.max_steps} steps...") + print(f"Effective batch size: {self.args.batch_size * self.args.accum_steps}") + + start_time = time.time() + + while self.global_step < self.args.max_steps: + # Get batch + try: + batch = next(train_iter) + except StopIteration: + train_iter = iter(self.train_dataloader) + batch = next(train_iter) + + # Move to device + input_ids = batch["input_ids"].to(self.device) + targets = batch["targets"].to(self.device) + + # Forward + with autocast(enabled=self.args.amp): + outputs = self.model(input_ids, targets=targets) + loss = outputs["loss"] / self.args.accum_steps + + # Backward + if self.scaler is not None: + self.scaler.scale(loss).backward() + else: + loss.backward() + + accum_loss += loss.item() * self.args.accum_steps + accum_steps += 1 + + # Optimizer step + if accum_steps == self.args.accum_steps: + if self.scaler is not None: + self.scaler.unscale_(self.optimizer) + + # Gradient clipping + torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.gradient_clip) + + if self.scaler is not None: + self.scaler.step(self.optimizer) + self.scaler.update() + else: + self.optimizer.step() + + self.scheduler.step() + self.optimizer.zero_grad() + + self.global_step += 1 + + # Logging + if self.global_step % 100 == 0: + elapsed = time.time() - start_time + steps_per_sec = self.global_step / elapsed + lr = self.scheduler.get_last_lr()[0] + + print(f"Step {self.global_step}/{self.args.max_steps} | " + f"Loss: {accum_loss:.4f} | " + f"LR: {lr:.2e} | " + f"Steps/s: {steps_per_sec:.2f}") + + # Log loss components if available + if "loss_components" in outputs: + components = outputs["loss_components"] + comp_str = " | ".join([f"{k}: {v.item():.4f}" for k, v in components.items() if k != "total_loss"]) + print(f" Components: {comp_str}") + + accum_loss = 0.0 + accum_steps = 0 + + # Evaluation + if self.global_step % self.args.eval_every == 0 and self.eval_dataloader is not None: + eval_loss = self.evaluate() + print(f"Eval loss: {eval_loss:.4f}") + + if eval_loss < self.best_eval_loss: + self.best_eval_loss = eval_loss + self.save_checkpoint("best") + + self.model.train() + + # Save checkpoint + if self.global_step % self.args.save_every == 0: + self.save_checkpoint(f"step_{self.global_step}") + + print("Training complete!") + self.save_checkpoint("final") + + @torch.no_grad() + def evaluate(self) -> float: + """Évalue le modèle.""" + self.model.eval() + total_loss = 0.0 + num_batches = 0 + + for batch in self.eval_dataloader: + input_ids = batch["input_ids"].to(self.device) + targets = batch["targets"].to(self.device) + + with autocast(enabled=self.args.amp): + outputs = self.model(input_ids, targets=targets) + total_loss += outputs["loss"].item() + + num_batches += 1 + + if num_batches >= 100: # Limiter l'évaluation + break + + return total_loss / num_batches + + def save_checkpoint(self, name: str): + """Sauvegarde un checkpoint.""" + path = os.path.join(self.args.output_dir, f"checkpoint_{name}.pt") + + torch.save({ + "model_state_dict": self.model.state_dict(), + "optimizer_state_dict": self.optimizer.state_dict(), + "scheduler_state_dict": self.scheduler.state_dict(), + "global_step": self.global_step, + "best_eval_loss": self.best_eval_loss, + "config": self.model.cfg, + }, path) + + print(f"Saved checkpoint to {path}") + + +# ============================================================================= +# Main +# ============================================================================= + +def main(): + parser = argparse.ArgumentParser(description="Train SLGA-Reasoning Model") + parser.add_argument("--config", type=str, default="config/config_reasoning.yaml") + parser.add_argument("--resume", type=str, default=None, help="Resume from checkpoint") + args = parser.parse_args() + + # Load config + with open(args.config) as f: + config = yaml.safe_load(f) + + print("=" * 60) + print("SLGA-Reasoning Training") + print("=" * 60) + + # Device + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"Using device: {device}") + + # Load tokenizer (using GPT-2 tokenizer) + try: + from transformers import GPT2Tokenizer + tokenizer = GPT2Tokenizer.from_pretrained("gpt2") + except ImportError: + print("transformers not installed. Using dummy tokenizer.") + tokenizer = None + + # Load data + print("Loading datasets...") + train_data = [] + + # GSM8K + gsm8k_data = load_gsm8k("train") + print(f" GSM8K: {len(gsm8k_data)} examples") + train_data.extend(gsm8k_data) + + # MATH + math_data = load_math_dataset("train") + print(f" MATH: {len(math_data)} examples") + train_data.extend(math_data) + + # CoT Collection + cot_data = load_cot_collection("train") + print(f" CoT-Collection: {len(cot_data)} examples") + train_data.extend(cot_data) + + print(f"Total training examples: {len(train_data)}") + + if len(train_data) == 0: + print("No training data available. Please install 'datasets' library.") + print(" pip install datasets") + return + + # Create datasets + train_dataset = ReasoningDataset(train_data, tokenizer, max_seq_len=config.get("model", {}).get("max_seq_len", 4096)) + train_dataloader = DataLoader( + train_dataset, + batch_size=config.get("training", {}).get("batch_size", 4), + shuffle=True, + collate_fn=collate_fn, + num_workers=4, + pin_memory=True, + ) + + # Create model + print("Creating model...") + model_config = config.get("model", {}) + reasoning_config = config.get("reasoning", {}) + + model = create_reasoning_model( + vocab_size=model_config.get("vocab_size", 50261), + embed_dim=model_config.get("embed_dim", 768), + n_layers=model_config.get("n_layers", 16), + max_seq_len=model_config.get("max_seq_len", 4096), + local_window=config.get("slga", {}).get("local_window", 256), + global_k=config.get("slga", {}).get("global_k", 48), + max_reasoning_steps=reasoning_config.get("max_reasoning_steps", 16), + use_prm=reasoning_config.get("use_prm", True), + ) + + num_params = model.get_num_params() + print(f"Model parameters: {num_params / 1e6:.1f}M") + + # Resume if specified + if args.resume: + print(f"Resuming from {args.resume}") + checkpoint = torch.load(args.resume) + model.load_state_dict(checkpoint["model_state_dict"]) + + # Create trainer + training_config = config.get("training", {}) + train_args = TrainingArgs( + batch_size=training_config.get("batch_size", 4), + accum_steps=training_config.get("accum_steps", 8), + lr=training_config.get("lr", 1e-4), + warmup_steps=training_config.get("warmup_steps", 1000), + max_steps=training_config.get("max_steps", 100000), + eval_every=training_config.get("eval_every_steps", 1000), + save_every=training_config.get("save_every_steps", 5000), + amp=training_config.get("amp", True), + ) + + trainer = ReasoningTrainer( + model=model, + train_dataloader=train_dataloader, + eval_dataloader=None, # TODO: Add eval dataloader + args=train_args, + device=device, + ) + + # Train + trainer.train() + + +if __name__ == "__main__": + main() diff --git a/src/__init__.py b/src/__init__.py index 48c648a..f1b5f09 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -2,17 +2,39 @@ SLGA-Plus: Sparse Local-Global Attention Transformer Main package for efficient long-context language modeling. +Now with Reasoning Model support! """ -__version__ = "0.1.0" +__version__ = "0.2.0" __author__ = "SLGA Team" +# Core SLGA modules from .slga import SLGAModule from .landmarks import LearnableLandmarkSelector, landmark_diversity_loss, landmark_sparsity_loss from .model import Config, LLMTransformer from .data import get_tokenizer, load_text_dataset, CollatorLocal, CollatorLocalGlobal +# Reasoning modules +from .reasoning import ( + ReasoningConfig, + ThoughtTokenEmbedding, + ReasoningController, + ProcessRewardModel, + ChainOfThoughtLoss, + ReasoningLandmarkSelector, + SelfConsistencyDecoder, + create_reasoning_tokens, + format_cot_example, + extract_reasoning_steps, +) +from .reasoning_model import ( + ReasoningModelConfig, + SLGAReasoningModel, + create_reasoning_model, +) + __all__ = [ + # Core "SLGAModule", "LearnableLandmarkSelector", "landmark_diversity_loss", @@ -22,5 +44,19 @@ "get_tokenizer", "load_text_dataset", "CollatorLocal", - "CollatorLocalGlobal" + "CollatorLocalGlobal", + # Reasoning + "ReasoningConfig", + "ThoughtTokenEmbedding", + "ReasoningController", + "ProcessRewardModel", + "ChainOfThoughtLoss", + "ReasoningLandmarkSelector", + "SelfConsistencyDecoder", + "create_reasoning_tokens", + "format_cot_example", + "extract_reasoning_steps", + "ReasoningModelConfig", + "SLGAReasoningModel", + "create_reasoning_model", ] diff --git a/src/reasoning.py b/src/reasoning.py new file mode 100644 index 0000000..ae53511 --- /dev/null +++ b/src/reasoning.py @@ -0,0 +1,640 @@ +# reasoning.py +""" +Reasoning Components for SLGA-Reasoning Model + +Composants pour transformer SLGA en un modèle de raisonnement: +1. ReasoningController - Contrôle la profondeur de raisonnement +2. ThoughtTokens - Gestion des tokens spéciaux , , +3. ProcessRewardModel - Récompense les étapes intermédiaires +4. ChainOfThoughtLoss - Loss pour l'entraînement CoT +""" + +from __future__ import annotations +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Optional, Tuple, Dict, List, Any +from dataclasses import dataclass + + +@dataclass +class ReasoningConfig: + """Configuration pour le module de raisonnement""" + # Tokens spéciaux + think_token_id: int = 50257 # + end_think_token_id: int = 50258 # + step_token_id: int = 50259 # + answer_token_id: int = 50260 # + + # Contrôle du raisonnement + max_reasoning_steps: int = 16 # Maximum d'étapes de raisonnement + min_reasoning_steps: int = 1 # Minimum d'étapes + reasoning_depth_penalty: float = 0.01 # Pénalité pour trop d'étapes + + # Récompenses + step_reward: float = 0.1 # Récompense par étape correcte + answer_reward: float = 1.0 # Récompense pour bonne réponse + + # Architecture + embed_dim: int = 512 + num_heads: int = 8 + + # Process Reward Model (PRM) + use_prm: bool = True # Activer le Process Reward Model + prm_hidden_dim: int = 256 + + +class ThoughtTokenEmbedding(nn.Module): + """ + Embeddings spéciaux pour les tokens de raisonnement. + + Apprend des embeddings distincts pour: + - : début de réflexion + - : fin de réflexion + - : étape intermédiaire + - : réponse finale + """ + + def __init__(self, embed_dim: int, num_special_tokens: int = 4): + super().__init__() + self.embed_dim = embed_dim + self.num_special_tokens = num_special_tokens + + # Embeddings appris pour tokens spéciaux + self.special_embeddings = nn.Embedding(num_special_tokens, embed_dim) + + # Projection pour fusionner avec embeddings existants + self.fusion = nn.Linear(embed_dim * 2, embed_dim) + + # Initialisation + nn.init.normal_(self.special_embeddings.weight, std=0.02) + + def forward( + self, + token_embeddings: torch.Tensor, + special_token_mask: torch.Tensor, + special_token_ids: torch.Tensor + ) -> torch.Tensor: + """ + Fusionne les embeddings de tokens spéciaux avec les embeddings standards. + + Args: + token_embeddings: (B, L, D) embeddings des tokens + special_token_mask: (B, L) bool, True si token spécial + special_token_ids: (B, L) IDs locaux (0-3) des tokens spéciaux + + Returns: + enriched_embeddings: (B, L, D) embeddings enrichis + """ + B, L, D = token_embeddings.shape + + # Récupérer embeddings spéciaux + # Clamp pour éviter index out of bounds + special_ids_clamped = torch.clamp(special_token_ids, 0, self.num_special_tokens - 1) + special_emb = self.special_embeddings(special_ids_clamped) # (B, L, D) + + # Fusionner seulement aux positions de tokens spéciaux + fused = self.fusion(torch.cat([token_embeddings, special_emb], dim=-1)) + + # Masquer: utiliser fusion seulement pour tokens spéciaux + mask_expanded = special_token_mask.unsqueeze(-1).expand_as(token_embeddings) + output = torch.where(mask_expanded, fused, token_embeddings) + + return output + + +class ReasoningController(nn.Module): + """ + Contrôleur de raisonnement adaptatif. + + Décide dynamiquement: + 1. Faut-il continuer à raisonner ou donner la réponse? + 2. Quelle profondeur de raisonnement est nécessaire? + 3. Le raisonnement actuel est-il correct? + + Utilise un mécanisme d'attention sur l'historique pour décider. + """ + + def __init__(self, config: ReasoningConfig): + super().__init__() + self.config = config + embed_dim = config.embed_dim + + # Encodeur du contexte de raisonnement + self.context_encoder = nn.Sequential( + nn.Linear(embed_dim, embed_dim), + nn.GELU(), + nn.Dropout(0.1), + nn.Linear(embed_dim, embed_dim), + ) + + # Tête de décision: continuer ou répondre + self.decision_head = nn.Sequential( + nn.Linear(embed_dim, embed_dim // 2), + nn.GELU(), + nn.Linear(embed_dim // 2, 2), # [continue, answer] + ) + + # Estimateur de confiance + self.confidence_head = nn.Sequential( + nn.Linear(embed_dim, embed_dim // 4), + nn.GELU(), + nn.Linear(embed_dim // 4, 1), + nn.Sigmoid(), + ) + + # Compteur d'étapes (buffer, pas de gradients) + self.register_buffer("step_count", torch.tensor(0)) + + def forward( + self, + hidden_states: torch.Tensor, + step_positions: Optional[torch.Tensor] = None, + ) -> Dict[str, torch.Tensor]: + """ + Évalue l'état du raisonnement et décide de la prochaine action. + + Args: + hidden_states: (B, L, D) états cachés actuels + step_positions: (B, num_steps) positions des tokens + + Returns: + dict avec: + - decision: (B, 2) logits [continuer, répondre] + - confidence: (B, 1) confiance dans le raisonnement actuel + - should_continue: (B,) bool, continuer ou non + """ + B, L, D = hidden_states.shape + + # Pooling sur la séquence (mean pooling) + context = hidden_states.mean(dim=1) # (B, D) + + # Si positions des étapes fournies, enrichir avec attention + if step_positions is not None and step_positions.numel() > 0: + # Gather les états aux positions des étapes + num_steps = step_positions.size(1) + step_positions_clamped = torch.clamp(step_positions, 0, L - 1) + step_pos_exp = step_positions_clamped.unsqueeze(-1).expand(B, num_steps, D) + step_states = torch.gather(hidden_states, dim=1, index=step_pos_exp) # (B, num_steps, D) + + # Attention sur les étapes précédentes + step_context = step_states.mean(dim=1) # (B, D) + context = context + step_context + + # Encoder le contexte + encoded = self.context_encoder(context) # (B, D) + + # Décision + decision_logits = self.decision_head(encoded) # (B, 2) + decision_probs = F.softmax(decision_logits, dim=-1) + + # Confiance + confidence = self.confidence_head(encoded) # (B, 1) + + # Seuil adaptatif basé sur le nombre d'étapes + # Plus on a d'étapes, plus on favorise "répondre" + step_bias = min(self.step_count.item() / self.config.max_reasoning_steps, 1.0) + + # Décision finale: continue si prob(continue) > 0.5 - step_bias * 0.3 + threshold = 0.5 - step_bias * 0.3 + should_continue = decision_probs[:, 0] > threshold + + return { + "decision_logits": decision_logits, + "decision_probs": decision_probs, + "confidence": confidence, + "should_continue": should_continue, + } + + +class ProcessRewardModel(nn.Module): + """ + Process Reward Model (PRM) pour le raisonnement. + + Évalue la qualité de chaque étape de raisonnement, pas seulement + la réponse finale. Inspiré de "Let's Verify Step by Step" (OpenAI). + + Avantages vs Outcome Reward Model (ORM): + - Feedback plus granulaire + - Détecte les erreurs tôt + - Meilleur crédit assignment + """ + + def __init__(self, config: ReasoningConfig): + super().__init__() + self.config = config + embed_dim = config.embed_dim + hidden_dim = config.prm_hidden_dim + + # Encodeur d'étape + self.step_encoder = nn.Sequential( + nn.Linear(embed_dim, hidden_dim), + nn.GELU(), + nn.Dropout(0.1), + nn.Linear(hidden_dim, hidden_dim), + ) + + # Tête de récompense par étape + self.reward_head = nn.Sequential( + nn.Linear(hidden_dim, hidden_dim // 2), + nn.GELU(), + nn.Linear(hidden_dim // 2, 1), + ) + + # Tête de classification: étape correcte/incorrecte + self.correctness_head = nn.Sequential( + nn.Linear(hidden_dim, hidden_dim // 2), + nn.GELU(), + nn.Linear(hidden_dim // 2, 2), # [incorrect, correct] + ) + + def forward( + self, + step_hidden_states: torch.Tensor, + ) -> Dict[str, torch.Tensor]: + """ + Évalue chaque étape de raisonnement. + + Args: + step_hidden_states: (B, num_steps, D) états des tokens + + Returns: + dict avec: + - step_rewards: (B, num_steps) récompense par étape + - step_correctness: (B, num_steps, 2) logits correct/incorrect + """ + B, S, D = step_hidden_states.shape + + # Encoder chaque étape + encoded = self.step_encoder(step_hidden_states) # (B, S, hidden) + + # Récompenses + rewards = self.reward_head(encoded).squeeze(-1) # (B, S) + + # Correctness + correctness_logits = self.correctness_head(encoded) # (B, S, 2) + + return { + "step_rewards": rewards, + "step_correctness": correctness_logits, + } + + def compute_process_reward( + self, + step_rewards: torch.Tensor, + step_labels: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Calcule la récompense totale du processus de raisonnement. + + Args: + step_rewards: (B, num_steps) récompenses prédites + step_labels: (B, num_steps) labels (1=correct, 0=incorrect) si supervisé + + Returns: + total_reward: (B,) récompense totale par séquence + """ + if step_labels is not None: + # Mode supervisé: récompense = sum des étapes correctes + masked_rewards = step_rewards * step_labels + return masked_rewards.sum(dim=-1) + else: + # Mode non-supervisé: somme simple + return step_rewards.sum(dim=-1) + + +class ChainOfThoughtLoss(nn.Module): + """ + Loss function pour l'entraînement Chain-of-Thought. + + Combine: + 1. Language modeling loss (cross-entropy) + 2. Process reward loss (récompenser les bonnes étapes) + 3. Reasoning depth loss (pénaliser trop/peu d'étapes) + 4. Consistency loss (étapes doivent être cohérentes) + """ + + def __init__(self, config: ReasoningConfig): + super().__init__() + self.config = config + + # Poids des différentes pertes + self.lm_weight = 1.0 + self.process_weight = 0.5 + self.depth_weight = 0.1 + self.consistency_weight = 0.1 + + def forward( + self, + logits: torch.Tensor, + targets: torch.Tensor, + step_rewards: Optional[torch.Tensor] = None, + step_labels: Optional[torch.Tensor] = None, + num_steps: Optional[torch.Tensor] = None, + step_hidden_states: Optional[torch.Tensor] = None, + ) -> Dict[str, torch.Tensor]: + """ + Calcule la loss totale pour l'entraînement CoT. + + Args: + logits: (B, L, V) logits du modèle + targets: (B, L) tokens cibles + step_rewards: (B, S) récompenses par étape (optionnel) + step_labels: (B, S) labels des étapes (optionnel) + num_steps: (B,) nombre d'étapes par séquence (optionnel) + step_hidden_states: (B, S, D) pour consistency loss (optionnel) + + Returns: + dict avec toutes les losses et la loss totale + """ + B, L, V = logits.shape + + # 1. Language modeling loss + lm_loss = F.cross_entropy( + logits.view(-1, V), + targets.view(-1), + ignore_index=-100, # Padding token + ) + + total_loss = self.lm_weight * lm_loss + losses = {"lm_loss": lm_loss} + + # 2. Process reward loss (si fourni) + if step_rewards is not None and step_labels is not None: + # Binary cross-entropy pour prédire si étape correcte + process_loss = F.binary_cross_entropy_with_logits( + step_rewards, + step_labels.float(), + ) + total_loss = total_loss + self.process_weight * process_loss + losses["process_loss"] = process_loss + + # 3. Reasoning depth loss + if num_steps is not None: + # Pénaliser si trop loin de la cible optimale + target_steps = (self.config.min_reasoning_steps + self.config.max_reasoning_steps) / 2 + depth_loss = ((num_steps.float() - target_steps) ** 2).mean() + total_loss = total_loss + self.depth_weight * depth_loss + losses["depth_loss"] = depth_loss + + # 4. Consistency loss (étapes successives doivent être cohérentes) + if step_hidden_states is not None and step_hidden_states.size(1) > 1: + # Cosine similarity entre étapes adjacentes + step_i = step_hidden_states[:, :-1, :] # (B, S-1, D) + step_j = step_hidden_states[:, 1:, :] # (B, S-1, D) + + # Normaliser + step_i_norm = F.normalize(step_i, dim=-1) + step_j_norm = F.normalize(step_j, dim=-1) + + # Similarité (voulons qu'elle soit élevée mais pas 1.0) + similarity = (step_i_norm * step_j_norm).sum(dim=-1) # (B, S-1) + + # Cible: similarité autour de 0.7 (cohérent mais pas identique) + target_sim = 0.7 + consistency_loss = ((similarity - target_sim) ** 2).mean() + total_loss = total_loss + self.consistency_weight * consistency_loss + losses["consistency_loss"] = consistency_loss + + losses["total_loss"] = total_loss + return losses + + +class ReasoningLandmarkSelector(nn.Module): + """ + Sélecteur de landmarks spécialisé pour le raisonnement. + + Apprend à identifier les positions importantes pour le raisonnement: + - Prémisses (données du problème) + - Définitions et contraintes + - Étapes intermédiaires clés + - Conclusions partielles + """ + + def __init__( + self, + embed_dim: int, + num_landmarks: int, + num_reasoning_types: int = 4, # prémisse, définition, étape, conclusion + ): + super().__init__() + + self.embed_dim = embed_dim + self.num_landmarks = num_landmarks + self.num_types = num_reasoning_types + + # Classifieur de type de token + self.type_classifier = nn.Sequential( + nn.Linear(embed_dim, embed_dim // 2), + nn.GELU(), + nn.Linear(embed_dim // 2, num_reasoning_types), + ) + + # Importance scorer par type + self.importance_scorers = nn.ModuleList([ + nn.Linear(embed_dim, 1) for _ in range(num_reasoning_types) + ]) + + # Pondération apprise des types + self.type_weights = nn.Parameter(torch.ones(num_reasoning_types)) + + def forward( + self, + x: torch.Tensor, + return_types: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """ + Sélectionne les landmarks importants pour le raisonnement. + + Args: + x: (B, L, D) séquence d'entrée + return_types: si True, retourne aussi les types prédits + + Returns: + landmark_indices: (B, G) indices des landmarks + landmark_states: (B, G, D) états correspondants + type_predictions: (B, L, num_types) si return_types=True + """ + B, L, D = x.shape + + # Classifier les types + type_logits = self.type_classifier(x) # (B, L, num_types) + type_probs = F.softmax(type_logits, dim=-1) # (B, L, num_types) + + # Scorer l'importance par type + scores_per_type = [] + for i, scorer in enumerate(self.importance_scorers): + score = scorer(x).squeeze(-1) # (B, L) + scores_per_type.append(score) + + scores_stacked = torch.stack(scores_per_type, dim=-1) # (B, L, num_types) + + # Score final = somme pondérée par probabilité de type et poids appris + weights = F.softmax(self.type_weights, dim=-1) # (num_types,) + final_scores = (scores_stacked * type_probs * weights).sum(dim=-1) # (B, L) + + # Top-K + k = min(self.num_landmarks, L) + _, landmark_indices = torch.topk(final_scores, k=k, dim=-1) # (B, k) + + # Gather states + landmark_indices_safe = torch.clamp(landmark_indices, 0, L - 1) + landmark_indices_exp = landmark_indices_safe.unsqueeze(-1).expand(B, k, D) + landmark_states = torch.gather(x, dim=1, index=landmark_indices_exp) + + if return_types: + return landmark_indices, landmark_states, type_probs + return landmark_indices, landmark_states, None + + +class SelfConsistencyDecoder(nn.Module): + """ + Décodeur avec Self-Consistency pour améliorer le raisonnement. + + Génère plusieurs chaînes de raisonnement et vote pour la réponse + finale. Implémente "Self-Consistency Improves Chain of Thought + Reasoning in Language Models" (Wang et al., 2022). + """ + + def __init__( + self, + config: ReasoningConfig, + num_samples: int = 5, + temperature: float = 0.7, + ): + super().__init__() + self.config = config + self.num_samples = num_samples + self.temperature = temperature + + # Agrégateur de votes + self.vote_aggregator = nn.Sequential( + nn.Linear(config.embed_dim * num_samples, config.embed_dim), + nn.GELU(), + nn.Linear(config.embed_dim, config.embed_dim), + ) + + def aggregate_answers( + self, + answer_embeddings: torch.Tensor, + answer_tokens: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Agrège plusieurs réponses par vote majoritaire. + + Args: + answer_embeddings: (B, num_samples, D) embeddings des réponses + answer_tokens: (B, num_samples, answer_len) tokens des réponses + + Returns: + best_answer_embedding: (B, D) + vote_distribution: (B, num_samples) poids de chaque réponse + """ + B, S, D = answer_embeddings.shape + + # Calculer similarité entre toutes les paires de réponses + # Réponses similaires = même "vote" + answer_norm = F.normalize(answer_embeddings, dim=-1) # (B, S, D) + similarity = torch.bmm(answer_norm, answer_norm.transpose(1, 2)) # (B, S, S) + + # Score de chaque réponse = somme des similarités (popularité) + vote_scores = similarity.sum(dim=-1) # (B, S) + vote_distribution = F.softmax(vote_scores, dim=-1) # (B, S) + + # Réponse finale = moyenne pondérée + best_answer = (answer_embeddings * vote_distribution.unsqueeze(-1)).sum(dim=1) # (B, D) + + return best_answer, vote_distribution + + +# === Fonctions utilitaires === + +def create_reasoning_tokens(tokenizer, vocab_size: int) -> Dict[str, int]: + """ + Crée les tokens spéciaux pour le raisonnement. + + Returns: + dict avec les IDs des tokens spéciaux + """ + special_tokens = { + "": vocab_size, + "": vocab_size + 1, + "": vocab_size + 2, + "": vocab_size + 3, + } + return special_tokens + + +def format_cot_example( + question: str, + reasoning_steps: List[str], + answer: str, +) -> str: + """ + Formate un exemple Chain-of-Thought. + + Args: + question: Question posée + reasoning_steps: Liste des étapes de raisonnement + answer: Réponse finale + + Returns: + formatted: Texte formaté avec tokens spéciaux + """ + parts = [question, ""] + + for i, step in enumerate(reasoning_steps): + parts.append(f" Step {i+1}: {step}") + + parts.append("") + parts.append(f"{answer}") + + return " ".join(parts) + + +def extract_reasoning_steps( + generated_text: str, +) -> Tuple[List[str], str]: + """ + Extrait les étapes de raisonnement et la réponse d'un texte généré. + + Args: + generated_text: Texte généré par le modèle + + Returns: + steps: Liste des étapes de raisonnement + answer: Réponse finale + """ + steps = [] + answer = "" + + # Extraire entre et + import re + think_match = re.search(r"(.*?)", generated_text, re.DOTALL) + if think_match: + think_content = think_match.group(1) + # Extraire chaque + step_matches = re.findall(r"\s*(.*?)(?=|$)", think_content, re.DOTALL) + steps = [s.strip() for s in step_matches if s.strip()] + + # Extraire la réponse + answer_match = re.search(r"(.*?)", generated_text, re.DOTALL) + if answer_match: + answer = answer_match.group(1).strip() + + return steps, answer + + +__all__ = [ + "ReasoningConfig", + "ThoughtTokenEmbedding", + "ReasoningController", + "ProcessRewardModel", + "ChainOfThoughtLoss", + "ReasoningLandmarkSelector", + "SelfConsistencyDecoder", + "create_reasoning_tokens", + "format_cot_example", + "extract_reasoning_steps", +] diff --git a/src/reasoning_model.py b/src/reasoning_model.py new file mode 100644 index 0000000..7f346ad --- /dev/null +++ b/src/reasoning_model.py @@ -0,0 +1,506 @@ +# reasoning_model.py +""" +SLGA-Reasoning: Modèle de Raisonnement Efficace + +Intègre les composants de raisonnement avec l'architecture SLGA +pour créer un modèle capable de raisonnement multi-étapes efficace. +""" + +from __future__ import annotations +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from dataclasses import dataclass, field +from typing import Optional, Tuple, Dict, List, Any + +from .model import Config, LLMTransformer, TransformerBlock, FeedForward +from .slga import SLGAModule +from .landmarks import LearnableLandmarkSelector +from .reasoning import ( + ReasoningConfig, + ThoughtTokenEmbedding, + ReasoningController, + ProcessRewardModel, + ChainOfThoughtLoss, + ReasoningLandmarkSelector, + SelfConsistencyDecoder, + extract_reasoning_steps, +) + + +@dataclass +class ReasoningModelConfig(Config): + """Configuration étendue pour le modèle de raisonnement""" + + # Hérite de Config, ajoute: + reasoning_enabled: bool = True + + # Tokens spéciaux + think_token_id: int = 50257 + end_think_token_id: int = 50258 + step_token_id: int = 50259 + answer_token_id: int = 50260 + + # Paramètres de raisonnement + max_reasoning_steps: int = 16 + min_reasoning_steps: int = 2 + reasoning_depth_penalty: float = 0.01 + + # Process Reward Model + use_prm: bool = True + prm_hidden_dim: int = 256 + + # Self-Consistency + use_self_consistency: bool = True + num_reasoning_samples: int = 5 + consistency_temperature: float = 0.7 + + +class SLGAReasoningModel(nn.Module): + """ + Modèle de raisonnement basé sur SLGA. + + Architecture: + 1. Token + Position + Thought Embeddings + 2. N × TransformerBlock avec SLGA (landmarks appris pour raisonnement) + 3. Reasoning Controller (décide quand répondre) + 4. Process Reward Model (évalue les étapes) + 5. LM Head + + Capacités: + - Raisonnement multi-étapes avec tokens , , + - Évaluation des étapes intermédiaires (PRM) + - Self-consistency pour améliorer la fiabilité + - Attention sparse efficace pour longues chaînes de raisonnement + """ + + def __init__(self, cfg: ReasoningModelConfig): + super().__init__() + + self.cfg = cfg + + # === Embeddings === + # Token embeddings (incluant tokens spéciaux) + self.token_emb = nn.Embedding(cfg.vocab_size, cfg.embed_dim) + self.pos_emb = nn.Embedding(cfg.max_seq_len, cfg.embed_dim) + self.emb_dropout = nn.Dropout(cfg.dropout_rate) + + # Embeddings spéciaux pour tokens de raisonnement + self.thought_emb = ThoughtTokenEmbedding( + embed_dim=cfg.embed_dim, + num_special_tokens=4, # , , , + ) + + # === Landmark Selector (spécialisé raisonnement) === + if cfg.learned_landmarks: + self.landmark_selector = ReasoningLandmarkSelector( + embed_dim=cfg.embed_dim, + num_landmarks=cfg.global_k * 2, + num_reasoning_types=4, + ) + else: + self.landmark_selector = None + + # === Transformer Blocks === + self.blocks = nn.ModuleList([ + TransformerBlock(cfg, layer_idx=i) for i in range(cfg.n_layers) + ]) + + # === Reasoning Controller === + reasoning_cfg = ReasoningConfig( + embed_dim=cfg.embed_dim, + num_heads=cfg.num_heads, + max_reasoning_steps=cfg.max_reasoning_steps, + min_reasoning_steps=cfg.min_reasoning_steps, + reasoning_depth_penalty=cfg.reasoning_depth_penalty, + use_prm=cfg.use_prm, + prm_hidden_dim=cfg.prm_hidden_dim, + ) + self.reasoning_controller = ReasoningController(reasoning_cfg) + + # === Process Reward Model === + if cfg.use_prm: + self.prm = ProcessRewardModel(reasoning_cfg) + else: + self.prm = None + + # === Self-Consistency Decoder === + if cfg.use_self_consistency: + self.self_consistency = SelfConsistencyDecoder( + reasoning_cfg, + num_samples=cfg.num_reasoning_samples, + temperature=cfg.consistency_temperature, + ) + else: + self.self_consistency = None + + # === Output === + self.final_norm = nn.LayerNorm(cfg.embed_dim) + self.lm_head = nn.Linear(cfg.embed_dim, cfg.vocab_size, bias=False) + + # Tie embeddings + self.lm_head.weight = self.token_emb.weight + + # === Loss === + self.cot_loss = ChainOfThoughtLoss(reasoning_cfg) + + # Initialize + self.apply(self._init_weights) + + # Cache pour les IDs des tokens spéciaux + self._special_token_ids = torch.tensor([ + cfg.think_token_id, + cfg.end_think_token_id, + cfg.step_token_id, + cfg.answer_token_id, + ]) + + def _init_weights(self, module: nn.Module): + """Initialisation GPT-2 style""" + if isinstance(module, nn.Linear): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + elif isinstance(module, nn.LayerNorm): + torch.nn.init.ones_(module.weight) + torch.nn.init.zeros_(module.bias) + + def _identify_special_tokens( + self, input_ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Identifie les tokens spéciaux dans la séquence. + + Returns: + special_mask: (B, L) bool, True si token spécial + special_local_ids: (B, L) IDs locaux (0-3) des tokens spéciaux + """ + B, L = input_ids.shape + device = input_ids.device + + special_ids = self._special_token_ids.to(device) + + # Créer masque pour chaque token spécial + special_mask = torch.zeros(B, L, dtype=torch.bool, device=device) + special_local_ids = torch.zeros(B, L, dtype=torch.long, device=device) + + for i, sid in enumerate(special_ids): + matches = input_ids == sid + special_mask = special_mask | matches + special_local_ids = torch.where(matches, torch.full_like(special_local_ids, i), special_local_ids) + + return special_mask, special_local_ids + + def _find_step_positions(self, input_ids: torch.Tensor) -> torch.Tensor: + """ + Trouve les positions des tokens . + + Returns: + step_positions: (B, max_steps) avec padding -1 + """ + B, L = input_ids.shape + step_token = self.cfg.step_token_id + max_steps = self.cfg.max_reasoning_steps + + step_positions = torch.full( + (B, max_steps), -1, dtype=torch.long, device=input_ids.device + ) + + for b in range(B): + step_mask = input_ids[b] == step_token + step_indices = step_mask.nonzero(as_tuple=True)[0] + num_steps = min(len(step_indices), max_steps) + if num_steps > 0: + step_positions[b, :num_steps] = step_indices[:num_steps] + + return step_positions + + def forward( + self, + input_ids: torch.Tensor, + targets: Optional[torch.Tensor] = None, + return_reasoning_info: bool = False, + global_weight: float = 1.0, + ) -> Dict[str, torch.Tensor]: + """ + Forward pass du modèle de raisonnement. + + Args: + input_ids: (B, L) tokens d'entrée + targets: (B, L) tokens cibles pour l'entraînement + return_reasoning_info: Si True, retourne les infos de raisonnement + global_weight: Poids de l'attention globale + + Returns: + dict avec: + - logits: (B, L, V) + - loss: scalar (si targets fourni) + - reasoning_info: dict (si return_reasoning_info) + """ + B, L = input_ids.shape + device = input_ids.device + + # === 1. Embeddings === + tok_emb = self.token_emb(input_ids) # (B, L, D) + pos = torch.arange(L, device=device).unsqueeze(0).expand(B, L) + pos_emb = self.pos_emb(pos) + + x = tok_emb + pos_emb + + # Enrichir avec thought embeddings + special_mask, special_local_ids = self._identify_special_tokens(input_ids) + x = self.thought_emb(x, special_mask, special_local_ids) + + x = self.emb_dropout(x) + + # === 2. Landmark Selection === + landmark_indices = None + landmark_type_probs = None + + if self.landmark_selector is not None: + landmark_indices, landmark_states, landmark_type_probs = self.landmark_selector( + x, return_types=True + ) + + # === 3. Transformer Blocks === + for block in self.blocks: + if landmark_indices is not None: + B_cur, L_cur, D = x.shape + G = landmark_indices.size(1) + landmark_indices_safe = torch.clamp(landmark_indices, 0, L_cur - 1) + landmark_indices_exp = landmark_indices_safe.unsqueeze(-1).expand(B_cur, G, D) + landmark_states = torch.gather(x, dim=1, index=landmark_indices_exp) + else: + landmark_states = None + + x = block(x, cache_global=landmark_states, global_weight=global_weight) + + # === 4. Final norm and LM head === + x = self.final_norm(x) + logits = self.lm_head(x) # (B, L, V) + + # === 5. Reasoning Controller === + step_positions = self._find_step_positions(input_ids) + valid_step_mask = step_positions >= 0 + + reasoning_decision = self.reasoning_controller( + x, + step_positions=step_positions[valid_step_mask.any(dim=1)] if valid_step_mask.any() else None, + ) + + # === 6. Process Reward Model (si entraînement) === + step_rewards = None + if self.prm is not None and valid_step_mask.any(): + # Gather états aux positions des étapes + valid_positions = step_positions.clamp(min=0) + step_pos_exp = valid_positions.unsqueeze(-1).expand(B, self.cfg.max_reasoning_steps, self.cfg.embed_dim) + step_states = torch.gather(x, dim=1, index=step_pos_exp) + + # Masquer les positions invalides + step_states = step_states * valid_step_mask.unsqueeze(-1).float() + + prm_output = self.prm(step_states) + step_rewards = prm_output["step_rewards"] + + # === 7. Compute Loss (si targets) === + output = {"logits": logits} + + if targets is not None: + # Compter le nombre d'étapes + num_steps = (input_ids == self.cfg.step_token_id).sum(dim=1) + + loss_dict = self.cot_loss( + logits=logits, + targets=targets, + step_rewards=step_rewards, + step_labels=None, # Non supervisé par défaut + num_steps=num_steps, + step_hidden_states=step_states if self.prm is not None and valid_step_mask.any() else None, + ) + output["loss"] = loss_dict["total_loss"] + output["loss_components"] = loss_dict + + # === 8. Reasoning Info === + if return_reasoning_info: + output["reasoning_info"] = { + "decision": reasoning_decision, + "step_positions": step_positions, + "step_rewards": step_rewards, + "landmark_indices": landmark_indices, + "landmark_types": landmark_type_probs, + "num_steps": (input_ids == self.cfg.step_token_id).sum(dim=1), + } + + return output + + @torch.no_grad() + def generate_reasoning( + self, + input_ids: torch.Tensor, + max_new_tokens: int = 512, + temperature: float = 0.0, + top_k: Optional[int] = 50, + top_p: Optional[float] = 0.9, + use_self_consistency: bool = False, + num_samples: int = 5, + ) -> Dict[str, Any]: + """ + Génère une chaîne de raisonnement complète. + + Args: + input_ids: (B, L) prompt (question) + max_new_tokens: Maximum de tokens à générer + temperature: Température (0 = déterministe) + top_k: Top-K sampling + top_p: Nucleus sampling + use_self_consistency: Si True, génère plusieurs chemins et vote + num_samples: Nombre de chemins pour self-consistency + + Returns: + dict avec: + - generated_ids: (B, L + generated) tokens générés + - reasoning_steps: List[str] étapes extraites + - answer: str réponse finale + - confidence: float confiance + """ + self.eval() + B = input_ids.size(0) + device = input_ids.device + + if use_self_consistency and self.self_consistency is not None: + # Générer plusieurs chemins + all_generations = [] + all_answers = [] + + for _ in range(num_samples): + gen = self._generate_single( + input_ids.clone(), + max_new_tokens, + temperature=self.cfg.consistency_temperature, + top_k=top_k, + top_p=top_p, + ) + all_generations.append(gen) + + # Voter pour la meilleure réponse + # TODO: Implémenter le vote basé sur les embeddings + best_gen = all_generations[0] # Simplification + + return { + "generated_ids": best_gen, + "all_generations": all_generations, + "num_samples": num_samples, + } + + else: + # Génération simple + generated = self._generate_single( + input_ids, + max_new_tokens, + temperature, + top_k, + top_p, + ) + + return { + "generated_ids": generated, + } + + def _generate_single( + self, + input_ids: torch.Tensor, + max_new_tokens: int, + temperature: float, + top_k: Optional[int], + top_p: Optional[float], + ) -> torch.Tensor: + """Génération d'une seule séquence.""" + + for _ in range(max_new_tokens): + # Truncate si nécessaire + if input_ids.size(1) > self.cfg.max_seq_len: + input_ids = input_ids[:, -self.cfg.max_seq_len:] + + # Forward + outputs = self(input_ids, return_reasoning_info=True) + logits = outputs["logits"][:, -1, :] # (B, V) + + # Vérifier si le controller dit de s'arrêter + decision = outputs["reasoning_info"]["decision"] + if not decision["should_continue"].any(): + # Forcer génération de + logits[:, self.cfg.answer_token_id] += 10.0 + + # Sampling + if temperature == 0.0: + next_token = torch.argmax(logits, dim=-1, keepdim=True) + else: + logits = logits / temperature + + if top_k is not None: + topk_vals, topk_idxs = torch.topk(logits, k=min(top_k, logits.size(-1))) + logits_filtered = torch.full_like(logits, float('-inf')) + logits_filtered.scatter_(1, topk_idxs, topk_vals) + logits = logits_filtered + + if top_p is not None: + sorted_logits, sorted_indices = torch.sort(logits, descending=True) + cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) + sorted_indices_to_remove = cumulative_probs > top_p + sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() + sorted_indices_to_remove[..., 0] = False + sorted_logits[sorted_indices_to_remove] = float('-inf') + logits = logits.scatter(1, sorted_indices, sorted_logits) + + probs = F.softmax(logits, dim=-1) + next_token = torch.multinomial(probs, num_samples=1) + + input_ids = torch.cat([input_ids, next_token], dim=1) + + # Arrêter si ou EOS + if next_token.item() in [self.cfg.answer_token_id, 50256]: + break + + return input_ids + + def get_num_params(self) -> int: + """Compte les paramètres du modèle.""" + return sum(p.numel() for p in self.parameters()) + + +def create_reasoning_model( + vocab_size: int = 50261, + embed_dim: int = 768, + n_layers: int = 16, + **kwargs +) -> SLGAReasoningModel: + """ + Factory pour créer un modèle de raisonnement. + + Args: + vocab_size: Taille du vocabulaire (incluant tokens spéciaux) + embed_dim: Dimension des embeddings + n_layers: Nombre de couches transformer + **kwargs: Arguments supplémentaires pour ReasoningModelConfig + + Returns: + SLGAReasoningModel configuré + """ + cfg = ReasoningModelConfig( + vocab_size=vocab_size, + embed_dim=embed_dim, + n_layers=n_layers, + num_heads=embed_dim // 64, # 64 dim per head + **kwargs + ) + return SLGAReasoningModel(cfg) + + +__all__ = [ + "ReasoningModelConfig", + "SLGAReasoningModel", + "create_reasoning_model", +] From 5dcb8f95394164d4bbf2dec22e91d39b9802cfa0 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 26 Nov 2025 20:09:57 +0000 Subject: [PATCH 2/5] Add Tree of Thought (ToT) for advanced multi-path reasoning Implements arborescent exploration of reasoning paths: - ThoughtNode: Tree structure with parent/children links - ThoughtGenerator: Creates diverse alternative thoughts - ThoughtEvaluator: Scores coherence, progress, validity - TreeOfThought: Orchestrates search through the tree Search strategies implemented: - BFS: Breadth-first, exhaustive level-by-level - DFS: Depth-first with backtracking - BEAM: Keeps top-K branches at each level - MCTS: Monte Carlo Tree Search with UCB selection - BEST_FIRST: A*-like heuristic search Key improvements over linear CoT: - Explores multiple reasoning paths in parallel - Allows backtracking when a path fails - Voting mechanism for answer confidence - Pruning of low-quality branches Based on "Tree of Thoughts" (Yao et al., 2023) --- src/tree_of_thought.py | 722 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 722 insertions(+) create mode 100644 src/tree_of_thought.py diff --git a/src/tree_of_thought.py b/src/tree_of_thought.py new file mode 100644 index 0000000..c8a72e4 --- /dev/null +++ b/src/tree_of_thought.py @@ -0,0 +1,722 @@ +# tree_of_thought.py +""" +Tree of Thought (ToT) pour SLGA-Reasoning + +Implémente une exploration arborescente du raisonnement: +- Génère plusieurs branches à chaque étape +- Évalue et élague les branches prometteuses +- Permet le backtracking +- Trouve le meilleur chemin de raisonnement + +Basé sur: "Tree of Thoughts: Deliberate Problem Solving with LLMs" (Yao et al., 2023) +""" + +from __future__ import annotations +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Optional, Tuple, Dict, List, Any, Callable +from dataclasses import dataclass, field +from enum import Enum +import heapq + + +class SearchStrategy(Enum): + """Stratégies de recherche dans l'arbre.""" + BFS = "bfs" # Breadth-First Search + DFS = "dfs" # Depth-First Search + BEAM = "beam" # Beam Search + MCTS = "mcts" # Monte Carlo Tree Search + BEST_FIRST = "best_first" # Best-First Search (A*) + + +@dataclass +class ThoughtNode: + """ + Noeud dans l'arbre de pensée. + + Représente une étape de raisonnement avec: + - Le contenu de la pensée + - Son score d'évaluation + - Ses enfants (branches alternatives) + - Son parent (pour backtracking) + """ + thought: str # Contenu de cette étape + hidden_state: Optional[torch.Tensor] # État caché du modèle + score: float = 0.0 # Score d'évaluation + depth: int = 0 # Profondeur dans l'arbre + parent: Optional['ThoughtNode'] = None + children: List['ThoughtNode'] = field(default_factory=list) + is_terminal: bool = False # True si c'est une réponse finale + visits: int = 0 # Pour MCTS + value_sum: float = 0.0 # Pour MCTS + + @property + def value(self) -> float: + """Valeur moyenne (pour MCTS).""" + return self.value_sum / max(1, self.visits) + + @property + def ucb_score(self) -> float: + """Upper Confidence Bound score (pour MCTS).""" + if self.visits == 0: + return float('inf') + exploitation = self.value + exploration = math.sqrt(2 * math.log(self.parent.visits + 1) / self.visits) + return exploitation + exploration + + def get_path(self) -> List['ThoughtNode']: + """Retourne le chemin depuis la racine.""" + path = [] + node = self + while node is not None: + path.append(node) + node = node.parent + return list(reversed(path)) + + def get_full_reasoning(self) -> str: + """Retourne le raisonnement complet depuis la racine.""" + path = self.get_path() + return " → ".join([n.thought for n in path if n.thought]) + + +@dataclass +class ToTConfig: + """Configuration pour Tree of Thought.""" + # Structure de l'arbre + max_depth: int = 8 # Profondeur maximale + branching_factor: int = 3 # Branches par noeud + min_branches: int = 2 # Minimum de branches à explorer + + # Recherche + search_strategy: SearchStrategy = SearchStrategy.BEAM + beam_width: int = 5 # Pour beam search + max_iterations: int = 100 # Limite d'itérations + + # Évaluation + use_value_network: bool = True # Utiliser un réseau de valeur + use_self_evaluation: bool = True # Le modèle s'auto-évalue + pruning_threshold: float = 0.3 # Seuil pour élaguer + + # MCTS spécifique + mcts_simulations: int = 50 # Simulations par noeud + mcts_temperature: float = 1.0 # Température d'exploration + + # Génération + thought_temperature: float = 0.7 # Température pour générer les pensées + thought_max_tokens: int = 100 # Tokens max par pensée + + +class ThoughtGenerator(nn.Module): + """ + Génère des pensées alternatives à chaque étape. + + Utilise le modèle de base avec différents sampling + pour produire des branches diverses. + """ + + def __init__(self, embed_dim: int, vocab_size: int): + super().__init__() + self.embed_dim = embed_dim + self.vocab_size = vocab_size + + # Projecteur de diversité: encourage des pensées différentes + self.diversity_proj = nn.Sequential( + nn.Linear(embed_dim, embed_dim), + nn.GELU(), + nn.Linear(embed_dim, embed_dim), + ) + + # Tête de type de pensée (déduction, hypothèse, vérification, etc.) + self.thought_type_head = nn.Linear(embed_dim, 5) + + def forward( + self, + hidden_states: torch.Tensor, + num_branches: int = 3, + existing_thoughts: Optional[List[torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Génère des représentations pour plusieurs pensées alternatives. + + Args: + hidden_states: (B, L, D) état actuel + num_branches: nombre de branches à générer + existing_thoughts: pensées existantes (pour diversité) + + Returns: + branch_states: (B, num_branches, D) états pour chaque branche + thought_types: (B, num_branches, 5) types de pensée + """ + B, L, D = hidden_states.shape + + # Pooling sur la séquence + pooled = hidden_states.mean(dim=1) # (B, D) + + # Générer des états diversifiés + branch_states = [] + for i in range(num_branches): + # Ajouter du bruit pour diversité + noise = torch.randn_like(pooled) * (0.1 * (i + 1)) + diverse_state = self.diversity_proj(pooled + noise) + + # Répulsion des pensées existantes + if existing_thoughts: + for existing in existing_thoughts: + sim = F.cosine_similarity(diverse_state, existing, dim=-1) + # Pousser loin des pensées similaires + diverse_state = diverse_state - 0.1 * sim.unsqueeze(-1) * existing + + branch_states.append(diverse_state) + + branch_states = torch.stack(branch_states, dim=1) # (B, num_branches, D) + + # Types de pensée + thought_types = self.thought_type_head(branch_states) # (B, num_branches, 5) + + return branch_states, thought_types + + +class ThoughtEvaluator(nn.Module): + """ + Évalue la qualité d'une pensée/étape de raisonnement. + + Critères d'évaluation: + 1. Cohérence avec le contexte + 2. Progrès vers la solution + 3. Validité logique + 4. Originalité (pas de répétition) + """ + + def __init__(self, embed_dim: int, hidden_dim: int = 256): + super().__init__() + + # Encodeur de pensée + self.thought_encoder = nn.Sequential( + nn.Linear(embed_dim, hidden_dim), + nn.GELU(), + nn.Dropout(0.1), + nn.Linear(hidden_dim, hidden_dim), + ) + + # Encodeur de contexte (question + pensées précédentes) + self.context_encoder = nn.Sequential( + nn.Linear(embed_dim, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, hidden_dim), + ) + + # Têtes d'évaluation + self.coherence_head = nn.Linear(hidden_dim * 2, 1) + self.progress_head = nn.Linear(hidden_dim * 2, 1) + self.validity_head = nn.Linear(hidden_dim * 2, 1) + self.final_score_head = nn.Linear(hidden_dim * 2, 1) + + def forward( + self, + thought_state: torch.Tensor, + context_state: torch.Tensor, + return_components: bool = False, + ) -> Dict[str, torch.Tensor]: + """ + Évalue une pensée dans son contexte. + + Args: + thought_state: (B, D) état de la pensée à évaluer + context_state: (B, D) état du contexte (question + historique) + return_components: si True, retourne les scores individuels + + Returns: + dict avec score final et optionnellement les composants + """ + # Encoder + thought_enc = self.thought_encoder(thought_state) # (B, hidden) + context_enc = self.context_encoder(context_state) # (B, hidden) + + # Concatener + combined = torch.cat([thought_enc, context_enc], dim=-1) # (B, 2*hidden) + + # Scores individuels + coherence = torch.sigmoid(self.coherence_head(combined)) + progress = torch.sigmoid(self.progress_head(combined)) + validity = torch.sigmoid(self.validity_head(combined)) + + # Score final (combinaison apprise) + final_score = torch.sigmoid(self.final_score_head(combined)) + + result = {"score": final_score.squeeze(-1)} + + if return_components: + result["coherence"] = coherence.squeeze(-1) + result["progress"] = progress.squeeze(-1) + result["validity"] = validity.squeeze(-1) + + return result + + +class TreeOfThought(nn.Module): + """ + Module principal Tree of Thought. + + Orchestre: + 1. La génération de pensées alternatives + 2. L'évaluation des branches + 3. La recherche dans l'arbre + 4. La sélection du meilleur chemin + """ + + def __init__( + self, + model: nn.Module, # Le modèle SLGA-Reasoning + config: ToTConfig, + embed_dim: int = 768, + vocab_size: int = 50261, + ): + super().__init__() + + self.model = model + self.config = config + self.embed_dim = embed_dim + + # Composants + self.thought_generator = ThoughtGenerator(embed_dim, vocab_size) + self.thought_evaluator = ThoughtEvaluator(embed_dim) + + # Value network (prédit la valeur d'un état) + if config.use_value_network: + self.value_network = nn.Sequential( + nn.Linear(embed_dim, embed_dim // 2), + nn.GELU(), + nn.Linear(embed_dim // 2, 1), + nn.Tanh(), # Valeur entre -1 et 1 + ) + + def generate_thoughts( + self, + node: ThoughtNode, + input_ids: torch.Tensor, + num_thoughts: int = 3, + ) -> List[ThoughtNode]: + """ + Génère des pensées alternatives à partir d'un noeud. + + Args: + node: Noeud parent + input_ids: Tokens du contexte actuel + num_thoughts: Nombre de pensées à générer + + Returns: + Liste de noeuds enfants + """ + device = input_ids.device + + # Forward pour obtenir les états cachés + with torch.no_grad(): + outputs = self.model(input_ids, return_reasoning_info=True) + hidden_states = outputs.get("hidden_states", outputs["logits"]) + + # Si hidden_states est logits, on fait un pooling simple + if hidden_states.dim() == 3 and hidden_states.size(-1) == self.model.cfg.vocab_size: + # C'est logits, on utilise le dernier état avant lm_head + context_state = hidden_states.mean(dim=1) + else: + context_state = hidden_states[:, -1, :] # Dernier état + + # Générer des états diversifiés + existing_states = [c.hidden_state for c in node.children if c.hidden_state is not None] + branch_states, thought_types = self.thought_generator( + hidden_states if hidden_states.dim() == 3 else hidden_states.unsqueeze(1), + num_branches=num_thoughts, + existing_thoughts=existing_states, + ) + + children = [] + for i in range(num_thoughts): + branch_state = branch_states[:, i, :] # (B, D) + + # Générer le texte de la pensée + thought_text = self._generate_thought_text( + input_ids, + branch_state, + thought_type=thought_types[:, i, :].argmax(dim=-1).item(), + ) + + # Évaluer la pensée + eval_result = self.thought_evaluator( + branch_state, + context_state, + return_components=True, + ) + + # Créer le noeud enfant + child = ThoughtNode( + thought=thought_text, + hidden_state=branch_state.detach(), + score=eval_result["score"].item(), + depth=node.depth + 1, + parent=node, + is_terminal=self._is_terminal(thought_text), + ) + + children.append(child) + node.children.append(child) + + return children + + def _generate_thought_text( + self, + input_ids: torch.Tensor, + branch_state: torch.Tensor, + thought_type: int = 0, + ) -> str: + """Génère le texte d'une pensée.""" + # Types de pensée + type_prompts = [ + "Let me deduce: ", # Déduction + "I hypothesize: ", # Hypothèse + "Checking: ", # Vérification + "Breaking down: ", # Décomposition + "Alternatively: ", # Alternative + ] + + prefix = type_prompts[thought_type % len(type_prompts)] + + # Génération simplifiée (en pratique, utiliserait le modèle) + # TODO: Implémenter génération complète avec le modèle + return f"{prefix}[Generated thought at depth {branch_state.mean().item():.2f}]" + + def _is_terminal(self, thought: str) -> bool: + """Vérifie si une pensée est terminale (réponse finale).""" + terminal_indicators = ["answer:", "therefore:", "conclusion:", "final answer:"] + thought_lower = thought.lower() + return any(ind in thought_lower for ind in terminal_indicators) + + def search( + self, + input_ids: torch.Tensor, + question: str = "", + ) -> Tuple[ThoughtNode, List[ThoughtNode]]: + """ + Recherche le meilleur chemin de raisonnement. + + Args: + input_ids: Tokens de la question + question: Texte de la question (optionnel) + + Returns: + best_node: Meilleur noeud terminal trouvé + all_nodes: Tous les noeuds explorés + """ + strategy = self.config.search_strategy + + if strategy == SearchStrategy.BEAM: + return self._beam_search(input_ids, question) + elif strategy == SearchStrategy.BFS: + return self._bfs_search(input_ids, question) + elif strategy == SearchStrategy.DFS: + return self._dfs_search(input_ids, question) + elif strategy == SearchStrategy.MCTS: + return self._mcts_search(input_ids, question) + elif strategy == SearchStrategy.BEST_FIRST: + return self._best_first_search(input_ids, question) + else: + raise ValueError(f"Unknown search strategy: {strategy}") + + def _beam_search( + self, + input_ids: torch.Tensor, + question: str, + ) -> Tuple[ThoughtNode, List[ThoughtNode]]: + """Beam Search: garde les K meilleures branches à chaque niveau.""" + + # Racine + root = ThoughtNode(thought=question, hidden_state=None, depth=0) + beam = [root] + all_nodes = [root] + terminals = [] + + for depth in range(self.config.max_depth): + candidates = [] + + for node in beam: + if node.is_terminal: + terminals.append(node) + continue + + # Générer des enfants + children = self.generate_thoughts( + node, input_ids, + num_thoughts=self.config.branching_factor + ) + candidates.extend(children) + all_nodes.extend(children) + + if not candidates: + break + + # Garder les K meilleurs + candidates.sort(key=lambda n: n.score, reverse=True) + beam = candidates[:self.config.beam_width] + + # Élaguer les branches faibles + beam = [n for n in beam if n.score >= self.config.pruning_threshold] + + # Ajouter les noeuds terminaux restants + terminals.extend([n for n in beam if n.is_terminal]) + + # Retourner le meilleur + if terminals: + best = max(terminals, key=lambda n: n.score) + else: + best = max(all_nodes, key=lambda n: n.score) + + return best, all_nodes + + def _bfs_search( + self, + input_ids: torch.Tensor, + question: str, + ) -> Tuple[ThoughtNode, List[ThoughtNode]]: + """Breadth-First Search: explore niveau par niveau.""" + from collections import deque + + root = ThoughtNode(thought=question, hidden_state=None, depth=0) + queue = deque([root]) + all_nodes = [root] + best_terminal = None + + iterations = 0 + while queue and iterations < self.config.max_iterations: + node = queue.popleft() + iterations += 1 + + if node.is_terminal: + if best_terminal is None or node.score > best_terminal.score: + best_terminal = node + continue + + if node.depth >= self.config.max_depth: + continue + + children = self.generate_thoughts(node, input_ids) + for child in children: + if child.score >= self.config.pruning_threshold: + queue.append(child) + all_nodes.append(child) + + return best_terminal or root, all_nodes + + def _dfs_search( + self, + input_ids: torch.Tensor, + question: str, + ) -> Tuple[ThoughtNode, List[ThoughtNode]]: + """Depth-First Search avec backtracking.""" + + root = ThoughtNode(thought=question, hidden_state=None, depth=0) + all_nodes = [root] + best_terminal = None + + def dfs(node: ThoughtNode, depth: int): + nonlocal best_terminal + + if node.is_terminal: + if best_terminal is None or node.score > best_terminal.score: + best_terminal = node + return + + if depth >= self.config.max_depth: + return + + children = self.generate_thoughts(node, input_ids) + # Trier par score (explorer les meilleurs d'abord) + children.sort(key=lambda n: n.score, reverse=True) + + for child in children: + if child.score >= self.config.pruning_threshold: + all_nodes.append(child) + dfs(child, depth + 1) + + # Backtrack si on a trouvé une solution + if best_terminal and best_terminal.score > 0.9: + return + + dfs(root, 0) + return best_terminal or root, all_nodes + + def _mcts_search( + self, + input_ids: torch.Tensor, + question: str, + ) -> Tuple[ThoughtNode, List[ThoughtNode]]: + """Monte Carlo Tree Search.""" + + root = ThoughtNode(thought=question, hidden_state=None, depth=0) + all_nodes = [root] + + for _ in range(self.config.mcts_simulations): + # 1. Selection: descendre jusqu'à une feuille + node = root + while node.children and not node.is_terminal: + # UCB selection + node = max(node.children, key=lambda n: n.ucb_score) + + # 2. Expansion: ajouter des enfants + if not node.is_terminal and node.depth < self.config.max_depth: + children = self.generate_thoughts(node, input_ids, num_thoughts=1) + if children: + node = children[0] + all_nodes.append(node) + + # 3. Simulation: évaluer + value = self._simulate(node, input_ids) + + # 4. Backpropagation + while node is not None: + node.visits += 1 + node.value_sum += value + node = node.parent + + # Sélectionner le meilleur enfant de la racine + if root.children: + best = max(root.children, key=lambda n: n.visits) + # Descendre jusqu'à un terminal + while best.children: + best = max(best.children, key=lambda n: n.visits) + else: + best = root + + return best, all_nodes + + def _simulate(self, node: ThoughtNode, input_ids: torch.Tensor) -> float: + """Simulation rapide pour MCTS.""" + if node.is_terminal: + return node.score + + # Simulation simple: utiliser le value network + if self.config.use_value_network and node.hidden_state is not None: + value = self.value_network(node.hidden_state) + return value.item() + + return node.score + + def _best_first_search( + self, + input_ids: torch.Tensor, + question: str, + ) -> Tuple[ThoughtNode, List[ThoughtNode]]: + """Best-First Search (A*-like).""" + + root = ThoughtNode(thought=question, hidden_state=None, depth=0) + # Priority queue: (-score, node) car heapq est min-heap + frontier = [(-root.score, id(root), root)] + all_nodes = [root] + best_terminal = None + + iterations = 0 + while frontier and iterations < self.config.max_iterations: + _, _, node = heapq.heappop(frontier) + iterations += 1 + + if node.is_terminal: + if best_terminal is None or node.score > best_terminal.score: + best_terminal = node + continue + + if node.depth >= self.config.max_depth: + continue + + children = self.generate_thoughts(node, input_ids) + for child in children: + if child.score >= self.config.pruning_threshold: + heapq.heappush(frontier, (-child.score, id(child), child)) + all_nodes.append(child) + + return best_terminal or root, all_nodes + + def get_best_reasoning_path( + self, + input_ids: torch.Tensor, + question: str = "", + ) -> Dict[str, Any]: + """ + Interface principale: trouve le meilleur chemin de raisonnement. + + Returns: + dict avec: + - path: Liste des pensées dans l'ordre + - answer: Réponse finale + - score: Score du chemin + - tree_stats: Statistiques de l'arbre + """ + best_node, all_nodes = self.search(input_ids, question) + + path = best_node.get_path() + + return { + "path": [n.thought for n in path], + "full_reasoning": best_node.get_full_reasoning(), + "answer": path[-1].thought if path else "", + "score": best_node.score, + "depth": best_node.depth, + "tree_stats": { + "total_nodes": len(all_nodes), + "max_depth_reached": max(n.depth for n in all_nodes), + "terminals_found": sum(1 for n in all_nodes if n.is_terminal), + "avg_branching": sum(len(n.children) for n in all_nodes) / max(1, len(all_nodes)), + } + } + + +def create_tree_of_thought( + model: nn.Module, + strategy: str = "beam", + beam_width: int = 5, + max_depth: int = 8, + **kwargs +) -> TreeOfThought: + """ + Factory pour créer un module Tree of Thought. + + Args: + model: Modèle SLGA-Reasoning + strategy: "beam", "bfs", "dfs", "mcts", "best_first" + beam_width: Largeur du beam search + max_depth: Profondeur maximale de l'arbre + + Returns: + TreeOfThought configuré + """ + strategy_map = { + "beam": SearchStrategy.BEAM, + "bfs": SearchStrategy.BFS, + "dfs": SearchStrategy.DFS, + "mcts": SearchStrategy.MCTS, + "best_first": SearchStrategy.BEST_FIRST, + } + + config = ToTConfig( + search_strategy=strategy_map.get(strategy, SearchStrategy.BEAM), + beam_width=beam_width, + max_depth=max_depth, + **kwargs + ) + + return TreeOfThought( + model=model, + config=config, + embed_dim=getattr(model.cfg, 'embed_dim', 768), + vocab_size=getattr(model.cfg, 'vocab_size', 50261), + ) + + +__all__ = [ + "SearchStrategy", + "ThoughtNode", + "ToTConfig", + "ThoughtGenerator", + "ThoughtEvaluator", + "TreeOfThought", + "create_tree_of_thought", +] From f5145dea9681be3f4963ec73a89793b2837af60a Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 26 Nov 2025 20:10:20 +0000 Subject: [PATCH 3/5] Export Tree of Thought modules in package init --- src/__init__.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/__init__.py b/src/__init__.py index f1b5f09..06e9dfa 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -32,6 +32,13 @@ SLGAReasoningModel, create_reasoning_model, ) +from .tree_of_thought import ( + SearchStrategy, + ThoughtNode, + ToTConfig, + TreeOfThought, + create_tree_of_thought, +) __all__ = [ # Core @@ -59,4 +66,10 @@ "ReasoningModelConfig", "SLGAReasoningModel", "create_reasoning_model", + # Tree of Thought + "SearchStrategy", + "ThoughtNode", + "ToTConfig", + "TreeOfThought", + "create_tree_of_thought", ] From 5ae11755c721b980b25ad194fd0e0e1d16825b24 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 26 Nov 2025 20:16:32 +0000 Subject: [PATCH 4/5] Add complete standalone training script for reasoning model Self-contained script with everything included: - ReasoningModel: Simplified SLGA with landmarks and PRM - SimpleTokenizer: GPT-2 based with special token handling - CoTDataset: Formats data as chain-of-thought - Data loaders: GSM8K + synthetic fallback - Training loop with AMP, gradient accumulation, scheduling - Evaluation and checkpointing Usage: python train_reasoning_simple.py # Full training python train_reasoning_simple.py --small # Quick test python train_reasoning_simple.py --dataset gsm8k No external dependencies beyond torch, datasets, transformers --- train_reasoning_simple.py | 831 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 831 insertions(+) create mode 100644 train_reasoning_simple.py diff --git a/train_reasoning_simple.py b/train_reasoning_simple.py new file mode 100644 index 0000000..5dfaa85 --- /dev/null +++ b/train_reasoning_simple.py @@ -0,0 +1,831 @@ +#!/usr/bin/env python3 +""" +🧠 SLGA-Reasoning: Script d'Entraînement Complet +================================================ + +Ce script est AUTONOME - tout est inclus pour entraîner un reasoning model. + +Usage: + python train_reasoning_simple.py # Entraînement par défaut + python train_reasoning_simple.py --small # Version légère (test) + python train_reasoning_simple.py --dataset gsm8k # Dataset spécifique + python train_reasoning_simple.py --resume ckpt.pt # Reprendre + +Prérequis: + pip install torch datasets transformers tqdm +""" + +import os +import sys +import argparse +import json +import math +import time +import random +from pathlib import Path +from typing import Dict, Any, Optional, List, Tuple +from dataclasses import dataclass, field + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import DataLoader, Dataset + +# ============================================================================== +# CONFIGURATION +# ============================================================================== + +@dataclass +class TrainConfig: + """Configuration complète d'entraînement.""" + # Modèle + vocab_size: int = 50261 # GPT-2 + 4 tokens spéciaux + embed_dim: int = 512 # 768 pour plus de capacité + num_heads: int = 8 + n_layers: int = 12 # 16 pour plus de capacité + max_seq_len: int = 1024 # 2048+ pour raisonnement long + dropout: float = 0.1 + + # SLGA + local_window: int = 128 + global_k: int = 32 + + # Reasoning + max_reasoning_steps: int = 12 + use_prm: bool = True + + # Entraînement + batch_size: int = 4 + accum_steps: int = 4 # Batch effectif = 16 + lr: float = 3e-4 + weight_decay: float = 0.1 + warmup_steps: int = 500 + max_steps: int = 10000 + eval_every: int = 500 + save_every: int = 1000 + + # Hardware + device: str = "cuda" if torch.cuda.is_available() else "cpu" + amp: bool = True + + # Chemins + output_dir: str = "checkpoints/reasoning" + + @classmethod + def small(cls) -> 'TrainConfig': + """Config légère pour test.""" + return cls( + embed_dim=256, + num_heads=4, + n_layers=4, + max_seq_len=512, + batch_size=2, + max_steps=1000, + eval_every=100, + ) + + +# ============================================================================== +# TOKENS SPÉCIAUX +# ============================================================================== + +SPECIAL_TOKENS = { + "": 50257, + "": 50258, + "": 50259, + "": 50260, +} + +# ============================================================================== +# MODÈLE SIMPLIFIÉ (tout-en-un) +# ============================================================================== + +class SLGAAttention(nn.Module): + """Sparse Local-Global Attention simplifié.""" + + def __init__(self, embed_dim: int, num_heads: int, local_window: int = 128, global_k: int = 32): + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.head_dim = embed_dim // num_heads + self.local_window = local_window + self.global_k = global_k + self.scale = self.head_dim ** -0.5 + + self.qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=False) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False) + self.dropout = nn.Dropout(0.1) + + # Gate pour fusion local/global + self.gate = nn.Linear(2 * self.head_dim, self.head_dim) + + def forward(self, x: torch.Tensor, landmarks: Optional[torch.Tensor] = None) -> torch.Tensor: + B, L, D = x.shape + + # QKV + qkv = self.qkv(x).reshape(B, L, 3, self.num_heads, self.head_dim) + qkv = qkv.permute(2, 0, 3, 1, 4) # (3, B, H, L, Dh) + q, k, v = qkv[0], qkv[1], qkv[2] + + # === Local Attention (fenêtre glissante) === + # Simplifié: attention causale standard pour les petites séquences + scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale + + # Masque causal + mask = torch.triu(torch.ones(L, L, device=x.device), diagonal=1).bool() + scores = scores.masked_fill(mask, float('-inf')) + + # Masque fenêtre locale (optionnel pour efficacité) + if L > self.local_window * 2: + window_mask = torch.ones(L, L, device=x.device).bool() + for i in range(L): + start = max(0, i - self.local_window) + window_mask[i, start:i+1] = False + scores = scores.masked_fill(window_mask, float('-inf')) + + attn = F.softmax(scores, dim=-1) + attn = self.dropout(attn) + + ctx_local = torch.matmul(attn, v) # (B, H, L, Dh) + + # === Global Attention (sur landmarks) === + if landmarks is not None: + G = landmarks.size(1) + # Project landmarks + lm_qkv = self.qkv(landmarks).reshape(B, G, 3, self.num_heads, self.head_dim) + lm_qkv = lm_qkv.permute(2, 0, 3, 1, 4) + lm_k, lm_v = lm_qkv[1], lm_qkv[2] + + # Global scores + global_scores = torch.matmul(q, lm_k.transpose(-2, -1)) * self.scale + global_attn = F.softmax(global_scores, dim=-1) + ctx_global = torch.matmul(global_attn, lm_v) + + # Fusion avec gate + gate_input = torch.cat([ctx_local, ctx_global], dim=-1) + gate_input = gate_input.permute(0, 2, 1, 3).reshape(B * L, self.num_heads, 2 * self.head_dim) + gate_weights = torch.sigmoid(self.gate(gate_input)) + gate_weights = gate_weights.reshape(B, L, self.num_heads, self.head_dim).permute(0, 2, 1, 3) + + ctx = gate_weights * ctx_local + (1 - gate_weights) * ctx_global + else: + ctx = ctx_local + + # Output + ctx = ctx.permute(0, 2, 1, 3).reshape(B, L, D) + return self.out_proj(ctx) + + +class TransformerBlock(nn.Module): + """Bloc Transformer avec SLGA.""" + + def __init__(self, embed_dim: int, num_heads: int, local_window: int, global_k: int, dropout: float = 0.1): + super().__init__() + self.norm1 = nn.LayerNorm(embed_dim) + self.attn = SLGAAttention(embed_dim, num_heads, local_window, global_k) + self.norm2 = nn.LayerNorm(embed_dim) + self.ffn = nn.Sequential( + nn.Linear(embed_dim, embed_dim * 4), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(embed_dim * 4, embed_dim), + nn.Dropout(dropout), + ) + + def forward(self, x: torch.Tensor, landmarks: Optional[torch.Tensor] = None) -> torch.Tensor: + x = x + self.attn(self.norm1(x), landmarks) + x = x + self.ffn(self.norm2(x)) + return x + + +class ReasoningModel(nn.Module): + """Modèle de raisonnement complet.""" + + def __init__(self, config: TrainConfig): + super().__init__() + self.config = config + + # Embeddings + self.token_emb = nn.Embedding(config.vocab_size, config.embed_dim) + self.pos_emb = nn.Embedding(config.max_seq_len, config.embed_dim) + self.dropout = nn.Dropout(config.dropout) + + # Transformer blocks + self.blocks = nn.ModuleList([ + TransformerBlock( + config.embed_dim, + config.num_heads, + config.local_window, + config.global_k, + config.dropout + ) for _ in range(config.n_layers) + ]) + + # Output + self.norm = nn.LayerNorm(config.embed_dim) + self.lm_head = nn.Linear(config.embed_dim, config.vocab_size, bias=False) + self.lm_head.weight = self.token_emb.weight # Tie weights + + # Landmark selector (simple) + self.landmark_scorer = nn.Sequential( + nn.Linear(config.embed_dim, config.embed_dim // 2), + nn.GELU(), + nn.Linear(config.embed_dim // 2, 1), + ) + + # Process Reward Model (optionnel) + if config.use_prm: + self.prm_head = nn.Sequential( + nn.Linear(config.embed_dim, config.embed_dim // 2), + nn.GELU(), + nn.Linear(config.embed_dim // 2, 1), + ) + + self._init_weights() + + def _init_weights(self): + for module in self.modules(): + if isinstance(module, nn.Linear): + torch.nn.init.normal_(module.weight, std=0.02) + if module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + torch.nn.init.normal_(module.weight, std=0.02) + + def _select_landmarks(self, x: torch.Tensor, k: int) -> torch.Tensor: + """Sélectionne les K positions les plus importantes.""" + B, L, D = x.shape + scores = self.landmark_scorer(x).squeeze(-1) # (B, L) + _, indices = torch.topk(scores, k=min(k, L), dim=-1) # (B, K) + + # Gather + indices_exp = indices.unsqueeze(-1).expand(-1, -1, D) + landmarks = torch.gather(x, dim=1, index=indices_exp) + return landmarks + + def forward( + self, + input_ids: torch.Tensor, + targets: Optional[torch.Tensor] = None + ) -> Dict[str, torch.Tensor]: + B, L = input_ids.shape + device = input_ids.device + + # Embeddings + tok_emb = self.token_emb(input_ids) + pos = torch.arange(L, device=device).unsqueeze(0) + pos_emb = self.pos_emb(pos) + x = self.dropout(tok_emb + pos_emb) + + # Transformer avec landmarks + landmarks = None + for i, block in enumerate(self.blocks): + # Sélectionner landmarks après quelques couches + if i == self.config.n_layers // 3: + landmarks = self._select_landmarks(x, self.config.global_k) + x = block(x, landmarks) + # Mettre à jour landmarks + if landmarks is not None and i % 2 == 0: + landmarks = self._select_landmarks(x, self.config.global_k) + + # Output + x = self.norm(x) + logits = self.lm_head(x) + + output = {"logits": logits} + + # Loss + if targets is not None: + loss = F.cross_entropy( + logits.view(-1, self.config.vocab_size), + targets.view(-1), + ignore_index=-100, + ) + output["loss"] = loss + + # PRM loss sur les tokens + if self.config.use_prm: + step_mask = (input_ids == SPECIAL_TOKENS[""]) + if step_mask.any(): + step_states = x[step_mask] + step_rewards = self.prm_head(step_states).squeeze(-1) + # Récompense positive par défaut (sera affinée avec labels) + prm_loss = F.binary_cross_entropy_with_logits( + step_rewards, + torch.ones_like(step_rewards) + ) + output["prm_loss"] = prm_loss + output["loss"] = output["loss"] + 0.1 * prm_loss + + return output + + @torch.no_grad() + def generate( + self, + input_ids: torch.Tensor, + max_new_tokens: int = 256, + temperature: float = 0.7, + top_p: float = 0.9, + ) -> torch.Tensor: + """Génère du texte avec raisonnement.""" + self.eval() + + for _ in range(max_new_tokens): + # Truncate si nécessaire + if input_ids.size(1) >= self.config.max_seq_len: + input_ids = input_ids[:, -self.config.max_seq_len:] + + # Forward + outputs = self(input_ids) + logits = outputs["logits"][:, -1, :] + + # Sampling + if temperature > 0: + logits = logits / temperature + + # Top-p + sorted_logits, sorted_indices = torch.sort(logits, descending=True) + cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) + sorted_indices_to_remove = cumulative_probs > top_p + sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() + sorted_indices_to_remove[..., 0] = False + sorted_logits[sorted_indices_to_remove] = float('-inf') + + probs = F.softmax(sorted_logits, dim=-1) + idx = torch.multinomial(probs, num_samples=1) + next_token = torch.gather(sorted_indices, -1, idx) + else: + next_token = torch.argmax(logits, dim=-1, keepdim=True) + + input_ids = torch.cat([input_ids, next_token], dim=1) + + # Stop sur ou + if next_token.item() in [SPECIAL_TOKENS[""], SPECIAL_TOKENS[""], 50256]: + break + + return input_ids + + +# ============================================================================== +# DATASET +# ============================================================================== + +class CoTDataset(Dataset): + """Dataset Chain-of-Thought.""" + + def __init__( + self, + data: List[Dict[str, Any]], + tokenizer, + max_length: int = 1024 + ): + self.data = data + self.tokenizer = tokenizer + self.max_length = max_length + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + item = self.data[idx] + + # Formater en CoT + text = self._format_cot(item) + + # Tokenizer + tokens = self.tokenizer.encode(text) + + # Truncate + if len(tokens) > self.max_length: + tokens = tokens[:self.max_length] + + # Créer input/target + input_ids = torch.tensor(tokens[:-1], dtype=torch.long) + targets = torch.tensor(tokens[1:], dtype=torch.long) + + return {"input_ids": input_ids, "targets": targets} + + def _format_cot(self, item: Dict[str, Any]) -> str: + """Formate un exemple en Chain-of-Thought.""" + question = item.get("question", item.get("problem", "")) + + # Si des étapes sont fournies + if "reasoning_steps" in item: + steps = item["reasoning_steps"] + answer = item.get("answer", "") + + steps_text = " ".join([f" {s}" for s in steps]) + return f"Question: {question}\n{steps_text}\n{answer}" + + # Si solution brute + elif "solution" in item: + solution = item["solution"] + answer = item.get("answer", solution.split("=")[-1].strip() if "=" in solution else "") + return f"Question: {question}\n {solution}\n{answer}" + + # Fallback + else: + answer = item.get("answer", "") + return f"Question: {question}\n Let me solve this.\n{answer}" + + +def collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]: + """Collate avec padding.""" + max_len = max(item["input_ids"].size(0) for item in batch) + + input_ids = [] + targets = [] + + for item in batch: + pad_len = max_len - item["input_ids"].size(0) + input_ids.append(F.pad(item["input_ids"], (0, pad_len), value=50256)) + targets.append(F.pad(item["targets"], (0, pad_len), value=-100)) + + return { + "input_ids": torch.stack(input_ids), + "targets": torch.stack(targets), + } + + +# ============================================================================== +# CHARGEMENT DES DONNÉES +# ============================================================================== + +def load_gsm8k() -> List[Dict[str, Any]]: + """Charge GSM8K.""" + try: + from datasets import load_dataset + ds = load_dataset("openai/gsm8k", "main", split="train") + + data = [] + for item in ds: + answer_text = item["answer"] + parts = answer_text.split("####") + + if len(parts) == 2: + steps = [s.strip() for s in parts[0].split("\n") if s.strip()] + final = parts[1].strip() + data.append({ + "question": item["question"], + "reasoning_steps": steps, + "answer": final, + }) + + print(f"✓ GSM8K: {len(data)} exemples chargés") + return data + + except Exception as e: + print(f"⚠ GSM8K non disponible: {e}") + return [] + + +def load_synthetic_data(n: int = 1000) -> List[Dict[str, Any]]: + """Génère des données synthétiques simples.""" + data = [] + + for _ in range(n): + # Arithmétique simple + a = random.randint(1, 100) + b = random.randint(1, 100) + op = random.choice(["+", "-", "*"]) + + if op == "+": + result = a + b + steps = [f"{a} + {b} = {result}"] + elif op == "-": + result = a - b + steps = [f"{a} - {b} = {result}"] + else: + result = a * b + steps = [f"{a} × {b} = {result}"] + + data.append({ + "question": f"What is {a} {op} {b}?", + "reasoning_steps": steps, + "answer": str(result), + }) + + print(f"✓ Synthétique: {len(data)} exemples générés") + return data + + +def load_data(dataset_name: str = "all") -> List[Dict[str, Any]]: + """Charge les données selon le choix.""" + data = [] + + if dataset_name in ["gsm8k", "all"]: + data.extend(load_gsm8k()) + + if dataset_name in ["synthetic", "all"] or len(data) == 0: + data.extend(load_synthetic_data(2000 if len(data) == 0 else 500)) + + random.shuffle(data) + return data + + +# ============================================================================== +# TOKENIZER SIMPLE +# ============================================================================== + +class SimpleTokenizer: + """Tokenizer simple basé sur GPT-2.""" + + def __init__(self): + try: + from transformers import GPT2Tokenizer + self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2") + self.tokenizer.pad_token = self.tokenizer.eos_token + except: + self.tokenizer = None + print("⚠ transformers non installé, utilisation d'un tokenizer basique") + + def encode(self, text: str) -> List[int]: + if self.tokenizer: + # Remplacer les tokens spéciaux avant encodage + for token, idx in SPECIAL_TOKENS.items(): + text = text.replace(token, f" {token} ") + + tokens = self.tokenizer.encode(text) + + # Post-process: remplacer les tokens spéciaux + result = [] + i = 0 + while i < len(tokens): + # Chercher les tokens spéciaux dans le texte décodé + found = False + for token, idx in SPECIAL_TOKENS.items(): + token_ids = self.tokenizer.encode(token, add_special_tokens=False) + if tokens[i:i+len(token_ids)] == token_ids: + result.append(idx) + i += len(token_ids) + found = True + break + + if not found: + result.append(tokens[i]) + i += 1 + + return result + else: + # Fallback très basique + return [ord(c) % 50000 for c in text] + + def decode(self, tokens: List[int]) -> str: + if self.tokenizer: + # Filtrer les tokens spéciaux pour le décodage + regular_tokens = [t for t in tokens if t < 50257] + text = self.tokenizer.decode(regular_tokens) + + # Réinsérer les tokens spéciaux + for token, idx in SPECIAL_TOKENS.items(): + if idx in tokens: + text = text # Simplification + + return text + else: + return "".join([chr(t) for t in tokens if t < 128]) + + +# ============================================================================== +# ENTRAÎNEMENT +# ============================================================================== + +def train(config: TrainConfig, dataset_name: str = "all", resume: Optional[str] = None): + """Boucle d'entraînement principale.""" + + print("=" * 60) + print("🧠 SLGA-Reasoning Training") + print("=" * 60) + print(f"Device: {config.device}") + print(f"Embed dim: {config.embed_dim}") + print(f"Layers: {config.n_layers}") + print(f"Batch size: {config.batch_size} x {config.accum_steps} = {config.batch_size * config.accum_steps}") + print("=" * 60) + + # Créer le dossier de sortie + os.makedirs(config.output_dir, exist_ok=True) + + # Tokenizer + tokenizer = SimpleTokenizer() + + # Données + print("\n📚 Chargement des données...") + data = load_data(dataset_name) + + if len(data) == 0: + print("❌ Aucune donnée disponible!") + return + + # Split train/val + split_idx = int(len(data) * 0.95) + train_data = data[:split_idx] + val_data = data[split_idx:] + + print(f"Train: {len(train_data)}, Val: {len(val_data)}") + + # Datasets + train_dataset = CoTDataset(train_data, tokenizer, config.max_seq_len) + val_dataset = CoTDataset(val_data, tokenizer, config.max_seq_len) + + train_loader = DataLoader( + train_dataset, + batch_size=config.batch_size, + shuffle=True, + collate_fn=collate_fn, + num_workers=0, + pin_memory=True, + ) + val_loader = DataLoader( + val_dataset, + batch_size=config.batch_size, + collate_fn=collate_fn, + ) + + # Modèle + print("\n🏗️ Création du modèle...") + model = ReasoningModel(config).to(config.device) + + num_params = sum(p.numel() for p in model.parameters()) + print(f"Paramètres: {num_params / 1e6:.1f}M") + + # Optimizer + optimizer = torch.optim.AdamW( + model.parameters(), + lr=config.lr, + weight_decay=config.weight_decay, + betas=(0.9, 0.95), + ) + + # Scheduler + def lr_lambda(step): + if step < config.warmup_steps: + return step / config.warmup_steps + progress = (step - config.warmup_steps) / (config.max_steps - config.warmup_steps) + return max(0.1, 0.5 * (1 + math.cos(math.pi * progress))) + + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) + + # AMP + scaler = torch.cuda.amp.GradScaler() if config.amp and config.device == "cuda" else None + + # Resume + start_step = 0 + if resume: + print(f"\n📂 Reprise depuis {resume}") + ckpt = torch.load(resume, map_location=config.device) + model.load_state_dict(ckpt["model"]) + optimizer.load_state_dict(ckpt["optimizer"]) + start_step = ckpt.get("step", 0) + + # Training loop + print("\n🚀 Démarrage de l'entraînement...") + model.train() + + train_iter = iter(train_loader) + accum_loss = 0.0 + accum_steps = 0 + best_val_loss = float('inf') + + start_time = time.time() + + for step in range(start_step, config.max_steps): + # Get batch + try: + batch = next(train_iter) + except StopIteration: + train_iter = iter(train_loader) + batch = next(train_iter) + + input_ids = batch["input_ids"].to(config.device) + targets = batch["targets"].to(config.device) + + # Forward + with torch.cuda.amp.autocast(enabled=config.amp and config.device == "cuda"): + outputs = model(input_ids, targets) + loss = outputs["loss"] / config.accum_steps + + # Backward + if scaler: + scaler.scale(loss).backward() + else: + loss.backward() + + accum_loss += loss.item() * config.accum_steps + accum_steps += 1 + + # Optimizer step + if accum_steps == config.accum_steps: + if scaler: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + scaler.step(optimizer) + scaler.update() + else: + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + + scheduler.step() + optimizer.zero_grad() + + # Log + if (step + 1) % 50 == 0: + elapsed = time.time() - start_time + steps_per_sec = (step + 1 - start_step) / elapsed + lr = scheduler.get_last_lr()[0] + + print(f"Step {step+1}/{config.max_steps} | " + f"Loss: {accum_loss:.4f} | " + f"LR: {lr:.2e} | " + f"Speed: {steps_per_sec:.1f} steps/s") + + accum_loss = 0.0 + accum_steps = 0 + + # Eval + if (step + 1) % config.eval_every == 0: + val_loss = evaluate(model, val_loader, config) + print(f" → Val loss: {val_loss:.4f}") + + if val_loss < best_val_loss: + best_val_loss = val_loss + save_checkpoint(model, optimizer, step, config, "best") + + model.train() + + # Save + if (step + 1) % config.save_every == 0: + save_checkpoint(model, optimizer, step, config, f"step_{step+1}") + + # Final save + save_checkpoint(model, optimizer, config.max_steps, config, "final") + print("\n✅ Entraînement terminé!") + + return model + + +@torch.no_grad() +def evaluate(model: nn.Module, dataloader: DataLoader, config: TrainConfig) -> float: + """Évalue le modèle.""" + model.eval() + total_loss = 0.0 + count = 0 + + for batch in dataloader: + input_ids = batch["input_ids"].to(config.device) + targets = batch["targets"].to(config.device) + + with torch.cuda.amp.autocast(enabled=config.amp and config.device == "cuda"): + outputs = model(input_ids, targets) + total_loss += outputs["loss"].item() + + count += 1 + if count >= 50: + break + + return total_loss / count + + +def save_checkpoint(model: nn.Module, optimizer, step: int, config: TrainConfig, name: str): + """Sauvegarde un checkpoint.""" + path = os.path.join(config.output_dir, f"checkpoint_{name}.pt") + torch.save({ + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "step": step, + "config": config, + }, path) + print(f" 💾 Saved: {path}") + + +# ============================================================================== +# MAIN +# ============================================================================== + +def main(): + parser = argparse.ArgumentParser(description="Train SLGA-Reasoning Model") + parser.add_argument("--small", action="store_true", help="Use small config for testing") + parser.add_argument("--dataset", type=str, default="all", choices=["gsm8k", "synthetic", "all"]) + parser.add_argument("--resume", type=str, default=None, help="Resume from checkpoint") + parser.add_argument("--steps", type=int, default=None, help="Override max_steps") + parser.add_argument("--batch-size", type=int, default=None, help="Override batch_size") + args = parser.parse_args() + + # Config + if args.small: + config = TrainConfig.small() + print("📦 Using SMALL config (for testing)") + else: + config = TrainConfig() + + # Overrides + if args.steps: + config.max_steps = args.steps + if args.batch_size: + config.batch_size = args.batch_size + + # Train + train(config, args.dataset, args.resume) + + +if __name__ == "__main__": + main() From ffc61d8037c18716650df460e804eb841e45fafb Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 26 Nov 2025 21:26:44 +0000 Subject: [PATCH 5/5] Add evaluation, chat interface, and tests for reasoning model New utilities: - eval_reasoning.py: Evaluate on GSM8K benchmark - Accuracy, exact match, step analysis - Error analysis with examples - chat_reasoning.py: Interactive chat interface - Real-time reasoning generation - Commands: /tot, /cot, /temp, /config - Colored output with step highlighting - test_reasoning_quick.py: Quick validation tests - Import check - Model creation - Forward/backward pass - Generation - Dataset loading Complete workflow now available: 1. python train_reasoning_simple.py --small # Train 2. python test_reasoning_quick.py # Verify 3. python eval_reasoning.py --checkpoint ... # Evaluate 4. python chat_reasoning.py --checkpoint ... # Interactive --- chat_reasoning.py | 218 +++++++++++++++++++++++++++++++++ eval_reasoning.py | 258 ++++++++++++++++++++++++++++++++++++++++ test_reasoning_quick.py | 205 +++++++++++++++++++++++++++++++ 3 files changed, 681 insertions(+) create mode 100644 chat_reasoning.py create mode 100644 eval_reasoning.py create mode 100644 test_reasoning_quick.py diff --git a/chat_reasoning.py b/chat_reasoning.py new file mode 100644 index 0000000..6ece5bd --- /dev/null +++ b/chat_reasoning.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python3 +""" +💬 Chat Interactif avec le Reasoning Model + +Interface en ligne de commande pour tester le modèle. + +Usage: + python chat_reasoning.py --checkpoint checkpoints/reasoning/checkpoint_best.pt + python chat_reasoning.py --checkpoint ckpt.pt --temperature 0.7 + +Commandes spéciales: + /quit, /exit - Quitter + /clear - Effacer l'historique + /config - Afficher la config + /tot - Activer Tree of Thought + /cot - Revenir à Chain of Thought + /temp 0.5 - Changer la température +""" + +import os +import sys +import argparse +import torch +from typing import Optional + +from train_reasoning_simple import ( + ReasoningModel, + TrainConfig, + SimpleTokenizer, + SPECIAL_TOKENS, +) + + +class ReasoningChat: + """Interface de chat interactive.""" + + def __init__( + self, + model: ReasoningModel, + tokenizer: SimpleTokenizer, + device: str = "cuda", + temperature: float = 0.7, + max_tokens: int = 512, + ): + self.model = model + self.tokenizer = tokenizer + self.device = device + self.temperature = temperature + self.max_tokens = max_tokens + self.use_tot = False + self.history = [] + + def generate_response(self, question: str) -> str: + """Génère une réponse avec raisonnement.""" + # Formater le prompt + prompt = f"Question: {question}\n" + input_ids = torch.tensor([self.tokenizer.encode(prompt)]).to(self.device) + + # Générer + with torch.no_grad(): + output_ids = self.model.generate( + input_ids, + max_new_tokens=self.max_tokens, + temperature=self.temperature, + ) + + # Décoder + full_text = self.tokenizer.decode(output_ids[0].tolist()) + + # Extraire la partie générée + generated = full_text[len(prompt):] + + return generated + + def format_response(self, response: str) -> str: + """Formate la réponse pour l'affichage.""" + # Coloriser les tokens spéciaux + formatted = response + formatted = formatted.replace('', '\n🤔 \033[94m\033[0m') + formatted = formatted.replace('', '\033[94m\033[0m') + formatted = formatted.replace('', '\n 📍 \033[93m') + formatted = formatted.replace('', '\n\n✅ \033[92m\033[0m ') + formatted = formatted.replace('', ' \033[92m\033[0m') + + return formatted + + def handle_command(self, cmd: str) -> Optional[str]: + """Gère les commandes spéciales.""" + parts = cmd.strip().split() + command = parts[0].lower() + + if command in ['/quit', '/exit', '/q']: + return "EXIT" + + elif command == '/clear': + self.history = [] + return "Historique effacé." + + elif command == '/config': + return f"""Configuration: + Temperature: {self.temperature} + Max tokens: {self.max_tokens} + Mode: {'Tree of Thought' if self.use_tot else 'Chain of Thought'} + Device: {self.device}""" + + elif command == '/tot': + self.use_tot = True + return "Mode Tree of Thought activé (non implémenté dans cette version)." + + elif command == '/cot': + self.use_tot = False + return "Mode Chain of Thought activé." + + elif command == '/temp' and len(parts) > 1: + try: + self.temperature = float(parts[1]) + return f"Température: {self.temperature}" + except: + return "Usage: /temp 0.7" + + elif command == '/help': + return """Commandes disponibles: + /quit, /exit - Quitter + /clear - Effacer l'historique + /config - Afficher la configuration + /tot - Activer Tree of Thought + /cot - Activer Chain of Thought + /temp - Changer la température + /help - Afficher cette aide""" + + return None + + def run(self): + """Boucle principale du chat.""" + print("\n" + "=" * 60) + print("🧠 SLGA Reasoning Model - Chat Interactif") + print("=" * 60) + print("Tapez votre question ou /help pour l'aide") + print("=" * 60 + "\n") + + while True: + try: + # Input + user_input = input("\033[96m❓ Vous:\033[0m ").strip() + + if not user_input: + continue + + # Commande? + if user_input.startswith('/'): + result = self.handle_command(user_input) + if result == "EXIT": + print("\n👋 Au revoir!") + break + if result: + print(f"\033[90m{result}\033[0m\n") + continue + + # Générer la réponse + print("\033[90m⏳ Réflexion en cours...\033[0m") + + response = self.generate_response(user_input) + formatted = self.format_response(response) + + print(f"\n🤖 Modèle:{formatted}\n") + + # Historique + self.history.append({ + "question": user_input, + "response": response, + }) + + except KeyboardInterrupt: + print("\n\n👋 Au revoir!") + break + except Exception as e: + print(f"\033[91m❌ Erreur: {e}\033[0m\n") + + +def main(): + parser = argparse.ArgumentParser(description="Chat with Reasoning Model") + parser.add_argument("--checkpoint", type=str, required=True, help="Path to checkpoint") + parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature") + parser.add_argument("--max-tokens", type=int, default=512, help="Max tokens to generate") + args = parser.parse_args() + + device = "cuda" if torch.cuda.is_available() else "cpu" + + # Charger le modèle + print(f"📂 Chargement: {args.checkpoint}") + ckpt = torch.load(args.checkpoint, map_location=device) + + config = ckpt.get("config", TrainConfig()) + if isinstance(config, dict): + config = TrainConfig(**config) + + model = ReasoningModel(config).to(device) + model.load_state_dict(ckpt["model"]) + model.eval() + + print(f"✓ Modèle chargé ({sum(p.numel() for p in model.parameters()) / 1e6:.1f}M params)") + + # Tokenizer + tokenizer = SimpleTokenizer() + + # Chat + chat = ReasoningChat( + model=model, + tokenizer=tokenizer, + device=device, + temperature=args.temperature, + max_tokens=args.max_tokens, + ) + chat.run() + + +if __name__ == "__main__": + main() diff --git a/eval_reasoning.py b/eval_reasoning.py new file mode 100644 index 0000000..1ba25c9 --- /dev/null +++ b/eval_reasoning.py @@ -0,0 +1,258 @@ +#!/usr/bin/env python3 +""" +🧪 Évaluation du Reasoning Model sur Benchmarks + +Évalue le modèle sur: +- GSM8K (mathématiques grade school) +- Accuracy, exact match, step analysis + +Usage: + python eval_reasoning.py --checkpoint checkpoints/reasoning/checkpoint_best.pt + python eval_reasoning.py --checkpoint ckpt.pt --samples 100 --verbose +""" + +import os +import sys +import argparse +import json +import re +import torch +from typing import Dict, List, Any, Optional, Tuple +from tqdm import tqdm + +# Import du modèle +from train_reasoning_simple import ( + ReasoningModel, + TrainConfig, + SimpleTokenizer, + SPECIAL_TOKENS, + load_gsm8k, +) + + +def extract_answer(text: str) -> str: + """Extrait la réponse finale du texte généré.""" + # Chercher ... + match = re.search(r'(.*?)', text, re.DOTALL) + if match: + return match.group(1).strip() + + # Chercher après "=" ou "is" + match = re.search(r'(?:=|is)\s*([\d,.\-]+)', text) + if match: + return match.group(1).replace(',', '') + + # Dernier nombre dans le texte + numbers = re.findall(r'[\d,]+(?:\.\d+)?', text) + if numbers: + return numbers[-1].replace(',', '') + + return text.strip() + + +def normalize_answer(answer: str) -> str: + """Normalise une réponse pour comparaison.""" + # Enlever les virgules, espaces, $, etc. + answer = answer.replace(',', '').replace('$', '').replace('%', '') + answer = answer.strip().lower() + + # Extraire le nombre si possible + match = re.search(r'([\d.\-]+)', answer) + if match: + try: + # Convertir en float puis en string normalisé + num = float(match.group(1)) + if num == int(num): + return str(int(num)) + return f"{num:.2f}" + except: + pass + + return answer + + +def count_reasoning_steps(text: str) -> int: + """Compte le nombre d'étapes de raisonnement.""" + return text.count('') + + +def evaluate_model( + model: ReasoningModel, + tokenizer: SimpleTokenizer, + data: List[Dict[str, Any]], + device: str = "cuda", + max_samples: int = -1, + verbose: bool = False, + use_tot: bool = False, +) -> Dict[str, Any]: + """ + Évalue le modèle sur un dataset. + + Returns: + dict avec métriques d'évaluation + """ + model.eval() + + results = { + "correct": 0, + "total": 0, + "exact_match": 0, + "avg_steps": 0, + "examples": [], + } + + samples = data[:max_samples] if max_samples > 0 else data + total_steps = 0 + + for item in tqdm(samples, desc="Evaluating"): + question = item["question"] + gold_answer = normalize_answer(item.get("answer", "")) + + # Préparer le prompt + prompt = f"Question: {question}\n" + input_ids = torch.tensor([tokenizer.encode(prompt)]).to(device) + + # Générer + with torch.no_grad(): + output_ids = model.generate( + input_ids, + max_new_tokens=256, + temperature=0.0, # Déterministe pour éval + ) + + generated_text = tokenizer.decode(output_ids[0].tolist()) + + # Extraire la réponse + pred_answer = normalize_answer(extract_answer(generated_text)) + + # Comparer + is_correct = pred_answer == gold_answer + is_exact = item.get("answer", "").strip() == extract_answer(generated_text) + + # Stats + num_steps = count_reasoning_steps(generated_text) + total_steps += num_steps + + results["total"] += 1 + if is_correct: + results["correct"] += 1 + if is_exact: + results["exact_match"] += 1 + + # Log exemple + example = { + "question": question, + "gold": gold_answer, + "pred": pred_answer, + "correct": is_correct, + "num_steps": num_steps, + } + + if verbose: + status = "✓" if is_correct else "✗" + print(f"\n{status} Q: {question[:60]}...") + print(f" Gold: {gold_answer}") + print(f" Pred: {pred_answer}") + print(f" Steps: {num_steps}") + + results["examples"].append(example) + + # Calculer métriques + results["accuracy"] = results["correct"] / max(1, results["total"]) + results["exact_match_rate"] = results["exact_match"] / max(1, results["total"]) + results["avg_steps"] = total_steps / max(1, results["total"]) + + return results + + +def print_results(results: Dict[str, Any]): + """Affiche les résultats.""" + print("\n" + "=" * 60) + print("📊 RÉSULTATS D'ÉVALUATION") + print("=" * 60) + print(f"Total samples: {results['total']}") + print(f"Correct: {results['correct']}") + print(f"Accuracy: {results['accuracy']:.1%}") + print(f"Exact match: {results['exact_match_rate']:.1%}") + print(f"Avg steps: {results['avg_steps']:.1f}") + print("=" * 60) + + # Analyse des erreurs + errors = [e for e in results["examples"] if not e["correct"]] + if errors: + print(f"\n❌ Exemples d'erreurs ({len(errors)} total):") + for e in errors[:5]: + print(f" Q: {e['question'][:50]}...") + print(f" Gold: {e['gold']} | Pred: {e['pred']}") + + +def main(): + parser = argparse.ArgumentParser(description="Evaluate Reasoning Model") + parser.add_argument("--checkpoint", type=str, required=True, help="Path to checkpoint") + parser.add_argument("--samples", type=int, default=-1, help="Max samples (-1 = all)") + parser.add_argument("--verbose", action="store_true", help="Print each example") + parser.add_argument("--output", type=str, default=None, help="Save results to JSON") + parser.add_argument("--dataset", type=str, default="gsm8k", choices=["gsm8k"]) + args = parser.parse_args() + + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"Device: {device}") + + # Charger le modèle + print(f"\n📂 Chargement du checkpoint: {args.checkpoint}") + ckpt = torch.load(args.checkpoint, map_location=device) + + config = ckpt.get("config", TrainConfig()) + if isinstance(config, dict): + config = TrainConfig(**config) + + model = ReasoningModel(config).to(device) + model.load_state_dict(ckpt["model"]) + model.eval() + + print(f"Modèle chargé: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M params") + + # Tokenizer + tokenizer = SimpleTokenizer() + + # Charger les données + print(f"\n📚 Chargement du dataset: {args.dataset}") + if args.dataset == "gsm8k": + try: + from datasets import load_dataset + ds = load_dataset("openai/gsm8k", "main", split="test") + data = [] + for item in ds: + parts = item["answer"].split("####") + if len(parts) == 2: + data.append({ + "question": item["question"], + "answer": parts[1].strip(), + }) + print(f"GSM8K test: {len(data)} exemples") + except Exception as e: + print(f"Erreur chargement GSM8K: {e}") + return + + # Évaluer + results = evaluate_model( + model, + tokenizer, + data, + device=device, + max_samples=args.samples, + verbose=args.verbose, + ) + + # Afficher + print_results(results) + + # Sauvegarder + if args.output: + with open(args.output, 'w') as f: + json.dump(results, f, indent=2) + print(f"\n💾 Résultats sauvegardés: {args.output}") + + +if __name__ == "__main__": + main() diff --git a/test_reasoning_quick.py b/test_reasoning_quick.py new file mode 100644 index 0000000..4310e80 --- /dev/null +++ b/test_reasoning_quick.py @@ -0,0 +1,205 @@ +#!/usr/bin/env python3 +""" +🧪 Test Rapide du Reasoning Model + +Vérifie que tout fonctionne correctement. + +Usage: + python test_reasoning_quick.py +""" + +import torch +import sys + + +def test_imports(): + """Test des imports.""" + print("1. Test des imports...") + try: + from train_reasoning_simple import ( + ReasoningModel, + TrainConfig, + SimpleTokenizer, + SPECIAL_TOKENS, + ) + print(" ✓ Imports OK") + return True + except ImportError as e: + print(f" ✗ Erreur d'import: {e}") + return False + + +def test_model_creation(): + """Test de création du modèle.""" + print("2. Test création du modèle...") + try: + from train_reasoning_simple import ReasoningModel, TrainConfig + + config = TrainConfig.small() # Version légère + model = ReasoningModel(config) + + num_params = sum(p.numel() for p in model.parameters()) + print(f" ✓ Modèle créé: {num_params / 1e6:.1f}M params") + return model, config + except Exception as e: + print(f" ✗ Erreur: {e}") + return None, None + + +def test_forward_pass(model, config): + """Test du forward pass.""" + print("3. Test forward pass...") + try: + device = "cuda" if torch.cuda.is_available() else "cpu" + model = model.to(device) + + # Input factice + batch_size = 2 + seq_len = 64 + input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len)).to(device) + targets = torch.randint(0, config.vocab_size, (batch_size, seq_len)).to(device) + + # Forward + outputs = model(input_ids, targets) + + assert "logits" in outputs + assert "loss" in outputs + assert outputs["logits"].shape == (batch_size, seq_len, config.vocab_size) + + print(f" ✓ Forward OK - Loss: {outputs['loss'].item():.4f}") + return True + except Exception as e: + print(f" ✗ Erreur: {e}") + return False + + +def test_generation(model, config): + """Test de génération.""" + print("4. Test génération...") + try: + from train_reasoning_simple import SimpleTokenizer + + device = "cuda" if torch.cuda.is_available() else "cpu" + model = model.to(device) + model.eval() + + tokenizer = SimpleTokenizer() + + prompt = "Question: What is 2 + 2?\n" + input_ids = torch.tensor([tokenizer.encode(prompt)]).to(device) + + with torch.no_grad(): + output_ids = model.generate( + input_ids, + max_new_tokens=50, + temperature=0.7, + ) + + generated = tokenizer.decode(output_ids[0].tolist()) + print(f" ✓ Génération OK") + print(f" → Output: {generated[:100]}...") + return True + except Exception as e: + print(f" ✗ Erreur: {e}") + return False + + +def test_dataset(): + """Test du chargement de données.""" + print("5. Test données synthétiques...") + try: + from train_reasoning_simple import load_synthetic_data, CoTDataset, SimpleTokenizer + + data = load_synthetic_data(n=10) + assert len(data) == 10 + + tokenizer = SimpleTokenizer() + dataset = CoTDataset(data, tokenizer, max_length=256) + + sample = dataset[0] + assert "input_ids" in sample + assert "targets" in sample + + print(f" ✓ Dataset OK - {len(data)} exemples") + return True + except Exception as e: + print(f" ✗ Erreur: {e}") + return False + + +def test_training_step(model, config): + """Test d'une étape d'entraînement.""" + print("6. Test étape d'entraînement...") + try: + device = "cuda" if torch.cuda.is_available() else "cpu" + model = model.to(device) + model.train() + + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) + + # Forward + input_ids = torch.randint(0, config.vocab_size, (2, 32)).to(device) + targets = torch.randint(0, config.vocab_size, (2, 32)).to(device) + + outputs = model(input_ids, targets) + loss = outputs["loss"] + + # Backward + loss.backward() + optimizer.step() + optimizer.zero_grad() + + print(f" ✓ Training step OK - Loss: {loss.item():.4f}") + return True + except Exception as e: + print(f" ✗ Erreur: {e}") + return False + + +def main(): + print("=" * 60) + print("🧪 SLGA Reasoning Model - Tests Rapides") + print("=" * 60) + print(f"PyTorch: {torch.__version__}") + print(f"CUDA: {torch.cuda.is_available()}") + print("=" * 60 + "\n") + + results = [] + + # Tests + results.append(("Imports", test_imports())) + + model, config = test_model_creation() + results.append(("Model Creation", model is not None)) + + if model: + results.append(("Forward Pass", test_forward_pass(model, config))) + results.append(("Generation", test_generation(model, config))) + results.append(("Training Step", test_training_step(model, config))) + + results.append(("Dataset", test_dataset())) + + # Résumé + print("\n" + "=" * 60) + print("📊 RÉSUMÉ") + print("=" * 60) + + passed = sum(1 for _, r in results if r) + total = len(results) + + for name, result in results: + status = "✓" if result else "✗" + print(f" {status} {name}") + + print(f"\nTotal: {passed}/{total} tests passés") + + if passed == total: + print("\n🎉 Tous les tests sont passés!") + return 0 + else: + print("\n⚠️ Certains tests ont échoué.") + return 1 + + +if __name__ == "__main__": + sys.exit(main())