cognitive_base.py

"""
COGNITIVE-CORE: Base Classes for Cognitive Architectures
=========================================================

This module provides the foundational classes for building cognitive AI models
that follow the Ame Web Studio standard. All cognitive models (vision, language,
world model, multimodal) should inherit from these base classes.

Copyright © 2026 Mike Amega (Logo) - Ame Web Studio
License: Proprietary - All Rights Reserved
"""

import torch
import torch.nn as nn
from typing import Dict, List, Optional, Any, Tuple
from abc import ABC, abstractmethod

from transformers import PreTrainedModel, PretrainedConfig


# ==============================================================================
# CONFIGURATION DE BASE
# ==============================================================================


class CognitiveConfig(PretrainedConfig):
    """
    Configuration de base pour tous les modèles cognitifs.

    Tous les modèles cognitifs (vision, language, world, multimodal) doivent
    hériter de cette configuration pour garantir la compatibilité.
    """

    model_type = "cognitive"

    def __init__(
        self,
        # Dimensions de base
        d_model: int = 512,
        d_ff: int = 2048,
        n_layers: int = 12,
        n_heads: int = 8,
        dropout: float = 0.1,
        # Modules cognitifs (peuvent être activés/désactivés)
        use_memory: bool = True,
        use_temporal: bool = True,
        use_synaptic: bool = True,
        use_dream: bool = True,
        use_world_model: bool = True,
        use_neurogenesis: bool = True,
        # Mémoire
        memory_size: int = 8192,
        short_term_dim: int = 512,
        long_term_dim: int = 256,
        # États internes
        internal_state_dim: int = 128,
        latent_state_dim: int = 768,
        # Meta
        version: str = "1.0",
        author: str = "Mike Amega",
        license: str = "Proprietary",
        **kwargs,
    ):
        super().__init__(**kwargs)

        # Dimensions
        self.d_model = d_model
        self.hidden_size = d_model  # Alias HuggingFace
        self.d_ff = d_ff
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.dropout = dropout

        # Modules cognitifs
        self.use_memory = use_memory
        self.use_temporal = use_temporal
        self.use_synaptic = use_synaptic
        self.use_dream = use_dream
        self.use_world_model = use_world_model
        self.use_neurogenesis = use_neurogenesis

        # Mémoire
        self.memory_size = memory_size
        self.short_term_dim = short_term_dim
        self.long_term_dim = long_term_dim

        # États
        self.internal_state_dim = internal_state_dim
        self.latent_state_dim = latent_state_dim

        # Meta
        self.version = version
        self.author = author
        self.license = license

    @property
    def head_dim(self) -> int:
        return self.d_model // self.n_heads


# ==============================================================================
# MODULES COGNITIFS ABSTRAITS
# ==============================================================================


class CognitiveModule(nn.Module, ABC):
    """
    Classe de base abstraite pour tous les modules cognitifs.

    Chaque module cognitif doit implémenter:
    - forward(): traitement principal
    - reset_state(): réinitialisation des états internes
    - get_state(): récupérer l'état courant
    """

    def __init__(self, config: CognitiveConfig):
        super().__init__()
        self.config = config

    @abstractmethod
    def forward(self, x: torch.Tensor, **kwargs) -> Dict[str, Any]:
        """Traitement principal du module."""
        pass

    @abstractmethod
    def reset_state(self):
        """Réinitialiser les états internes du module."""
        pass

    def get_state(self) -> Dict[str, torch.Tensor]:
        """Récupérer l'état courant (pour sauvegarde/debug)."""
        return {}


class MemoryModule(CognitiveModule):
    """Interface pour les modules de mémoire."""

    @abstractmethod
    def store(self, key: torch.Tensor, value: torch.Tensor):
        """Stocker une information en mémoire."""
        pass

    @abstractmethod
    def retrieve(self, query: torch.Tensor, k: int = 1) -> torch.Tensor:
        """Récupérer les k informations les plus pertinentes."""
        pass


class TemporalModule(CognitiveModule):
    """Interface pour les modules temporels/prédictifs."""

    @abstractmethod
    def predict(self, state: torch.Tensor, horizon: int = 1) -> torch.Tensor:
        """Prédire l'état futur à l'horizon donné."""
        pass


class WorldModelModule(CognitiveModule):
    """Interface pour les modèles du monde."""

    @abstractmethod
    def update(self, observation: torch.Tensor) -> Dict[str, float]:
        """Mettre à jour le modèle du monde avec une observation."""
        pass

    @abstractmethod
    def imagine(self, action: torch.Tensor) -> torch.Tensor:
        """Imaginer l'effet d'une action."""
        pass


# ==============================================================================
# MODÈLE COGNITIF DE BASE
# ==============================================================================


class CognitivePreTrainedModel(PreTrainedModel):
    """
    Classe de base pour tous les modèles cognitifs HuggingFace-compatibles.

    Fournit:
    - Remappage automatique des clés de checkpoint
    - Gestion des modules cognitifs optionnels
    - Méthodes d'initialisation standardisées
    """

    config_class = CognitiveConfig
    base_model_prefix = "cognitive"
    supports_gradient_checkpointing = False  # Incompatible avec architecture cognitive

    # Clés à ignorer lors du chargement (buffers dynamiques)
    _keys_to_ignore_on_load_missing = [
        r".*\.state$",
        r".*\.history$",
        r".*\.buffer$",
        r".*rope\..*_cache",
        r".*rope\.inv_freq",
    ]

    def _init_weights(self, module):
        """Initialisation standard des poids."""
        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)

    def _load_from_state_dict(
        self,
        state_dict,
        prefix,
        local_metadata,
        strict,
        missing_keys,
        unexpected_keys,
        error_msgs,
    ):
        """
        Remappage automatique des clés de checkpoint.

        Gère les différences de préfixes entre formats de checkpoint
        (ex: avec/sans 'model.' prefix).
        """
        from .cognitive_checkpoint import remap_checkpoint_keys

        # Remapper les clés si nécessaire
        remapped = remap_checkpoint_keys(state_dict, self.state_dict())

        # Appeler l'implémentation parent
        super()._load_from_state_dict(
            remapped,
            prefix,
            local_metadata,
            strict,
            missing_keys,
            unexpected_keys,
            error_msgs,
        )

    def get_cognitive_state(self) -> Dict[str, Any]:
        """Récupérer l'état de tous les modules cognitifs."""
        state = {}
        for name, module in self.named_modules():
            if isinstance(module, CognitiveModule):
                state[name] = module.get_state()
        return state

    def reset_cognitive_state(self):
        """Réinitialiser l'état de tous les modules cognitifs."""
        for module in self.modules():
            if isinstance(module, CognitiveModule):
                module.reset_state()


# ==============================================================================
# UTILITAIRES D'ENREGISTREMENT AUTO
# ==============================================================================


def register_cognitive_model(config_class, model_class):
    """
    Enregistrer un modèle cognitif pour utilisation avec AutoModel.

    Usage:
        register_cognitive_model(MyConfig, MyModel)
        # Puis: AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True)
    """
    from transformers import AutoConfig, AutoModel

    AutoConfig.register(config_class.model_type, config_class)
    AutoModel.register(config_class, model_class)