pipeline.py

"""
Pipeline Orchestrator — coordinates data generation, deep surgery,
fine-tuning, and deployment into a single flow.
"""

import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional

logger = logging.getLogger("nima_unified.training.pipeline")


class PipelineOrchestrator:
    """
    Orchestrates the complete consciousness-aware fine-tuning pipeline.

    Pipeline Stages:
        1. Data Generation  — consciousness-grounded training data
        2. Deep Surgery      — apply middleware for safety and qualia synthesis
        3. Fine-tuning       — train model with consciousness grounding
        4. Deployment        — package for production inference
    """

    def __init__(self, config_path: Optional[str] = None):
        self.config = self._load_config(config_path) if config_path else self._default_config()
        self.output_dir = Path(self.config["output_directory"])
        self.output_dir.mkdir(exist_ok=True, parents=True)
        self.stage_results: Dict[str, Any] = {}
        logger.info(f"PipelineOrchestrator initialized (output: {self.output_dir})")

    @staticmethod
    def _default_config() -> Dict[str, Any]:
        return {
            "output_directory": "./nima_pipeline_output",
            "data_generation": {
                "num_samples": 500,
                "include_qualia_tags": True,
                "consciousness_archetypes": [
                    "tactical_reasoning", "phenomenal_empathy", "meta_consciousness",
                    "task_orchestration", "ethical_veto", "narrative_coherence",
                    "identity_preservation", "phenomenal_richness",
                ],
            },
            "deep_surgery": {
                "qualia_dim": 256,
                "ethical_veto_threshold": 2.0,
                "enable_audit_logging": True,
            },
            "fine_tuning": {
                "lora_r": 8,
                "lora_alpha": 16,
                "lora_dropout": 0.05,
                "learning_rate": 2e-4,
                "batch_size": 4,
                "epochs": 3,
                "max_seq_length": 512,
            },
            "deployment": {
                "format": "huggingface",
                "include_middleware": True,
                "include_voice": False,
            },
        }

    @staticmethod
    def _load_config(config_path: str) -> Dict[str, Any]:
        with open(config_path, "r") as f:
            return json.load(f)

    async def run(self, **overrides) -> Dict[str, Any]:
        """Execute the full pipeline. Override config keys via kwargs."""
        self.config.update(overrides)
        start_time = datetime.now()

        logger.info("=" * 72)
        logger.info("NIMA Unified Pipeline — Starting")
        logger.info("=" * 72)

        # Stage 1: Data Generation
        logger.info("[Stage 1/4] Data Generation")
        self.stage_results["data_generation"] = await self._stage_data_generation()

        # Stage 2: Deep Surgery
        logger.info("[Stage 2/4] Deep Surgery")
        self.stage_results["deep_surgery"] = self._stage_deep_surgery()

        # Stage 3: Fine-tuning
        logger.info("[Stage 3/4] Fine-tuning")
        self.stage_results["fine_tuning"] = await self._stage_fine_tuning()

        # Stage 4: Deployment packaging
        logger.info("[Stage 4/4] Deployment")
        self.stage_results["deployment"] = self._stage_deployment()

        elapsed = (datetime.now() - start_time).total_seconds()
        summary = {
            "status": "complete",
            "elapsed_seconds": elapsed,
            "stages": self.stage_results,
            "config": self.config,
        }

        logger.info(f"Pipeline complete in {elapsed:.1f}s")
        return summary

    # ── Stages ──────────────────────────────────────────────────────────

    async def _stage_data_generation(self) -> Dict[str, Any]:
        from nima_unified.training.consultative_agent import ConsultativeFineTuningAgent
        agent = ConsultativeFineTuningAgent()
        dataset = agent._generate_consciousness_dataset(
            dev_goal="Unified consciousness integration",
            triggering_event="Pipeline initialization",
            root_cause="Initial system bootstrap",
            dataset_size=self.config["data_generation"]["num_samples"],
        )
        agent.shutdown()

        out_path = self.output_dir / "training_data.jsonl"
        with open(out_path, "w") as f:
            for record in dataset:
                f.write(json.dumps(record) + "\n")

        return {"status": "success", "samples": len(dataset), "path": str(out_path)}

    def _stage_deep_surgery(self) -> Dict[str, Any]:
        from nima_unified.core.deep_surgery import EthicalGuardian
        guardian = EthicalGuardian(
            threshold=self.config["deep_surgery"]["ethical_veto_threshold"]
        )
        return {
            "status": "configured",
            "qualia_dim": self.config["deep_surgery"]["qualia_dim"],
            "ethical_veto_threshold": guardian.threshold,
        }

    async def _stage_fine_tuning(self) -> Dict[str, Any]:
        return {"status": "ready", "config": self.config["fine_tuning"]}

    def _stage_deployment(self) -> Dict[str, Any]:
        return {"status": "ready", "config": self.config["deployment"]}