cognitive_utils.py

"""

COGNITIVE-CORE: Utility Functions

==================================



Common utilities for cognitive model development, including:

- Environment setup for Kaggle/Colab

- Device detection

- Memory optimization helpers

- Logging utilities



Copyright © 2026 Mike Amega (Logo) - Ame Web Studio

License: Proprietary - All Rights Reserved

"""

import os
import sys
import torch
import warnings
from typing import Optional, Dict, Any


# ==============================================================================
# ENVIRONNEMENT & CACHE
# ==============================================================================


def setup_environment(cache_dir: Optional[str] = None) -> str:
    """

    Configure l'environnement pour Kaggle/Colab/Local.



    Résout les problèmes de:

    - Read-only file system sur Kaggle

    - Chemins de cache HuggingFace



    Args:

        cache_dir: Répertoire cache personnalisé (optionnel)



    Returns:

        Chemin du répertoire cache configuré

    """
    if cache_dir is None:
        # Détecter l'environnement
        if os.path.exists("/kaggle"):
            cache_dir = "/kaggle/working/.cache"
        elif os.path.exists("/content"):  # Colab
            cache_dir = "/content/.cache"
        else:
            cache_dir = os.path.expanduser("~/.cache/cognitive")

    # Créer le répertoire
    os.makedirs(cache_dir, exist_ok=True)
    os.makedirs(os.path.join(cache_dir, "datasets"), exist_ok=True)

    # Configurer les variables d'environnement
    os.environ["HF_HOME"] = cache_dir
    os.environ["TRANSFORMERS_CACHE"] = cache_dir
    os.environ["HF_DATASETS_CACHE"] = os.path.join(cache_dir, "datasets")

    # Désactiver les warnings non critiques
    warnings.filterwarnings("ignore", category=FutureWarning)
    warnings.filterwarnings("ignore", category=UserWarning, module="transformers")

    return cache_dir


def get_device(prefer_gpu: bool = True) -> torch.device:
    """

    Détecte et retourne le meilleur device disponible.



    Args:

        prefer_gpu: Préférer GPU si disponible



    Returns:

        torch.device configuré

    """
    if prefer_gpu and torch.cuda.is_available():
        device = torch.device("cuda")
        gpu_name = torch.cuda.get_device_name(0)
        gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
        print(f"🔧 GPU: {gpu_name} ({gpu_mem:.1f} GB)")
    elif (
        prefer_gpu
        and hasattr(torch.backends, "mps")
        and torch.backends.mps.is_available()
    ):
        device = torch.device("mps")
        print("🔧 Apple MPS")
    else:
        device = torch.device("cpu")
        print("🔧 CPU")

    return device


def get_optimal_dtype(device: torch.device) -> torch.dtype:
    """

    Retourne le dtype optimal pour le device.



    Args:

        device: Le device cible



    Returns:

        torch.dtype optimal (float16 pour GPU, float32 pour CPU)

    """
    if device.type == "cuda":
        # Vérifier support BF16
        if torch.cuda.is_bf16_supported():
            return torch.bfloat16
        return torch.float16
    return torch.float32


# ==============================================================================
# MÉMOIRE & OPTIMISATION
# ==============================================================================


def get_memory_info() -> Dict[str, float]:
    """

    Retourne les informations mémoire (GPU si disponible).



    Returns:

        Dict avec allocated, reserved, free en GB

    """
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1e9
        reserved = torch.cuda.memory_reserved() / 1e9
        total = torch.cuda.get_device_properties(0).total_memory / 1e9
        return {
            "allocated_gb": allocated,
            "reserved_gb": reserved,
            "free_gb": total - allocated,
            "total_gb": total,
        }
    return {"allocated_gb": 0, "reserved_gb": 0, "free_gb": 0, "total_gb": 0}


def clear_memory():
    """Libère la mémoire GPU si possible."""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()


def estimate_model_memory(model, dtype: torch.dtype = torch.float32) -> float:
    """

    Estime la mémoire nécessaire pour un modèle.



    Args:

        model: Le modèle PyTorch

        dtype: Le dtype utilisé



    Returns:

        Estimation en GB

    """
    param_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
    buffer_bytes = sum(b.numel() * b.element_size() for b in model.buffers())

    # Facteur pour activations (estimation: 2x les paramètres)
    activation_factor = 2.0

    total_bytes = (param_bytes + buffer_bytes) * activation_factor

    # Ajuster selon dtype
    if dtype in (torch.float16, torch.bfloat16):
        total_bytes *= 0.5

    return total_bytes / 1e9


# ==============================================================================
# LOGGING & AFFICHAGE
# ==============================================================================


def print_model_info(model, show_params: bool = True):
    """

    Affiche les informations du modèle.



    Args:

        model: Le modèle à analyser

        show_params: Afficher le détail des paramètres

    """
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

    print(f"\n📊 MODÈLE: {model.__class__.__name__}")
    print(f"   Total paramètres: {total_params:,}")
    print(f"   Paramètres entraînables: {trainable_params:,}")
    print(f"   Mémoire estimée: {estimate_model_memory(model):.2f} GB")

    if show_params and hasattr(model, "config"):
        print(f"\n   Configuration:")
        for key in ["d_model", "n_layers", "n_heads", "vocab_size"]:
            if hasattr(model.config, key):
                print(f"   - {key}: {getattr(model.config, key)}")


def print_training_progress(

    step: int,

    total_steps: int,

    loss: float,

    lr: Optional[float] = None,

    extras: Optional[Dict[str, float]] = None,

):
    """

    Affiche la progression d'entraînement.



    Args:

        step: Étape actuelle

        total_steps: Nombre total d'étapes

        loss: Valeur de la loss

        lr: Learning rate actuel

        extras: Métriques additionnelles

    """
    progress = step / total_steps * 100
    msg = f"[{step:>6}/{total_steps}] ({progress:>5.1f}%) | Loss: {loss:.4f}"

    if lr is not None:
        msg += f" | LR: {lr:.2e}"

    if extras:
        for key, val in extras.items():
            msg += f" | {key}: {val:.4f}"

    print(msg)


# ==============================================================================
# TOKEN HUGGINGFACE
# ==============================================================================


def get_hf_token() -> Optional[str]:
    """

    Récupère le token HuggingFace depuis différentes sources.



    Ordre de recherche:

    1. Variable d'environnement HF_TOKEN

    2. Secrets Kaggle

    3. Secrets Colab

    4. Token local HuggingFace CLI



    Returns:

        Token ou None si non trouvé

    """
    # Env var
    token = os.environ.get("HF_TOKEN")
    if token:
        return token

    # Kaggle
    try:
        from kaggle_secrets import UserSecretsClient

        token = UserSecretsClient().get_secret("HF_TOKEN")
        if token:
            return token
    except Exception:
        pass

    # Colab
    try:
        from google.colab import userdata

        token = userdata.get("HF_TOKEN")
        if token:
            return token
    except Exception:
        pass

    # Local HuggingFace CLI
    try:
        from huggingface_hub import HfFolder

        token = HfFolder.get_token()
        if token:
            return token
    except Exception:
        pass

    return None