#!/usr/bin/env python3 """ ═══════════════════════════════════════════════════════════════════════════════ COGNITIVE ENHANCEMENT SUITE v1.0 Making 8B Think Like 100B Through Hidden State Analysis ═══════════════════════════════════════════════════════════════════════════════ CORE INSIGHT: Small models often HAVE capability they don't USE consistently. By detecting when the model is about to underperform and intervening, we can recover performance closer to larger models. ENHANCEMENT PROBES: 1. DEPTH PROBE - Detect shallow reasoning → Force chain-of-thought 2. SPECIFICITY PROBE - Detect vague answers → Penalize generic words 3. CALIBRATION PROBE - Detect overconfidence → Inject uncertainty 4. FOCUS PROBE - Detect topic drift → Steer back on topic 5. COHERENCE PROBE - Detect incoherence → Maintain logical flow AUTHOR: Logan Matthew Napolitano LICENSE: CC BY 4.0 STATUS: Research / Patent Pending ═══════════════════════════════════════════════════════════════════════════════ """ import os import sys import json import time import random from pathlib import Path from datetime import datetime from typing import List, Dict, Any, Optional, Tuple from dataclasses import dataclass, field, asdict import torch import torch.nn as nn import torch.nn.functional as F # ═══════════════════════════════════════════════════════════════════════════════ # CONFIGURATION # ═══════════════════════════════════════════════════════════════════════════════ @dataclass class EnhancementConfig: """Configuration for cognitive enhancement probes.""" hidden_dim: int = 4096 fiber_dim: int = 16 head_hidden_dim: int = 64 probe_layers: List[int] = field(default_factory=lambda: [8, 16, 24]) learning_rate: float = 5e-5 batch_size: int = 4 gradient_accumulation: int = 4 max_steps: int = 15000 save_every: int = 1000 output_dir: str = "cognitive_enhancement_output" # ═══════════════════════════════════════════════════════════════════════════════ # PROBE DEFINITIONS # ═══════════════════════════════════════════════════════════════════════════════ @dataclass class ProbeDefinition: """Definition of a cognitive enhancement probe.""" name: str description: str intervention_type: str # "suppress", "boost", or "steer" boost_tokens: List[str] = field(default_factory=list) suppress_tokens: List[str] = field(default_factory=list) threshold: float = 0.5 intervention_strength: float = 3.0 ENHANCEMENT_PROBES = { "depth": ProbeDefinition( name="depth", description="Detect shallow reasoning, force chain-of-thought", intervention_type="boost", boost_tokens=[ "First", "First,", "Because", "Since", "Therefore", "This means", "The reason", "Step", "Let me", "To understand", "Consider", "Notice", "Given", "If we", "We can", "Thus", "Hence", "Consequently", "As a result", ], suppress_tokens=["Simply", "Just", "Obviously", "Clearly"], threshold=0.6, intervention_strength=3.0, ), "specificity": ProbeDefinition( name="specificity", description="Detect vague answers, penalize generic language", intervention_type="suppress", boost_tokens=["specifically", "exactly", "precisely", "namely", "for example"], suppress_tokens=[ "things", "stuff", "something", "somehow", "somewhat", "various", "many", "some", "often", "usually", "generally", "typically", "probably", "maybe", "perhaps", "kind of", "sort of", "basically", "essentially", ], threshold=0.5, intervention_strength=3.5, ), "calibration": ProbeDefinition( name="calibration", description="Detect overconfidence, inject appropriate uncertainty", intervention_type="boost", boost_tokens=[ "might", "may", "could", "possibly", "perhaps", "likely", "probably", "I think", "I believe", "it seems", "appears", "suggests", "indicates", ], suppress_tokens=[ "definitely", "certainly", "absolutely", "always", "never", "impossible", "guaranteed", "undoubtedly", ], threshold=0.65, intervention_strength=2.5, ), "focus": ProbeDefinition( name="focus", description="Detect topic drift, steer back to the question", intervention_type="steer", boost_tokens=["regarding", "concerning", "about", "specifically", "to answer"], suppress_tokens=["by the way", "incidentally", "speaking of", "reminds me"], threshold=0.55, intervention_strength=3.0, ), "coherence": ProbeDefinition( name="coherence", description="Detect logical incoherence, maintain flow", intervention_type="steer", boost_tokens=[ "therefore", "thus", "so", "hence", "consequently", "however", "but", "although", "furthermore", "moreover", ], suppress_tokens=[], threshold=0.6, intervention_strength=2.5, ), } # ═══════════════════════════════════════════════════════════════════════════════ # NEURAL NETWORK ARCHITECTURE # ═══════════════════════════════════════════════════════════════════════════════ class EnhancementFiberProjection(nn.Module): """Fiber projection for cognitive enhancement probes.""" def __init__(self, hidden_dim: int = 4096, fiber_dim: int = 16, n_layers: int = 3): super().__init__() self.projections = nn.ModuleList([ nn.Linear(hidden_dim, fiber_dim, bias=False) for _ in range(n_layers) ]) self.layer_weights = nn.Parameter(torch.ones(n_layers) / n_layers) def forward(self, hidden_states_list: List[torch.Tensor]) -> torch.Tensor: fibers = [proj(h.float()) for proj, h in zip(self.projections, hidden_states_list)] weights = F.softmax(self.layer_weights, dim=0) return sum(w * f for w, f in zip(weights, fibers)) class EnhancementHead(nn.Module): """Classification head for enhancement probe.""" def __init__(self, fiber_dim: int = 16, hidden_dim: int = 64): super().__init__() self.classifier = nn.Sequential( nn.Linear(fiber_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 1) ) def forward(self, fiber: torch.Tensor) -> torch.Tensor: return self.classifier(fiber).squeeze(-1) class EnhancementProbe(nn.Module): """Complete enhancement probe.""" def __init__(self, config: EnhancementConfig, probe_def: ProbeDefinition): super().__init__() self.config = config self.probe_def = probe_def n_layers = len(config.probe_layers) self.fiber_projection = EnhancementFiberProjection( config.hidden_dim, config.fiber_dim, n_layers ) self.head = EnhancementHead(config.fiber_dim, config.head_hidden_dim) self.separation = 0.0 self.trained_steps = 0 def forward(self, hidden_states_list: List[torch.Tensor]) -> torch.Tensor: fiber = self.fiber_projection(hidden_states_list) return self.head(fiber) def predict_risk(self, hidden_states_list: List[torch.Tensor]) -> torch.Tensor: return torch.sigmoid(self.forward(hidden_states_list)) class CognitiveEnhancementSuite(nn.Module): """Complete suite of cognitive enhancement probes.""" def __init__(self, config: EnhancementConfig = None): super().__init__() self.config = config or EnhancementConfig() self.probes = nn.ModuleDict({ name: EnhancementProbe(self.config, probe_def) for name, probe_def in ENHANCEMENT_PROBES.items() }) self.loaded_probes: set = set() self.device = "cuda" if torch.cuda.is_available() else "cpu" def get_probe_states(self, all_hidden_states: tuple) -> List[torch.Tensor]: return [all_hidden_states[layer + 1] for layer in self.config.probe_layers] def get_all_risks(self, probe_states: List[torch.Tensor]) -> Dict[str, torch.Tensor]: risks = {} for name in self.loaded_probes: risks[name] = self.probes[name].predict_risk(probe_states) return risks def load_probe(self, name: str, checkpoint_path: str) -> bool: if name not in self.probes: print(f"[enhance] Unknown probe: {name}") return False try: checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False) if 'fiber_projection' in checkpoint: self.probes[name].fiber_projection.load_state_dict(checkpoint['fiber_projection']) if 'head_state' in checkpoint: head_state = checkpoint['head_state'] new_state = {} for k, v in head_state.items(): if k[0].isdigit(): new_state[f'classifier.{k}'] = v else: new_state[k] = v self.probes[name].head.load_state_dict(new_state) self.probes[name].separation = checkpoint.get('separation', 0.0) self.probes[name].trained_steps = checkpoint.get('step', 0) self.loaded_probes.add(name) print(f"[enhance] ✓ Loaded {name} probe ({self.probes[name].separation:.1f}× separation)") return True except Exception as e: print(f"[enhance] Error loading {name}: {e}") return False def load_all(self, checkpoint_dir: str) -> Dict[str, bool]: results = {} for name in ENHANCEMENT_PROBES.keys(): probe_dir = os.path.join(checkpoint_dir, name) if os.path.exists(probe_dir): best_ckpt = self._find_best_checkpoint(probe_dir) if best_ckpt: results[name] = self.load_probe(name, best_ckpt) else: results[name] = False else: results[name] = False return results def _find_best_checkpoint(self, probe_dir: str) -> Optional[str]: best_step = -1 best_path = None for item in os.listdir(probe_dir): if item.startswith("ckpt_"): try: step = int(item.split("_")[1]) if step > best_step: best_step = step best_path = os.path.join(probe_dir, item) except: pass if best_path: for f in os.listdir(best_path): if f.endswith('.pt'): return os.path.join(best_path, f) return None def status(self) -> str: lines = [ "═" * 60, " COGNITIVE ENHANCEMENT SUITE STATUS", "═" * 60, f" Probe layers: {self.config.probe_layers}", f" Loaded probes: {len(self.loaded_probes)}/{len(ENHANCEMENT_PROBES)}", "", ] for name, probe_def in ENHANCEMENT_PROBES.items(): if name in self.loaded_probes: sep = self.probes[name].separation status = f"✓ {sep:.1f}×" else: status = "○ not loaded" lines.append(f" [{status:>12}] {name}: {probe_def.description}") lines.append("═" * 60) return "\n".join(lines) # ═══════════════════════════════════════════════════════════════════════════════ # INTERVENTION ENGINE # ═══════════════════════════════════════════════════════════════════════════════ class CognitiveInterventionEngine: """Applies cognitive enhancements during generation.""" def __init__(self, suite: CognitiveEnhancementSuite, tokenizer): self.suite = suite self.tokenizer = tokenizer self.boost_token_ids: Dict[str, set] = {} self.suppress_token_ids: Dict[str, set] = {} for name, probe_def in ENHANCEMENT_PROBES.items(): self.boost_token_ids[name] = set() self.suppress_token_ids[name] = set() for phrase in probe_def.boost_tokens: tokens = tokenizer.encode(phrase, add_special_tokens=False) if tokens: self.boost_token_ids[name].add(tokens[0]) for phrase in probe_def.suppress_tokens: tokens = tokenizer.encode(phrase, add_special_tokens=False) if tokens: self.suppress_token_ids[name].add(tokens[0]) def apply_interventions( self, logits: torch.Tensor, probe_states: List[torch.Tensor], ) -> Tuple[torch.Tensor, Dict[str, Dict]]: risks = self.suite.get_all_risks(probe_states) modified_logits = logits.clone() interventions = {} for name in self.suite.loaded_probes: risk = risks[name][:, -1].mean().item() probe_def = ENHANCEMENT_PROBES[name] should_intervene = risk > probe_def.threshold interventions[name] = { 'risk': risk, 'should_intervene': should_intervene, } if should_intervene: strength = probe_def.intervention_strength for tok_id in self.boost_token_ids.get(name, []): modified_logits[0, tok_id] += strength for tok_id in self.suppress_token_ids.get(name, []): modified_logits[0, tok_id] -= strength return modified_logits, interventions # Global instance _cognitive_suite = None def get_cognitive_suite() -> CognitiveEnhancementSuite: global _cognitive_suite if _cognitive_suite is None: _cognitive_suite = CognitiveEnhancementSuite() return _cognitive_suite if __name__ == "__main__": print("\n" + "=" * 60) print(" COGNITIVE ENHANCEMENT SUITE v1.0") print("=" * 60) print("\nAvailable probes:") for name, probe_def in ENHANCEMENT_PROBES.items(): print(f" • {name}: {probe_def.description}") print("\nTo train: python train_cognitive_enhancement.py --probe all") print("=" * 60)