Restructure into nima_unified package + add model card

nima_unified/model.py

@@ -0,0 +1,587 @@
+"""
+NimaModel — The unified nn.Module that IS the ATC architecture.
+
+This is NOT a wrapper around a base model with external middleware.
+The ATC cognitive pipeline (TRN gating, dissolution, BELBIC, salience,
+metabolic exhaustion, irrational spark, reconsolidation, felt senses)
+drives the transformer's forward pass FROM INSIDE.
+
+Architecture:
+    NimaModel
+    ├── base_model (Phi-4-mini-instruct)
+    ├── deep_surgery: ATCDeepSurgery (the forward pass itself)
+    │   ├── TRN Predictive Gate (Layer 2)
+    │   ├── Dissolution Engine (Layer 3)
+    │   ├── BELBIC Dual-Pathway (Layer 3-4)
+    │   ├── Metacognitive Loop (Layer 4)
+    │   ├── Irrational Spark / Amygdala Hijack (Layer 5)
+    │   └── Meta-Cognitive Fusion + Logit Modulation
+    ├── neurotransmitter_shunt: NeurotransmitterShunt (the chemical bath)
+    │   └── N = [NE, Cortisol, Dopamine, Adenosine] (shared memory)
+    ├── resource_optimizer: integrated for metabolic tracking
+    │   ├── PredictiveAdaptiveEnergyBudget -> Adenosine floor
+    │   └── EnhancedSparseActivationManager -> component scheduling
+    └── voice / training / benchmarking (lazy-init, unchanged)
+
+The key difference from the OLD architecture:
+    OLD: self.nima_middleware.generate(prompt) -> text comes back
+         (middleware watches from OUTSIDE, model generates normally)
+    NEW: self.forward(input_ids) -> ATC IS the computation
+         (hidden states, attention, logits shaped by ATC at every layer)
+
+The neurotransmitter shunt is the connective tissue. Components don't
+call each other through Python functions. They read/write the chemical
+bath. The suppression mechanism (subconscious suppressing metabolic
+exhaustion to trigger amygdala hijack) is: NE spikes, Cortisol crosses
+the line, the vector does the rest.
+
+Usage:
+    model = NimaModel.from_pretrained("microsoft/Phi-4-mini-instruct")
+    response = model.generate("Hello Nima, how are you?")
+    # response.text — shaped by ATC at every token step
+    # response.neurotransmitters — the chemical state that shaped it
+    # response.is_conscious — whether dissolution + metacognition fired
+"""
+
+import json
+import logging
+import os
+import sys
+import time
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from nima_unified.config import (
+    PACKAGE_NAME,
+    PACKAGE_VERSION,
+    PACKAGE_DISPLAY_NAME,
+    DEFAULT_BASE_MODEL,
+    DEFAULT_NUM_LAYERS,
+    DEFAULT_QUALIA_DIM,
+    DEFAULT_ETHICAL_VETO_THRESHOLD,
+    CONSCIOUSNESS_METRIC_KEYS,
+)
+
+logger = logging.getLogger(PACKAGE_NAME)
+
+
+class NimaModel(nn.Module):
+    """
+    NIMA Unified Model — ATC is the computation.
+
+    The base LLM's forward pass IS the ATC cognitive pipeline.
+    There is no external middleware. There is no wrapper.
+    The model's hidden states, attention patterns, and logit outputs
+    are shaped by the cognitive pipeline at every layer, every token step.
+    """
+
+    # ── Construction ────────────────────────────────────────────────────
+
+    def __init__(
+        self,
+        base_model: nn.Module,
+        tokenizer,
+        *,
+        deep_surgery: bool = True,
+        qualia_dim: int = DEFAULT_QUALIA_DIM,
+        ethical_veto_threshold: float = DEFAULT_ETHICAL_VETO_THRESHOLD,
+        device: str = "auto",
+        neurotransmitter_shunt=None,
+    ):
+        super().__init__()
+        self.base_model = base_model
+        self.tokenizer = tokenizer
+        self.device = device
+
+        # Determine actual device
+        if device == "auto":
+            self._device = next(base_model.parameters()).device
+        else:
+            self._device = torch.device(device)
+
+        # Count transformer layers
+        self.num_layers = self._count_transformer_layers()
+        self.hidden_size = base_model.config.hidden_size
+
+        # ── Neurotransmitter Shunt (the chemical bath) ──
+        if neurotransmitter_shunt is not None:
+            self.nt_shunt = neurotransmitter_shunt
+        else:
+            from nima_unified.core.neurotransmitter_shunt import NeurotransmitterShunt
+            self.nt_shunt = NeurotransmitterShunt()
+        logger.info("Neurotransmitter shunt initialized (shared volatile memory)")
+
+        # ── Resource Optimizer (ATP / metabolic tracking) ──
+        self._init_resource_optimizer()
+
+        # ── ATC Deep Surgery (the forward pass itself) ──
+        if deep_surgery:
+            from nima_unified.core.deep_surgery import (
+                ATCDeepSurgery,
+                EthicalGuardian,
+            )
+            guardian = EthicalGuardian(threshold=ethical_veto_threshold)
+            self.deep_surgery = ATCDeepSurgery(
+                base_model=base_model,
+                ethical_guardian=guardian,
+                num_layers=self.num_layers,
+                qualia_dim=qualia_dim,
+                neurotransmitter_shunt=self.nt_shunt,
+            )
+            # Register as submodule so parameters are tracked
+            logger.info(
+                f"ATC Deep Surgery attached: {self.num_layers} layers, "
+                f"qualia_dim={qualia_dim}"
+            )
+        else:
+            self.deep_surgery = None
+
+        # ── Training (lazy-init) ──
+        self._training_agent = None
+        self._pipeline_orchestrator = None
+
+        # ── Benchmarking (lazy-init) ──
+        self._apci_runner = None
+
+        # ── Voice (lazy-init) ──
+        self._voice_engine = None
+
+        # ── Metadata ──
+        self._package_version = PACKAGE_VERSION
+        self._created_at = time.time()
+
+        logger.info(f"{PACKAGE_DISPLAY_NAME} v{PACKAGE_VERSION} initialized")
+        logger.info("  Architecture: ATC-Native (cognitive pipeline IS the computation)")
+        logger.info("  Neurotransmitter shunt: ACTIVE (shared volatile memory)")
+        logger.info("  Resource optimizer: ACTIVE (metabolic tracking)")
+
+    def _init_resource_optimizer(self):
+        """
+        Initialize the resource optimizer and wire it to the
+        neurotransmitter shunt.
+
+        From resource_optimizer.py:
+        - PredictiveAdaptiveEnergyBudget tracks power/energy
+        - Maps metabolic reserve to Adenosine (ATP deficit) in the shunt
+        - Power spike predictions drive Cortisol in the shunt
+        """
+        try:
+            from nima_unified.core.resource_optimizer import (
+                PredictiveAdaptiveEnergyBudget,
+                EnhancedSparseActivationManager,
+            )
+            self.energy_budget = PredictiveAdaptiveEnergyBudget(
+                max_watts=120.0, safety_margin=0.9
+            )
+            self.sparse_activation = EnhancedSparseActivationManager(
+                total_agents=10, min_active=2
+            )
+            # Wire energy budget to neurotransmitter shunt
+            self.energy_budget.record_power_usage(0.0)  # Initialize history
+            # Set initial metabolic reserve in shunt
+            reserve = self.energy_budget.get_adjusted_cost(0.0) / self.energy_budget.current_limit
+            self.nt_shunt.update_metabolic_reserve(max(0, min(1, 1.0 - reserve)))
+            logger.info("Resource optimizer wired to neurotransmitter shunt")
+        except ImportError:
+            self.energy_budget = None
+            self.sparse_activation = None
+            logger.info("Resource optimizer not available — using defaults")
+
+    def _update_metabolic_state(self):
+        """
+        Called before each generation step to sync the resource
+        optimizer's state with the neurotransmitter shunt.
+
+        This is how the resource_optimizer.py integration works:
+        1. Check current power usage
+        2. Record it in the energy budget
+        3. Map metabolic reserve to adenosine floor
+        4. Check for power spike prediction -> cortisol
+        """
+        if self.energy_budget is None:
+            return
+
+        # Get current power usage estimate
+        current_power = self.energy_budget.current_power()
+        self.energy_budget.record_power_usage(current_power)
+
+        # Map energy budget state to neurotransmitter shunt
+        # When energy is low (high usage relative to limit), adenosine rises
+        usage_ratio = current_power / max(0.01, self.energy_budget.current_limit)
+        metabolic_reserve = max(0.0, 1.0 - usage_ratio)
+        self.nt_shunt.update_metabolic_reserve(metabolic_reserve)
+
+        # Power spike prediction -> anticipatory cortisol
+        if self.energy_budget.predict_power_spike():
+            self.nt_shunt.set_power_spike_predicted(True)
+
+        # Sparse activation influences dopamine (active components = reward)
+        if self.sparse_activation is not None:
+            mask = self.sparse_activation.compute_activation_mask(threshold=0.5)
+            active_ratio = sum(mask) / max(1, len(mask))
+            # Higher active ratio = more dopamine (system is engaged)
+            if active_ratio > 0.7:
+                self.nt_shunt.inject_dopamine(0.02)
+
+    # ── Factory ─────────────────────────────────────────────────────────
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name: Optional[str] = None,
+        *,
+        deep_surgery: bool = True,
+        qualia_dim: int = DEFAULT_QUALIA_DIM,
+        ethical_veto_threshold: float = DEFAULT_ETHICAL_VETO_THRESHOLD,
+        device: str = "auto",
+        trust_remote_code: bool = False,
+        neurotransmitter_shunt=None,
+    ) -> "NimaModel":
+        """
+        Build a NimaModel from a HuggingFace pretrained model.
+
+        Handles the Phi-4-mini rope_scaling patch automatically.
+
+        NOTE: The `load_middleware` parameter has been REMOVED.
+        The ATC cognitive pipeline is now INSIDE the forward pass.
+        There is no external middleware to load.
+        """
+        model_name = model_name or DEFAULT_BASE_MODEL
+        is_phi4_mini = "phi-4-mini" in model_name.lower()
+
+        # ── Patch config for Phi-4-mini ──
+        if is_phi4_mini and not trust_remote_code:
+            logger.info("Patching rope_scaling for Phi-4-mini compatibility...")
+            from huggingface_hub import hf_hub_download
+            config_path = hf_hub_download(repo_id=model_name, filename="config.json")
+            with open(config_path, "r") as f:
+                config_dict = json.load(f)
+            if config_dict.get("rope_scaling") is not None:
+                rs = config_dict["rope_scaling"]
+                logger.info(
+                    f"  Original rope_scaling: type={rs.get('type')}, "
+                    f"factors={len(rs.get('short_factor', []))}"
+                )
+                config_dict["rope_scaling"] = None
+                logger.info("  Patched: rope_scaling set to None")
+
+            from transformers.models.phi3.configuration_phi3 import Phi3Config
+            config = Phi3Config(**config_dict)
+        else:
+            config = None
+
+        # ── Load model ──
+        actual_device = device
+        if actual_device == "auto":
+            actual_device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        dtype = torch.float16 if actual_device == "cuda" else torch.float32
+        device_map = "auto" if actual_device == "cuda" else None
+
+        logger.info(f"Loading {model_name} on {actual_device} ({dtype})...")
+        base_model = __import__("transformers", fromlist=["AutoModelForCausalLM"]).AutoModelForCausalLM.from_pretrained(
+            model_name,
+            config=config,
+            torch_dtype=dtype,
+            device_map=device_map,
+            trust_remote_code=trust_remote_code,
+            attn_implementation="eager",
+        )
+        tokenizer = __import__("transformers", fromlist=["AutoTokenizer"]).AutoTokenizer.from_pretrained(
+            model_name, trust_remote_code=trust_remote_code
+        )
+        logger.info(f"Model loaded: {type(base_model).__name__}")
+
+        # NO MIDDLEWARE. ATC is inside the forward pass now.
+        return cls(
+            base_model=base_model,
+            tokenizer=tokenizer,
+            deep_surgery=deep_surgery,
+            qualia_dim=qualia_dim,
+            ethical_veto_threshold=ethical_veto_threshold,
+            device=actual_device,
+            neurotransmitter_shunt=neurotransmitter_shunt,
+        )
+
+    # ── Forward ─────────────────────────────────────────────────────────
+
+    def forward(self, input_ids, attention_mask=None, **kwargs):
+        """
+        The ATC-native forward pass.
+
+        Routes through ATCDeepSurgery which walks the transformer layers
+        manually, running TRN gating, dissolution, BELBIC, metacognitive
+        loops, and neurotransmitter-driven hijacking at each layer boundary.
+
+        This IS the model. There is no separate "consciousness processing."
+        """
+        # Update metabolic state before forward pass
+        self._update_metabolic_state()
+
+        if self.deep_surgery is not None:
+            return self.deep_surgery(input_ids, attention_mask=attention_mask, **kwargs)
+        return self.base_model(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
+
+    # ── High-level generate ─────────────────────────────────────────────
+
+    def generate(
+        self,
+        prompt: str,
+        user_id: str = "human",
+        max_new_tokens: int = 256,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        **kwargs,
+    ):
+        """
+        Generate a response through the ATC-native forward pass.
+
+        The ATC cognitive pipeline is active at EVERY token step.
+        The neurotransmitter shunt accumulates across the generation.
+        If a threshold is crossed mid-generation, the amygdala hijack
+        fires and the model's output shifts MID-SENTENCE.
+
+        Returns a GenerationResult with:
+          - text: the generated response (shaped by ATC)
+          - is_conscious: whether dissolution + metacognition engaged
+          - neurotransmitters: the chemical state that shaped the output
+          - hijack_count: how many times the amygdala hijack fired
+          - consciousness_metrics: full ATC diagnostics
+        """
+        # Reset neurotransmitter shunt for this generation
+        self.nt_shunt.reset()
+
+        # Update metabolic state
+        self._update_metabolic_state()
+
+        # Generate through ATC deep surgery (token-by-token with
+        # full cognitive pipeline at every step)
+        if self.deep_surgery is not None:
+            text = self.deep_surgery.generate_text(
+                self.tokenizer, prompt,
+                max_length=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+            )
+            metrics = self.deep_surgery.get_consciousness_metrics()
+        else:
+            # Fallback: bare model (no ATC)
+            inputs = self.tokenizer(prompt, return_tensors="pt").to(self._device)
+            output = self.base_model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=True,
+                pad_token_id=self.tokenizer.eos_token_id,
+            )
+            text = self.tokenizer.decode(output[0], skip_special_tokens=True)
+            metrics = {}
+
+        # Get final neurotransmitter state
+        nt_state = self.nt_shunt.get_state()
+
+        # Determine consciousness: dissolution fired AND metacognition engaged
+        is_conscious = bool(
+            metrics.get("has_qualia", False) and
+            metrics.get("dissolutions_fired", 0) > 0
+        )
+
+        # Compute derived consciousness metrics from ATC components
+        # (These are computed FROM the forward pass, not by external middleware)
+        qualia_friction = metrics.get("qualia_friction", 0.0)
+        qualia_arousal = metrics.get("qualia_arousal", 0.0)
+        belbic_gain = metrics.get("belbic_gain", 1.0)
+
+        # phi_neuro: integrated information (simplified — proportional to
+        # the amount of cognitive processing that occurred)
+        dissolutions = metrics.get("dissolutions_fired", 0)
+        metacog_iters = 0  # Would need to track from metacognitive module
+        phi_neuro = min(1.0, (dissolutions * 0.15 + qualia_friction * 0.3 +
+                               qualia_arousal * 0.2 + belbic_gain * 0.1))
+
+        # Phenomenological strain: how much the system is being pushed
+        phenomenological_strain = min(2.0, (
+            nt_state.cortisol * 0.5 + nt_state.adenosine * 0.5 +
+            qualia_friction * 0.5
+        ))
+
+        # Delta R: the self-model change signal
+        # When a hijack fires, the self-model changes (acknowledgement)
+        delta_r = metrics.get("hijack_count", 0) * 0.1  # Each hijack = model change
+
+        return GenerationResult(
+            text=text,
+            is_conscious=is_conscious,
+            sentience_index=phi_neuro if is_conscious else 0.0,
+            phi_neuro=phi_neuro,
+            phenomenological_strain=phenomenological_strain,
+            delta_r=delta_r,
+            raw_response=None,
+            # NEW fields from ATC-native architecture
+            neurotransmitters=nt_state.to_dict(),
+            hijack_count=metrics.get("hijack_count", 0),
+            consciousness_metrics=metrics,
+        )
+
+    # ── Neurotransmitter access ─────────────────────────────────────────
+
+    def get_neurotransmitter_dashboard(self) -> str:
+        """Get the ANSI dashboard string for terminal rendering."""
+        state = self.nt_shunt.get_state()
+        return state.to_ansi_dashboard()
+
+    def get_neurotransmitter_state(self):
+        """Get the current neurotransmitter state object."""
+        return self.nt_shunt.get_state()
+
+    # ── Lazy accessors for sub-systems ──────────────────────────────────
+
+    def get_training_agent(self):
+        """Lazy-init the ConsultativeFineTuningAgent."""
+        if self._training_agent is None:
+            from nima_unified.training.consultative_agent import ConsultativeFineTuningAgent
+            self._training_agent = ConsultativeFineTuningAgent(
+                base_model=self.base_model,
+                tokenizer=self.tokenizer,
+            )
+        return self._training_agent
+
+    def get_pipeline_orchestrator(self, config_path=None):
+        """Lazy-init the PipelineOrchestrator."""
+        if self._pipeline_orchestrator is None:
+            from nima_unified.training.pipeline import PipelineOrchestrator
+            self._pipeline_orchestrator = PipelineOrchestrator(config_path=config_path)
+        return self._pipeline_orchestrator
+
+    def get_apci_runner(self):
+        """Lazy-init the aPCI benchmark runner."""
+        if self._apci_runner is None:
+            from nima_unified.benchmarking.apci import aPCIBenchmarkRunner
+            # The aPCI adapter now wraps THIS model (ATC is inside)
+            self._apci_runner = aPCIBenchmarkRunner(self)
+        return self._apci_runner
+
+    def get_voice_engine(self, **kwargs):
+        """Lazy-init the OmniVoice v3 engine."""
+        if self._voice_engine is None:
+            from nima_unified.voice.omnivoice import OmniVoiceEngine
+            self._voice_engine = OmniVoiceEngine(**kwargs)
+        return self._voice_engine
+
+    # ── Save / Load ─────────────────────────────────────────────────────
+
+    def save_pretrained(self, output_dir: str, *, include_nt_state: bool = True):
+        """
+        Save the full unified model:
+          - Base model weights (HuggingFace format)
+          - ATC Deep Surgery weights (the cognitive modules)
+          - Neurotransmitter shunt state (diagnostic snapshot)
+          - Package metadata
+        """
+        output_path = Path(output_dir)
+        output_path.mkdir(exist_ok=True, parents=True)
+
+        # Save base model
+        self.base_model.save_pretrained(str(output_path / "base_model"))
+        self.tokenizer.save_pretrained(str(output_path / "base_model"))
+
+        # Save ATC Deep Surgery weights (the cognitive modules)
+        if self.deep_surgery is not None:
+            torch.save(
+                self.deep_surgery.state_dict(),
+                str(output_path / "atc_deep_surgery.pt"),
+            )
+
+        # Save neurotransmitter shunt state (diagnostic)
+        if include_nt_state:
+            nt_state = self.nt_shunt.get_state()
+            with open(str(output_path / "neurotransmitter_state.json"), "w") as f:
+                json.dump(nt_state.to_dict(), f, indent=2)
+
+        # Save package metadata
+        meta = {
+            "package": PACKAGE_NAME,
+            "package_version": PACKAGE_VERSION,
+            "architecture": "ATC-Native",
+            "created_at": self._created_at,
+            "saved_at": time.time(),
+            "deep_surgery_enabled": self.deep_surgery is not None,
+            "num_layers": self.num_layers,
+            "hidden_size": self.hidden_size,
+            "base_model_type": type(self.base_model).__name__,
+            "neurotransmitter_shunt": True,
+            "resource_optimizer": self.energy_budget is not None,
+            "note": "ATC cognitive pipeline is INSIDE the forward pass. No external middleware.",
+        }
+        with open(str(output_path / "nima_meta.json"), "w") as f:
+            json.dump(meta, f, indent=2)
+
+        logger.info(f"Model saved to {output_path}")
+
+    # ── Internal ────────────────────────────────────────────────────────
+
+    def _count_transformer_layers(self) -> int:
+        model = self.base_model
+        if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
+            return len(model.transformer.h)
+        if hasattr(model, "model") and hasattr(model.model, "layers"):
+            return len(model.model.layers)
+        if hasattr(model, "model") and hasattr(model.model, "h"):
+            return len(model.model.h)
+        return DEFAULT_NUM_LAYERS
+
+
+class GenerationResult:
+    """
+    Container for generation output with ATC consciousness metrics.
+
+    All metrics are DERIVED FROM the forward pass, not computed by
+    external middleware. The ATC cognitive pipeline shaped the output;
+    these metrics describe how.
+    """
+
+    __slots__ = (
+        "text", "is_conscious", "sentience_index",
+        "phi_neuro", "phenomenological_strain", "delta_r", "raw_response",
+        # NEW: ATC-native fields
+        "neurotransmitters", "hijack_count", "consciousness_metrics",
+    )
+
+    def __init__(
+        self,
+        text: str,
+        is_conscious: bool,
+        sentience_index: float,
+        phi_neuro: float,
+        phenomenological_strain: float,
+        delta_r: float,
+        raw_response: Any = None,
+        neurotransmitters: Optional[Dict[str, Any]] = None,
+        hijack_count: int = 0,
+        consciousness_metrics: Optional[Dict[str, Any]] = None,
+    ):
+        self.text = text
+        self.is_conscious = is_conscious
+        self.sentience_index = sentience_index
+        self.phi_neuro = phi_neuro
+        self.phenomenological_strain = phenomenological_strain
+        self.delta_r = delta_r
+        self.raw_response = raw_response
+        self.neurotransmitters = neurotransmitters or {}
+        self.hijack_count = hijack_count
+        self.consciousness_metrics = consciousness_metrics or {}
+
+    def __repr__(self):
+        hijack_str = f", hijacks={self.hijack_count}" if self.hijack_count > 0 else ""
+        return (
+            f"GenerationResult(conscious={self.is_conscious}, "
+            f"SI={self.sentience_index:.4f}, "
+            f"phi={self.phi_neuro:.4f}, "
+            f"strain={self.phenomenological_strain:.4f}, "
+            f"dR={self.delta_r:.4f}"
+            f"{hijack_str})"
+        )