feat: Phi-3.5-MoE multi-agent model repository

- Core Phi-3.5-MoE integration and multi-agent training system
- Multi-agent dataset loading and tokenization
- CPU-optimized LoRA training with Hugging Face
- Configuration and deployment scripts
- Git LFS configuration for model files

.gitattributes

@@ -0,0 +1,8 @@
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.json filter=lfs diff=lfs merge=lfs -text
+*.md filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,37 @@
+# Large files that shouldn't be in model repository
+imported_aws_backup/
+artifacts/
+models/gradient_descent_ensemble/
+models/improved_ensemble/
+models/expanded_agents_1gb/
+
+# Python cache
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+.Python
+*.so
+
+# Virtual environments
+venv/
+env/
+ENV/
+
+# IDE files
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# OS files
+.DS_Store
+Thumbs.db
+
+# Logs
+*.log
+logs/
+
+# Temporary files
+*.tmp
+*.temp

MULTI_AGENT_TRAINING_README.md

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:291d3f54b6a84723e0fe0b2bfa96a8d135d67546c2e2c278d5d3ea8f8a308ed7
+size 13132

deploy_multi_agent_training.py

@@ -0,0 +1,384 @@
+#!/usr/bin/env python3
+"""
+Multi-Agent Training Deployment Script
+
+This script provides comprehensive deployment capabilities for the multi-agent
+training system, including Docker container management, environment setup,
+and training execution.
+"""
+
+import os
+import sys
+import json
+import yaml
+import argparse
+import subprocess
+import logging
+from pathlib import Path
+from typing import Dict, List, Optional, Any
+from dataclasses import dataclass
+
+# Add src to path for imports
+sys.path.append('src')
+
+from training.multi_agent_trainer import MultiAgentTrainingConfig, MultiAgentTrainingPipeline
+from datasets.multi_agent_loader import MultiAgentDatasetConfig
+
+@dataclass
+class DeploymentConfig:
+    """Configuration for deployment"""
+    project_root: str
+    dataset_path: str
+    model_repo_id: str
+    dataset_repo_id: Optional[str] = None
+    agents_file: Optional[str] = None
+    config_file: Optional[str] = None
+    docker_image_name: str = "phi35moe-cpu:latest"
+    output_dir: str = "./outputs"
+    logs_dir: str = "./logs"
+    max_steps: int = 50
+    balance_agents: bool = True
+    push_to_hub: bool = True
+    dry_run: bool = False
+
+class MultiAgentTrainingDeployment:
+    """
+    Comprehensive deployment manager for multi-agent training
+    """
+    
+    def __init__(self, config: DeploymentConfig):
+        self.config = config
+        self.setup_logging()
+        
+    def setup_logging(self):
+        """Setup logging configuration"""
+        logging.basicConfig(
+            level=logging.INFO,
+            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+            handlers=[
+                logging.StreamHandler(),
+                logging.FileHandler(os.path.join(self.config.logs_dir, 'deployment.log'))
+            ]
+        )
+        self.logger = logging.getLogger(__name__)
+        
+        # Create logs directory
+        os.makedirs(self.config.logs_dir, exist_ok=True)
+    
+    def validate_environment(self) -> bool:
+        """Validate deployment environment"""
+        self.logger.info("Validating deployment environment")
+        
+        # Check required environment variables
+        required_env_vars = ["HF_TOKEN"]
+        missing_vars = []
+        
+        for var in required_env_vars:
+            if not os.getenv(var):
+                missing_vars.append(var)
+        
+        if missing_vars:
+            self.logger.error(f"Missing required environment variables: {missing_vars}")
+            return False
+        
+        # Check Docker availability
+        try:
+            result = subprocess.run(["docker", "--version"], capture_output=True, text=True)
+            if result.returncode != 0:
+                self.logger.error("Docker not available")
+                return False
+        except FileNotFoundError:
+            self.logger.error("Docker not installed")
+            return False
+        
+        # Check dataset path
+        if not os.path.exists(self.config.dataset_path):
+            self.logger.error(f"Dataset path not found: {self.config.dataset_path}")
+            return False
+        
+        # Check agents file if specified
+        if self.config.agents_file and not os.path.exists(self.config.agents_file):
+            self.logger.error(f"Agents file not found: {self.config.agents_file}")
+            return False
+        
+        self.logger.info("Environment validation passed")
+        return True
+    
+    def build_docker_image(self) -> bool:
+        """Build Docker image for training"""
+        self.logger.info("Building Docker image")
+        
+        dockerfile_path = "docker/multi_agent_training/Dockerfile.cpu"
+        if not os.path.exists(dockerfile_path):
+            self.logger.error(f"Dockerfile not found: {dockerfile_path}")
+            return False
+        
+        try:
+            cmd = [
+                "docker", "build",
+                "-f", dockerfile_path,
+                "-t", self.config.docker_image_name,
+                "docker/multi_agent_training/"
+            ]
+            
+            self.logger.info(f"Running command: {' '.join(cmd)}")
+            
+            if not self.config.dry_run:
+                result = subprocess.run(cmd, check=True, capture_output=True, text=True)
+                self.logger.info("Docker image built successfully")
+            else:
+                self.logger.info("Dry run: Would build Docker image")
+            
+            return True
+            
+        except subprocess.CalledProcessError as e:
+            self.logger.error(f"Docker build failed: {e}")
+            self.logger.error(f"Error output: {e.stderr}")
+            return False
+    
+    def create_training_command(self) -> List[str]:
+        """Create training command for Docker execution"""
+        cmd = [
+            "python", "/app/train_lora_cpu_multiagent.py",
+            "--dataset_path", self.config.dataset_path,
+            "--hub_repo_id", self.config.model_repo_id,
+            "--output_dir", "/app/outputs",
+            "--max_steps", str(self.config.max_steps),
+            "--logging_steps", "5",
+            "--save_steps", "50",
+            "--eval_steps", "25"
+        ]
+        
+        if self.config.balance_agents:
+            cmd.append("--balance_agents")
+        
+        if self.config.push_to_hub:
+            cmd.append("--push_to_hub")
+        
+        if self.config.agents_file:
+            cmd.extend(["--agents_file", self.config.agents_file])
+        
+        if self.config.dataset_repo_id:
+            cmd.extend(["--push_dataset_repo", self.config.dataset_repo_id])
+        
+        return cmd
+    
+    def run_training(self) -> bool:
+        """Run training in Docker container"""
+        self.logger.info("Starting training in Docker container")
+        
+        # Create training command
+        training_cmd = self.create_training_command()
+        
+        # Prepare Docker run command
+        docker_cmd = [
+            "docker", "run", "--rm", "-it",
+            "-e", f"HF_TOKEN={os.getenv('HF_TOKEN')}",
+            "-v", f"{os.path.abspath(self.config.dataset_path)}:{self.config.dataset_path}:ro",
+            "-v", f"{os.path.abspath(self.config.output_dir)}:/app/outputs",
+            "-v", f"{os.path.abspath(self.config.logs_dir)}:/app/logs",
+            self.config.docker_image_name,
+            "bash", "-lc"
+        ]
+        
+        # Create full command with training script
+        full_cmd = " ".join(training_cmd)
+        docker_cmd.append(full_cmd)
+        
+        self.logger.info(f"Running command: {' '.join(docker_cmd[:-1])} '{full_cmd}'")
+        
+        try:
+            if not self.config.dry_run:
+                result = subprocess.run(docker_cmd, check=True)
+                self.logger.info("Training completed successfully")
+            else:
+                self.logger.info("Dry run: Would execute training")
+            
+            return True
+            
+        except subprocess.CalledProcessError as e:
+            self.logger.error(f"Training failed: {e}")
+            return False
+    
+    def generate_deployment_report(self) -> Dict[str, Any]:
+        """Generate deployment report"""
+        report = {
+            "deployment_config": {
+                "project_root": self.config.project_root,
+                "dataset_path": self.config.dataset_path,
+                "model_repo_id": self.config.model_repo_id,
+                "dataset_repo_id": self.config.dataset_repo_id,
+                "docker_image_name": self.config.docker_image_name,
+                "max_steps": self.config.max_steps,
+                "balance_agents": self.config.balance_agents,
+                "push_to_hub": self.config.push_to_hub
+            },
+            "environment": {
+                "hf_token_set": bool(os.getenv("HF_TOKEN")),
+                "docker_available": self._check_docker_available(),
+                "dataset_exists": os.path.exists(self.config.dataset_path),
+                "agents_file_exists": os.path.exists(self.config.agents_file) if self.config.agents_file else None
+            },
+            "deployment_status": "completed" if not self.config.dry_run else "dry_run"
+        }
+        
+        return report
+    
+    def _check_docker_available(self) -> bool:
+        """Check if Docker is available"""
+        try:
+            result = subprocess.run(["docker", "--version"], capture_output=True, text=True)
+            return result.returncode == 0
+        except FileNotFoundError:
+            return False
+    
+    def save_deployment_report(self, report: Dict[str, Any]):
+        """Save deployment report to file"""
+        report_file = os.path.join(self.config.logs_dir, "deployment_report.json")
+        with open(report_file, 'w') as f:
+            json.dump(report, f, indent=2)
+        
+        self.logger.info(f"Deployment report saved to {report_file}")
+    
+    def deploy(self) -> bool:
+        """Execute complete deployment process"""
+        self.logger.info("Starting multi-agent training deployment")
+        
+        try:
+            # Validate environment
+            if not self.validate_environment():
+                return False
+            
+            # Build Docker image
+            if not self.build_docker_image():
+                return False
+            
+            # Run training
+            if not self.run_training():
+                return False
+            
+            # Generate and save report
+            report = self.generate_deployment_report()
+            self.save_deployment_report(report)
+            
+            self.logger.info("Deployment completed successfully")
+            return True
+            
+        except Exception as e:
+            self.logger.error(f"Deployment failed: {e}")
+            return False
+
+def load_config_from_file(config_file: str) -> Dict[str, Any]:
+    """Load configuration from YAML file"""
+    with open(config_file, 'r') as f:
+        return yaml.safe_load(f)
+
+def create_deployment_config_from_yaml(config_data: Dict[str, Any], 
+                                     dataset_path: str, 
+                                     model_repo_id: str) -> DeploymentConfig:
+    """Create deployment configuration from YAML data"""
+    training_config = config_data.get("training", {})
+    multi_agent_config = config_data.get("multi_agent", {})
+    huggingface_config = config_data.get("huggingface", {})
+    
+    return DeploymentConfig(
+        project_root=os.getcwd(),
+        dataset_path=dataset_path,
+        model_repo_id=model_repo_id,
+        dataset_repo_id=huggingface_config.get("push_dataset_repo"),
+        agents_file=config_data.get("dataset", {}).get("agents_file"),
+        config_file=None,  # Will be set if provided
+        docker_image_name="phi35moe-cpu:latest",
+        output_dir=training_config.get("output_dir", "./outputs"),
+        logs_dir=training_config.get("logging", {}).get("logging_dir", "./logs"),
+        max_steps=training_config.get("max_steps", 50),
+        balance_agents=multi_agent_config.get("balance_agents", True),
+        push_to_hub=huggingface_config.get("push_to_hub", True)
+    )
+
+def main():
+    """Main deployment function"""
+    parser = argparse.ArgumentParser(description="Deploy Multi-Agent Training System")
+    
+    # Required arguments
+    parser.add_argument("--dataset_path", required=True,
+                       help="Path to multi-agent dataset")
+    parser.add_argument("--model_repo_id", required=True,
+                       help="Hugging Face model repository ID")
+    
+    # Optional arguments
+    parser.add_argument("--dataset_repo_id", default="",
+                       help="Optional dataset repository ID")
+    parser.add_argument("--agents_file", default="",
+                       help="Optional agents YAML file")
+    parser.add_argument("--config_file", default="",
+                       help="Optional configuration YAML file")
+    parser.add_argument("--docker_image_name", default="phi35moe-cpu:latest",
+                       help="Docker image name")
+    parser.add_argument("--output_dir", default="./outputs",
+                       help="Output directory")
+    parser.add_argument("--logs_dir", default="./logs",
+                       help="Logs directory")
+    parser.add_argument("--max_steps", type=int, default=50,
+                       help="Maximum training steps")
+    parser.add_argument("--balance_agents", action="store_true",
+                       help="Balance dataset across agents")
+    parser.add_argument("--push_to_hub", action="store_true",
+                       help="Push model to Hugging Face Hub")
+    parser.add_argument("--dry_run", action="store_true",
+                       help="Perform dry run without actual execution")
+    parser.add_argument("--log_level", default="INFO",
+                       help="Logging level")
+    
+    args = parser.parse_args()
+    
+    # Setup logging
+    logging.basicConfig(level=getattr(logging, args.log_level.upper()))
+    
+    try:
+        # Load configuration from file if provided
+        if args.config_file and os.path.exists(args.config_file):
+            config_data = load_config_from_file(args.config_file)
+            deployment_config = create_deployment_config_from_yaml(
+                config_data, args.dataset_path, args.model_repo_id
+            )
+        else:
+            # Create configuration from command line arguments
+            deployment_config = DeploymentConfig(
+                project_root=os.getcwd(),
+                dataset_path=args.dataset_path,
+                model_repo_id=args.model_repo_id,
+                dataset_repo_id=args.dataset_repo_id if args.dataset_repo_id else None,
+                agents_file=args.agents_file if args.agents_file else None,
+                config_file=args.config_file if args.config_file else None,
+                docker_image_name=args.docker_image_name,
+                output_dir=args.output_dir,
+                logs_dir=args.logs_dir,
+                max_steps=args.max_steps,
+                balance_agents=args.balance_agents,
+                push_to_hub=args.push_to_hub,
+                dry_run=args.dry_run
+            )
+        
+        # Create deployment manager
+        deployment = MultiAgentTrainingDeployment(deployment_config)
+        
+        # Execute deployment
+        success = deployment.deploy()
+        
+        if success:
+            print("✅ Deployment completed successfully!")
+            print(f"📁 Outputs: {deployment_config.output_dir}")
+            print(f"📊 Logs: {deployment_config.logs_dir}")
+            if deployment_config.push_to_hub:
+                print(f"🚀 Model: https://huggingface.co/{deployment_config.model_repo_id}")
+        else:
+            print("❌ Deployment failed!")
+            sys.exit(1)
+            
+    except Exception as e:
+        print(f"❌ Deployment error: {e}")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main()

multi_agent_training_config.yaml

@@ -0,0 +1,181 @@
+# Multi-Agent Training Configuration
+# Comprehensive configuration for multi-agent LoRA training
+
+# Model Configuration
+model:
+  base_model: "microsoft/Phi-3.5-MoE-instruct"
+  model_cache_dir: "/app/cache"
+  trust_remote_code: true
+  max_seq_length: 2048
+
+# Training Configuration
+training:
+  output_dir: "/app/outputs"
+  max_steps: 50
+  num_train_epochs: 1
+  per_device_train_batch_size: 1
+  per_device_eval_batch_size: 1
+  gradient_accumulation_steps: 8
+  learning_rate: 2e-5
+  lr_scheduler_type: "cosine"
+  warmup_steps: 0
+  
+  # CPU Optimization
+  use_cpu: true
+  bf16: false
+  fp16: false
+  gradient_checkpointing: true
+  dataloader_num_workers: 0
+  remove_unused_columns: false
+
+# LoRA Configuration
+lora:
+  r: 8
+  alpha: 16
+  dropout: 0.05
+  target_modules: "all-linear"
+  bias: "none"
+
+# Multi-Agent Configuration
+multi_agent:
+  agent_prefix: "<|agent:"
+  agent_suffix: "|>"
+  balance_agents: true
+  balance_cap: null  # null means no cap
+  
+  # Agent-Expert Mapping
+  agent_expert_mapping:
+    SWE: "code"
+    SQE: "code"
+    DevOps: "code"
+    Architect: "reasoning"
+    Security: "reasoning"
+    Product_Manager: "general"
+    VP_Product: "general"
+    Philosopher: "reasoning"
+    EvalOps: "general"
+    CostOptimizer: "general"
+    Discovery: "general"
+
+# Dataset Configuration
+dataset:
+  dataset_path: "/data/datasets/phi_moe_chat"
+  agents_file: "/data/datasets/phi_moe_chat/agents.yaml"
+  validation_split: 0.1
+  seed: 42
+  
+  # Dataset patterns supported
+  patterns:
+    - "single_folder"  # Single folder with agent field in rows
+    - "subfolders"     # Per-agent subfolders
+
+# Logging and Monitoring
+logging:
+  logging_steps: 5
+  save_steps: 50
+  eval_steps: 25
+  save_total_limit: 1
+  logging_dir: "/app/logs"
+  report_to: "none"  # Options: wandb, tensorboard, none
+  
+  # Log levels
+  log_level: "INFO"
+
+# Hugging Face Hub Configuration
+huggingface:
+  hub_repo_id: null  # Set to your repo ID
+  push_to_hub: false
+  push_dataset_repo: null  # Optional dataset repo
+  hub_token: null  # Set via environment variable HF_TOKEN
+
+# Performance Configuration
+performance:
+  # Memory optimization
+  gradient_checkpointing: true
+  use_cache: false
+  
+  # CPU optimization
+  attn_implementation: "eager"
+  torch_dtype: "float32"
+  
+  # Training optimization
+  optim: "adamw_torch"
+  weight_decay: 0.01
+  max_grad_norm: 1.0
+
+# Validation Configuration
+validation:
+  validate_dataset: true
+  dry_run: false
+  
+  # Dataset validation
+  check_agent_field: true
+  check_text_content: true
+  validate_jsonl_format: true
+
+# Deployment Configuration
+deployment:
+  # Docker configuration
+  docker:
+    base_image: "python:3.12-slim-bookworm"
+    working_dir: "/app"
+    cache_dir: "/app/cache"
+    outputs_dir: "/app/outputs"
+    logs_dir: "/app/logs"
+  
+  # Model serving
+  serving:
+    enable_gradio: true
+    gradio_port: 7860
+    max_concurrent_requests: 10
+
+# Environment Configuration
+environment:
+  # Python environment
+  python_version: "3.12"
+  pip_index_url: "https://download.pytorch.org/whl/cpu"
+  
+  # System requirements
+  system_deps:
+    - "git"
+    - "git-lfs"
+    - "build-essential"
+    - "curl"
+    - "wget"
+    - "ca-certificates"
+  
+  # Environment variables
+  env_vars:
+    PYTHONDONTWRITEBYTECODE: "1"
+    PYTHONUNBUFFERED: "1"
+    HF_HOME: "/app/cache"
+    TRANSFORMERS_CACHE: "/app/cache/transformers"
+    HF_DATASETS_CACHE: "/app/cache/datasets"
+
+# Development Configuration
+development:
+  # Testing
+  run_tests: true
+  test_coverage: true
+  
+  # Code quality
+  lint_code: true
+  format_code: true
+  
+  # Documentation
+  generate_docs: true
+  update_readme: true
+
+# Security Configuration
+security:
+  # Token management
+  secure_token_storage: true
+  token_rotation: false
+  
+  # Model security
+  trust_remote_code: true
+  verify_checksums: true
+  
+  # Data security
+  encrypt_sensitive_data: false
+  secure_data_transmission: true

src/models/multi_agent_moe_integration.py

@@ -0,0 +1,393 @@
+#!/usr/bin/env python3
+"""
+Multi-Agent MoE Integration
+
+This module provides integration between the multi-agent training system and the existing
+MoE framework, allowing for seamless combination of agent-specific conditioning and
+expert specialization.
+"""
+
+import os
+import json
+import logging
+from typing import Dict, List, Optional, Tuple, Any, Union
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+from .phi35_moe_integration import EnhancedMoEFramework, Phi35MoEConfig, Phi35MoEExpert
+from ..multi_agent_tokenization.agent_tokenizer import AgentTokenManager, AgentTokenConfig
+from ..multi_agent_training.multi_agent_trainer import MultiAgentTrainingConfig
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class MultiAgentMoEConfig:
+    """Configuration for multi-agent MoE integration"""
+    # Base MoE configuration
+    moe_config: Phi35MoEConfig
+    
+    # Multi-agent configuration
+    agent_prefix: str = "<|agent:"
+    agent_suffix: str = "|>"
+    agents_file: Optional[str] = None
+    
+    # Integration settings
+    enable_agent_conditioning: bool = True
+    enable_expert_routing: bool = True
+    hybrid_mode: bool = True  # Use both agent tokens and expert routing
+    
+    # Model paths
+    base_model_path: str = "microsoft/Phi-3.5-MoE-instruct"
+    lora_adapter_path: Optional[str] = None
+    
+    # Agent-Expert mapping
+    agent_expert_mapping: Optional[Dict[str, str]] = None
+
+class MultiAgentMoEExpert(Phi35MoEExpert):
+    """
+    Enhanced MoE expert with multi-agent support
+    """
+    
+    def __init__(self, expert_id: str, specialization: str, config: Phi35MoEConfig, 
+                 agent_manager: Optional[AgentTokenManager] = None):
+        super().__init__(expert_id, specialization, config)
+        self.agent_manager = agent_manager
+        self.supported_agents: List[str] = []
+        
+    def add_agent_support(self, agent: str):
+        """Add agent support to this expert"""
+        if agent not in self.supported_agents:
+            self.supported_agents.append(agent)
+            logger.info(f"Added agent '{agent}' support to expert '{self.expert_id}'")
+    
+    def format_agent_prompt(self, agent: str, messages: List[Dict[str, str]]) -> List[Dict[str, str]]:
+        """Format messages with agent conditioning"""
+        if not self.agent_manager or agent not in self.supported_agents:
+            return messages
+        
+        # Add agent token to system message or create one
+        formatted_messages = messages.copy()
+        
+        # Check if there's a system message
+        has_system = any(msg.get("role") == "system" for msg in formatted_messages)
+        
+        if not has_system:
+            # Add system message with agent token
+            agent_token = f"{self.agent_manager.config.agent_prefix}{agent}{self.agent_manager.config.agent_suffix}"
+            system_message = {
+                "role": "system",
+                "content": f"You are a {agent} agent specialized in {self.specialization}."
+            }
+            formatted_messages.insert(0, system_message)
+        else:
+            # Update existing system message
+            for msg in formatted_messages:
+                if msg.get("role") == "system":
+                    agent_token = f"{self.agent_manager.config.agent_prefix}{agent}{self.agent_manager.config.agent_suffix}"
+                    msg["content"] = f"{agent_token}\n{msg['content']}"
+                    break
+        
+        return formatted_messages
+    
+    async def generate_response(self, messages: List[Dict[str, str]], 
+                              agent: Optional[str] = None, **kwargs) -> Dict[str, Any]:
+        """Generate response with optional agent conditioning"""
+        if agent and agent in self.supported_agents:
+            messages = self.format_agent_prompt(agent, messages)
+        
+        # Call parent method
+        return await super().generate_response(messages, **kwargs)
+
+class MultiAgentMoERouter:
+    """
+    Enhanced router that considers both agent and content for expert selection
+    """
+    
+    def __init__(self, agent_manager: Optional[AgentTokenManager] = None):
+        self.agent_manager = agent_manager
+        self.agent_expert_mapping: Dict[str, str] = {}
+        self.expert_specializations = {
+            "code": ["programming", "software", "development", "coding", "algorithm", "python", "javascript", "java", "function", "code"],
+            "math": ["mathematics", "calculation", "equation", "formula", "statistics", "derivative", "integral", "algebra", "calculus", "math", "solve", "calculate"],
+            "reasoning": ["logic", "analysis", "reasoning", "problem-solving", "critical", "explain", "why", "how", "because"],
+            "multilingual": ["translation", "language", "multilingual", "localization", "translate", "spanish", "french", "german"],
+            "general": ["general", "conversation", "assistance", "help", "hello", "hi", "what", "who", "when", "where"]
+        }
+    
+    def set_agent_expert_mapping(self, mapping: Dict[str, str]):
+        """Set mapping from agents to preferred experts"""
+        self.agent_expert_mapping = mapping
+        logger.info(f"Set agent-expert mapping: {mapping}")
+    
+    def extract_agent_from_messages(self, messages: List[Dict[str, str]]) -> Optional[str]:
+        """Extract agent from messages"""
+        if not self.agent_manager:
+            return None
+        
+        # Look for agent token in system message
+        for msg in messages:
+            if msg.get("role") == "system":
+                content = msg.get("content", "")
+                agent = self.agent_manager.extract_agent_from_text(content)
+                if agent:
+                    return agent
+        
+        return None
+    
+    def route_experts(self, messages: List[Dict[str, str]], available_experts: List[MultiAgentMoEExpert]) -> List[MultiAgentMoEExpert]:
+        """Route to appropriate experts considering both agent and content"""
+        # Extract agent
+        agent = self.extract_agent_from_messages(messages)
+        
+        # Get content for analysis
+        content = ""
+        for msg in messages:
+            if msg.get("role") in ["user", "assistant"]:
+                content += " " + msg.get("content", "")
+        
+        content_lower = content.lower()
+        
+        # First, try agent-based routing
+        if agent and agent in self.agent_expert_mapping:
+            preferred_expert_type = self.agent_expert_mapping[agent]
+            agent_experts = [exp for exp in available_experts if exp.specialization == preferred_expert_type]
+            if agent_experts:
+                logger.debug(f"Routing agent '{agent}' to {preferred_expert_type} expert")
+                return agent_experts
+        
+        # Fall back to content-based routing
+        for specialization, keywords in self.expert_specializations.items():
+            if any(keyword in content_lower for keyword in keywords):
+                content_experts = [exp for exp in available_experts if exp.specialization == specialization]
+                if content_experts:
+                    logger.debug(f"Routing based on content to {specialization} expert")
+                    return content_experts
+        
+        # Default to general expert
+        general_experts = [exp for exp in available_experts if exp.specialization == "general"]
+        if general_experts:
+            logger.debug("Routing to general expert")
+            return general_experts
+        
+        # Return all experts if no specific routing
+        return available_experts
+
+class MultiAgentMoEFramework(EnhancedMoEFramework):
+    """
+    Enhanced MoE framework with multi-agent support
+    """
+    
+    def __init__(self, config: MultiAgentMoEConfig):
+        super().__init__(config.moe_config)
+        self.multi_agent_config = config
+        self.agent_manager: Optional[AgentTokenManager] = None
+        self.agent_expert_mapping: Dict[str, str] = {}
+        self.agents: List[str] = []
+        
+    async def initialize_agents(self, agents: List[str], agent_expert_mapping: Optional[Dict[str, str]] = None):
+        """Initialize multi-agent support"""
+        self.agents = agents
+        
+        # Create agent token manager
+        agent_config = AgentTokenConfig(
+            agent_prefix=self.multi_agent_config.agent_prefix,
+            agent_suffix=self.multi_agent_config.agent_suffix
+        )
+        self.agent_manager = AgentTokenManager(agent_config)
+        
+        # Set agent-expert mapping
+        if agent_expert_mapping:
+            self.agent_expert_mapping = agent_expert_mapping
+        else:
+            # Default mapping based on agent names
+            self.agent_expert_mapping = self._create_default_mapping(agents)
+        
+        # Update router with agent manager
+        if hasattr(self, 'router') and isinstance(self.router, MultiAgentMoERouter):
+            self.router.agent_manager = self.agent_manager
+            self.router.set_agent_expert_mapping(self.agent_expert_mapping)
+        
+        logger.info(f"Initialized multi-agent support for {len(agents)} agents")
+        logger.info(f"Agent-expert mapping: {self.agent_expert_mapping}")
+    
+    def _create_default_mapping(self, agents: List[str]) -> Dict[str, str]:
+        """Create default agent-expert mapping"""
+        mapping = {}
+        
+        for agent in agents:
+            agent_lower = agent.lower()
+            if any(keyword in agent_lower for keyword in ["swe", "developer", "programmer", "engineer"]):
+                mapping[agent] = "code"
+            elif any(keyword in agent_lower for keyword in ["sqa", "tester", "qa", "quality"]):
+                mapping[agent] = "code"
+            elif any(keyword in agent_lower for keyword in ["devops", "ops", "deployment"]):
+                mapping[agent] = "code"
+            elif any(keyword in agent_lower for keyword in ["architect", "design", "system"]):
+                mapping[agent] = "reasoning"
+            elif any(keyword in agent_lower for keyword in ["security", "sec", "cyber"]):
+                mapping[agent] = "reasoning"
+            elif any(keyword in agent_lower for keyword in ["math", "analyst", "data"]):
+                mapping[agent] = "math"
+            elif any(keyword in agent_lower for keyword in ["translate", "localization", "lang"]):
+                mapping[agent] = "multilingual"
+            else:
+                mapping[agent] = "general"
+        
+        return mapping
+    
+    async def initialize_experts(self, expert_configs: List[Dict[str, str]]):
+        """Initialize experts with multi-agent support"""
+        # Call parent method
+        await super().initialize_experts(expert_configs)
+        
+        # Enhance experts with agent support
+        for expert_id, expert in self.experts.items():
+            if isinstance(expert, MultiAgentMoEExpert):
+                # Add agent support based on mapping
+                for agent, expert_type in self.agent_expert_mapping.items():
+                    if expert.specialization == expert_type:
+                        expert.add_agent_support(agent)
+    
+    async def process_query(self, query: str, agent: Optional[str] = None, 
+                          system_message: Optional[str] = None, 
+                          use_multiple_experts: bool = True, **kwargs) -> Dict[str, Any]:
+        """Process query with optional agent conditioning"""
+        # Format messages with agent context
+        messages = [{"role": "user", "content": query}]
+        
+        if system_message:
+            messages.insert(0, {"role": "system", "content": system_message})
+        
+        # Add agent conditioning if specified
+        if agent and self.agent_manager:
+            agent_token = f"{self.agent_manager.config.agent_prefix}{agent}{self.agent_manager.config.agent_suffix}"
+            if messages[0].get("role") == "system":
+                messages[0]["content"] = f"{agent_token}\n{messages[0]['content']}"
+            else:
+                messages.insert(0, {"role": "system", "content": f"{agent_token}\nYou are a {agent} agent."})
+        
+        # Process with enhanced framework
+        return await super().process_query(
+            query=query,
+            system_message=system_message,
+            use_multiple_experts=use_multiple_experts,
+            **kwargs
+        )
+    
+    def get_agent_statistics(self) -> Dict[str, Any]:
+        """Get statistics about agent support"""
+        stats = {
+            "total_agents": len(self.agents),
+            "agents": self.agents,
+            "agent_expert_mapping": self.agent_expert_mapping,
+            "expert_agent_support": {}
+        }
+        
+        for expert_id, expert in self.experts.items():
+            if isinstance(expert, MultiAgentMoEExpert):
+                stats["expert_agent_support"][expert_id] = {
+                    "specialization": expert.specialization,
+                    "supported_agents": expert.supported_agents
+                }
+        
+        return stats
+    
+    def save_agent_configuration(self, output_dir: str):
+        """Save agent configuration for deployment"""
+        os.makedirs(output_dir, exist_ok=True)
+        
+        config_data = {
+            "agents": self.agents,
+            "agent_expert_mapping": self.agent_expert_mapping,
+            "agent_token_config": {
+                "agent_prefix": self.multi_agent_config.agent_prefix,
+                "agent_suffix": self.multi_agent_config.agent_suffix
+            },
+            "expert_agent_support": self.get_agent_statistics()["expert_agent_support"]
+        }
+        
+        config_file = os.path.join(output_dir, "agent_config.json")
+        with open(config_file, 'w') as f:
+            json.dump(config_data, f, indent=2)
+        
+        logger.info(f"Saved agent configuration to {config_file}")
+        return config_file
+
+class MultiAgentMoEAdapter:
+    """
+    Adapter for integrating multi-agent MoE with existing systems
+    """
+    
+    def __init__(self, multi_agent_moe: MultiAgentMoEFramework):
+        self.multi_agent_moe = multi_agent_moe
+    
+    async def process_agent_query(self, agent: str, query: str, **kwargs) -> Dict[str, Any]:
+        """Process query for specific agent"""
+        return await self.multi_agent_moe.process_query(
+            query=query,
+            agent=agent,
+            **kwargs
+        )
+    
+    def get_agent_capabilities(self, agent: str) -> Dict[str, Any]:
+        """Get capabilities for specific agent"""
+        if agent not in self.multi_agent_moe.agents:
+            return {"error": f"Agent '{agent}' not found"}
+        
+        expert_type = self.multi_agent_moe.agent_expert_mapping.get(agent, "general")
+        
+        return {
+            "agent": agent,
+            "expert_type": expert_type,
+            "supported": True,
+            "capabilities": self.multi_agent_moe.expert_specializations.get(expert_type, [])
+        }
+    
+    def list_available_agents(self) -> List[Dict[str, Any]]:
+        """List all available agents and their capabilities"""
+        agents_info = []
+        
+        for agent in self.multi_agent_moe.agents:
+            expert_type = self.multi_agent_moe.agent_expert_mapping.get(agent, "general")
+            agents_info.append({
+                "agent": agent,
+                "expert_type": expert_type,
+                "capabilities": self.multi_agent_moe.expert_specializations.get(expert_type, [])
+            })
+        
+        return agents_info
+
+# Example usage and testing
+if __name__ == "__main__":
+    # Configure logging
+    logging.basicConfig(level=logging.INFO)
+    
+    # Example configuration
+    moe_config = Phi35MoEConfig()
+    multi_agent_config = MultiAgentMoEConfig(
+        moe_config=moe_config,
+        agent_prefix="<|agent:",
+        agent_suffix="|>"
+    )
+    
+    # Create framework
+    framework = MultiAgentMoEFramework(multi_agent_config)
+    
+    # Example agents
+    agents = ["SWE", "SQE", "DevOps", "Architect", "Security"]
+    
+    # Example agent-expert mapping
+    agent_expert_mapping = {
+        "SWE": "code",
+        "SQE": "code", 
+        "DevOps": "code",
+        "Architect": "reasoning",
+        "Security": "reasoning"
+    }
+    
+    print("Multi-agent MoE framework ready")
+    print(f"Agents: {agents}")
+    print(f"Agent-expert mapping: {agent_expert_mapping}")

src/models/phi35_moe_integration.py

@@ -0,0 +1,397 @@
+#!/usr/bin/env python3
+"""
+Phi-3.5-MoE Integration Module
+
+This module integrates Microsoft's Phi-3.5-MoE-instruct model as a baseline
+for expanding our current MoE (Mixture of Experts) framework.
+
+Based on: https://huggingface.co/microsoft/Phi-3.5-MoE-instruct
+- 41.9B parameters
+- 128K context length
+- Multilingual support
+- Strong reasoning capabilities (code, math, logic)
+"""
+
+import torch
+import logging
+from typing import Dict, List, Optional, Any, Union
+from transformers import (
+    AutoModelForCausalLM, 
+    AutoTokenizer, 
+    pipeline,
+    Pipeline
+)
+from pathlib import Path
+import json
+import asyncio
+from dataclasses import dataclass
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class Phi35MoEConfig:
+    """Configuration for Phi-3.5-MoE integration"""
+    model_name: str = "microsoft/Phi-3.5-MoE-instruct"
+    device_map: str = "auto"
+    torch_dtype: str = "auto"
+    trust_remote_code: bool = False
+    max_new_tokens: int = 500
+    temperature: float = 0.7
+    top_p: float = 0.9
+    do_sample: bool = True
+    context_length: int = 128000  # 128K context length
+    vocabulary_size: int = 32064
+
+class Phi35MoEExpert:
+    """Individual expert using Phi-3.5-MoE model"""
+    
+    def __init__(self, config: Phi35MoEConfig, expert_id: str, specialization: str):
+        self.config = config
+        self.expert_id = expert_id
+        self.specialization = specialization
+        self.model = None
+        self.tokenizer = None
+        self.pipeline = None
+        self.is_loaded = False
+        
+    async def load_model(self):
+        """Load the Phi-3.5-MoE model and tokenizer"""
+        try:
+            logger.info(f"Loading Phi-3.5-MoE model for expert {self.expert_id}")
+            
+            # Load model
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.config.model_name,
+                device_map=self.config.device_map,
+                torch_dtype=self.config.torch_dtype,
+                trust_remote_code=self.config.trust_remote_code,
+            )
+            
+            # Load tokenizer
+            self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_name)
+            
+            # Create pipeline
+            self.pipeline = pipeline(
+                "text-generation",
+                model=self.model,
+                tokenizer=self.tokenizer,
+                device_map=self.config.device_map,
+            )
+            
+            self.is_loaded = True
+            logger.info(f"Successfully loaded Phi-3.5-MoE model for expert {self.expert_id}")
+            
+        except Exception as e:
+            logger.error(f"Failed to load Phi-3.5-MoE model for expert {self.expert_id}: {e}")
+            raise
+    
+    async def generate_response(
+        self, 
+        messages: List[Dict[str, str]], 
+        **generation_kwargs
+    ) -> Dict[str, Any]:
+        """Generate response using Phi-3.5-MoE model"""
+        if not self.is_loaded:
+            await self.load_model()
+        
+        try:
+            # Set default generation parameters
+            generation_args = {
+                "max_new_tokens": self.config.max_new_tokens,
+                "temperature": self.config.temperature,
+                "top_p": self.config.top_p,
+                "do_sample": self.config.do_sample,
+                "return_full_text": False,
+                **generation_kwargs
+            }
+            
+            # Generate response
+            response = self.pipeline(messages, **generation_args)
+            
+            return {
+                "expert_id": self.expert_id,
+                "specialization": self.specialization,
+                "response": response[0]["generated_text"] if response else "",
+                "model": self.config.model_name,
+                "success": True
+            }
+            
+        except Exception as e:
+            logger.error(f"Error generating response for expert {self.expert_id}: {e}")
+            return {
+                "expert_id": self.expert_id,
+                "specialization": self.specialization,
+                "response": "",
+                "error": str(e),
+                "success": False
+            }
+
+class Phi35MoERouter:
+    """Router for selecting appropriate Phi-3.5-MoE experts"""
+    
+    def __init__(self):
+        self.expert_specializations = {
+            "code": ["programming", "software", "development", "coding", "algorithm", "python", "javascript", "java", "function", "code"],
+            "math": ["mathematics", "calculation", "equation", "formula", "statistics", "derivative", "integral", "algebra", "calculus", "math", "solve", "calculate"],
+            "reasoning": ["logic", "analysis", "reasoning", "problem-solving", "critical", "explain", "why", "how", "because"],
+            "multilingual": ["translation", "language", "multilingual", "localization", "translate", "spanish", "french", "german"],
+            "general": ["general", "conversation", "assistance", "help", "hello", "hi", "what", "who", "when", "where"]
+        }
+    
+    def route_query(self, query: str, available_experts: List[Phi35MoEExpert]) -> List[Phi35MoEExpert]:
+        """Route query to appropriate experts based on content analysis"""
+        query_lower = query.lower()
+        selected_experts = []
+        
+        # Score each expert based on query content
+        expert_scores = {}
+        for expert in available_experts:
+            score = 0
+            for keyword in self.expert_specializations.get(expert.specialization, []):
+                if keyword in query_lower:
+                    score += 1
+            expert_scores[expert.expert_id] = score
+        
+        # Select experts with highest scores
+        sorted_experts = sorted(expert_scores.items(), key=lambda x: x[1], reverse=True)
+        
+        # Always include at least one expert (preferably general)
+        if sorted_experts and sorted_experts[0][1] > 0:
+            # Select experts with score > 0
+            for expert_id, score in sorted_experts:
+                if score > 0:
+                    expert = next((e for e in available_experts if e.expert_id == expert_id), None)
+                    if expert:
+                        selected_experts.append(expert)
+        else:
+            # Fallback to general expert
+            general_expert = next((e for e in available_experts if e.specialization == "general"), None)
+            if general_expert:
+                selected_experts.append(general_expert)
+        
+        return selected_experts[:3]  # Limit to top 3 experts
+
+class EnhancedMoEFramework:
+    """Enhanced MoE framework using Phi-3.5-MoE as baseline"""
+    
+    def __init__(self, config: Optional[Phi35MoEConfig] = None):
+        self.config = config or Phi35MoEConfig()
+        self.experts: Dict[str, Phi35MoEExpert] = {}
+        self.router = Phi35MoERouter()
+        self.is_initialized = False
+        
+    async def initialize_experts(self, expert_configs: List[Dict[str, str]]):
+        """Initialize multiple Phi-3.5-MoE experts with different specializations"""
+        logger.info("Initializing Phi-3.5-MoE experts...")
+        
+        for expert_config in expert_configs:
+            expert_id = expert_config["expert_id"]
+            specialization = expert_config["specialization"]
+            
+            expert = Phi35MoEExpert(
+                config=self.config,
+                expert_id=expert_id,
+                specialization=specialization
+            )
+            
+            self.experts[expert_id] = expert
+        
+        self.is_initialized = True
+        logger.info(f"Initialized {len(self.experts)} Phi-3.5-MoE experts")
+    
+    async def process_query(
+        self, 
+        query: str, 
+        system_message: Optional[str] = None,
+        use_multiple_experts: bool = True
+    ) -> Dict[str, Any]:
+        """Process query using the enhanced MoE framework"""
+        if not self.is_initialized:
+            raise RuntimeError("MoE framework not initialized. Call initialize_experts() first.")
+        
+        # Prepare messages
+        messages = []
+        if system_message:
+            messages.append({"role": "system", "content": system_message})
+        messages.append({"role": "user", "content": query})
+        
+        if use_multiple_experts:
+            # Route to multiple experts
+            selected_experts = self.router.route_query(query, list(self.experts.values()))
+            
+            # Generate responses from selected experts
+            tasks = []
+            for expert in selected_experts:
+                task = expert.generate_response(messages)
+                tasks.append(task)
+            
+            # Wait for all responses
+            responses = await asyncio.gather(*tasks, return_exceptions=True)
+            
+            # Process responses
+            successful_responses = []
+            for response in responses:
+                if isinstance(response, dict) and response.get("success", False):
+                    successful_responses.append(response)
+            
+            # Combine responses
+            combined_response = self._combine_responses(successful_responses)
+            
+            return {
+                "query": query,
+                "responses": successful_responses,
+                "combined_response": combined_response,
+                "num_experts_used": len(successful_responses),
+                "success": len(successful_responses) > 0
+            }
+        else:
+            # Use single expert (general)
+            general_expert = self.experts.get("general")
+            if not general_expert:
+                general_expert = list(self.experts.values())[0]  # Fallback to first expert
+            
+            response = await general_expert.generate_response(messages)
+            return {
+                "query": query,
+                "response": response,
+                "success": response.get("success", False)
+            }
+    
+    def _combine_responses(self, responses: List[Dict[str, Any]]) -> str:
+        """Combine multiple expert responses into a coherent answer"""
+        if not responses:
+            return "No responses available."
+        
+        if len(responses) == 1:
+            return responses[0]["response"]
+        
+        # Combine responses with expert attribution
+        combined = "Based on analysis from multiple experts:\n\n"
+        for i, response in enumerate(responses, 1):
+            expert_id = response.get("expert_id", f"Expert {i}")
+            specialization = response.get("specialization", "general")
+            expert_response = response.get("response", "")
+            
+            combined += f"**{expert_id} ({specialization}):**\n{expert_response}\n\n"
+        
+        return combined.strip()
+    
+    async def benchmark_performance(self, test_queries: List[str]) -> Dict[str, Any]:
+        """Benchmark the enhanced MoE framework performance"""
+        logger.info("Starting Phi-3.5-MoE framework benchmark...")
+        
+        results = {
+            "total_queries": len(test_queries),
+            "successful_queries": 0,
+            "failed_queries": 0,
+            "average_response_time": 0,
+            "expert_usage_stats": {},
+            "detailed_results": []
+        }
+        
+        total_time = 0
+        
+        for i, query in enumerate(test_queries):
+            start_time = asyncio.get_event_loop().time()
+            
+            try:
+                result = await self.process_query(query)
+                end_time = asyncio.get_event_loop().time()
+                response_time = end_time - start_time
+                total_time += response_time
+                
+                if result.get("success", False):
+                    results["successful_queries"] += 1
+                else:
+                    results["failed_queries"] += 1
+                
+                # Track expert usage
+                if "responses" in result:
+                    for response in result["responses"]:
+                        expert_id = response.get("expert_id", "unknown")
+                        results["expert_usage_stats"][expert_id] = results["expert_usage_stats"].get(expert_id, 0) + 1
+                
+                results["detailed_results"].append({
+                    "query": query,
+                    "success": result.get("success", False),
+                    "response_time": response_time,
+                    "num_experts_used": result.get("num_experts_used", 1)
+                })
+                
+            except Exception as e:
+                logger.error(f"Error processing query {i}: {e}")
+                results["failed_queries"] += 1
+                results["detailed_results"].append({
+                    "query": query,
+                    "success": False,
+                    "error": str(e),
+                    "response_time": 0
+                })
+        
+        results["average_response_time"] = total_time / len(test_queries) if test_queries else 0
+        results["success_rate"] = results["successful_queries"] / results["total_queries"] if test_queries else 0
+        
+        logger.info(f"Benchmark completed. Success rate: {results['success_rate']:.2%}")
+        return results
+
+# Example usage and testing
+async def main():
+    """Example usage of the enhanced MoE framework"""
+    
+    # Initialize configuration
+    config = Phi35MoEConfig(
+        max_new_tokens=300,
+        temperature=0.7
+    )
+    
+    # Initialize framework
+    moe_framework = EnhancedMoEFramework(config)
+    
+    # Define expert configurations
+    expert_configs = [
+        {"expert_id": "code_expert", "specialization": "code"},
+        {"expert_id": "math_expert", "specialization": "math"},
+        {"expert_id": "reasoning_expert", "specialization": "reasoning"},
+        {"expert_id": "general_expert", "specialization": "general"}
+    ]
+    
+    # Initialize experts
+    await moe_framework.initialize_experts(expert_configs)
+    
+    # Test queries
+    test_queries = [
+        "How do I implement a binary search algorithm in Python?",
+        "What is the derivative of x^2 + 3x + 1?",
+        "Explain the logical reasoning behind the Monty Hall problem",
+        "Hello, how are you today?"
+    ]
+    
+    # Process queries
+    for query in test_queries:
+        print(f"\n{'='*50}")
+        print(f"Query: {query}")
+        print(f"{'='*50}")
+        
+        result = await moe_framework.process_query(query)
+        
+        if result.get("success", False):
+            if "combined_response" in result:
+                print(f"Combined Response:\n{result['combined_response']}")
+            else:
+                print(f"Response:\n{result['response']['response']}")
+        else:
+            print("Failed to generate response")
+    
+    # Run benchmark
+    print(f"\n{'='*50}")
+    print("Running Performance Benchmark...")
+    print(f"{'='*50}")
+    
+    benchmark_results = await moe_framework.benchmark_performance(test_queries)
+    
+    print(f"Success Rate: {benchmark_results['success_rate']:.2%}")
+    print(f"Average Response Time: {benchmark_results['average_response_time']:.2f}s")
+    print(f"Expert Usage Stats: {benchmark_results['expert_usage_stats']}")
+
+if __name__ == "__main__":
+    asyncio.run(main())

src/multi_agent_datasets/__init__.py

@@ -0,0 +1 @@
+# Multi-agent dataset loading package

src/multi_agent_datasets/multi_agent_loader.py

@@ -0,0 +1,515 @@
+#!/usr/bin/env python3
+"""
+Multi-Agent Dataset Loader
+
+This module provides comprehensive support for loading and processing multi-agent datasets
+with two supported patterns:
+A) Single folder with JSONLs that include an "agent" field
+B) Per-agent subfolders (agent name == folder name)
+
+Supports agent balancing, dataset validation, and integration with existing training pipelines.
+"""
+
+import os
+import json
+import yaml
+import logging
+from pathlib import Path
+from typing import Dict, List, Optional, Union, Tuple, Any
+from collections import Counter, defaultdict
+from dataclasses import dataclass
+
+import torch
+from datasets import load_dataset, DatasetDict, Dataset, concatenate_datasets
+from transformers import AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class MultiAgentDatasetConfig:
+    """Configuration for multi-agent dataset loading"""
+    dataset_path: str
+    agents_file: Optional[str] = None
+    agent_prefix: str = "<|agent:"
+    agent_suffix: str = "|>"
+    balance_agents: bool = False
+    balance_cap: Optional[int] = None
+    max_seq_length: int = 2048
+    validation_split: float = 0.1
+    seed: int = 42
+
+class MultiAgentDatasetLoader:
+    """
+    Multi-agent dataset loader supporting two patterns:
+    1. Single folder with JSONLs containing 'agent' field
+    2. Per-agent subfolders with agent name == folder name
+    """
+    
+    def __init__(self, config: MultiAgentDatasetConfig):
+        self.config = config
+        self.agents = []
+        self.dataset_stats = {}
+        
+    def read_agents_yaml(self, path: str) -> List[str]:
+        """Read agents list from YAML file"""
+        yml_path = os.path.join(path, "agents.yaml")
+        if os.path.isfile(yml_path):
+            try:
+                with open(yml_path, "r") as f:
+                    obj = yaml.safe_load(f) or {}
+                    agents = [str(a) for a in obj.get("agents", [])]
+                    logger.info(f"Loaded {len(agents)} agents from YAML: {agents}")
+                    return agents
+            except Exception as e:
+                logger.warning(f"Failed to read agents.yaml: {e}")
+        return []
+    
+    def list_agent_subdirs(self, path: str) -> List[Tuple[str, str]]:
+        """List agent subdirectories with train/test.jsonl files"""
+        items = []
+        if not os.path.isdir(path):
+            return items
+            
+        for name in sorted(os.listdir(path)):
+            subdir_path = os.path.join(path, name)
+            if os.path.isdir(subdir_path):
+                train_file = os.path.join(subdir_path, "train.jsonl")
+                test_file = os.path.join(subdir_path, "test.jsonl")
+                if os.path.isfile(train_file) or os.path.isfile(test_file):
+                    items.append((name, subdir_path))
+                    logger.debug(f"Found agent subdirectory: {name}")
+        
+        return items
+    
+    def load_single_folder_dataset(self, dataset_path: str) -> DatasetDict:
+        """Load dataset from single folder with agent field in rows"""
+        data_files = {}
+        
+        train_file = os.path.join(dataset_path, "train.jsonl")
+        test_file = os.path.join(dataset_path, "test.jsonl")
+        
+        if os.path.isfile(train_file):
+            data_files["train"] = train_file
+        if os.path.isfile(test_file):
+            data_files["test"] = test_file
+            
+        if not data_files:
+            raise FileNotFoundError(f"No dataset files found in {dataset_path}")
+        
+        logger.info(f"Loading single folder dataset from {data_files}")
+        dataset = load_dataset("json", data_files=data_files)
+        
+        # Validate that agent field exists
+        for split_name, split_data in dataset.items():
+            if "agent" not in split_data.column_names:
+                raise ValueError(f"Agent field not found in {split_name} split")
+        
+        return dataset
+    
+    def load_subfolder_dataset(self, dataset_path: str) -> DatasetDict:
+        """Load dataset from per-agent subfolders"""
+        subdirs = self.list_agent_subdirs(dataset_path)
+        if not subdirs:
+            raise FileNotFoundError(f"No agent subdirectories found in {dataset_path}")
+        
+        parts_train, parts_test = [], []
+        
+        for agent_name, agent_dir in subdirs:
+            train_file = os.path.join(agent_dir, "train.jsonl")
+            test_file = os.path.join(agent_dir, "test.jsonl")
+            
+            def add_agent_field(example):
+                example["agent"] = agent_name
+                return example
+            
+            if os.path.isfile(train_file):
+                logger.debug(f"Loading train data for agent: {agent_name}")
+                train_data = load_dataset("json", data_files={"train": train_file})["train"]
+                train_data = train_data.map(add_agent_field)
+                parts_train.append(train_data)
+            
+            if os.path.isfile(test_file):
+                logger.debug(f"Loading test data for agent: {agent_name}")
+                test_data = load_dataset("json", data_files={"test": test_file})["test"]
+                test_data = test_data.map(add_agent_field)
+                parts_test.append(test_data)
+        
+        dataset_dict = {}
+        if parts_train:
+            dataset_dict["train"] = concatenate_datasets(parts_train)
+        if parts_test:
+            dataset_dict["test"] = concatenate_datasets(parts_test)
+        
+        if not dataset_dict:
+            raise ValueError("No data splits found in agent subdirectories")
+        
+        return DatasetDict(dataset_dict)
+    
+    def load_multiagent_dataset(self) -> DatasetDict:
+        """
+        Load multi-agent dataset supporting both patterns:
+        - Single folder with 'agent' field in rows
+        - Per-agent subfolders
+        """
+        dataset_path = self.config.dataset_path
+        
+        # Try single folder pattern first
+        if os.path.isfile(os.path.join(dataset_path, "train.jsonl")):
+            logger.info("Loading dataset using single folder pattern")
+            return self.load_single_folder_dataset(dataset_path)
+        
+        # Try subfolder pattern
+        logger.info("Loading dataset using subfolder pattern")
+        return self.load_subfolder_dataset(dataset_path)
+    
+    def infer_agents_from_dataset(self, dataset: DatasetDict) -> List[str]:
+        """Infer agent list from dataset"""
+        agents = set()
+        
+        for split_name, split_data in dataset.items():
+            if "agent" in split_data.column_names:
+                agent_values = [a for a in set(split_data["agent"]) if a is not None]
+                agents.update(agent_values)
+                logger.debug(f"Found agents in {split_name}: {agent_values}")
+        
+        agents_list = sorted(list(agents))
+        logger.info(f"Inferred {len(agents_list)} agents from dataset: {agents_list}")
+        return agents_list
+    
+    def resolve_agents_list(self, dataset: DatasetDict) -> List[str]:
+        """Resolve agents list from YAML file or dataset inference"""
+        agents = []
+        
+        # Try to load from agents file first
+        if self.config.agents_file and os.path.isfile(self.config.agents_file):
+            try:
+                with open(self.config.agents_file, "r") as f:
+                    obj = yaml.safe_load(f) or {}
+                    agents = [str(a) for a in obj.get("agents", [])]
+                    logger.info(f"Loaded agents from file: {agents}")
+            except Exception as e:
+                logger.warning(f"Failed to load agents from file: {e}")
+        
+        # Fall back to dataset inference
+        if not agents:
+            agents = self.infer_agents_from_dataset(dataset)
+        
+        self.agents = agents
+        return agents
+    
+    def balance_by_agent(self, dataset: Dataset, agent_col: str = "agent") -> Dataset:
+        """
+        Balance dataset by upsampling minority agents to the max count
+        """
+        if agent_col not in dataset.column_names:
+            logger.warning(f"Agent column '{agent_col}' not found, skipping balancing")
+            return dataset
+        
+        counts = Counter(dataset[agent_col])
+        if not counts:
+            logger.warning("No agent counts found, skipping balancing")
+            return dataset
+        
+        max_count = max(counts.values())
+        if self.config.balance_cap:
+            max_count = min(max_count, self.config.balance_cap)
+        
+        logger.info(f"Balancing agents. Current counts: {dict(counts)}")
+        logger.info(f"Target count per agent: {max_count}")
+        
+        parts = []
+        for agent, count in counts.items():
+            agent_subset = dataset.filter(lambda x: x[agent_col] == agent)
+            parts.append(agent_subset)
+            
+            # Calculate how many additional samples we need
+            needed = max_count - count
+            if needed > 0:
+                agent_subset_len = len(agent_subset)
+                if agent_subset_len == 0:
+                    logger.warning(f"Agent '{agent}' has zero samples, cannot upsample.")
+                    continue
+                # Calculate repetitions needed
+                reps = needed // agent_subset_len
+                remainder = needed % agent_subset_len
+                
+                # Add full repetitions
+                for _ in range(reps):
+                    parts.append(agent_subset)
+                
+                # Add remainder samples
+                if remainder > 0:
+                    remainder_subset = agent_subset.shuffle(seed=self.config.seed).select(range(remainder))
+                    parts.append(remainder_subset)
+        
+        balanced_dataset = concatenate_datasets(parts).shuffle(seed=self.config.seed)
+        
+        # Log final counts
+        final_counts = Counter(balanced_dataset[agent_col])
+        logger.info(f"Balanced dataset counts: {dict(final_counts)}")
+        
+        return balanced_dataset
+    
+    def apply_agent_prefix(self, dataset: Dataset, tokenizer: AutoTokenizer) -> Dataset:
+        """
+        Apply agent prefix to dataset text using chat template or direct text
+        """
+        def add_agent_prefix(example):
+            agent = example.get("agent", None)
+            prefix = f"{self.config.agent_prefix}{agent}{self.config.agent_suffix}\n" if agent else ""
+            
+            # Handle different text formats
+            if "messages" in example and example["messages"] is not None:
+                # Use chat template if available
+                try:
+                    text = tokenizer.apply_chat_template(
+                        example["messages"], 
+                        tokenize=False, 
+                        add_generation_prompt=False
+                    )
+                    example["text"] = prefix + text
+                except Exception as e:
+                    logger.warning(f"Failed to apply chat template: {e}")
+                    # Fallback to simple concatenation
+                    text = str(example["messages"])
+                    example["text"] = prefix + text
+                    
+            elif "text" in example and example["text"] is not None:
+                example["text"] = prefix + example["text"]
+                
+            else:
+                # Handle prompt/response format
+                prompt = example.get("prompt", "")
+                response = example.get("response", "")
+                example["text"] = prefix + prompt + ("\n" if response else "") + response
+            
+            return example
+        
+        # Get original features to preserve
+        original_features = list(dataset.features)
+        features_to_remove = [f for f in original_features if f not in ["text", "agent"]]
+        
+        logger.info("Applying agent prefixes to dataset")
+        processed_dataset = dataset.map(
+            add_agent_prefix,
+            remove_columns=features_to_remove,
+            desc="Adding agent prefixes"
+        )
+        
+        return processed_dataset
+    
+    def validate_dataset(self, dataset: DatasetDict) -> Dict[str, Any]:
+        """Validate dataset and return statistics"""
+        stats = {
+            "total_samples": 0,
+            "agents": {},
+            "splits": {},
+            "validation_errors": []
+        }
+        
+        for split_name, split_data in dataset.items():
+            split_stats = {
+                "samples": len(split_data),
+                "agents": {},
+                "columns": split_data.column_names
+            }
+            
+            stats["total_samples"] += len(split_data)
+            
+            # Validate required columns
+            if "agent" not in split_data.column_names:
+                stats["validation_errors"].append(f"Missing 'agent' column in {split_name}")
+            
+            # Count agents
+            if "agent" in split_data.column_names:
+                agent_counts = Counter(split_data["agent"])
+                split_stats["agents"] = dict(agent_counts)
+                
+                # Update global agent counts
+                for agent, count in agent_counts.items():
+                    if agent not in stats["agents"]:
+                        stats["agents"][agent] = 0
+                    stats["agents"][agent] += count
+            
+            stats["splits"][split_name] = split_stats
+        
+        self.dataset_stats = stats
+        logger.info(f"Dataset validation complete. Stats: {stats}")
+        
+        return stats
+    
+    def load_and_process(self, tokenizer: AutoTokenizer) -> Tuple[DatasetDict, List[str], Dict[str, Any]]:
+        """
+        Complete dataset loading and processing pipeline
+        """
+        logger.info(f"Loading multi-agent dataset from {self.config.dataset_path}")
+        
+        # Load dataset
+        dataset = self.load_multiagent_dataset()
+        
+        # Resolve agents list
+        agents = self.resolve_agents_list(dataset)
+        
+        # Validate dataset
+        stats = self.validate_dataset(dataset)
+        
+        # Apply agent prefixes
+        if "train" in dataset:
+            dataset["train"] = self.apply_agent_prefix(dataset["train"], tokenizer)
+        if "test" in dataset:
+            dataset["test"] = self.apply_agent_prefix(dataset["test"], tokenizer)
+        
+        # Balance agents if requested
+        if self.config.balance_agents and "train" in dataset:
+            dataset["train"] = self.balance_by_agent(dataset["train"])
+        
+        logger.info(f"Dataset processing complete. Loaded {len(agents)} agents with {stats['total_samples']} total samples")
+        
+        return dataset, agents, stats
+
+class MultiAgentDatasetValidator:
+    """Validator for multi-agent datasets"""
+    
+    @staticmethod
+    def validate_jsonl_file(file_path: str) -> List[str]:
+        """Validate JSONL file format and content"""
+        errors = []
+        
+        if not os.path.isfile(file_path):
+            errors.append(f"File not found: {file_path}")
+            return errors
+        
+        try:
+            with open(file_path, 'r') as f:
+                for line_num, line in enumerate(f, 1):
+                    line = line.strip()
+                    if not line:
+                        continue
+                    
+                    try:
+                        data = json.loads(line)
+                        
+                        # Check required fields
+                        if not isinstance(data, dict):
+                            errors.append(f"Line {line_num}: Not a JSON object")
+                            continue
+                        
+                        # Check for agent field
+                        if "agent" not in data:
+                            errors.append(f"Line {line_num}: Missing 'agent' field")
+                        
+                        # Check for text content
+                        has_text = any(field in data for field in ["text", "messages", "prompt"])
+                        if not has_text:
+                            errors.append(f"Line {line_num}: No text content found")
+                        
+                    except json.JSONDecodeError as e:
+                        errors.append(f"Line {line_num}: JSON decode error - {e}")
+        
+        except Exception as e:
+            errors.append(f"File read error: {e}")
+        
+        return errors
+    
+    @staticmethod
+    def validate_dataset_structure(dataset_path: str) -> Dict[str, Any]:
+        """Validate complete dataset structure"""
+        validation_result = {
+            "valid": True,
+            "errors": [],
+            "warnings": [],
+            "structure": {}
+        }
+        
+        if not os.path.isdir(dataset_path):
+            validation_result["valid"] = False
+            validation_result["errors"].append(f"Dataset path is not a directory: {dataset_path}")
+            return validation_result
+        
+        # Check for single folder pattern
+        train_file = os.path.join(dataset_path, "train.jsonl")
+        test_file = os.path.join(dataset_path, "test.jsonl")
+        
+        if os.path.isfile(train_file):
+            validation_result["structure"]["pattern"] = "single_folder"
+            validation_result["structure"]["files"] = []
+            
+            if os.path.isfile(train_file):
+                validation_result["structure"]["files"].append("train.jsonl")
+                errors = MultiAgentDatasetValidator.validate_jsonl_file(train_file)
+                validation_result["errors"].extend(errors)
+            
+            if os.path.isfile(test_file):
+                validation_result["structure"]["files"].append("test.jsonl")
+                errors = MultiAgentDatasetValidator.validate_jsonl_file(test_file)
+                validation_result["errors"].extend(errors)
+        
+        else:
+            # Check for subfolder pattern
+            validation_result["structure"]["pattern"] = "subfolders"
+            validation_result["structure"]["agents"] = []
+            
+            for item in os.listdir(dataset_path):
+                item_path = os.path.join(dataset_path, item)
+                if os.path.isdir(item_path):
+                    agent_train = os.path.join(item_path, "train.jsonl")
+                    agent_test = os.path.join(item_path, "test.jsonl")
+                    
+                    if os.path.isfile(agent_train) or os.path.isfile(agent_test):
+                        validation_result["structure"]["agents"].append(item)
+                        
+                        if os.path.isfile(agent_train):
+                            errors = MultiAgentDatasetValidator.validate_jsonl_file(agent_train)
+                            validation_result["errors"].extend([f"{item}/train.jsonl: {e}" for e in errors])
+                        
+                        if os.path.isfile(agent_test):
+                            errors = MultiAgentDatasetValidator.validate_jsonl_file(agent_test)
+                            validation_result["errors"].extend([f"{item}/test.jsonl: {e}" for e in errors])
+        
+        # Check for agents.yaml
+        agents_yaml = os.path.join(dataset_path, "agents.yaml")
+        if os.path.isfile(agents_yaml):
+            validation_result["structure"]["has_agents_yaml"] = True
+            try:
+                with open(agents_yaml, 'r') as f:
+                    yaml.safe_load(f)
+            except Exception as e:
+                validation_result["warnings"].append(f"Invalid agents.yaml: {e}")
+        else:
+            validation_result["structure"]["has_agents_yaml"] = False
+        
+        validation_result["valid"] = len(validation_result["errors"]) == 0
+        
+        return validation_result
+
+# Example usage and testing
+if __name__ == "__main__":
+    # Configure logging
+    logging.basicConfig(level=logging.INFO)
+    
+    # Example configuration
+    config = MultiAgentDatasetConfig(
+        dataset_path="/path/to/dataset",
+        balance_agents=True,
+        balance_cap=1000
+    )
+    
+    # Create loader
+    loader = MultiAgentDatasetLoader(config)
+    
+    # Example tokenizer (would be loaded from actual model)
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3.5-MoE-instruct")
+    
+    try:
+        # Load and process dataset
+        dataset, agents, stats = loader.load_and_process(tokenizer)
+        
+        print(f"Loaded dataset with {len(agents)} agents:")
+        for agent in agents:
+            print(f"  - {agent}")
+        
+        print(f"Dataset stats: {stats}")
+        
+    except Exception as e:
+        print(f"Error loading dataset: {e}")

src/multi_agent_tokenization/__init__.py

@@ -0,0 +1 @@
+# Agent tokenization package

src/multi_agent_tokenization/agent_tokenizer.py

@@ -0,0 +1,402 @@
+#!/usr/bin/env python3
+"""
+Agent Token Management System
+
+This module provides comprehensive agent token management for multi-agent training,
+including special token handling, embedding management, and integration with
+existing tokenization systems.
+"""
+
+import os
+import json
+import logging
+from typing import Dict, List, Optional, Tuple, Any, Union
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from transformers import AutoTokenizer, PreTrainedTokenizer
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class AgentTokenConfig:
+    """Configuration for agent token management"""
+    agent_prefix: str = "<|agent:"
+    agent_suffix: str = "|>"
+    special_tokens: Optional[Dict[str, str]] = None
+    resize_embeddings: bool = True
+    save_tokens: bool = True
+    tokens_file: str = "agent_tokens.json"
+
+class AgentTokenManager:
+    """
+    Manages agent-specific tokens and their integration with tokenizers
+    """
+    
+    def __init__(self, config: AgentTokenConfig):
+        self.config = config
+        self.agent_tokens: Dict[str, str] = {}
+        self.token_ids: Dict[str, int] = {}
+        self.original_vocab_size: Optional[int] = None
+        
+    def generate_agent_tokens(self, agents: List[str]) -> List[str]:
+        """Generate agent tokens for given agent list"""
+        tokens = []
+        for agent in agents:
+            token = f"{self.config.agent_prefix}{agent}{self.config.agent_suffix}"
+            tokens.append(token)
+            self.agent_tokens[agent] = token
+        
+        logger.info(f"Generated {len(tokens)} agent tokens: {tokens}")
+        return tokens
+    
+    def add_agent_tokens_to_tokenizer(self, tokenizer: PreTrainedTokenizer, agents: List[str]) -> Tuple[PreTrainedTokenizer, List[str]]:
+        """
+        Add agent tokens to tokenizer and return updated tokenizer with token list
+        """
+        if not agents:
+            logger.warning("No agents provided, skipping token addition")
+            return tokenizer, []
+        
+        # Generate agent tokens
+        agent_tokens = self.generate_agent_tokens(agents)
+        
+        # Check which tokens need to be added
+        existing_tokens = set(tokenizer.get_vocab().keys())
+        tokens_to_add = [token for token in agent_tokens if token not in existing_tokens]
+        
+        if not tokens_to_add:
+            logger.info("All agent tokens already exist in tokenizer")
+            return tokenizer, agent_tokens
+        
+        # Store original vocab size
+        self.original_vocab_size = len(tokenizer)
+        
+        # Add special tokens
+        logger.info(f"Adding {len(tokens_to_add)} new agent tokens to tokenizer")
+        tokenizer.add_special_tokens({
+            "additional_special_tokens": tokens_to_add
+        })
+        
+        # Update token IDs mapping
+        for agent, token in self.agent_tokens.items():
+            if token in tokenizer.get_vocab():
+                self.token_ids[agent] = tokenizer.convert_tokens_to_ids(token)
+        
+        logger.info(f"Added agent tokens. New vocab size: {len(tokenizer)}")
+        return tokenizer, agent_tokens
+    
+    def resize_model_embeddings(self, model: torch.nn.Module, tokenizer: PreTrainedTokenizer) -> torch.nn.Module:
+        """
+        Resize model embeddings to accommodate new agent tokens
+        """
+        if not self.config.resize_embeddings:
+            logger.info("Embedding resize disabled, skipping")
+            return model
+        
+        if self.original_vocab_size is None:
+            logger.warning("Original vocab size not set, cannot resize embeddings")
+            return model
+        
+        new_vocab_size = len(tokenizer)
+        if new_vocab_size == self.original_vocab_size:
+            logger.info("Vocab size unchanged, no embedding resize needed")
+            return model
+        
+        logger.info(f"Resizing model embeddings from {self.original_vocab_size} to {new_vocab_size}")
+        
+        # Resize embeddings
+        model.resize_token_embeddings(new_vocab_size)
+        
+        # Initialize new embeddings (copy from unk token or use random initialization)
+        if hasattr(model, 'get_input_embeddings'):
+            embeddings = model.get_input_embeddings()
+            if hasattr(embeddings, 'weight'):
+                with torch.no_grad():
+                    # Initialize new embeddings with small random values
+                    new_embeddings = embeddings.weight[self.original_vocab_size:]
+                    torch.nn.init.normal_(new_embeddings, mean=0.0, std=0.02)
+        
+        logger.info("Model embeddings resized successfully")
+        return model
+    
+    def format_agent_prompt(self, agent: str, text: str) -> str:
+        """Format text with agent token prefix"""
+        if agent not in self.agent_tokens:
+            logger.warning(f"Agent '{agent}' not found in token mappings")
+            return text
+        
+        agent_token = self.agent_tokens[agent]
+        return f"{agent_token}\n{text}"
+    
+    def extract_agent_from_text(self, text: str) -> Optional[str]:
+        """Extract agent name from text if it starts with agent token"""
+        for agent, token in self.agent_tokens.items():
+            if text.startswith(token):
+                return agent
+        return None
+    
+    def get_agent_token_id(self, agent: str) -> Optional[int]:
+        """Get token ID for agent token"""
+        return self.token_ids.get(agent)
+    
+    def save_agent_tokens(self, output_dir: str) -> str:
+        """Save agent tokens to file"""
+        if not self.config.save_tokens:
+            return ""
+        
+        os.makedirs(output_dir, exist_ok=True)
+        tokens_file = os.path.join(output_dir, self.config.tokens_file)
+        
+        tokens_data = {
+            "agent_tokens": self.agent_tokens,
+            "token_ids": self.token_ids,
+            "config": {
+                "agent_prefix": self.config.agent_prefix,
+                "agent_suffix": self.config.agent_suffix,
+                "original_vocab_size": self.original_vocab_size
+            }
+        }
+        
+        with open(tokens_file, 'w') as f:
+            json.dump(tokens_data, f, indent=2)
+        
+        logger.info(f"Saved agent tokens to {tokens_file}")
+        return tokens_file
+    
+    def load_agent_tokens(self, tokens_file: str) -> bool:
+        """Load agent tokens from file"""
+        if not os.path.isfile(tokens_file):
+            logger.warning(f"Agent tokens file not found: {tokens_file}")
+            return False
+        
+        try:
+            with open(tokens_file, 'r') as f:
+                tokens_data = json.load(f)
+            
+            self.agent_tokens = tokens_data.get("agent_tokens", {})
+            self.token_ids = tokens_data.get("token_ids", {})
+            
+            config_data = tokens_data.get("config", {})
+            self.original_vocab_size = config_data.get("original_vocab_size")
+            
+            logger.info(f"Loaded {len(self.agent_tokens)} agent tokens from {tokens_file}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to load agent tokens: {e}")
+            return False
+    
+    def get_agent_statistics(self) -> Dict[str, Any]:
+        """Get statistics about agent tokens"""
+        return {
+            "total_agents": len(self.agent_tokens),
+            "agents": list(self.agent_tokens.keys()),
+            "token_ids": self.token_ids,
+            "original_vocab_size": self.original_vocab_size,
+            "config": {
+                "agent_prefix": self.config.agent_prefix,
+                "agent_suffix": self.config.agent_suffix
+            }
+        }
+
+class AgentTokenizer:
+    """
+    Enhanced tokenizer wrapper that integrates agent token management
+    """
+    
+    def __init__(self, tokenizer: PreTrainedTokenizer, agent_manager: AgentTokenManager):
+        self.tokenizer = tokenizer
+        self.agent_manager = agent_manager
+    
+    def tokenize_agent_text(self, agent: str, text: str, **kwargs) -> Dict[str, Any]:
+        """Tokenize text with agent prefix"""
+        formatted_text = self.agent_manager.format_agent_prompt(agent, text)
+        return self.tokenizer(formatted_text, **kwargs)
+    
+    def decode_agent_tokens(self, token_ids: Union[List[int], torch.Tensor], **kwargs) -> str:
+        """Decode token IDs back to text"""
+        return self.tokenizer.decode(token_ids, **kwargs)
+    
+    def get_agent_attention_mask(self, input_ids: torch.Tensor, agent: str) -> torch.Tensor:
+        """Get attention mask with special handling for agent tokens"""
+        attention_mask = torch.ones_like(input_ids)
+        
+        # Find agent token position
+        agent_token_id = self.agent_manager.get_agent_token_id(agent)
+        if agent_token_id is not None:
+            # Ensure agent token is attended to
+            agent_positions = (input_ids == agent_token_id)
+            attention_mask[agent_positions] = 1
+        
+        return attention_mask
+    
+    def __getattr__(self, name):
+        """Delegate unknown attributes to underlying tokenizer"""
+        return getattr(self.tokenizer, name)
+
+class AgentTokenValidator:
+    """Validator for agent token configurations"""
+    
+    @staticmethod
+    def validate_agent_tokens(agents: List[str], config: AgentTokenConfig) -> Dict[str, Any]:
+        """Validate agent token configuration"""
+        validation_result = {
+            "valid": True,
+            "errors": [],
+            "warnings": [],
+            "tokens": {}
+        }
+        
+        if not agents:
+            validation_result["warnings"].append("No agents provided")
+            return validation_result
+        
+        # Check for duplicate agents
+        if len(agents) != len(set(agents)):
+            validation_result["errors"].append("Duplicate agents found")
+            validation_result["valid"] = False
+        
+        # Generate and validate tokens
+        manager = AgentTokenManager(config)
+        tokens = manager.generate_agent_tokens(agents)
+        
+        # Check for token conflicts
+        token_set = set(tokens)
+        if len(token_set) != len(tokens):
+            validation_result["errors"].append("Duplicate tokens generated")
+            validation_result["valid"] = False
+        
+        # Check token length
+        for agent, token in zip(agents, tokens):
+            if len(token) > 50:  # Reasonable limit
+                validation_result["warnings"].append(f"Long token for agent '{agent}': {token}")
+        
+        validation_result["tokens"] = dict(zip(agents, tokens))
+        
+        return validation_result
+    
+    @staticmethod
+    def validate_tokenizer_compatibility(tokenizer: PreTrainedTokenizer, agents: List[str], config: AgentTokenConfig) -> Dict[str, Any]:
+        """Validate tokenizer compatibility with agent tokens"""
+        validation_result = {
+            "compatible": True,
+            "errors": [],
+            "warnings": [],
+            "existing_tokens": [],
+            "new_tokens": []
+        }
+        
+        if not agents:
+            return validation_result
+        
+        # Generate tokens
+        manager = AgentTokenManager(config)
+        tokens = manager.generate_agent_tokens(agents)
+        
+        # Check existing vocabulary
+        vocab = tokenizer.get_vocab()
+        for agent, token in zip(agents, tokens):
+            if token in vocab:
+                validation_result["existing_tokens"].append(agent)
+            else:
+                validation_result["new_tokens"].append(agent)
+        
+        # Check for potential conflicts
+        for token in tokens:
+            if token in vocab:
+                # Check if it's already a special token
+                if hasattr(tokenizer, 'special_tokens_map'):
+                    special_tokens = tokenizer.special_tokens_map
+                    if token not in special_tokens.values():
+                        validation_result["warnings"].append(f"Token '{token}' exists in vocab but not as special token")
+        
+        return validation_result
+
+# Integration with existing MoE framework
+class MoEAgentTokenIntegration:
+    """
+    Integration layer between agent tokens and MoE framework
+    """
+    
+    def __init__(self, agent_manager: AgentTokenManager):
+        self.agent_manager = agent_manager
+        self.agent_to_expert_mapping: Dict[str, str] = {}
+    
+    def map_agent_to_expert(self, agent: str, expert: str):
+        """Map agent to MoE expert specialization"""
+        self.agent_to_expert_mapping[agent] = expert
+        logger.info(f"Mapped agent '{agent}' to expert '{expert}'")
+    
+    def get_expert_for_agent(self, agent: str) -> Optional[str]:
+        """Get expert specialization for agent"""
+        return self.agent_to_expert_mapping.get(agent)
+    
+    def format_moe_prompt(self, agent: str, text: str, expert: Optional[str] = None) -> str:
+        """Format prompt for MoE framework with agent and expert context"""
+        # Start with agent token
+        formatted_text = self.agent_manager.format_agent_prompt(agent, text)
+        
+        # Add expert context if available
+        if expert:
+            expert_context = f"\n<|expert:{expert}|>\n"
+            formatted_text = formatted_text.replace("\n", expert_context, 1)
+        
+        return formatted_text
+    
+    def extract_agent_and_expert(self, text: str) -> Tuple[Optional[str], Optional[str]]:
+        """Extract both agent and expert from formatted text"""
+        agent = self.agent_manager.extract_agent_from_text(text)
+        
+        # Extract expert if present
+        expert = None
+        if "<|expert:" in text and "|>" in text:
+            start = text.find("<|expert:") + 9
+            end = text.find("|>", start)
+            if end > start:
+                expert = text[start:end]
+        
+        return agent, expert
+
+# Example usage and testing
+if __name__ == "__main__":
+    # Configure logging
+    logging.basicConfig(level=logging.INFO)
+    
+    # Example configuration
+    config = AgentTokenConfig(
+        agent_prefix="<|agent:",
+        agent_suffix="|>",
+        resize_embeddings=True
+    )
+    
+    # Example agents
+    agents = ["SWE", "SQE", "DevOps", "Architect", "Security"]
+    
+    # Create agent manager
+    manager = AgentTokenManager(config)
+    
+    # Generate tokens
+    tokens = manager.generate_agent_tokens(agents)
+    print(f"Generated tokens: {tokens}")
+    
+    # Example tokenizer (would be loaded from actual model)
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3.5-MoE-instruct")
+    
+    # Add tokens to tokenizer
+    updated_tokenizer, agent_tokens = manager.add_agent_tokens_to_tokenizer(tokenizer, agents)
+    
+    print(f"Updated tokenizer vocab size: {len(updated_tokenizer)}")
+    print(f"Agent token IDs: {manager.token_ids}")
+    
+    # Test formatting
+    test_text = "How do I implement a binary search?"
+    formatted = manager.format_agent_prompt("SWE", test_text)
+    print(f"Formatted prompt: {formatted}")
+    
+    # Test extraction
+    extracted_agent = manager.extract_agent_from_text(formatted)
+    print(f"Extracted agent: {extracted_agent}")

src/multi_agent_training/__init__.py

@@ -0,0 +1 @@
+# Multi-agent training package

src/multi_agent_training/core_framework.py

@@ -0,0 +1,232 @@
+"""
+Core Training Framework for MangoMAS Local
+
+This module provides the foundation for specialized training modules,
+allowing for modular training of different cognitive capabilities.
+"""
+
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+import torch
+import yaml
+
+from .lora_trainer import LoRADistillationTrainer
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TrainingModuleConfig:
+    """Configuration for a specialized training module."""
+
+    name: str
+    module_type: str
+    enabled: bool = True
+    loss_weight: float = 1.0
+    learning_rate: Optional[float] = None
+    batch_size: Optional[int] = None
+    data_path: Optional[str] = None
+    module_config: Dict[str, Any] = field(default_factory=dict)
+
+
+class SpecializedTrainingModule(ABC):
+    """
+    Abstract base class for specialized training modules.
+    Each cognitive capability (reasoning, memory, etc.) should implement this interface.
+    """
+
+    def __init__(self, config: TrainingModuleConfig, tokenizer):
+        """
+        Initialize the specialized training module.
+
+        Args:
+            config: Module configuration
+            tokenizer: Tokenizer for text processing
+        """
+        self.config = config
+        self.tokenizer = tokenizer
+        self.name = config.name
+        self.enabled = config.enabled
+        self.loss_weight = config.loss_weight
+        self.device = torch.device(
+            "cuda"
+            if torch.cuda.is_available()
+            else "mps" if torch.backends.mps.is_available() else "cpu"
+        )
+
+        logger.info(f"Initialized {self.name} training module")
+        logger.info(f"Module config: {self.config}")
+
+    @abstractmethod
+    def prepare_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """
+        Prepare a batch of data for this specific training module.
+
+        Args:
+            batch: The input batch from the dataloader
+
+        Returns:
+            Processed batch ready for the module
+        """
+        pass
+
+    @abstractmethod
+    def compute_loss(
+        self, student_outputs: Any, teacher_outputs: Any, batch: Dict[str, torch.Tensor]
+    ) -> torch.Tensor:
+        """
+        Compute the specialized loss for this module.
+
+        Args:
+            student_outputs: Outputs from the student model
+            teacher_outputs: Outputs from the teacher model
+            batch: The processed input batch
+
+        Returns:
+            Loss tensor for this module
+        """
+        pass
+
+    @abstractmethod
+    def get_metrics(self) -> Dict[str, float]:
+        """
+        Get metrics specific to this training module.
+
+        Returns:
+            Dictionary of metric names and values
+        """
+        pass
+
+
+class ModularTrainingManager:
+    """
+    Training manager that orchestrates multiple specialized training modules.
+    """
+
+    def __init__(self, config_path: str):
+        """
+        Initialize the modular training manager.
+
+        Args:
+            config_path: Path to the training configuration file
+        """
+        with open(config_path, "r") as f:
+            self.config = yaml.safe_load(f)
+
+        # Set up core components
+        self.base_trainer = LoRADistillationTrainer(config_path)
+        self.tokenizer = self.base_trainer.tokenizer
+        self.student_model = self.base_trainer.student_model
+        self.teacher_model = self.base_trainer.teacher_manager.model
+
+        # Initialize modules
+        self.modules = self._initialize_modules()
+
+        logger.info(
+            f"Initialized ModularTrainingManager with {len(self.modules)} modules"
+        )
+
+    def _initialize_modules(self) -> List[SpecializedTrainingModule]:
+        """
+        Initialize all specialized training modules based on configuration.
+
+        Returns:
+            List of initialized training modules
+        """
+        modules = []
+        module_configs = self.config.get("specialized_modules", [])
+
+        for module_config in module_configs:
+            if not module_config.get("enabled", True):
+                logger.info(f"Skipping disabled module: {module_config.get('name')}")
+                continue
+
+            try:
+                # Convert to proper config object
+                config_obj = TrainingModuleConfig(**module_config)
+
+                # Import the module dynamically
+                module_type = config_obj.module_type
+                module_class = self._import_module_class(module_type)
+
+                # Initialize the module
+                module = module_class(config_obj, self.tokenizer)
+                modules.append(module)
+
+                logger.info(f"Successfully loaded module: {config_obj.name}")
+            except Exception as e:
+                logger.error(
+                    f"Failed to load module {module_config.get('name')}: {str(e)}"
+                )
+
+        return modules
+
+    def _import_module_class(self, module_type: str) -> type:
+        """
+        Dynamically import a module class based on its type.
+
+        Args:
+            module_type: The module type identifier
+
+        Returns:
+            The module class
+        """
+        if module_type == "reasoning":
+            from .specialized.reasoning_module import ReasoningTrainingModule
+
+            return ReasoningTrainingModule
+        elif module_type == "memory":
+            from .specialized.memory_module import MemoryTrainingModule
+
+            return MemoryTrainingModule
+        elif module_type == "ethics":
+            from .specialized.ethics_module import EthicsTrainingModule
+
+            return EthicsTrainingModule
+        elif module_type == "empathy":
+            from .specialized.empathy_module import EmpathyTrainingModule
+
+            return EmpathyTrainingModule
+        elif module_type == "curiosity":
+            from .specialized.curiosity_module import CuriosityTrainingModule
+
+            return CuriosityTrainingModule
+        else:
+            raise ValueError(f"Unknown module type: {module_type}")
+
+    def train(self, agent_type: str = None) -> Dict[str, Any]:
+        """
+        Train the model using all enabled specialized modules.
+
+        Args:
+            agent_type: Optional agent type for specialized training
+
+        Returns:
+            Training metrics and results
+        """
+        # Delegate to base trainer for core training functionality
+        # but integrate specialized module losses
+        logger.info(f"Starting modular training for agent: {agent_type or 'all'}")
+
+        # TODO: Implement the full training loop integrating all modules
+
+        # This is a placeholder until we implement the full integration
+        return self.base_trainer.train(agent_type)
+
+    def evaluate(self, agent_type: str = None) -> Dict[str, Any]:
+        """
+        Evaluate the model using all enabled specialized modules.
+
+        Args:
+            agent_type: Optional agent type for specialized evaluation
+
+        Returns:
+            Evaluation metrics and results
+        """
+        # TODO: Implement evaluation using specialized modules
+
+        # This is a placeholder until we implement the full integration
+        return self.base_trainer.evaluate(agent_type)

src/multi_agent_training/data_processor.py

@@ -0,0 +1,582 @@
+"""
+Data Processor for MangoMAS Local Training
+
+This module processes the original MangoMAS datasets from JSONL format
+into training-ready datasets with proper splits and preprocessing.
+"""
+
+import argparse
+import json
+import logging
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import yaml
+from sklearn.model_selection import train_test_split
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+
+class MangoMASDataProcessor:
+    """Process MangoMAS datasets for local training."""
+
+    def __init__(
+        self,
+        input_dir,
+        output_dir=None,
+        min_length: int = 10,
+        max_length: int = 2048,
+        config_path: str = None,
+    ):
+        """Initialize with input/output directories or config path for flexibility."""
+        # Support both test interface and config-driven approach
+        if config_path is not None:
+            # Config-driven initialization (original functionality)
+            with open(config_path, "r") as f:
+                self.config = yaml.safe_load(f)
+            self.data_config = self.config["data"]
+            self.agents_config = self.config["agents"]
+            self.input_dir = Path(input_dir) if input_dir else None
+            self.output_dir = Path(output_dir) if output_dir else Path("/Volumes/Mango_MAS/data/processed")
+            self.min_length = self.data_config.get("preprocessing", {}).get(
+                "min_length", min_length
+            )
+            self.max_length = self.data_config.get("preprocessing", {}).get(
+                "max_length", max_length
+            )
+        else:
+            # Direct initialization (test interface)
+            self.input_dir = Path(input_dir)
+            self.output_dir = Path(output_dir) if output_dir else Path("/Volumes/Mango_MAS/data/processed")
+            self.min_length = min_length
+            self.max_length = max_length
+            self.config = None
+            self.data_config = None
+            self.agents_config = None
+
+        logging.basicConfig(level=logging.INFO)
+
+    def process_datasets(
+        self, input_dir: str, output_dir: str = "/Volumes/Mango_MAS/data/processed"
+    ) -> None:
+        """
+        Process all agent datasets from input directory.
+
+        Args:
+            input_dir: Directory containing original JSONL files
+            output_dir: Directory to save processed datasets
+        """
+        input_path = Path(input_dir)
+        output_path = Path(output_dir)
+        output_path.mkdir(parents=True, exist_ok=True)
+
+        # Define dataset mappings
+        datasets = {
+            "infrastructure": input_path
+            / "infrastructure_agent_synthetic_prompts.jsonl",
+            "devsecops": input_path / "devsecops_agent_synthetic_prompts.jsonl",
+            "risk_assessment": input_path
+            / "risk_assessment_agent_synthetic_prompts.jsonl",
+        }
+
+        for agent_type, file_path in datasets.items():
+            if file_path.exists():
+                logger.info(f"Processing {agent_type} dataset from {file_path}")
+                self._process_single_dataset(file_path, output_path, agent_type)
+            else:
+                logger.warning(f"Dataset file not found: {file_path}")
+
+    def _process_single_dataset(
+        self, input_file: Path, output_dir: Path, agent_type: str
+    ) -> None:
+        """Process a single agent dataset."""
+        # Load data
+        data = self._load_jsonl(input_file)
+        logger.info(f"Loaded {len(data)} samples for {agent_type}")
+
+        # Clean and preprocess
+        cleaned_data = self._clean_data(data)
+        logger.info(f"After cleaning: {len(cleaned_data)} samples")
+
+        # Convert to training format
+        training_data = self._convert_to_training_format(cleaned_data, agent_type)
+
+        # Create splits
+        train_data, val_data, test_data = self._create_splits(training_data)
+
+        # Save processed datasets
+        self._save_datasets(train_data, val_data, test_data, output_dir, agent_type)
+
+        logger.info(
+            f"Saved {agent_type} dataset: "
+            f"{len(train_data)} train, {len(val_data)} val, {len(test_data)} test"
+        )
+
+    def _load_jsonl(self, file_path: Path) -> List[Dict]:
+        """Load data from JSONL file."""
+        data = []
+        with open(file_path, "r", encoding="utf-8") as f:
+            for line_num, line in enumerate(f, 1):
+                try:
+                    data.append(json.loads(line.strip()))
+                except json.JSONDecodeError as e:
+                    logger.warning(f"Skipping invalid JSON on line {line_num}: {e}")
+        return data
+
+    def _clean_data(self, data: List[Dict]) -> List[Dict]:
+        """Clean and validate the data."""
+        cleaned = []
+
+        for item in tqdm(data, desc="Cleaning data"):
+            # Check required fields
+            if not all(key in item for key in ["instruction", "input", "output"]):
+                continue
+
+            # Check text lengths
+            input_text = f"{item['instruction']} {item['input']}"
+            output_text = item["output"]
+
+            if (
+                len(input_text) < self.data_config["preprocessing"]["min_length"]
+                or len(input_text) > self.data_config["preprocessing"]["max_length"]
+            ):
+                continue
+
+            if (
+                len(output_text) < self.data_config["preprocessing"]["min_length"]
+                or len(output_text) > self.data_config["preprocessing"]["max_length"]
+            ):
+                continue
+
+            cleaned.append(item)
+
+        # Remove duplicates if configured
+        if self.data_config["preprocessing"]["remove_duplicates"]:
+            cleaned = self._remove_duplicates(cleaned)
+
+        return cleaned
+
+    def _remove_duplicates(self, data: List[Dict]) -> List[Dict]:
+        """Remove duplicate entries based on input text."""
+        seen_inputs = set()
+        unique_data = []
+
+        for item in data:
+            input_text = f"{item['instruction']} {item['input']}"
+            if input_text not in seen_inputs:
+                seen_inputs.add(input_text)
+                unique_data.append(item)
+
+        logger.info(f"Removed {len(data) - len(unique_data)} duplicates")
+        return unique_data
+
+    def _validate_sample(self, sample: Dict) -> bool:
+        """Validate a single sample for required fields and length constraints."""
+        # Check required fields
+        required_fields = ["instruction", "input", "output", "agent_type"]
+        if not all(key in sample for key in required_fields):
+            return False
+
+        # Check text lengths
+        combined_text = f"{sample['instruction']} {sample['input']} {sample['output']}"
+        if len(combined_text) < self.min_length or len(combined_text) > self.max_length:
+            return False
+
+        return True
+
+    def _clean_text(self, text: str) -> str:
+        """Clean text by normalizing whitespace and removing extra spaces."""
+        import re
+
+        # Remove extra whitespace and normalize
+        cleaned = re.sub(r"\s+", " ", text.strip())
+        return cleaned
+
+    def _format_conversation(self, sample: Dict) -> Dict:
+        """Format sample into conversation format suitable for training."""
+        # Create conversation text
+        if sample.get("input", "").strip():
+            conversation_text = f"Human: {sample['instruction']}\n{sample['input']}\n\nAssistant: {sample['output']}"
+        else:
+            conversation_text = (
+                f"Human: {sample['instruction']}\n\nAssistant: {sample['output']}"
+            )
+
+        return {
+            "text": conversation_text,
+            "agent_type": sample["agent_type"],
+            "instruction": sample["instruction"],
+            "input": sample["input"],
+            "output": sample["output"],
+        }
+
+    def _split_dataset(
+        self,
+        data: List[Dict],
+        train_ratio: float = 0.8,
+        val_ratio: float = 0.1,
+        test_ratio: float = 0.1,
+    ) -> Tuple[List[Dict], List[Dict], List[Dict]]:
+        """Split dataset into train/validation/test sets."""
+        if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-6:
+            raise ValueError(
+                f"Split ratios must sum to 1.0, got {train_ratio + val_ratio + test_ratio}"
+            )
+
+        if not data:
+            return [], [], []
+
+        # Use sklearn for consistent splitting
+        from sklearn.model_selection import train_test_split
+
+        # First split: train vs (val + test)
+        if len(data) == 1:
+            return data, [], []
+
+        train_data, temp_data = train_test_split(
+            data, test_size=(val_ratio + test_ratio), random_state=42, shuffle=True
+        )
+
+        # Second split: val vs test
+        if temp_data and val_ratio > 0 and test_ratio > 0:
+            val_ratio_normalized = val_ratio / (val_ratio + test_ratio)
+            val_data, test_data = train_test_split(
+                temp_data,
+                test_size=(1 - val_ratio_normalized),
+                random_state=42,
+                shuffle=True,
+            )
+        elif val_ratio > 0:
+            val_data, test_data = temp_data, []
+        else:
+            val_data, test_data = [], temp_data
+
+        return train_data, val_data, test_data
+
+    def _calculate_stats(self, data: List[Dict]) -> Dict:
+        """Calculate statistics for the dataset."""
+        if not data:
+            return {
+                "total_samples": 0,
+                "avg_length": 0,
+                "min_length": 0,
+                "max_length": 0,
+                "agent_distribution": {},
+            }
+
+        lengths = [len(item.get("text", "")) for item in data]
+        agent_counts = {}
+
+        for item in data:
+            agent = item.get("agent_type", "unknown")
+            agent_counts[agent] = agent_counts.get(agent, 0) + 1
+
+        return {
+            "total_samples": len(data),
+            "avg_length": sum(lengths) / len(lengths),
+            "min_length": min(lengths),
+            "max_length": max(lengths),
+            "agent_distribution": agent_counts,
+        }
+
+    def _load_agent_data(self, agent_type: str) -> List[Dict]:
+        """Load data for a specific agent type."""
+        if not self.input_dir:
+            return []
+
+        # Look for files matching the agent type. We intentionally call glob even
+        # if the directory may not exist in test environments, since tests patch
+        # pathlib.Path.glob.
+        pattern = f"*{agent_type}*.jsonl"
+        matching_files = list(self.input_dir.glob(pattern))
+
+        data = []
+        for file_path in matching_files:
+            file_data = self._load_jsonl(file_path)
+            data.extend(file_data)
+
+        return data
+
+    def _save_jsonl(self, data: List[Dict], output_path: Path) -> None:
+        """Save data to JSONL file."""
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(output_path, "w", encoding="utf-8") as f:
+            for item in data:
+                f.write(json.dumps(item, ensure_ascii=False) + "\n")
+
+    def _save_stats(self, stats: Dict, output_path: Path) -> None:
+        """Save statistics to JSON file."""
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(output_path, "w", encoding="utf-8") as f:
+            json.dump(stats, f, indent=2, ensure_ascii=False)
+
+    def process_agent(
+        self,
+        agent_type: str,
+        train_ratio: float = 0.8,
+        val_ratio: float = 0.1,
+        test_ratio: float = 0.1,
+    ) -> None:
+        """Process data for a single agent type."""
+        # Load data
+        data = self._load_agent_data(agent_type)
+
+        if not data:
+            raise ValueError(f"No valid data found for agent type: {agent_type}")
+
+        # Validate and clean data
+        valid_data = []
+        for sample in data:
+            if self._validate_sample(sample):
+                formatted = self._format_conversation(sample)
+                valid_data.append(formatted)
+
+        if not valid_data:
+            raise ValueError(
+                f"No valid data found after processing for agent type: {agent_type}"
+            )
+
+        # Remove duplicates
+        unique_data = self._remove_duplicates(valid_data)
+
+        # Split dataset
+        train_data, val_data, test_data = self._split_dataset(
+            unique_data, train_ratio, val_ratio, test_ratio
+        )
+
+        # Save datasets
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self._save_jsonl(train_data, self.output_dir / f"{agent_type}_train.jsonl")
+        self._save_jsonl(val_data, self.output_dir / f"{agent_type}_val.jsonl")
+        self._save_jsonl(test_data, self.output_dir / f"{agent_type}_test.jsonl")
+
+        # Save statistics
+        stats = self._calculate_stats(unique_data)
+        self._save_stats(stats, self.output_dir / f"{agent_type}_stats.json")
+
+        logger.info(
+            f"Processed {agent_type}: {len(train_data)} train, {len(val_data)} val, {len(test_data)} test samples"
+        )
+
+    def _convert_to_training_format(
+        self, data: List[Dict], agent_type: str
+    ) -> List[Dict]:
+        """Convert to format suitable for training."""
+        training_data = []
+
+        for item in data:
+            # Create conversation format suitable for language modeling
+            conversation = {
+                "messages": [
+                    {
+                        "role": "system",
+                        "content": f"You are a {agent_type.replace('_', ' ')} specialist. "
+                        f"Provide expert recommendations and analysis.",
+                    },
+                    {
+                        "role": "user",
+                        "content": f"{item['instruction']}\n\n{item['input']}",
+                    },
+                    {"role": "assistant", "content": item["output"]},
+                ],
+                "metadata": item.get("metadata", {}),
+                "agent_type": agent_type,
+            }
+            training_data.append(conversation)
+
+        return training_data
+
+    def _create_splits(
+        self, data: List[Dict]
+    ) -> Tuple[List[Dict], List[Dict], List[Dict]]:
+        """Create train/validation/test splits."""
+        train_size = self.data_config["train_split"]
+        val_size = self.data_config["validation_split"]
+        test_size = self.data_config["test_split"]
+
+        # Normalize splits to sum to 1
+        total = train_size + val_size + test_size
+        train_size /= total
+        val_size /= total
+        test_size /= total
+
+        # First split: train vs (val + test)
+        train_data, temp_data = train_test_split(
+            data, test_size=(val_size + test_size), random_state=42, shuffle=True
+        )
+
+        # Second split: val vs test
+        val_ratio = val_size / (val_size + test_size)
+        val_data, test_data = train_test_split(
+            temp_data, test_size=(1 - val_ratio), random_state=42, shuffle=True
+        )
+
+        return train_data, val_data, test_data
+
+    def _save_datasets(
+        self,
+        train_data: List[Dict],
+        val_data: List[Dict],
+        test_data: List[Dict],
+        output_dir: Path,
+        agent_type: str,
+    ) -> None:
+        """Save processed datasets to files."""
+        datasets = {"train": train_data, "validation": val_data, "test": test_data}
+
+        for split_name, split_data in datasets.items():
+            output_file = output_dir / f"{agent_type}_{split_name}.jsonl"
+
+            with open(output_file, "w", encoding="utf-8") as f:
+                for item in split_data:
+                    f.write(json.dumps(item, ensure_ascii=False) + "\n")
+
+            logger.info(f"Saved {len(split_data)} samples to {output_file}")
+
+    def create_combined_dataset(self, output_dir: str = "/Volumes/Mango_MAS/data/processed") -> None:
+        """Create combined dataset with all agent types for multi-task training."""
+        output_path = Path(output_dir)
+
+        # Collect all processed data
+        all_train_data = []
+        all_val_data = []
+        all_test_data = []
+
+        for agent_type in self.agents_config.keys():
+            for split in ["train", "validation", "test"]:
+                file_path = output_path / f"{agent_type}_{split}.jsonl"
+                if file_path.exists():
+                    data = self._load_jsonl(file_path)
+
+                    if split == "train":
+                        all_train_data.extend(data)
+                    elif split == "validation":
+                        all_val_data.extend(data)
+                    else:
+                        all_test_data.extend(data)
+
+        # Shuffle combined datasets
+        import random
+
+        random.seed(42)
+        random.shuffle(all_train_data)
+        random.shuffle(all_val_data)
+        random.shuffle(all_test_data)
+
+        # Save combined datasets
+        combined_datasets = {
+            "train": all_train_data,
+            "validation": all_val_data,
+            "test": all_test_data,
+        }
+
+        for split_name, split_data in combined_datasets.items():
+            output_file = output_path / f"combined_{split_name}.jsonl"
+
+            with open(output_file, "w", encoding="utf-8") as f:
+                for item in split_data:
+                    f.write(json.dumps(item, ensure_ascii=False) + "\n")
+
+            logger.info(
+                f"Saved combined {split_name} dataset: {len(split_data)} samples"
+            )
+
+    def generate_statistics(self, output_dir: str = "/Volumes/Mango_MAS/data/processed") -> Dict:
+        """Generate statistics about the processed datasets."""
+        output_path = Path(output_dir)
+        stats = {}
+
+        for agent_type in list(self.agents_config.keys()) + ["combined"]:
+            agent_stats = {}
+
+            for split in ["train", "validation", "test"]:
+                file_path = output_path / f"{agent_type}_{split}.jsonl"
+
+                if file_path.exists():
+                    data = self._load_jsonl(file_path)
+
+                    # Calculate statistics
+                    lengths = []
+                    for item in data:
+                        if "messages" in item:
+                            # Calculate total text length
+                            total_length = sum(
+                                len(msg["content"]) for msg in item["messages"]
+                            )
+                            lengths.append(total_length)
+
+                    agent_stats[split] = {
+                        "count": len(data),
+                        "avg_length": sum(lengths) / len(lengths) if lengths else 0,
+                        "min_length": min(lengths) if lengths else 0,
+                        "max_length": max(lengths) if lengths else 0,
+                    }
+
+            stats[agent_type] = agent_stats
+
+        # Save statistics
+        stats_file = output_path / "dataset_statistics.json"
+        with open(stats_file, "w") as f:
+            json.dump(stats, f, indent=2)
+
+        logger.info(f"Generated dataset statistics: {stats_file}")
+        return stats
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Process MangoMAS datasets for local training"
+    )
+    parser.add_argument(
+        "--input_dir",
+        type=str,
+        default="/Users/iancruickshank/Documents/Model/mangomas-datasets/agents/",
+        help="Directory containing original JSONL files",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="/Volumes/Mango_MAS/data/processed",
+        help="Directory to save processed datasets",
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        default="config/training/distillation.yaml",
+        help="Path to configuration file",
+    )
+    parser.add_argument(
+        "--create_combined",
+        action="store_true",
+        help="Create combined multi-agent dataset",
+    )
+
+    args = parser.parse_args()
+
+    # Initialize processor
+    processor = MangoMASDataProcessor(args.config)
+
+    # Process datasets
+    processor.process_datasets(args.input_dir, args.output_dir)
+
+    # Create combined dataset if requested
+    if args.create_combined:
+        processor.create_combined_dataset(args.output_dir)
+
+    # Generate statistics
+    stats = processor.generate_statistics(args.output_dir)
+
+    print("\nDataset Statistics:")
+    print("=" * 50)
+    for agent_type, agent_stats in stats.items():
+        print(f"\n{agent_type.upper()}:")
+        for split, split_stats in agent_stats.items():
+            print(
+                f"  {split}: {split_stats['count']} samples, "
+                f"avg length: {split_stats['avg_length']:.0f} chars"
+            )
+
+
+if __name__ == "__main__":
+    main()

src/multi_agent_training/distillation_loss.py

@@ -0,0 +1,232 @@
+"""
+Knowledge Distillation Loss Implementation for MangoMAS Local
+
+This module implements custom loss functions for knowledge distillation,
+balancing task-specific loss with knowledge transfer from teacher models.
+"""
+
+import logging
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+class DistillationLoss:
+    """
+    Custom loss function for knowledge distillation combining task loss
+    and distillation loss with configurable temperature and alpha parameters.
+    """
+
+    def __init__(self, alpha: float = 0.5, temperature: float = 3.0):
+        """
+        Initialize distillation loss with parameters.
+
+        Args:
+            alpha: Balance between task loss and distillation loss (0.0-1.0)
+            temperature: Temperature for softmax in knowledge distillation
+        """
+        self.alpha = alpha
+        self.temperature = temperature
+
+        if not 0.0 <= alpha <= 1.0:
+            raise ValueError(f"Alpha must be between 0 and 1, got {alpha}")
+        if temperature <= 0.0:
+            raise ValueError(f"Temperature must be positive, got {temperature}")
+
+        # Add required loss attributes for tests
+        import torch.nn as nn
+
+        self.task_loss = nn.CrossEntropyLoss(ignore_index=-100)
+        self.kl_loss = nn.KLDivLoss(reduction="batchmean")
+
+        logger.info(
+            f"Initialized DistillationLoss with alpha={alpha}, temperature={temperature}"
+        )
+
+    def compute_loss(
+        self,
+        student_logits: torch.Tensor,
+        teacher_logits: torch.Tensor,
+        labels: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the combined distillation loss.
+
+        Args:
+            student_logits: Logits from student model [batch_size, seq_len, vocab_size]
+            teacher_logits: Logits from teacher model [batch_size, seq_len, vocab_size]
+            labels: Target labels [batch_size, seq_len]
+            attention_mask: Attention mask for padding tokens [batch_size, seq_len]
+
+        Returns:
+            Tuple of (total_loss, loss_dict) where loss_dict contains individual losses
+        """
+        # Task-specific loss (standard cross-entropy)
+        task_loss = self._compute_task_loss(student_logits, labels, attention_mask)
+
+        # Knowledge distillation loss
+        distill_loss = self._compute_distillation_loss(
+            student_logits, teacher_logits, attention_mask
+        )
+
+        # Combined loss
+        # NOTE: alpha is treated as the weight for the distillation loss to match
+        # unit-test expectations (alpha=1.0 => pure distillation loss).
+        total_loss = self.alpha * distill_loss + (1.0 - self.alpha) * task_loss
+
+        loss_dict = {
+            "total_loss": total_loss.item(),
+            "task_loss": task_loss.item(),
+            "distillation_loss": distill_loss.item(),
+            "alpha": self.alpha,
+            "temperature": self.temperature,
+        }
+
+        return total_loss, loss_dict
+
+    def _compute_task_loss(
+        self,
+        logits: torch.Tensor,
+        labels: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Compute standard cross-entropy task loss."""
+        # Reshape for cross entropy: [batch_size * seq_len, vocab_size]
+        shift_logits = logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+
+        # Flatten the tokens
+        flat_logits = shift_logits.view(-1, shift_logits.size(-1))
+        flat_labels = shift_labels.view(-1)
+
+        # Apply attention mask if provided
+        if attention_mask is not None:
+            shift_mask = attention_mask[..., 1:].contiguous()
+            flat_mask = shift_mask.view(-1)
+            # Only compute loss for non-padded tokens
+            valid_indices = flat_mask.bool()
+            if valid_indices.sum() == 0:
+                return torch.tensor(0.0, device=logits.device)
+            flat_logits = flat_logits[valid_indices]
+            flat_labels = flat_labels[valid_indices]
+
+        # Compute cross entropy loss, ignoring padding tokens
+        task_loss = F.cross_entropy(flat_logits, flat_labels, ignore_index=-100)
+
+        return task_loss
+
+    def _compute_distillation_loss(
+        self,
+        student_logits: torch.Tensor,
+        teacher_logits: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        eps = 1e-8
+        # Apply temperature scaling
+        student_log_probs = F.log_softmax(student_logits / self.temperature, dim=-1)
+        teacher_log_probs = F.log_softmax(teacher_logits / self.temperature, dim=-1)
+        # Optionally mask out padding tokens
+        if attention_mask is not None:
+            mask = attention_mask[..., 1:].contiguous().view(-1)
+            student_log_probs = (
+                student_log_probs[..., :-1, :]
+                .contiguous()
+                .view(-1, student_log_probs.size(-1))[mask.bool()]
+            )
+            teacher_log_probs = (
+                teacher_log_probs[..., :-1, :]
+                .contiguous()
+                .view(-1, teacher_log_probs.size(-1))[mask.bool()]
+            )
+        else:
+            student_log_probs = (
+                student_log_probs[..., :-1, :]
+                .contiguous()
+                .view(-1, student_log_probs.size(-1))
+            )
+            teacher_log_probs = (
+                teacher_log_probs[..., :-1, :]
+                .contiguous()
+                .view(-1, teacher_log_probs.size(-1))
+            )
+        # KLDivLoss expects log-probabilities for both input and target if log_target=True
+        if student_log_probs.shape[0] == 0:
+            return torch.tensor(0.0, device=student_logits.device)
+        distill_loss = F.kl_div(
+            student_log_probs, teacher_log_probs, reduction="batchmean", log_target=True
+        ) * (self.temperature**2)
+        return distill_loss
+
+    def update_alpha(self, new_alpha: float) -> None:
+        """Update the alpha parameter during training."""
+        if not 0.0 <= new_alpha <= 1.0:
+            raise ValueError(f"Alpha must be between 0.0 and 1.0, got {new_alpha}")
+        self.alpha = new_alpha
+        logger.info(f"Updated alpha to {new_alpha}")
+
+    def update_temperature(self, new_temperature: float) -> None:
+        """Update the temperature parameter during training."""
+        if new_temperature <= 0.0:
+            raise ValueError(f"Temperature must be positive, got {new_temperature}")
+        self.temperature = new_temperature
+        logger.info(f"Updated temperature to {new_temperature}")
+
+    def __call__(self, student_logits, teacher_logits, labels, attention_mask=None):
+        total_loss, _ = self.compute_loss(
+            student_logits, teacher_logits, labels, attention_mask
+        )
+        return total_loss
+
+
+class AdaptiveDistillationLoss(DistillationLoss):
+    """
+    Adaptive distillation loss that adjusts alpha based on training progress.
+    Starts with more focus on distillation, gradually shifting to task loss.
+    """
+
+    def __init__(
+        self,
+        initial_alpha: float = 0.5,
+        final_alpha: float = 0.1,
+        temperature: float = 3.0,
+        warmup_steps: int = 1000,
+    ):
+        if not 0.0 <= initial_alpha <= 1.0:
+            raise ValueError("initial_alpha must be between 0 and 1")
+        if not 0.0 <= final_alpha <= 1.0:
+            raise ValueError("final_alpha must be between 0 and 1")
+        if initial_alpha < final_alpha:
+            raise ValueError(
+                f"initial_alpha must be >= final_alpha, got {initial_alpha} < {final_alpha}"
+            )
+        super().__init__(alpha=initial_alpha, temperature=temperature)
+        self.initial_alpha = initial_alpha
+        self.final_alpha = final_alpha
+        self.current_alpha = initial_alpha
+        self.warmup_steps = warmup_steps
+        self.current_step = 0
+
+    def update_alpha(self, current_epoch: int, total_epochs: int):
+        """Update alpha based on current epoch and total epochs."""
+        if total_epochs <= 0:
+            raise ValueError("total_epochs must be positive")
+        if current_epoch < 0:
+            raise ValueError("current_epoch must be non-negative")
+        if current_epoch >= total_epochs:
+            self.current_alpha = self.final_alpha
+        elif current_epoch <= 0:
+            self.current_alpha = self.initial_alpha
+        else:
+            progress = current_epoch / total_epochs
+            self.current_alpha = (
+                self.initial_alpha - (self.initial_alpha - self.final_alpha) * progress
+            )
+        self.alpha = self.current_alpha
+
+    def get_alpha(self):
+        """Return the current alpha value."""
+        return self.current_alpha

src/multi_agent_training/gradient_descent/__init__.py

@@ -0,0 +1,49 @@
+"""
+Gradient Descent and Backpropagation Training System
+==================================================
+
+This module provides a comprehensive implementation of gradient descent optimization
+algorithms and backpropagation for neural network training, specifically designed
+for the MangoMAS multi-agent system.
+
+Key Components:
+- Optimizers: SGD, Adam, AdamW with proper mathematical implementations
+- Backpropagation: Chain rule-based gradient computation
+- Training Loop: Complete training orchestration with monitoring
+- Loss Functions: Various loss implementations for different tasks
+- Monitoring: Comprehensive gradient and training metrics tracking
+
+Usage:
+    from src.training.gradient_descent import GradientDescentTrainer
+    trainer = GradientDescentTrainer()
+    results = trainer.train_agent(agent_spec)
+"""
+
+from .optimizers import SGD, Adam, AdamW, Optimizer
+from .backpropagation import BackpropagationEngine
+from .training_loop import GradientDescentTrainer
+from .loss_functions import CrossEntropyLoss, KLDivergenceLoss, LossFunction
+from .monitoring import GradientMonitor, TrainingMonitor
+from .model_wrapper import ModelWrapper
+from .schedulers import LinearScheduler, CosineScheduler, StepScheduler
+
+__version__ = "1.0.0"
+__author__ = "MangoMAS Team"
+
+__all__ = [
+    "SGD",
+    "Adam", 
+    "AdamW",
+    "Optimizer",
+    "BackpropagationEngine",
+    "GradientDescentTrainer",
+    "CrossEntropyLoss",
+    "KLDivergenceLoss", 
+    "LossFunction",
+    "GradientMonitor",
+    "TrainingMonitor",
+    "ModelWrapper",
+    "LinearScheduler",
+    "CosineScheduler",
+    "StepScheduler"
+]

src/multi_agent_training/gradient_descent/backpropagation.py

@@ -0,0 +1,337 @@
+"""
+Backpropagation Engine
+=====================
+
+This module implements the backpropagation algorithm for neural network training,
+including support for transformer architectures and LoRA adapters.
+
+The implementation follows the mathematical chain rule for gradient computation
+and includes proper handling of various activation functions and layer types.
+"""
+
+import logging
+import math
+from typing import Dict, List, Any
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+
+class BackpropagationEngine:
+    """
+    Backpropagation engine for neural network training
+    
+    Implements the chain rule for gradient computation:
+    ∂L/∂w = ∂L/∂y * ∂y/∂w
+    
+    Supports various architectures including transformers and LoRA adapters.
+    """
+    
+    def __init__(self, model: nn.Module, device: torch.device = None):
+        self.model = model
+        self.device = device or torch.device('cpu')
+        self.gradient_norms = []
+        self.gradient_stats = {}
+        
+        # Move model to device
+        self.model.to(self.device)
+        
+        logger.info(f"Initialized backpropagation engine on device: {self.device}")
+    
+    def compute_gradients(self, loss: torch.Tensor, retain_graph: bool = False) -> Dict[str, torch.Tensor]:
+        """
+        Compute gradients using backpropagation
+        
+        Args:
+            loss: The computed loss tensor
+            retain_graph: Whether to retain the computational graph
+            
+        Returns:
+            Dictionary of parameter gradients
+        """
+        # Zero existing gradients
+        self.model.zero_grad()
+        
+        # Backward pass
+        loss.backward(retain_graph=retain_graph)
+        
+        # Collect gradients
+        gradients = {}
+        for name, param in self.model.named_parameters():
+            if param.grad is not None:
+                gradients[name] = param.grad.clone()
+                
+        # Compute gradient statistics
+        self._compute_gradient_stats(gradients)
+        
+        return gradients
+    
+    def _compute_gradient_stats(self, gradients: Dict[str, torch.Tensor]):
+        """Compute gradient statistics for monitoring"""
+        total_norm = 0.0
+        param_count = 0
+        
+        for name, grad in gradients.items():
+            if grad is not None:
+                param_norm = grad.data.norm(2)
+                total_norm += param_norm.item() ** 2
+                param_count += 1
+                
+                # Store individual gradient stats
+                self.gradient_stats[name] = {
+                    'norm': param_norm.item(),
+                    'mean': grad.data.mean().item(),
+                    'std': grad.data.std().item(),
+                    'max': grad.data.max().item(),
+                    'min': grad.data.min().item()
+                }
+        
+        # Compute total gradient norm
+        total_norm = total_norm ** (1. / 2)
+        self.gradient_norms.append(total_norm)
+        
+        logger.debug(f"Total gradient norm: {total_norm:.6f}")
+    
+    def apply_gradient_clipping(self, max_norm: float = 1.0, norm_type: float = 2.0) -> float:
+        """
+        Apply gradient clipping to prevent exploding gradients
+        
+        Args:
+            max_norm: Maximum gradient norm
+            norm_type: Type of norm to use (2.0 for L2 norm)
+            
+        Returns:
+            Actual gradient norm before clipping
+        """
+        # Compute total gradient norm
+        total_norm = 0.0
+        for param in self.model.parameters():
+            if param.grad is not None:
+                param_norm = param.grad.data.norm(norm_type)
+                total_norm += param_norm.item() ** norm_type
+        
+        total_norm = total_norm ** (1. / norm_type)
+        
+        # Apply clipping if necessary
+        if total_norm > max_norm:
+            clip_coef = max_norm / (total_norm + 1e-6)
+            for param in self.model.parameters():
+                if param.grad is not None:
+                    param.grad.data.mul_(clip_coef)
+            
+            logger.debug(f"Gradient clipping applied: {total_norm:.6f} -> {max_norm:.6f}")
+        
+        return total_norm
+    
+    def compute_parameter_updates(self, gradients: Dict[str, torch.Tensor], 
+                                 learning_rate: float) -> Dict[str, torch.Tensor]:
+        """
+        Compute parameter updates from gradients
+        
+        Args:
+            gradients: Dictionary of parameter gradients
+            learning_rate: Learning rate for updates
+            
+        Returns:
+            Dictionary of parameter updates
+        """
+        updates = {}
+        
+        for name, param in self.model.named_parameters():
+            if name in gradients and gradients[name] is not None:
+                # Simple gradient descent update
+                update = -learning_rate * gradients[name]
+                updates[name] = update
+                
+        return updates
+    
+    def apply_parameter_updates(self, updates: Dict[str, torch.Tensor]):
+        """
+        Apply parameter updates to the model
+        
+        Args:
+            updates: Dictionary of parameter updates
+        """
+        for name, param in self.model.named_parameters():
+            if name in updates:
+                param.data.add_(updates[name])
+    
+    def compute_hessian_diagonal(self, loss: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """
+        Compute diagonal elements of the Hessian matrix for second-order optimization
+        
+        Args:
+            loss: The computed loss tensor
+            
+        Returns:
+            Dictionary of Hessian diagonal elements
+        """
+        hessian_diag = {}
+        
+        # First backward pass
+        self.model.zero_grad()
+        loss.backward(retain_graph=True)
+        
+        # Compute second-order gradients
+        for name, param in self.model.named_parameters():
+            if param.grad is not None:
+                # Compute second derivatives
+                grad_grad = torch.autograd.grad(
+                    param.grad, param, retain_graph=True, create_graph=True
+                )[0]
+                hessian_diag[name] = grad_grad.diag()
+        
+        return hessian_diag
+    
+    def compute_gradient_flow(self) -> Dict[str, float]:
+        """
+        Compute gradient flow statistics for monitoring training health
+        
+        Returns:
+            Dictionary of gradient flow statistics
+        """
+        flow_stats = {}
+        
+        for name, param in self.model.named_parameters():
+            if param.grad is not None:
+                # Compute gradient-to-parameter ratio
+                grad_norm = param.grad.data.norm(2)
+                param_norm = param.data.norm(2)
+                
+                if param_norm > 0:
+                    flow_ratio = grad_norm / param_norm
+                    flow_stats[name] = flow_ratio.item()
+        
+        return flow_stats
+    
+    def detect_gradient_anomalies(self) -> Dict[str, Any]:
+        """
+        Detect gradient anomalies that might indicate training issues
+        
+        Returns:
+            Dictionary of detected anomalies
+        """
+        anomalies = {
+            'exploding_gradients': False,
+            'vanishing_gradients': False,
+            'gradient_imbalance': False,
+            'nan_gradients': False
+        }
+        
+        if not self.gradient_norms:
+            return anomalies
+        
+        current_norm = self.gradient_norms[-1]
+        
+        # Check for exploding gradients
+        if current_norm > 10.0:
+            anomalies['exploding_gradients'] = True
+            logger.warning(f"Exploding gradients detected: norm={current_norm:.6f}")
+        
+        # Check for vanishing gradients
+        if current_norm < 1e-6:
+            anomalies['vanishing_gradients'] = True
+            logger.warning(f"Vanishing gradients detected: norm={current_norm:.6f}")
+        
+        # Check for NaN gradients
+        for name, stats in self.gradient_stats.items():
+            if math.isnan(stats['norm']) or math.isnan(stats['mean']):
+                anomalies['nan_gradients'] = True
+                logger.warning(f"NaN gradients detected in parameter: {name}")
+                break
+        
+        # Check for gradient imbalance
+        if len(self.gradient_stats) > 1:
+            norms = [stats['norm'] for stats in self.gradient_stats.values()]
+            if max(norms) / min(norms) > 1000:
+                anomalies['gradient_imbalance'] = True
+                logger.warning("Gradient imbalance detected between parameters")
+        
+        return anomalies
+    
+    def get_gradient_summary(self) -> Dict[str, Any]:
+        """
+        Get comprehensive gradient summary for monitoring
+        
+        Returns:
+            Dictionary of gradient summary statistics
+        """
+        if not self.gradient_norms:
+            return {}
+        
+        return {
+            'total_norm': self.gradient_norms[-1],
+            'norm_history': self.gradient_norms[-10:],  # Last 10 steps
+            'parameter_stats': self.gradient_stats,
+            'anomalies': self.detect_gradient_anomalies(),
+            'flow_stats': self.compute_gradient_flow()
+        }
+
+
+class LoRABackpropagationEngine(BackpropagationEngine):
+    """
+    Specialized backpropagation engine for LoRA (Low-Rank Adaptation) models
+    
+    Only computes gradients for LoRA adapter parameters, not the base model.
+    """
+    
+    def __init__(self, model: nn.Module, lora_params: List[str], device: torch.device = None):
+        super().__init__(model, device)
+        self.lora_params = lora_params
+        
+        # Freeze base model parameters
+        for name, param in self.model.named_parameters():
+            if name not in self.lora_params:
+                param.requires_grad = False
+        
+        logger.info(f"Initialized LoRA backpropagation engine with {len(lora_params)} LoRA parameters")
+    
+    def compute_gradients(self, loss: torch.Tensor, retain_graph: bool = False) -> Dict[str, torch.Tensor]:
+        """
+        Compute gradients only for LoRA parameters
+        
+        Args:
+            loss: The computed loss tensor
+            retain_graph: Whether to retain the computational graph
+            
+        Returns:
+            Dictionary of LoRA parameter gradients
+        """
+        # Zero existing gradients
+        self.model.zero_grad()
+        
+        # Backward pass
+        loss.backward(retain_graph=retain_graph)
+        
+        # Collect only LoRA gradients
+        gradients = {}
+        for name, param in self.model.named_parameters():
+            if name in self.lora_params and param.grad is not None:
+                gradients[name] = param.grad.clone()
+                
+        # Compute gradient statistics
+        self._compute_gradient_stats(gradients)
+        
+        return gradients
+    
+    def get_lora_parameter_info(self) -> Dict[str, Any]:
+        """
+        Get information about LoRA parameters
+        
+        Returns:
+            Dictionary of LoRA parameter information
+        """
+        info = {}
+        
+        for name, param in self.model.named_parameters():
+            if name in self.lora_params:
+                info[name] = {
+                    'shape': list(param.shape),
+                    'numel': param.numel(),
+                    'requires_grad': param.requires_grad,
+                    'device': str(param.device),
+                    'dtype': str(param.dtype)
+                }
+        
+        return info

src/multi_agent_training/gradient_descent/loss_functions.py

@@ -0,0 +1,379 @@
+"""
+Loss Functions
+=============
+
+This module implements various loss functions for neural network training,
+including cross-entropy, KL divergence, and custom loss functions for
+the MangoMAS multi-agent system.
+"""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Dict, Optional, Any
+import torch
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+class LossFunction(ABC):
+    """Abstract base class for all loss functions"""
+    
+    def __init__(self, reduction: str = 'mean'):
+        self.reduction = reduction
+        
+    @abstractmethod
+    def forward(self, predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        """Compute the loss"""
+        pass
+    
+    def __call__(self, predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        """Call the loss function"""
+        return self.forward(predictions, targets)
+
+
+class CrossEntropyLoss(LossFunction):
+    """
+    Cross-entropy loss for classification tasks
+    
+    Mathematical formulation:
+    L = -∑(y_i * log(ŷ_i))
+    
+    Where y_i is the true label and ŷ_i is the predicted probability.
+    """
+    
+    def __init__(self, reduction: str = 'mean', label_smoothing: float = 0.0, 
+                 weight: Optional[torch.Tensor] = None):
+        super().__init__(reduction)
+        self.label_smoothing = label_smoothing
+        self.weight = weight
+        
+    def forward(self, predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        """
+        Compute cross-entropy loss
+        
+        Args:
+            predictions: Model predictions (logits)
+            targets: True labels
+            
+        Returns:
+            Computed loss
+        """
+        # Apply softmax to get probabilities
+        if predictions.dim() > 1 and predictions.size(1) > 1:
+            # Multi-class classification
+            log_probs = F.log_softmax(predictions, dim=1)
+            
+            if targets.dim() == 1:
+                # Targets are class indices
+                loss = F.nll_loss(log_probs, targets, weight=self.weight, 
+                                reduction=self.reduction, label_smoothing=self.label_smoothing)
+            else:
+                # Targets are one-hot encoded
+                loss = -(targets * log_probs).sum(dim=1)
+                if self.reduction == 'mean':
+                    loss = loss.mean()
+                elif self.reduction == 'sum':
+                    loss = loss.sum()
+        else:
+            # Binary classification
+            loss = F.binary_cross_entropy_with_logits(predictions, targets.float(), 
+                                                    weight=self.weight, reduction=self.reduction)
+        
+        return loss
+
+
+class KLDivergenceLoss(LossFunction):
+    """
+    Kullback-Leibler divergence loss for distribution matching
+    
+    Mathematical formulation:
+    KL(P||Q) = ∑ P(x) * log(P(x)/Q(x))
+    """
+    
+    def __init__(self, reduction: str = 'mean', log_target: bool = False):
+        super().__init__(reduction)
+        self.log_target = log_target
+        
+    def forward(self, predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        """
+        Compute KL divergence loss
+        
+        Args:
+            predictions: Predicted distribution (log probabilities)
+            targets: Target distribution (probabilities or log probabilities)
+            
+        Returns:
+            Computed KL divergence loss
+        """
+        if self.log_target:
+            # Both predictions and targets are in log space
+            loss = F.kl_div(predictions, targets, reduction=self.reduction, log_target=True)
+        else:
+            # Convert predictions to log space, targets are probabilities
+            log_predictions = F.log_softmax(predictions, dim=-1)
+            loss = F.kl_div(log_predictions, targets, reduction=self.reduction, log_target=False)
+        
+        return loss
+
+
+class MSELoss(LossFunction):
+    """
+    Mean Squared Error loss for regression tasks
+    
+    Mathematical formulation:
+    L = (1/n) * ∑(y_i - ŷ_i)²
+    """
+    
+    def forward(self, predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        """
+        Compute MSE loss
+        
+        Args:
+            predictions: Model predictions
+            targets: True values
+            
+        Returns:
+            Computed MSE loss
+        """
+        loss = F.mse_loss(predictions, targets, reduction=self.reduction)
+        return loss
+
+
+class HuberLoss(LossFunction):
+    """
+    Huber loss (smooth L1 loss) for robust regression
+    
+    Mathematical formulation:
+    L = { 0.5 * (y - ŷ)², if |y - ŷ| < δ
+        { δ * (|y - ŷ| - 0.5 * δ), otherwise
+    """
+    
+    def __init__(self, reduction: str = 'mean', delta: float = 1.0):
+        super().__init__(reduction)
+        self.delta = delta
+        
+    def forward(self, predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        """
+        Compute Huber loss
+        
+        Args:
+            predictions: Model predictions
+            targets: True values
+            
+        Returns:
+            Computed Huber loss
+        """
+        loss = F.huber_loss(predictions, targets, reduction=self.reduction, delta=self.delta)
+        return loss
+
+
+class FocalLoss(LossFunction):
+    """
+    Focal loss for addressing class imbalance
+    
+    Mathematical formulation:
+    FL = -α(1-p_t)^γ * log(p_t)
+    
+    Where p_t is the predicted probability for the true class.
+    """
+    
+    def __init__(self, alpha: float = 1.0, gamma: float = 2.0, reduction: str = 'mean'):
+        super().__init__(reduction)
+        self.alpha = alpha
+        self.gamma = gamma
+        
+    def forward(self, predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        """
+        Compute focal loss
+        
+        Args:
+            predictions: Model predictions (logits)
+            targets: True labels
+            
+        Returns:
+            Computed focal loss
+        """
+        # Compute probabilities
+        probs = F.softmax(predictions, dim=1)
+        
+        # Get probabilities for true classes
+        if targets.dim() == 1:
+            # Targets are class indices
+            pt = probs.gather(1, targets.unsqueeze(1)).squeeze(1)
+        else:
+            # Targets are one-hot encoded
+            pt = (probs * targets).sum(dim=1)
+        
+        # Compute focal loss
+        focal_weight = self.alpha * (1 - pt) ** self.gamma
+        loss = -focal_weight * torch.log(pt + 1e-8)
+        
+        if self.reduction == 'mean':
+            loss = loss.mean()
+        elif self.reduction == 'sum':
+            loss = loss.sum()
+        
+        return loss
+
+
+class ContrastiveLoss(LossFunction):
+    """
+    Contrastive loss for learning representations
+    
+    Mathematical formulation:
+    L = (1-y) * d² + y * max(0, margin - d)²
+    
+    Where d is the distance between embeddings and y is the similarity label.
+    """
+    
+    def __init__(self, margin: float = 1.0, reduction: str = 'mean'):
+        super().__init__(reduction)
+        self.margin = margin
+        
+    def forward(self, embeddings1: torch.Tensor, embeddings2: torch.Tensor, 
+                labels: torch.Tensor) -> torch.Tensor:
+        """
+        Compute contrastive loss
+        
+        Args:
+            embeddings1: First set of embeddings
+            embeddings2: Second set of embeddings
+            labels: Similarity labels (1 for similar, 0 for dissimilar)
+            
+        Returns:
+            Computed contrastive loss
+        """
+        # Compute Euclidean distance
+        distance = F.pairwise_distance(embeddings1, embeddings2)
+        
+        # Compute contrastive loss
+        positive_loss = labels.float() * distance.pow(2)
+        negative_loss = (1 - labels.float()) * F.relu(self.margin - distance).pow(2)
+        
+        loss = positive_loss + negative_loss
+        
+        if self.reduction == 'mean':
+            loss = loss.mean()
+        elif self.reduction == 'sum':
+            loss = loss.sum()
+        
+        return loss
+
+
+class CustomAgentLoss(LossFunction):
+    """
+    Custom loss function for MangoMAS agents
+    
+    Combines multiple loss components to optimize agent performance.
+    """
+    
+    def __init__(self, task_loss_weight: float = 1.0, 
+                 consistency_loss_weight: float = 0.1,
+                 regularization_weight: float = 0.01):
+        super().__init__('mean')
+        self.task_loss_weight = task_loss_weight
+        self.consistency_loss_weight = consistency_loss_weight
+        self.regularization_weight = regularization_weight
+        
+        # Initialize component losses
+        self.task_loss = CrossEntropyLoss()
+        self.consistency_loss = MSELoss()
+        
+    def forward(self, predictions: torch.Tensor, targets: torch.Tensor,
+                model_outputs: Optional[Dict[str, torch.Tensor]] = None) -> torch.Tensor:
+        """
+        Compute custom agent loss
+        
+        Args:
+            predictions: Model predictions
+            targets: True labels
+            model_outputs: Additional model outputs for consistency loss
+            
+        Returns:
+            Computed custom loss
+        """
+        # Task-specific loss
+        task_loss = self.task_loss(predictions, targets)
+        
+        # Consistency loss (if model outputs provided)
+        consistency_loss = torch.tensor(0.0, device=predictions.device)
+        if model_outputs is not None and 'hidden_states' in model_outputs:
+            # Encourage consistent hidden representations
+            hidden_states = model_outputs['hidden_states']
+            if len(hidden_states) > 1:
+                # Compute consistency between consecutive hidden states
+                for i in range(len(hidden_states) - 1):
+                    consistency_loss += self.consistency_loss(hidden_states[i], hidden_states[i+1])
+                consistency_loss /= (len(hidden_states) - 1)
+        
+        # Regularization loss (L2 penalty)
+        regularization_loss = torch.tensor(0.0, device=predictions.device)
+        # This would be computed from model parameters in practice
+        
+        # Combine losses
+        total_loss = (self.task_loss_weight * task_loss + 
+                     self.consistency_loss_weight * consistency_loss + 
+                     self.regularization_weight * regularization_loss)
+        
+        return total_loss
+
+
+class LossFunctionFactory:
+    """Factory class for creating loss functions"""
+    
+    @staticmethod
+    def create_loss_function(loss_type: str, **kwargs) -> LossFunction:
+        """Create a loss function instance"""
+        loss_functions = {
+            'cross_entropy': CrossEntropyLoss,
+            'kl_divergence': KLDivergenceLoss,
+            'mse': MSELoss,
+            'huber': HuberLoss,
+            'focal': FocalLoss,
+            'contrastive': ContrastiveLoss,
+            'custom_agent': CustomAgentLoss
+        }
+        
+        if loss_type.lower() not in loss_functions:
+            raise ValueError(f"Unknown loss function type: {loss_type}")
+            
+        loss_class = loss_functions[loss_type.lower()]
+        return loss_class(**kwargs)
+    
+    @staticmethod
+    def get_default_config(loss_type: str) -> Dict[str, Any]:
+        """Get default configuration for loss function"""
+        configs = {
+            'cross_entropy': {
+                'reduction': 'mean',
+                'label_smoothing': 0.0
+            },
+            'kl_divergence': {
+                'reduction': 'mean',
+                'log_target': False
+            },
+            'mse': {
+                'reduction': 'mean'
+            },
+            'huber': {
+                'reduction': 'mean',
+                'delta': 1.0
+            },
+            'focal': {
+                'alpha': 1.0,
+                'gamma': 2.0,
+                'reduction': 'mean'
+            },
+            'contrastive': {
+                'margin': 1.0,
+                'reduction': 'mean'
+            },
+            'custom_agent': {
+                'task_loss_weight': 1.0,
+                'consistency_loss_weight': 0.1,
+                'regularization_weight': 0.01
+            }
+        }
+        
+        return configs.get(loss_type.lower(), {})

src/multi_agent_training/gradient_descent/model_wrapper.py

@@ -0,0 +1,342 @@
+"""
+Model Wrapper
+=============
+
+This module provides a wrapper for neural network models to integrate
+with the gradient descent training system, including support for LoRA
+adapters and the MangoMAS agent system.
+"""
+
+import logging
+import torch
+import torch.nn as nn
+from typing import Dict, List, Optional, Any
+from pathlib import Path
+import json
+
+logger = logging.getLogger(__name__)
+
+
+class ModelWrapper:
+    """
+    Wrapper for neural network models to integrate with gradient descent training
+    
+    Provides a unified interface for different model types and handles
+    LoRA adapter integration for the MangoMAS system.
+    """
+    
+    def __init__(self, model: nn.Module, model_type: str = 'transformer', 
+                 lora_config: Optional[Dict[str, Any]] = None):
+        self.model = model
+        self.model_type = model_type
+        self.lora_config = lora_config or {}
+        self.lora_params = []
+        
+        # Initialize LoRA if configured
+        if lora_config:
+            self._setup_lora()
+        
+        logger.info(f"Initialized ModelWrapper for {model_type} model")
+    
+    def _setup_lora(self):
+        """Setup LoRA adapters for the model"""
+        if not self.lora_config:
+            return
+        
+        # Extract LoRA parameters
+        for name, param in self.model.named_parameters():
+            if 'lora' in name.lower() or 'adapter' in name.lower():
+                self.lora_params.append(name)
+                param.requires_grad = True
+            else:
+                param.requires_grad = False
+        
+        logger.info(f"Setup LoRA with {len(self.lora_params)} adapter parameters")
+    
+    def forward(self, inputs: torch.Tensor, **kwargs) -> torch.Tensor:
+        """
+        Forward pass through the model
+        
+        Args:
+            inputs: Input tensor
+            **kwargs: Additional arguments
+            
+        Returns:
+            Model output tensor
+        """
+        return self.model(inputs, **kwargs)
+    
+    def get_trainable_parameters(self) -> List[torch.Tensor]:
+        """
+        Get list of trainable parameters
+        
+        Returns:
+            List of trainable parameter tensors
+        """
+        if self.lora_params:
+            # Return only LoRA parameters
+            return [param for name, param in self.model.named_parameters() 
+                   if name in self.lora_params and param.requires_grad]
+        else:
+            # Return all trainable parameters
+            return [param for param in self.model.parameters() if param.requires_grad]
+    
+    def get_parameter_info(self) -> Dict[str, Any]:
+        """
+        Get information about model parameters
+        
+        Returns:
+            Dictionary of parameter information
+        """
+        info = {
+            'total_parameters': sum(p.numel() for p in self.model.parameters()),
+            'trainable_parameters': sum(p.numel() for p in self.get_trainable_parameters()),
+            'lora_parameters': len(self.lora_params),
+            'parameter_details': {}
+        }
+        
+        for name, param in self.model.named_parameters():
+            info['parameter_details'][name] = {
+                'shape': list(param.shape),
+                'numel': param.numel(),
+                'requires_grad': param.requires_grad,
+                'is_lora': name in self.lora_params
+            }
+        
+        return info
+    
+    def save_model(self, save_path: str, metadata: Optional[Dict[str, Any]] = None):
+        """
+        Save the model and metadata
+        
+        Args:
+            save_path: Path to save the model
+            metadata: Additional metadata to save
+        """
+        save_path = Path(save_path)
+        save_path.mkdir(parents=True, exist_ok=True)
+        
+        # Save model state
+        model_path = save_path / 'model.pt'
+        torch.save(self.model.state_dict(), model_path)
+        
+        # Save metadata
+        if metadata is None:
+            metadata = {}
+        
+        metadata.update({
+            'model_type': self.model_type,
+            'lora_config': self.lora_config,
+            'lora_params': self.lora_params,
+            'parameter_info': self.get_parameter_info()
+        })
+        
+        metadata_path = save_path / 'metadata.json'
+        with open(metadata_path, 'w') as f:
+            json.dump(metadata, f, indent=2)
+        
+        logger.info(f"Model saved to {save_path}")
+    
+    def load_model(self, load_path: str):
+        """
+        Load the model from saved files
+        
+        Args:
+            load_path: Path to load the model from
+        """
+        load_path = Path(load_path)
+        
+        # Load model state
+        model_path = load_path / 'model.pt'
+        if model_path.exists():
+            state_dict = torch.load(model_path, map_location='cpu')
+            self.model.load_state_dict(state_dict)
+            logger.info(f"Model loaded from {model_path}")
+        
+        # Load metadata
+        metadata_path = load_path / 'metadata.json'
+        if metadata_path.exists():
+            with open(metadata_path, 'r') as f:
+                metadata = json.load(f)
+            
+            self.model_type = metadata.get('model_type', self.model_type)
+            self.lora_config = metadata.get('lora_config', self.lora_config)
+            self.lora_params = metadata.get('lora_params', self.lora_params)
+            
+            logger.info(f"Metadata loaded from {metadata_path}")
+    
+    def to(self, device: torch.device):
+        """Move model to device"""
+        self.model.to(device)
+        return self
+    
+    def train(self):
+        """Set model to training mode"""
+        self.model.train()
+        return self
+    
+    def eval(self):
+        """Set model to evaluation mode"""
+        self.model.eval()
+        return self
+    
+    def __call__(self, *args, **kwargs):
+        """Call the model"""
+        return self.forward(*args, **kwargs)
+
+
+class LoRAModelWrapper(ModelWrapper):
+    """
+    Specialized wrapper for LoRA (Low-Rank Adaptation) models
+    
+    Provides enhanced functionality for LoRA adapter management
+    and integration with the MangoMAS system.
+    """
+    
+    def __init__(self, base_model: nn.Module, lora_config: Dict[str, Any]):
+        super().__init__(base_model, 'lora_transformer', lora_config)
+        self.base_model = base_model
+        self.adapters = {}
+        
+        # Initialize LoRA adapters
+        self._initialize_lora_adapters()
+    
+    def _initialize_lora_adapters(self):
+        """Initialize LoRA adapters based on configuration"""
+        rank = self.lora_config.get('rank', 16)
+        alpha = self.lora_config.get('alpha', 32)
+        dropout = self.lora_config.get('dropout', 0.1)
+        target_modules = self.lora_config.get('target_modules', ['c_attn', 'c_proj'])
+        
+        # Add LoRA adapters to target modules
+        for name, module in self.base_model.named_modules():
+            if any(target in name for target in target_modules):
+                if isinstance(module, (nn.Linear, nn.Conv2d)):
+                    # Add LoRA adapter
+                    adapter = LoRAAdapter(module, rank, alpha, dropout)
+                    self.adapters[name] = adapter
+                    
+                    # Replace original module
+                    self._replace_module(name, adapter)
+        
+        logger.info(f"Initialized {len(self.adapters)} LoRA adapters")
+    
+    def _replace_module(self, module_name: str, new_module: nn.Module):
+        """Replace a module in the model"""
+        parts = module_name.split('.')
+        parent = self.base_model
+        
+        for part in parts[:-1]:
+            parent = getattr(parent, part)
+        
+        setattr(parent, parts[-1], new_module)
+    
+    def get_lora_parameters(self) -> List[torch.Tensor]:
+        """Get LoRA adapter parameters"""
+        lora_params = []
+        for adapter in self.adapters.values():
+            lora_params.extend(adapter.parameters())
+        return lora_params
+    
+    def merge_adapters(self):
+        """Merge LoRA adapters into base model"""
+        for adapter in self.adapters.values():
+            adapter.merge()
+        logger.info("LoRA adapters merged into base model")
+    
+    def unmerge_adapters(self):
+        """Unmerge LoRA adapters from base model"""
+        for adapter in self.adapters.values():
+            adapter.unmerge()
+        logger.info("LoRA adapters unmerged from base model")
+
+
+class LoRAAdapter(nn.Module):
+    """
+    LoRA (Low-Rank Adaptation) adapter module
+    
+    Implements the LoRA technique for efficient fine-tuning of large models.
+    """
+    
+    def __init__(self, original_module: nn.Module, rank: int = 16, 
+                 alpha: float = 32, dropout: float = 0.1):
+        super().__init__()
+        self.original_module = original_module
+        self.rank = rank
+        self.alpha = alpha
+        self.dropout = dropout
+        
+        # Get original module dimensions
+        if isinstance(original_module, nn.Linear):
+            in_features = original_module.in_features
+            out_features = original_module.out_features
+        elif isinstance(original_module, nn.Conv2d):
+            in_features = original_module.in_channels
+            out_features = original_module.out_channels
+        else:
+            raise ValueError(f"Unsupported module type: {type(original_module)}")
+        
+        # Initialize LoRA matrices
+        self.lora_A = nn.Linear(in_features, rank, bias=False)
+        self.lora_B = nn.Linear(rank, out_features, bias=False)
+        self.dropout_layer = nn.Dropout(dropout)
+        
+        # Initialize weights
+        nn.init.kaiming_uniform_(self.lora_A.weight)
+        nn.init.zeros_(self.lora_B.weight)
+        
+        # Store original weights
+        self.original_weight = original_module.weight.data.clone()
+        self.merged = False
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through LoRA adapter"""
+        if self.merged:
+            # Use merged weights
+            return F.linear(x, self.original_weight, self.original_module.bias)
+        else:
+            # Use LoRA adaptation
+            lora_output = self.lora_B(self.dropout_layer(self.lora_A(x)))
+            original_output = F.linear(x, self.original_weight, self.original_module.bias)
+            return original_output + (self.alpha / self.rank) * lora_output
+    
+    def merge(self):
+        """Merge LoRA weights into original weights"""
+        if not self.merged:
+            lora_weight = (self.alpha / self.rank) * torch.mm(
+                self.lora_B.weight, self.lora_A.weight
+            )
+            self.original_weight += lora_weight
+            self.merged = True
+    
+    def unmerge(self):
+        """Unmerge LoRA weights from original weights"""
+        if self.merged:
+            lora_weight = (self.alpha / self.rank) * torch.mm(
+                self.lora_B.weight, self.lora_A.weight
+            )
+            self.original_weight -= lora_weight
+            self.merged = False
+
+
+class ModelFactory:
+    """Factory class for creating model wrappers"""
+    
+    @staticmethod
+    def create_model_wrapper(model_type: str, model: nn.Module, 
+                           **kwargs) -> ModelWrapper:
+        """Create a model wrapper instance"""
+        if model_type.lower() == 'lora':
+            return LoRAModelWrapper(model, kwargs.get('lora_config', {}))
+        else:
+            return ModelWrapper(model, model_type, kwargs.get('lora_config'))
+    
+    @staticmethod
+    def get_default_lora_config() -> Dict[str, Any]:
+        """Get default LoRA configuration"""
+        return {
+            'rank': 16,
+            'alpha': 32,
+            'dropout': 0.1,
+            'target_modules': ['c_attn', 'c_proj']
+        }

src/multi_agent_training/gradient_descent/monitoring.py

@@ -0,0 +1,548 @@
+"""
+Training Monitoring
+==================
+
+This module provides comprehensive monitoring capabilities for gradient descent
+training, including gradient tracking, loss monitoring, and performance metrics.
+"""
+
+import logging
+import math
+import time
+from typing import Dict, Optional, Any
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from collections import defaultdict, deque
+import json
+
+logger = logging.getLogger(__name__)
+
+
+class GradientMonitor:
+    """
+    Monitor gradient statistics during training
+    
+    Tracks gradient norms, distributions, and anomalies to ensure
+    stable training and detect potential issues.
+    """
+    
+    def __init__(self, max_history: int = 1000):
+        self.max_history = max_history
+        self.gradient_norms = deque(maxlen=max_history)
+        self.gradient_means = deque(maxlen=max_history)
+        self.gradient_stds = deque(maxlen=max_history)
+        self.gradient_maxs = deque(maxlen=max_history)
+        self.gradient_mins = deque(maxlen=max_history)
+        
+        self.parameter_stats = defaultdict(lambda: {
+            'norms': deque(maxlen=max_history),
+            'means': deque(maxlen=max_history),
+            'stds': deque(maxlen=max_history)
+        })
+        
+        self.anomaly_count = 0
+        self.last_anomaly_time = None
+        
+        logger.info("Initialized GradientMonitor")
+    
+    def update(self, gradients: Dict[str, torch.Tensor]):
+        """
+        Update gradient statistics
+        
+        Args:
+            gradients: Dictionary of parameter gradients
+        """
+        total_norm = 0.0
+        total_mean = 0.0
+        total_std = 0.0
+        total_max = float('-inf')
+        total_min = float('inf')
+        param_count = 0
+        
+        for name, grad in gradients.items():
+            if grad is not None:
+                # Compute statistics for this parameter
+                param_norm = grad.data.norm(2).item()
+                param_mean = grad.data.mean().item()
+                param_std = grad.data.std().item()
+                param_max = grad.data.max().item()
+                param_min = grad.data.min().item()
+                
+                # Update parameter-specific stats
+                self.parameter_stats[name]['norms'].append(param_norm)
+                self.parameter_stats[name]['means'].append(param_mean)
+                self.parameter_stats[name]['stds'].append(param_std)
+                
+                # Accumulate global stats
+                total_norm += param_norm ** 2
+                total_mean += param_mean
+                total_std += param_std ** 2
+                total_max = max(total_max, param_max)
+                total_min = min(total_min, param_min)
+                param_count += 1
+        
+        # Compute global statistics
+        if param_count > 0:
+            total_norm = math.sqrt(total_norm)
+            total_mean /= param_count
+            total_std = math.sqrt(total_std / param_count)
+            
+            # Store global stats
+            self.gradient_norms.append(total_norm)
+            self.gradient_means.append(total_mean)
+            self.gradient_stds.append(total_std)
+            self.gradient_maxs.append(total_max)
+            self.gradient_mins.append(total_min)
+    
+    def detect_anomalies(self) -> Dict[str, Any]:
+        """
+        Detect gradient anomalies
+        
+        Returns:
+            Dictionary of detected anomalies
+        """
+        anomalies = {
+            'exploding_gradients': False,
+            'vanishing_gradients': False,
+            'gradient_imbalance': False,
+            'nan_gradients': False,
+            'gradient_spikes': False
+        }
+        
+        if len(self.gradient_norms) < 2:
+            return anomalies
+        
+        current_norm = self.gradient_norms[-1]
+        
+        # Check for exploding gradients
+        if current_norm > 10.0:
+            anomalies['exploding_gradients'] = True
+            self.anomaly_count += 1
+            self.last_anomaly_time = time.time()
+            logger.warning(f"Exploding gradients detected: norm={current_norm:.6f}")
+        
+        # Check for vanishing gradients
+        if current_norm < 1e-6:
+            anomalies['vanishing_gradients'] = True
+            self.anomaly_count += 1
+            self.last_anomaly_time = time.time()
+            logger.warning(f"Vanishing gradients detected: norm={current_norm:.6f}")
+        
+        # Check for gradient spikes
+        if len(self.gradient_norms) >= 10:
+            recent_norms = list(self.gradient_norms)[-10:]
+            avg_norm = np.mean(recent_norms[:-1])
+            if current_norm > 3 * avg_norm:
+                anomalies['gradient_spikes'] = True
+                logger.warning(f"Gradient spike detected: {current_norm:.6f} vs avg {avg_norm:.6f}")
+        
+        # Check for NaN gradients
+        if math.isnan(current_norm) or math.isnan(self.gradient_means[-1]):
+            anomalies['nan_gradients'] = True
+            self.anomaly_count += 1
+            self.last_anomaly_time = time.time()
+            logger.warning("NaN gradients detected")
+        
+        # Check for gradient imbalance between parameters
+        if len(self.parameter_stats) > 1:
+            param_norms = [stats['norms'][-1] for stats in self.parameter_stats.values() 
+                          if len(stats['norms']) > 0]
+            if param_norms and max(param_norms) / min(param_norms) > 1000:
+                anomalies['gradient_imbalance'] = True
+                logger.warning("Gradient imbalance detected between parameters")
+        
+        return anomalies
+    
+    def get_statistics(self) -> Dict[str, Any]:
+        """
+        Get comprehensive gradient statistics
+        
+        Returns:
+            Dictionary of gradient statistics
+        """
+        if not self.gradient_norms:
+            return {}
+        
+        stats = {
+            'current_norm': self.gradient_norms[-1],
+            'mean_norm': np.mean(self.gradient_norms),
+            'std_norm': np.std(self.gradient_norms),
+            'min_norm': min(self.gradient_norms),
+            'max_norm': max(self.gradient_norms),
+            'current_mean': self.gradient_means[-1],
+            'current_std': self.gradient_stds[-1],
+            'current_max': self.gradient_maxs[-1],
+            'current_min': self.gradient_mins[-1],
+            'anomaly_count': self.anomaly_count,
+            'parameter_count': len(self.parameter_stats)
+        }
+        
+        # Add parameter-specific statistics
+        param_stats = {}
+        for name, stats_dict in self.parameter_stats.items():
+            if stats_dict['norms']:
+                param_stats[name] = {
+                    'current_norm': stats_dict['norms'][-1],
+                    'mean_norm': np.mean(stats_dict['norms']),
+                    'std_norm': np.std(stats_dict['norms'])
+                }
+        
+        stats['parameter_stats'] = param_stats
+        
+        return stats
+    
+    def plot_gradients(self, save_path: Optional[str] = None):
+        """
+        Plot gradient statistics
+        
+        Args:
+            save_path: Path to save the plot
+        """
+        if not self.gradient_norms:
+            logger.warning("No gradient data to plot")
+            return
+        
+        fig, axes = plt.subplots(2, 2, figsize=(12, 8))
+        
+        # Plot gradient norms
+        axes[0, 0].plot(self.gradient_norms)
+        axes[0, 0].set_title('Gradient Norms')
+        axes[0, 0].set_xlabel('Step')
+        axes[0, 0].set_ylabel('Norm')
+        axes[0, 0].grid(True)
+        
+        # Plot gradient means
+        axes[0, 1].plot(self.gradient_means)
+        axes[0, 1].set_title('Gradient Means')
+        axes[0, 1].set_xlabel('Step')
+        axes[0, 1].set_ylabel('Mean')
+        axes[0, 1].grid(True)
+        
+        # Plot gradient stds
+        axes[1, 0].plot(self.gradient_stds)
+        axes[1, 0].set_title('Gradient Standard Deviations')
+        axes[1, 0].set_xlabel('Step')
+        axes[1, 0].set_ylabel('Std')
+        axes[1, 0].grid(True)
+        
+        # Plot gradient range
+        axes[1, 1].plot(self.gradient_maxs, label='Max')
+        axes[1, 1].plot(self.gradient_mins, label='Min')
+        axes[1, 1].set_title('Gradient Range')
+        axes[1, 1].set_xlabel('Step')
+        axes[1, 1].set_ylabel('Value')
+        axes[1, 1].legend()
+        axes[1, 1].grid(True)
+        
+        plt.tight_layout()
+        
+        if save_path:
+            plt.savefig(save_path, dpi=300, bbox_inches='tight')
+            logger.info(f"Gradient plots saved to {save_path}")
+        
+        plt.show()
+
+
+class TrainingMonitor:
+    """
+    Monitor training progress and performance
+    
+    Tracks loss, accuracy, learning rates, and other training metrics
+    to provide comprehensive training insights.
+    """
+    
+    def __init__(self, max_history: int = 1000):
+        self.max_history = max_history
+        self.losses = deque(maxlen=max_history)
+        self.accuracies = deque(maxlen=max_history)
+        self.learning_rates = deque(maxlen=max_history)
+        self.training_times = deque(maxlen=max_history)
+        
+        self.epoch_metrics = defaultdict(lambda: {
+            'loss': deque(maxlen=max_history),
+            'accuracy': deque(maxlen=max_history),
+            'learning_rate': deque(maxlen=max_history)
+        })
+        
+        self.best_loss = float('inf')
+        self.best_accuracy = 0.0
+        self.training_start_time = time.time()
+        
+        logger.info("Initialized TrainingMonitor")
+    
+    def update(self, loss: float, accuracy: Optional[float] = None, 
+               learning_rate: Optional[float] = None, epoch: Optional[int] = None):
+        """
+        Update training metrics
+        
+        Args:
+            loss: Current loss value
+            accuracy: Current accuracy (optional)
+            learning_rate: Current learning rate (optional)
+            epoch: Current epoch (optional)
+        """
+        current_time = time.time()
+        
+        # Update global metrics
+        self.losses.append(loss)
+        if accuracy is not None:
+            self.accuracies.append(accuracy)
+        if learning_rate is not None:
+            self.learning_rates.append(learning_rate)
+        self.training_times.append(current_time - self.training_start_time)
+        
+        # Update epoch-specific metrics
+        if epoch is not None:
+            self.epoch_metrics[epoch]['loss'].append(loss)
+            if accuracy is not None:
+                self.epoch_metrics[epoch]['accuracy'].append(accuracy)
+            if learning_rate is not None:
+                self.epoch_metrics[epoch]['learning_rate'].append(learning_rate)
+        
+        # Update best metrics
+        if loss < self.best_loss:
+            self.best_loss = loss
+        if accuracy is not None and accuracy > self.best_accuracy:
+            self.best_accuracy = accuracy
+    
+    def get_statistics(self) -> Dict[str, Any]:
+        """
+        Get comprehensive training statistics
+        
+        Returns:
+            Dictionary of training statistics
+        """
+        if not self.losses:
+            return {}
+        
+        stats = {
+            'current_loss': self.losses[-1],
+            'best_loss': self.best_loss,
+            'mean_loss': np.mean(self.losses),
+            'std_loss': np.std(self.losses),
+            'min_loss': min(self.losses),
+            'max_loss': max(self.losses),
+            'best_accuracy': self.best_accuracy,
+            'total_steps': len(self.losses),
+            'training_time': self.training_times[-1] if self.training_times else 0
+        }
+        
+        if self.accuracies:
+            stats.update({
+                'current_accuracy': self.accuracies[-1],
+                'mean_accuracy': np.mean(self.accuracies),
+                'std_accuracy': np.std(self.accuracies)
+            })
+        
+        if self.learning_rates:
+            stats.update({
+                'current_learning_rate': self.learning_rates[-1],
+                'mean_learning_rate': np.mean(self.learning_rates),
+                'min_learning_rate': min(self.learning_rates),
+                'max_learning_rate': max(self.learning_rates)
+            })
+        
+        return stats
+    
+    def detect_convergence(self, patience: int = 10, threshold: float = 1e-4) -> bool:
+        """
+        Detect if training has converged
+        
+        Args:
+            patience: Number of steps to wait for improvement
+            threshold: Minimum improvement threshold
+            
+        Returns:
+            True if training has converged
+        """
+        if len(self.losses) < patience:
+            return False
+        
+        recent_losses = list(self.losses)[-patience:]
+        best_recent = min(recent_losses)
+        improvement = self.best_loss - best_recent
+        
+        return improvement < threshold
+    
+    def plot_training_curves(self, save_path: Optional[str] = None):
+        """
+        Plot training curves
+        
+        Args:
+            save_path: Path to save the plot
+        """
+        if not self.losses:
+            logger.warning("No training data to plot")
+            return
+        
+        fig, axes = plt.subplots(2, 2, figsize=(12, 8))
+        
+        # Plot loss curve
+        axes[0, 0].plot(self.losses)
+        axes[0, 0].set_title('Training Loss')
+        axes[0, 0].set_xlabel('Step')
+        axes[0, 0].set_ylabel('Loss')
+        axes[0, 0].grid(True)
+        
+        # Plot accuracy curve
+        if self.accuracies:
+            axes[0, 1].plot(self.accuracies)
+            axes[0, 1].set_title('Training Accuracy')
+            axes[0, 1].set_xlabel('Step')
+            axes[0, 1].set_ylabel('Accuracy')
+            axes[0, 1].grid(True)
+        
+        # Plot learning rate curve
+        if self.learning_rates:
+            axes[1, 0].plot(self.learning_rates)
+            axes[1, 0].set_title('Learning Rate')
+            axes[1, 0].set_xlabel('Step')
+            axes[1, 0].set_ylabel('Learning Rate')
+            axes[1, 0].grid(True)
+        
+        # Plot training time
+        if self.training_times:
+            axes[1, 1].plot(self.training_times)
+            axes[1, 1].set_title('Training Time')
+            axes[1, 1].set_xlabel('Step')
+            axes[1, 1].set_ylabel('Time (seconds)')
+            axes[1, 1].grid(True)
+        
+        plt.tight_layout()
+        
+        if save_path:
+            plt.savefig(save_path, dpi=300, bbox_inches='tight')
+            logger.info(f"Training curves saved to {save_path}")
+        
+        plt.show()
+    
+    def save_metrics(self, file_path: str):
+        """
+        Save training metrics to file
+        
+        Args:
+            file_path: Path to save the metrics
+        """
+        metrics = {
+            'losses': list(self.losses),
+            'accuracies': list(self.accuracies),
+            'learning_rates': list(self.learning_rates),
+            'training_times': list(self.training_times),
+            'best_loss': self.best_loss,
+            'best_accuracy': self.best_accuracy,
+            'statistics': self.get_statistics()
+        }
+        
+        with open(file_path, 'w') as f:
+            json.dump(metrics, f, indent=2)
+        
+        logger.info(f"Training metrics saved to {file_path}")
+    
+    def load_metrics(self, file_path: str):
+        """
+        Load training metrics from file
+        
+        Args:
+            file_path: Path to load the metrics from
+        """
+        with open(file_path, 'r') as f:
+            metrics = json.load(f)
+        
+        self.losses = deque(metrics['losses'], maxlen=self.max_history)
+        self.accuracies = deque(metrics['accuracies'], maxlen=self.max_history)
+        self.learning_rates = deque(metrics['learning_rates'], maxlen=self.max_history)
+        self.training_times = deque(metrics['training_times'], maxlen=self.max_history)
+        self.best_loss = metrics['best_loss']
+        self.best_accuracy = metrics['best_accuracy']
+        
+        logger.info(f"Training metrics loaded from {file_path}")
+
+
+class PerformanceMonitor:
+    """
+    Monitor system performance during training
+    
+    Tracks memory usage, compute time, and other system metrics
+    to optimize training efficiency.
+    """
+    
+    def __init__(self):
+        self.memory_usage = deque(maxlen=1000)
+        self.compute_times = deque(maxlen=1000)
+        self.gpu_usage = deque(maxlen=1000)
+        
+        self.step_times = []
+        self.forward_times = []
+        self.backward_times = []
+        self.optimizer_times = []
+        
+        logger.info("Initialized PerformanceMonitor")
+    
+    def update_memory(self, memory_mb: float):
+        """Update memory usage"""
+        self.memory_usage.append(memory_mb)
+    
+    def update_compute_time(self, time_seconds: float):
+        """Update compute time"""
+        self.compute_times.append(time_seconds)
+    
+    def update_gpu_usage(self, gpu_percent: float):
+        """Update GPU usage"""
+        self.gpu_usage.append(gpu_percent)
+    
+    def time_step(self, step_name: str):
+        """Context manager for timing steps"""
+        return StepTimer(self, step_name)
+    
+    def get_statistics(self) -> Dict[str, Any]:
+        """Get performance statistics"""
+        stats = {}
+        
+        if self.memory_usage:
+            stats['memory'] = {
+                'current_mb': self.memory_usage[-1],
+                'mean_mb': np.mean(self.memory_usage),
+                'max_mb': max(self.memory_usage)
+            }
+        
+        if self.compute_times:
+            stats['compute'] = {
+                'current_seconds': self.compute_times[-1],
+                'mean_seconds': np.mean(self.compute_times),
+                'total_seconds': sum(self.compute_times)
+            }
+        
+        if self.gpu_usage:
+            stats['gpu'] = {
+                'current_percent': self.gpu_usage[-1],
+                'mean_percent': np.mean(self.gpu_usage),
+                'max_percent': max(self.gpu_usage)
+            }
+        
+        return stats
+
+
+class StepTimer:
+    """Context manager for timing training steps"""
+    
+    def __init__(self, monitor: PerformanceMonitor, step_name: str):
+        self.monitor = monitor
+        self.step_name = step_name
+        self.start_time = None
+    
+    def __enter__(self):
+        self.start_time = time.time()
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        elapsed = time.time() - self.start_time
+        
+        if self.step_name == 'forward':
+            self.monitor.forward_times.append(elapsed)
+        elif self.step_name == 'backward':
+            self.monitor.backward_times.append(elapsed)
+        elif self.step_name == 'optimizer':
+            self.monitor.optimizer_times.append(elapsed)
+        else:
+            self.monitor.step_times.append(elapsed)

src/multi_agent_training/gradient_descent/optimizers.py

@@ -0,0 +1,319 @@
+"""
+Gradient Descent Optimizers
+==========================
+
+This module implements various gradient descent optimization algorithms
+with proper mathematical formulations and PyTorch compatibility.
+
+Algorithms implemented:
+- SGD (Stochastic Gradient Descent)
+- Adam (Adaptive Moment Estimation)
+- AdamW (Adam with Decoupled Weight Decay)
+
+All optimizers follow the mathematical principles and include proper
+momentum, learning rate scheduling, and weight decay handling.
+"""
+
+import math
+import logging
+from abc import ABC, abstractmethod
+from typing import Dict, List, Tuple, Any
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+class Optimizer(ABC):
+    """Abstract base class for all optimizers"""
+    
+    def __init__(self, params: List[torch.Tensor], lr: float = 1e-3, **kwargs):
+        self.params = params
+        self.lr = lr
+        self.step_count = 0
+        self.state = {}
+        
+    @abstractmethod
+    def step(self, closure=None):
+        """Perform a single optimization step"""
+        pass
+        
+    def zero_grad(self):
+        """Zero the gradients of all parameters"""
+        for param in self.params:
+            if param.grad is not None:
+                param.grad.zero_()
+                
+    def state_dict(self):
+        """Return the state of the optimizer"""
+        return {
+            'state': self.state,
+            'param_groups': [{'params': self.params, 'lr': self.lr}],
+            'step_count': self.step_count
+        }
+        
+    def load_state_dict(self, state_dict):
+        """Load the state of the optimizer"""
+        self.state = state_dict['state']
+        self.lr = state_dict['param_groups'][0]['lr']
+        self.step_count = state_dict['step_count']
+
+
+class SGD(Optimizer):
+    """
+    Stochastic Gradient Descent optimizer
+    
+    Mathematical formulation:
+    θ_{t+1} = θ_t - α * ∇_θ J(θ_t)
+    
+    With momentum:
+    v_{t+1} = μ * v_t + ∇_θ J(θ_t)
+    θ_{t+1} = θ_t - α * v_{t+1}
+    """
+    
+    def __init__(self, params: List[torch.Tensor], lr: float = 1e-3, 
+                 momentum: float = 0.0, weight_decay: float = 0.0, 
+                 dampening: float = 0.0, nesterov: bool = False):
+        super().__init__(params, lr)
+        self.momentum = momentum
+        self.weight_decay = weight_decay
+        self.dampening = dampening
+        self.nesterov = nesterov
+        
+        # Initialize momentum buffers
+        for param in self.params:
+            if momentum > 0:
+                self.state[param] = {'momentum_buffer': torch.zeros_like(param)}
+                
+        logger.info(f"Initialized SGD optimizer: lr={lr}, momentum={momentum}, weight_decay={weight_decay}")
+    
+    def step(self, closure=None):
+        """Perform SGD optimization step"""
+        loss = None
+        if closure is not None:
+            loss = closure()
+            
+        for param in self.params:
+            if param.grad is None:
+                continue
+                
+            grad = param.grad.data
+            
+            # Apply weight decay
+            if self.weight_decay != 0:
+                grad = grad.add(param.data, alpha=self.weight_decay)
+            
+            # Apply momentum
+            if self.momentum != 0:
+                param_state = self.state[param]
+                if 'momentum_buffer' not in param_state:
+                    param_state['momentum_buffer'] = torch.zeros_like(param.data)
+                
+                momentum_buffer = param_state['momentum_buffer']
+                momentum_buffer.mul_(self.momentum).add_(grad, alpha=1 - self.dampening)
+                
+                if self.nesterov:
+                    grad = grad.add(momentum_buffer, alpha=self.momentum)
+                else:
+                    grad = momentum_buffer
+            
+            # Update parameters
+            param.data.add_(grad, alpha=-self.lr)
+            
+        self.step_count += 1
+        return loss
+
+
+class Adam(Optimizer):
+    """
+    Adam (Adaptive Moment Estimation) optimizer
+    
+    Mathematical formulation:
+    m_t = β₁ * m_{t-1} + (1 - β₁) * g_t
+    v_t = β₂ * v_{t-1} + (1 - β₂) * g_t²
+    m̂_t = m_t / (1 - β₁ᵗ)
+    v̂_t = v_t / (1 - β₂ᵗ)
+    θ_{t+1} = θ_t - α * m̂_t / (√v̂_t + ε)
+    """
+    
+    def __init__(self, params: List[torch.Tensor], lr: float = 1e-3,
+                 betas: Tuple[float, float] = (0.9, 0.999), eps: float = 1e-8,
+                 weight_decay: float = 0.0, amsgrad: bool = False):
+        super().__init__(params, lr)
+        self.betas = betas
+        self.eps = eps
+        self.weight_decay = weight_decay
+        self.amsgrad = amsgrad
+        
+        # Initialize moment estimates
+        for param in self.params:
+            self.state[param] = {
+                'step': 0,
+                'exp_avg': torch.zeros_like(param.data),
+                'exp_avg_sq': torch.zeros_like(param.data)
+            }
+            if amsgrad:
+                self.state[param]['max_exp_avg_sq'] = torch.zeros_like(param.data)
+                
+        logger.info(f"Initialized Adam optimizer: lr={lr}, betas={betas}, eps={eps}")
+    
+    def step(self, closure=None):
+        """Perform Adam optimization step"""
+        loss = None
+        if closure is not None:
+            loss = closure()
+            
+        for param in self.params:
+            if param.grad is None:
+                continue
+                
+            grad = param.grad.data
+            
+            # Apply weight decay
+            if self.weight_decay != 0:
+                grad = grad.add(param.data, alpha=self.weight_decay)
+            
+            param_state = self.state[param]
+            exp_avg, exp_avg_sq = param_state['exp_avg'], param_state['exp_avg_sq']
+            beta1, beta2 = self.betas
+            
+            param_state['step'] += 1
+            bias_correction1 = 1 - beta1 ** param_state['step']
+            bias_correction2 = 1 - beta2 ** param_state['step']
+            
+            # Update biased first moment estimate
+            exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+            
+            # Update biased second raw moment estimate
+            exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+            
+            if self.amsgrad:
+                # Maintains the maximum of all 2nd moment running avg. of squared gradients
+                torch.max(param_state['max_exp_avg_sq'], exp_avg_sq, out=param_state['max_exp_avg_sq'])
+                # Use the max. for normalizing running avg. of squared gradients
+                denom = (param_state['max_exp_avg_sq'].sqrt() / math.sqrt(bias_correction2)).add_(self.eps)
+            else:
+                denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(self.eps)
+            
+            # Update parameters
+            step_size = self.lr / bias_correction1
+            param.data.addcdiv_(exp_avg, denom, value=-step_size)
+            
+        self.step_count += 1
+        return loss
+
+
+class AdamW(Optimizer):
+    """
+    AdamW (Adam with Decoupled Weight Decay) optimizer
+    
+    Mathematical formulation:
+    θ_t = θ_{t-1} - α * (m̂_t / (√v̂_t + ε) + λ * θ_{t-1})
+    
+    Where weight decay is applied directly to parameters, not gradients.
+    """
+    
+    def __init__(self, params: List[torch.Tensor], lr: float = 1e-3,
+                 betas: Tuple[float, float] = (0.9, 0.999), eps: float = 1e-8,
+                 weight_decay: float = 0.01, amsgrad: bool = False):
+        super().__init__(params, lr)
+        self.betas = betas
+        self.eps = eps
+        self.weight_decay = weight_decay
+        self.amsgrad = amsgrad
+        
+        # Initialize moment estimates
+        for param in self.params:
+            self.state[param] = {
+                'step': 0,
+                'exp_avg': torch.zeros_like(param.data),
+                'exp_avg_sq': torch.zeros_like(param.data)
+            }
+            if amsgrad:
+                self.state[param]['max_exp_avg_sq'] = torch.zeros_like(param.data)
+                
+        logger.info(f"Initialized AdamW optimizer: lr={lr}, betas={betas}, weight_decay={weight_decay}")
+    
+    def step(self, closure=None):
+        """Perform AdamW optimization step"""
+        loss = None
+        if closure is not None:
+            loss = closure()
+            
+        for param in self.params:
+            if param.grad is None:
+                continue
+                
+            grad = param.grad.data
+            param_state = self.state[param]
+            exp_avg, exp_avg_sq = param_state['exp_avg'], param_state['exp_avg_sq']
+            beta1, beta2 = self.betas
+            
+            param_state['step'] += 1
+            bias_correction1 = 1 - beta1 ** param_state['step']
+            bias_correction2 = 1 - beta2 ** param_state['step']
+            
+            # Update biased first moment estimate
+            exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+            
+            # Update biased second raw moment estimate
+            exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+            
+            if self.amsgrad:
+                # Maintains the maximum of all 2nd moment running avg. of squared gradients
+                torch.max(param_state['max_exp_avg_sq'], exp_avg_sq, out=param_state['max_exp_avg_sq'])
+                # Use the max. for normalizing running avg. of squared gradients
+                denom = (param_state['max_exp_avg_sq'].sqrt() / math.sqrt(bias_correction2)).add_(self.eps)
+            else:
+                denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(self.eps)
+            
+            # Update parameters with decoupled weight decay
+            step_size = self.lr / bias_correction1
+            param.data.mul_(1 - self.lr * self.weight_decay)
+            param.data.addcdiv_(exp_avg, denom, value=-step_size)
+            
+        self.step_count += 1
+        return loss
+
+
+class OptimizerFactory:
+    """Factory class for creating optimizers"""
+    
+    @staticmethod
+    def create_optimizer(optimizer_type: str, params: List[torch.Tensor], **kwargs) -> Optimizer:
+        """Create an optimizer instance"""
+        optimizers = {
+            'sgd': SGD,
+            'adam': Adam,
+            'adamw': AdamW
+        }
+        
+        if optimizer_type.lower() not in optimizers:
+            raise ValueError(f"Unknown optimizer type: {optimizer_type}")
+            
+        optimizer_class = optimizers[optimizer_type.lower()]
+        return optimizer_class(params, **kwargs)
+    
+    @staticmethod
+    def get_default_config(optimizer_type: str) -> Dict[str, Any]:
+        """Get default configuration for optimizer"""
+        configs = {
+            'sgd': {
+                'lr': 1e-3,
+                'momentum': 0.9,
+                'weight_decay': 1e-4
+            },
+            'adam': {
+                'lr': 1e-3,
+                'betas': (0.9, 0.999),
+                'eps': 1e-8,
+                'weight_decay': 1e-4
+            },
+            'adamw': {
+                'lr': 1e-3,
+                'betas': (0.9, 0.999),
+                'eps': 1e-8,
+                'weight_decay': 0.01
+            }
+        }
+        
+        return configs.get(optimizer_type.lower(), {})

src/multi_agent_training/gradient_descent/schedulers.py

@@ -0,0 +1,338 @@
+"""
+Learning Rate Schedulers
+=======================
+
+This module implements various learning rate scheduling strategies for
+gradient descent optimization, including linear, cosine, and step schedulers.
+
+Learning rate scheduling is crucial for stable training and convergence.
+"""
+
+import logging
+import math
+from abc import ABC, abstractmethod
+from typing import Dict, Optional, Any
+
+logger = logging.getLogger(__name__)
+
+
+class Scheduler(ABC):
+    """Abstract base class for all learning rate schedulers"""
+    
+    def __init__(self, optimizer, **kwargs):
+        self.optimizer = optimizer
+        self.base_lr = optimizer.lr
+        self.current_lr = self.base_lr
+        
+    @abstractmethod
+    def step(self, epoch: int = None, metrics: Dict[str, float] = None):
+        """Update the learning rate"""
+        pass
+    
+    def get_lr(self) -> float:
+        """Get the current learning rate"""
+        return self.current_lr
+    
+    def state_dict(self) -> Dict[str, Any]:
+        """Return the state of the scheduler"""
+        return {
+            'current_lr': self.current_lr,
+            'base_lr': self.base_lr
+        }
+    
+    def load_state_dict(self, state_dict: Dict[str, Any]):
+        """Load the state of the scheduler"""
+        self.current_lr = state_dict['current_lr']
+        self.base_lr = state_dict['base_lr']
+
+
+class LinearScheduler(Scheduler):
+    """
+    Linear learning rate scheduler
+    
+    Mathematical formulation:
+    lr_t = lr_0 * (1 - t/T)
+    
+    Where lr_0 is the initial learning rate, t is the current step,
+    and T is the total number of steps.
+    """
+    
+    def __init__(self, optimizer, total_steps: int, min_lr: float = 0.0):
+        super().__init__(optimizer)
+        self.total_steps = total_steps
+        self.min_lr = min_lr
+        self.current_step = 0
+        
+        logger.info(f"Initialized LinearScheduler: base_lr={self.base_lr}, total_steps={total_steps}, min_lr={min_lr}")
+    
+    def step(self, epoch: int = None, metrics: Dict[str, float] = None):
+        """Update learning rate linearly"""
+        if self.current_step < self.total_steps:
+            # Linear decay
+            decay_factor = 1.0 - (self.current_step / self.total_steps)
+            self.current_lr = max(self.min_lr, self.base_lr * decay_factor)
+            
+            # Update optimizer learning rate
+            self.optimizer.lr = self.current_lr
+            
+            self.current_step += 1
+            
+            logger.debug(f"LinearScheduler step {self.current_step}: lr={self.current_lr:.6f}")
+
+
+class CosineScheduler(Scheduler):
+    """
+    Cosine annealing learning rate scheduler
+    
+    Mathematical formulation:
+    lr_t = lr_min + (lr_max - lr_min) * (1 + cos(π * t/T)) / 2
+    
+    Where lr_max is the maximum learning rate, lr_min is the minimum,
+    t is the current step, and T is the total number of steps.
+    """
+    
+    def __init__(self, optimizer, total_steps: int, min_lr: float = 0.0, 
+                 max_lr: Optional[float] = None, warmup_steps: int = 0):
+        super().__init__(optimizer)
+        self.total_steps = total_steps
+        self.min_lr = min_lr
+        self.max_lr = max_lr or self.base_lr
+        self.warmup_steps = warmup_steps
+        self.current_step = 0
+        
+        logger.info(f"Initialized CosineScheduler: max_lr={self.max_lr}, min_lr={min_lr}, total_steps={total_steps}")
+    
+    def step(self, epoch: int = None, metrics: Dict[str, float] = None):
+        """Update learning rate with cosine annealing"""
+        if self.current_step < self.total_steps:
+            if self.current_step < self.warmup_steps:
+                # Warmup phase
+                self.current_lr = self.min_lr + (self.max_lr - self.min_lr) * (self.current_step / self.warmup_steps)
+            else:
+                # Cosine annealing phase
+                progress = (self.current_step - self.warmup_steps) / (self.total_steps - self.warmup_steps)
+                cosine_factor = (1 + math.cos(math.pi * progress)) / 2
+                self.current_lr = self.min_lr + (self.max_lr - self.min_lr) * cosine_factor
+            
+            # Update optimizer learning rate
+            self.optimizer.lr = self.current_lr
+            
+            self.current_step += 1
+            
+            logger.debug(f"CosineScheduler step {self.current_step}: lr={self.current_lr:.6f}")
+
+
+class StepScheduler(Scheduler):
+    """
+    Step learning rate scheduler
+    
+    Mathematical formulation:
+    lr_t = lr_0 * γ^⌊t/step_size⌋
+    
+    Where γ is the decay factor and step_size is the number of steps
+    between each decay.
+    """
+    
+    def __init__(self, optimizer, step_size: int, gamma: float = 0.1, 
+                 min_lr: float = 0.0):
+        super().__init__(optimizer)
+        self.step_size = step_size
+        self.gamma = gamma
+        self.min_lr = min_lr
+        self.current_step = 0
+        
+        logger.info(f"Initialized StepScheduler: step_size={step_size}, gamma={gamma}, min_lr={min_lr}")
+    
+    def step(self, epoch: int = None, metrics: Dict[str, float] = None):
+        """Update learning rate with step decay"""
+        if self.current_step % self.step_size == 0 and self.current_step > 0:
+            # Decay learning rate
+            decay_factor = self.gamma ** (self.current_step // self.step_size)
+            self.current_lr = max(self.min_lr, self.base_lr * decay_factor)
+            
+            # Update optimizer learning rate
+            self.optimizer.lr = self.current_lr
+            
+            logger.debug(f"StepScheduler decay at step {self.current_step}: lr={self.current_lr:.6f}")
+        
+        self.current_step += 1
+
+
+class ExponentialScheduler(Scheduler):
+    """
+    Exponential learning rate scheduler
+    
+    Mathematical formulation:
+    lr_t = lr_0 * γ^t
+    
+    Where γ is the decay factor.
+    """
+    
+    def __init__(self, optimizer, gamma: float = 0.95, min_lr: float = 0.0):
+        super().__init__(optimizer)
+        self.gamma = gamma
+        self.min_lr = min_lr
+        self.current_step = 0
+        
+        logger.info(f"Initialized ExponentialScheduler: gamma={gamma}, min_lr={min_lr}")
+    
+    def step(self, epoch: int = None, metrics: Dict[str, float] = None):
+        """Update learning rate exponentially"""
+        self.current_lr = max(self.min_lr, self.base_lr * (self.gamma ** self.current_step))
+        
+        # Update optimizer learning rate
+        self.optimizer.lr = self.current_lr
+        
+        self.current_step += 1
+        
+        logger.debug(f"ExponentialScheduler step {self.current_step}: lr={self.current_lr:.6f}")
+
+
+class ReduceLROnPlateauScheduler(Scheduler):
+    """
+    Reduce learning rate when a metric has stopped improving
+    
+    This scheduler monitors a metric and reduces the learning rate
+    when the metric stops improving for a specified number of epochs.
+    """
+    
+    def __init__(self, optimizer, mode: str = 'min', factor: float = 0.1, 
+                 patience: int = 10, threshold: float = 1e-4, 
+                 min_lr: float = 0.0, metric_name: str = 'loss'):
+        super().__init__(optimizer)
+        self.mode = mode
+        self.factor = factor
+        self.patience = patience
+        self.threshold = threshold
+        self.min_lr = min_lr
+        self.metric_name = metric_name
+        
+        self.best_metric = float('inf') if mode == 'min' else float('-inf')
+        self.patience_counter = 0
+        self.current_epoch = 0
+        
+        logger.info(f"Initialized ReduceLROnPlateauScheduler: mode={mode}, factor={factor}, patience={patience}")
+    
+    def step(self, epoch: int = None, metrics: Dict[str, float] = None):
+        """Update learning rate based on metric improvement"""
+        if metrics is None or self.metric_name not in metrics:
+            return
+        
+        current_metric = metrics[self.metric_name]
+        self.current_epoch = epoch or self.current_epoch
+        
+        # Check if metric improved
+        if self.mode == 'min':
+            improved = current_metric < (self.best_metric - self.threshold)
+        else:
+            improved = current_metric > (self.best_metric + self.threshold)
+        
+        if improved:
+            self.best_metric = current_metric
+            self.patience_counter = 0
+        else:
+            self.patience_counter += 1
+        
+        # Reduce learning rate if patience exceeded
+        if self.patience_counter >= self.patience:
+            old_lr = self.current_lr
+            self.current_lr = max(self.min_lr, self.current_lr * self.factor)
+            
+            # Update optimizer learning rate
+            self.optimizer.lr = self.current_lr
+            
+            self.patience_counter = 0
+            
+            logger.info(f"ReduceLROnPlateauScheduler reduced LR: {old_lr:.6f} -> {self.current_lr:.6f}")
+
+
+class WarmupScheduler(Scheduler):
+    """
+    Warmup learning rate scheduler
+    
+    Gradually increases learning rate from 0 to the target learning rate
+    over a specified number of warmup steps.
+    """
+    
+    def __init__(self, optimizer, warmup_steps: int, target_lr: Optional[float] = None):
+        super().__init__(optimizer)
+        self.warmup_steps = warmup_steps
+        self.target_lr = target_lr or self.base_lr
+        self.current_step = 0
+        
+        # Start with zero learning rate
+        self.current_lr = 0.0
+        self.optimizer.lr = self.current_lr
+        
+        logger.info(f"Initialized WarmupScheduler: warmup_steps={warmup_steps}, target_lr={self.target_lr}")
+    
+    def step(self, epoch: int = None, metrics: Dict[str, float] = None):
+        """Update learning rate with warmup"""
+        if self.current_step < self.warmup_steps:
+            # Linear warmup
+            self.current_lr = self.target_lr * (self.current_step / self.warmup_steps)
+            
+            # Update optimizer learning rate
+            self.optimizer.lr = self.current_lr
+            
+            logger.debug(f"WarmupScheduler step {self.current_step}: lr={self.current_lr:.6f}")
+        
+        self.current_step += 1
+
+
+class SchedulerFactory:
+    """Factory class for creating learning rate schedulers"""
+    
+    @staticmethod
+    def create_scheduler(scheduler_type: str, optimizer, **kwargs) -> Scheduler:
+        """Create a scheduler instance"""
+        schedulers = {
+            'linear': LinearScheduler,
+            'cosine': CosineScheduler,
+            'step': StepScheduler,
+            'exponential': ExponentialScheduler,
+            'plateau': ReduceLROnPlateauScheduler,
+            'warmup': WarmupScheduler
+        }
+        
+        if scheduler_type.lower() not in schedulers:
+            raise ValueError(f"Unknown scheduler type: {scheduler_type}")
+            
+        scheduler_class = schedulers[scheduler_type.lower()]
+        return scheduler_class(optimizer, **kwargs)
+    
+    @staticmethod
+    def get_default_config(scheduler_type: str) -> Dict[str, Any]:
+        """Get default configuration for scheduler"""
+        configs = {
+            'linear': {
+                'total_steps': 1000,
+                'min_lr': 0.0
+            },
+            'cosine': {
+                'total_steps': 1000,
+                'min_lr': 0.0,
+                'warmup_steps': 100
+            },
+            'step': {
+                'step_size': 100,
+                'gamma': 0.1,
+                'min_lr': 0.0
+            },
+            'exponential': {
+                'gamma': 0.95,
+                'min_lr': 0.0
+            },
+            'plateau': {
+                'mode': 'min',
+                'factor': 0.1,
+                'patience': 10,
+                'threshold': 1e-4,
+                'min_lr': 0.0
+            },
+            'warmup': {
+                'warmup_steps': 100
+            }
+        }
+        
+        return configs.get(scheduler_type.lower(), {})

src/multi_agent_training/gradient_descent/training_loop.py

@@ -0,0 +1,570 @@
+"""
+Gradient Descent Training Loop
+=============================
+
+This module implements the main training loop that orchestrates gradient descent
+optimization with backpropagation for the MangoMAS multi-agent system.
+
+The training loop includes:
+- Forward and backward passes
+- Gradient computation and optimization
+- Learning rate scheduling
+- Comprehensive monitoring and logging
+- Model checkpointing and validation
+"""
+
+import logging
+import time
+import math
+from typing import Dict, List, Optional, Tuple, Any
+import torch
+import torch.nn as nn
+from pathlib import Path
+
+from .optimizers import OptimizerFactory
+from .backpropagation import BackpropagationEngine, LoRABackpropagationEngine
+from .loss_functions import LossFunctionFactory
+from .schedulers import SchedulerFactory
+from .monitoring import GradientMonitor, TrainingMonitor, PerformanceMonitor
+
+logger = logging.getLogger(__name__)
+
+
+class GradientDescentTrainer:
+    """
+    Main training class that orchestrates gradient descent optimization
+    
+    This class provides a complete training pipeline with:
+    - Real gradient descent and backpropagation
+    - Comprehensive monitoring and logging
+    - Model checkpointing and validation
+    - Integration with MangoMAS agent system
+    """
+    
+    def __init__(self, 
+                 optimizer_type: str = 'adam',
+                 learning_rate: float = 1e-3,
+                 scheduler_type: str = 'cosine',
+                 loss_function_type: str = 'cross_entropy',
+                 device: torch.device = None,
+                 max_grad_norm: float = 1.0,
+                 gradient_accumulation_steps: int = 1,
+                 mixed_precision: bool = False,
+                 **kwargs):
+        
+        self.optimizer_type = optimizer_type
+        self.learning_rate = learning_rate
+        self.scheduler_type = scheduler_type
+        self.loss_function_type = loss_function_type
+        self.device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.max_grad_norm = max_grad_norm
+        self.gradient_accumulation_steps = gradient_accumulation_steps
+        self.mixed_precision = mixed_precision
+        
+        # Initialize components
+        self.optimizer = None
+        self.scheduler = None
+        self.loss_function = None
+        self.backprop_engine = None
+        
+        # Monitoring
+        self.gradient_monitor = GradientMonitor()
+        self.training_monitor = TrainingMonitor()
+        self.performance_monitor = PerformanceMonitor()
+        
+        # Training state
+        self.current_epoch = 0
+        self.current_step = 0
+        self.best_loss = float('inf')
+        self.training_start_time = None
+        
+        # Configuration
+        self.config = {
+            'optimizer_type': optimizer_type,
+            'learning_rate': learning_rate,
+            'scheduler_type': scheduler_type,
+            'loss_function_type': loss_function_type,
+            'max_grad_norm': max_grad_norm,
+            'gradient_accumulation_steps': gradient_accumulation_steps,
+            'mixed_precision': mixed_precision,
+            **kwargs
+        }
+        
+        logger.info(f"Initialized GradientDescentTrainer with config: {self.config}")
+    
+    def setup_training(self, model: nn.Module, training_data: List[Dict[str, Any]]):
+        """
+        Setup training components
+        
+        Args:
+            model: The neural network model to train
+            training_data: Training dataset
+        """
+        logger.info("Setting up training components...")
+        
+        # Move model to device
+        model.to(self.device)
+        
+        # Get trainable parameters
+        trainable_params = [p for p in model.parameters() if p.requires_grad]
+        logger.info(f"Found {len(trainable_params)} trainable parameters")
+        
+        # Initialize optimizer
+        optimizer_config = OptimizerFactory.get_default_config(self.optimizer_type)
+        optimizer_config.update({'lr': self.learning_rate})
+        
+        self.optimizer = OptimizerFactory.create_optimizer(
+            self.optimizer_type, trainable_params, **optimizer_config
+        )
+        
+        # Initialize scheduler
+        scheduler_config = SchedulerFactory.get_default_config(self.scheduler_type)
+        scheduler_config.update({'total_steps': len(training_data)})
+        
+        self.scheduler = SchedulerFactory.create_scheduler(
+            self.scheduler_type, self.optimizer, **scheduler_config
+        )
+        
+        # Initialize loss function
+        loss_config = LossFunctionFactory.get_default_config(self.loss_function_type)
+        self.loss_function = LossFunctionFactory.create_loss_function(
+            self.loss_function_type, **loss_config
+        )
+        
+        # Initialize backpropagation engine
+        if hasattr(model, 'lora_params'):
+            # LoRA model
+            self.backprop_engine = LoRABackpropagationEngine(
+                model, model.lora_params, self.device
+            )
+        else:
+            # Standard model
+            self.backprop_engine = BackpropagationEngine(model, self.device)
+        
+        logger.info("Training setup complete")
+    
+    def train_epoch(self, model: nn.Module, training_data: List[Dict[str, Any]], 
+                   epoch: int) -> Dict[str, float]:
+        """
+        Train for one epoch using gradient descent and backpropagation
+        
+        Args:
+            model: The neural network model
+            training_data: Training dataset
+            epoch: Current epoch number
+            
+        Returns:
+            Dictionary of training metrics
+        """
+        logger.info(f"Starting epoch {epoch}")
+        
+        model.train()
+        epoch_loss = 0.0
+        epoch_accuracy = 0.0
+        num_batches = 0
+        
+        # Process training data in batches
+        batch_size = 32  # Default batch size
+        num_batches = math.ceil(len(training_data) / batch_size)
+        
+        for batch_idx in range(num_batches):
+            start_idx = batch_idx * batch_size
+            end_idx = min(start_idx + batch_size, len(training_data))
+            batch_data = training_data[start_idx:end_idx]
+            
+            # Process batch
+            batch_metrics = self.train_batch(model, batch_data, epoch, batch_idx)
+            
+            epoch_loss += batch_metrics['loss']
+            epoch_accuracy += batch_metrics.get('accuracy', 0.0)
+            
+            # Update step counter
+            self.current_step += 1
+            
+            # Log progress
+            if batch_idx % 10 == 0:
+                logger.info(f"Epoch {epoch}, Batch {batch_idx}/{num_batches}, "
+                          f"Loss: {batch_metrics['loss']:.4f}")
+        
+        # Compute epoch averages
+        avg_loss = epoch_loss / num_batches
+        avg_accuracy = epoch_accuracy / num_batches if num_batches > 0 else 0.0
+        
+        # Update monitors
+        self.training_monitor.update(
+            loss=avg_loss,
+            accuracy=avg_accuracy,
+            learning_rate=self.optimizer.lr,
+            epoch=epoch
+        )
+        
+        # Update scheduler
+        self.scheduler.step(epoch=epoch, metrics={'loss': avg_loss})
+        
+        logger.info(f"Epoch {epoch} complete - Loss: {avg_loss:.4f}, "
+                   f"Accuracy: {avg_accuracy:.4f}, LR: {self.optimizer.lr:.6f}")
+        
+        return {
+            'loss': avg_loss,
+            'accuracy': avg_accuracy,
+            'learning_rate': self.optimizer.lr,
+            'num_batches': num_batches
+        }
+    
+    def train_batch(self, model: nn.Module, batch_data: List[Dict[str, Any]], 
+                   epoch: int, batch_idx: int) -> Dict[str, float]:
+        """
+        Train on a single batch using gradient descent and backpropagation
+        
+        Args:
+            model: The neural network model
+            batch_data: Batch of training data
+            epoch: Current epoch number
+            batch_idx: Current batch index
+            
+        Returns:
+            Dictionary of batch metrics
+        """
+        # Prepare batch data
+        inputs, targets = self._prepare_batch(batch_data)
+        
+        # Forward pass
+        with self.performance_monitor.time_step('forward'):
+            outputs = model(inputs)
+        
+        # Compute loss
+        loss = self.loss_function(outputs, targets)
+        
+        # Scale loss for gradient accumulation
+        if self.gradient_accumulation_steps > 1:
+            loss = loss / self.gradient_accumulation_steps
+        
+        # Backward pass
+        with self.performance_monitor.time_step('backward'):
+            loss.backward()
+        
+        # Gradient accumulation
+        if (batch_idx + 1) % self.gradient_accumulation_steps == 0:
+            # Apply gradient clipping
+            grad_norm = self.backprop_engine.apply_gradient_clipping(self.max_grad_norm)
+            
+            # Get gradients for monitoring
+            gradients = self.backprop_engine.compute_gradients(loss, retain_graph=False)
+            self.gradient_monitor.update(gradients)
+            
+            # Optimizer step
+            with self.performance_monitor.time_step('optimizer'):
+                self.optimizer.step()
+            
+            # Zero gradients
+            self.optimizer.zero_grad()
+            
+            # Update performance monitoring
+            self.performance_monitor.update_compute_time(time.time() - self.training_start_time)
+        
+        # Compute accuracy (if applicable)
+        accuracy = self._compute_accuracy(outputs, targets)
+        
+        return {
+            'loss': loss.item() * self.gradient_accumulation_steps,
+            'accuracy': accuracy,
+            'grad_norm': grad_norm if 'grad_norm' in locals() else 0.0
+        }
+    
+    def _prepare_batch(self, batch_data: List[Dict[str, Any]]) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Prepare batch data for training
+        
+        Args:
+            batch_data: Raw batch data
+            
+        Returns:
+            Tuple of (inputs, targets) tensors
+        """
+        # Extract inputs and targets
+        inputs = []
+        targets = []
+        
+        for item in batch_data:
+            # Convert text to tokens (simplified)
+            if 'instruction' in item and 'response' in item:
+                # For text generation tasks
+                input_text = item['instruction']
+                target_text = item['response']
+                
+                # Simple tokenization (in practice, use proper tokenizer)
+                input_tokens = self._simple_tokenize(input_text)
+                target_tokens = self._simple_tokenize(target_text)
+                
+                inputs.append(input_tokens)
+                targets.append(target_tokens)
+        
+        # Convert to tensors
+        if inputs and targets:
+            # Pad sequences to same length
+            max_len = max(len(seq) for seq in inputs + targets)
+            inputs = [seq + [0] * (max_len - len(seq)) for seq in inputs]
+            targets = [seq + [0] * (max_len - len(seq)) for seq in targets]
+            
+            inputs_tensor = torch.tensor(inputs, dtype=torch.long, device=self.device)
+            targets_tensor = torch.tensor(targets, dtype=torch.long, device=self.device)
+        else:
+            # Fallback: create dummy data
+            batch_size = len(batch_data)
+            seq_len = 128
+            inputs_tensor = torch.randint(0, 1000, (batch_size, seq_len), device=self.device)
+            targets_tensor = torch.randint(0, 1000, (batch_size, seq_len), device=self.device)
+        
+        return inputs_tensor, targets_tensor
+    
+    def _simple_tokenize(self, text: str) -> List[int]:
+        """
+        Simple tokenization for demonstration
+        
+        Args:
+            text: Input text
+            
+        Returns:
+            List of token IDs
+        """
+        # Simple character-based tokenization
+        tokens = []
+        for char in text[:100]:  # Limit length
+            tokens.append(ord(char) % 1000)  # Map to vocabulary
+        return tokens
+    
+    def _compute_accuracy(self, outputs: torch.Tensor, targets: torch.Tensor) -> float:
+        """
+        Compute accuracy for the batch
+        
+        Args:
+            outputs: Model outputs
+            targets: Target values
+            
+        Returns:
+            Accuracy score
+        """
+        if outputs.dim() > 1 and outputs.size(1) > 1:
+            # Classification task
+            predictions = torch.argmax(outputs, dim=1)
+            if targets.dim() == 1:
+                correct = (predictions == targets).float().sum()
+                accuracy = correct / targets.size(0)
+            else:
+                # Multi-label case
+                accuracy = 0.0
+        else:
+            # Regression task - use a simple threshold
+            accuracy = 0.0
+        
+        return accuracy.item() if isinstance(accuracy, torch.Tensor) else accuracy
+    
+    def validate(self, model: nn.Module, validation_data: List[Dict[str, Any]]) -> Dict[str, float]:
+        """
+        Validate the model
+        
+        Args:
+            model: The neural network model
+            validation_data: Validation dataset
+            
+        Returns:
+            Dictionary of validation metrics
+        """
+        logger.info("Running validation...")
+        
+        model.eval()
+        total_loss = 0.0
+        total_accuracy = 0.0
+        num_batches = 0
+        
+        with torch.no_grad():
+            batch_size = 32
+            num_batches = math.ceil(len(validation_data) / batch_size)
+            
+            for batch_idx in range(num_batches):
+                start_idx = batch_idx * batch_size
+                end_idx = min(start_idx + batch_size, len(validation_data))
+                batch_data = validation_data[start_idx:end_idx]
+                
+                # Prepare batch
+                inputs, targets = self._prepare_batch(batch_data)
+                
+                # Forward pass
+                outputs = model(inputs)
+                
+                # Compute loss
+                loss = self.loss_function(outputs, targets)
+                total_loss += loss.item()
+                
+                # Compute accuracy
+                accuracy = self._compute_accuracy(outputs, targets)
+                total_accuracy += accuracy
+        
+        # Compute averages
+        avg_loss = total_loss / num_batches if num_batches > 0 else 0.0
+        avg_accuracy = total_accuracy / num_batches if num_batches > 0 else 0.0
+        
+        logger.info(f"Validation - Loss: {avg_loss:.4f}, Accuracy: {avg_accuracy:.4f}")
+        
+        return {
+            'val_loss': avg_loss,
+            'val_accuracy': avg_accuracy
+        }
+    
+    def train(self, model: nn.Module, training_data: List[Dict[str, Any]], 
+              validation_data: Optional[List[Dict[str, Any]]] = None,
+              num_epochs: int = 10, save_dir: Optional[str] = None) -> Dict[str, Any]:
+        """
+        Complete training loop with gradient descent and backpropagation
+        
+        Args:
+            model: The neural network model to train
+            training_data: Training dataset
+            validation_data: Validation dataset (optional)
+            num_epochs: Number of training epochs
+            save_dir: Directory to save checkpoints
+            
+        Returns:
+            Dictionary of training results
+        """
+        logger.info(f"Starting training for {num_epochs} epochs")
+        
+        # Setup training
+        self.setup_training(model, training_data)
+        
+        # Initialize training state
+        self.training_start_time = time.time()
+        self.current_epoch = 0
+        self.current_step = 0
+        
+        # Training history
+        training_history = []
+        validation_history = []
+        
+        # Main training loop
+        for epoch in range(num_epochs):
+            self.current_epoch = epoch
+            
+            # Train epoch
+            epoch_metrics = self.train_epoch(model, training_data, epoch)
+            training_history.append(epoch_metrics)
+            
+            # Validation
+            if validation_data:
+                val_metrics = self.validate(model, validation_data)
+                validation_history.append(val_metrics)
+                
+                # Update best loss
+                if val_metrics['val_loss'] < self.best_loss:
+                    self.best_loss = val_metrics['val_loss']
+                    
+                    # Save best model
+                    if save_dir:
+                        self.save_checkpoint(model, save_dir, epoch, val_metrics)
+            
+            # Check for convergence
+            if self.training_monitor.detect_convergence():
+                logger.info("Training converged, stopping early")
+                break
+            
+            # Log epoch summary
+            logger.info(f"Epoch {epoch} Summary:")
+            logger.info(f"  Training Loss: {epoch_metrics['loss']:.4f}")
+            logger.info(f"  Training Accuracy: {epoch_metrics['accuracy']:.4f}")
+            if validation_data:
+                logger.info(f"  Validation Loss: {val_metrics['val_loss']:.4f}")
+                logger.info(f"  Validation Accuracy: {val_metrics['val_accuracy']:.4f}")
+            logger.info(f"  Learning Rate: {self.optimizer.lr:.6f}")
+        
+        # Training complete
+        training_time = time.time() - self.training_start_time
+        
+        # Get final statistics
+        gradient_stats = self.gradient_monitor.get_statistics()
+        training_stats = self.training_monitor.get_statistics()
+        performance_stats = self.performance_monitor.get_statistics()
+        
+        results = {
+            'training_history': training_history,
+            'validation_history': validation_history,
+            'final_metrics': {
+                'best_loss': self.best_loss,
+                'final_loss': training_history[-1]['loss'] if training_history else 0.0,
+                'final_accuracy': training_history[-1]['accuracy'] if training_history else 0.0,
+                'training_time': training_time,
+                'total_steps': self.current_step,
+                'total_epochs': self.current_epoch + 1
+            },
+            'gradient_stats': gradient_stats,
+            'training_stats': training_stats,
+            'performance_stats': performance_stats,
+            'config': self.config
+        }
+        
+        logger.info("Training complete!")
+        logger.info(f"Final Loss: {results['final_metrics']['final_loss']:.4f}")
+        logger.info(f"Best Loss: {results['final_metrics']['best_loss']:.4f}")
+        logger.info(f"Training Time: {training_time:.2f} seconds")
+        
+        return results
+    
+    def save_checkpoint(self, model: nn.Module, save_dir: str, epoch: int, 
+                       metrics: Dict[str, float]):
+        """
+        Save model checkpoint
+        
+        Args:
+            model: The neural network model
+            save_dir: Directory to save checkpoint
+            epoch: Current epoch
+            metrics: Training metrics
+        """
+        save_path = Path(save_dir)
+        save_path.mkdir(parents=True, exist_ok=True)
+        
+        checkpoint = {
+            'epoch': epoch,
+            'model_state_dict': model.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'scheduler_state_dict': self.scheduler.state_dict(),
+            'best_loss': self.best_loss,
+            'metrics': metrics,
+            'config': self.config
+        }
+        
+        checkpoint_path = save_path / f'checkpoint_epoch_{epoch}.pt'
+        torch.save(checkpoint, checkpoint_path)
+        
+        logger.info(f"Checkpoint saved to {checkpoint_path}")
+    
+    def load_checkpoint(self, model: nn.Module, checkpoint_path: str):
+        """
+        Load model checkpoint
+        
+        Args:
+            model: The neural network model
+            checkpoint_path: Path to checkpoint file
+        """
+        checkpoint = torch.load(checkpoint_path, map_location=self.device)
+        
+        model.load_state_dict(checkpoint['model_state_dict'])
+        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+        self.best_loss = checkpoint['best_loss']
+        
+        logger.info(f"Checkpoint loaded from {checkpoint_path}")
+    
+    def get_training_summary(self) -> Dict[str, Any]:
+        """
+        Get comprehensive training summary
+        
+        Returns:
+            Dictionary of training summary
+        """
+        return {
+            'gradient_stats': self.gradient_monitor.get_statistics(),
+            'training_stats': self.training_monitor.get_statistics(),
+            'performance_stats': self.performance_monitor.get_statistics(),
+            'anomalies': self.gradient_monitor.detect_anomalies(),
+            'convergence': self.training_monitor.detect_convergence()
+        }

src/multi_agent_training/lora_trainer.py

@@ -0,0 +1,699 @@
+"""
+LoRA Knowledge Distillation Trainer for MangoMAS Local
+
+This module implements the main training loop for knowledge distillation
+with LoRA fine-tuning optimized for Mac Mini hardware constraints.
+"""
+
+import argparse
+import json
+import logging
+import os
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List
+
+import torch
+import torch.nn as nn
+import yaml
+from peft import LoraConfig, TaskType, get_peft_model
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          get_linear_schedule_with_warmup)
+
+# Try to import context7 for enhanced training
+try:
+    from context7 import Context7
+
+    CONTEXT7_AVAILABLE = True
+except ImportError:
+    CONTEXT7_AVAILABLE = False
+    Context7 = None
+
+# Try to import MLflow for experiment tracking
+try:
+    import mlflow
+
+    MLFLOW_AVAILABLE = True
+except ImportError:
+    MLFLOW_AVAILABLE = False
+    mlflow = None
+
+# Fix import path issues for distillation loss
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+try:
+    from distillation_loss import AdaptiveDistillationLoss, DistillationLoss
+except ImportError:
+    try:
+        from training.distillation_loss import (AdaptiveDistillationLoss,
+                                                DistillationLoss)
+    except ImportError:
+        # Fallback: create minimal distillation loss if not available
+        class DistillationLoss:
+            def __init__(self, alpha=0.5, temperature=2.0):
+                self.alpha = alpha
+                self.temperature = temperature
+                self.task_loss = nn.CrossEntropyLoss()
+
+            def compute_loss(
+                self, student_logits, teacher_logits, labels, attention_mask=None
+            ):
+                # Task loss (standard cross-entropy)
+                shift_logits = student_logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+                task_loss = self.task_loss(
+                    shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
+                )
+
+                # Distillation loss (KL divergence)
+                if teacher_logits is not None:
+                    student_probs = nn.functional.log_softmax(
+                        student_logits / self.temperature, dim=-1
+                    )
+                    teacher_probs = nn.functional.softmax(
+                        teacher_logits / self.temperature, dim=-1
+                    )
+                    distill_loss = nn.functional.kl_div(
+                        student_probs, teacher_probs, reduction="batchmean"
+                    )
+                    distill_loss *= self.temperature**2
+                else:
+                    distill_loss = torch.tensor(0.0)
+
+                # Combined loss
+                total_loss = (1 - self.alpha) * task_loss + self.alpha * distill_loss
+
+                return total_loss, {
+                    "total_loss": total_loss.item(),
+                    "task_loss": task_loss.item(),
+                    "distillation_loss": (
+                        distill_loss.item()
+                        if isinstance(distill_loss, torch.Tensor)
+                        else 0.0
+                    ),
+                }
+
+        AdaptiveDistillationLoss = DistillationLoss  # Fallback
+
+logger = logging.getLogger(__name__)
+
+
+class ConversationDataset:
+    """Dataset class for conversation-based training data."""
+
+    def __init__(self, data_path: str, tokenizer, max_length: int = 512):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.data = self._load_data(data_path)
+
+    def _load_data(self, data_path: str) -> List[Dict]:
+        """Load conversation data from JSONL file."""
+        data = []
+        with open(data_path, "r", encoding="utf-8") as f:
+            for line in f:
+                data.append(json.loads(line.strip()))
+        return data
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        """Get tokenized conversation item."""
+        item = self.data[idx]
+
+        # Handle different data formats
+        if "messages" in item:
+            # Chat format with messages
+            conversation_text = ""
+            for message in item["messages"]:
+                role = message["role"]
+                content = message["content"]
+                conversation_text += f"<{role}>\n{content}\n</{role}>\n\n"
+        elif "instruction" in item and "response" in item:
+            # Instruction-response format
+            instruction = item["instruction"]
+            response = item["response"]
+            conversation_text = f"<user>\n{instruction}\n</user>\n\n<assistant>\n{response}\n</assistant>\n\n"
+        elif "prompt" in item and "completion" in item:
+            # Prompt-completion format
+            prompt = item["prompt"]
+            completion = item["completion"]
+            conversation_text = f"<user>\n{prompt}\n</user>\n\n<assistant>\n{completion}\n</assistant>\n\n"
+        else:
+            # Fallback - try to extract any text
+            conversation_text = str(item)
+
+        # Tokenize
+        encoding = self.tokenizer(
+            conversation_text,
+            truncation=True,
+            padding="max_length",
+            max_length=self.max_length,
+            return_tensors="pt",
+        )
+
+        return {
+            "input_ids": encoding["input_ids"].squeeze(),
+            "attention_mask": encoding["attention_mask"].squeeze(),
+            "labels": encoding["input_ids"].squeeze().clone(),
+            "agent_type": item.get("agent_type", "unknown"),
+        }
+
+
+class LoRADistillationTrainer:
+    """Main trainer class for LoRA knowledge distillation."""
+
+    def __init__(self, config_path: str):
+        """Initialize trainer with configuration."""
+        with open(config_path, "r") as f:
+            self.config = yaml.safe_load(f)
+
+        self.setup_logging()
+        self.setup_device()
+        self.setup_monitoring()
+
+        logger.info("Initialized LoRA Distillation Trainer")
+        logger.info(f"Device: {self.device}")
+        logger.info(f"Config: {config_path}")
+
+    def setup_logging(self) -> None:
+        """Set up logging configuration."""
+        log_dir = Path("logs")
+        log_dir.mkdir(exist_ok=True)
+
+        logging.basicConfig(
+            level=logging.INFO,
+            format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+            handlers=[
+                logging.FileHandler(log_dir / "training.log"),
+                logging.StreamHandler(),
+            ],
+        )
+
+    def setup_device(self) -> None:
+        """Set up compute device (MPS for Mac Mini)."""
+        device_config = self.config["hardware"]["device"]
+
+        if device_config == "mps" and torch.backends.mps.is_available():
+            self.device = torch.device("mps")
+            logger.info("Using Apple Metal Performance Shaders (MPS)")
+        elif device_config == "cuda" and torch.cuda.is_available():
+            self.device = torch.device("cuda")
+            logger.info(f"Using CUDA: {torch.cuda.get_device_name()}")
+        else:
+            self.device = torch.device("cpu")
+            logger.warning("Using CPU - training will be slow")
+
+    def setup_monitoring(self) -> None:
+        """Set up experiment tracking and monitoring."""
+        self.use_tensorboard = self.config["monitoring"]["use_tensorboard"]
+        self.use_mlflow = self.config["monitoring"]["use_mlflow"]
+
+        if self.use_tensorboard:
+            log_dir = self.config["monitoring"]["log_dir"]
+            Path(log_dir).mkdir(parents=True, exist_ok=True)
+            self.tb_writer = SummaryWriter(log_dir)
+            logger.info(f"TensorBoard logging to: {log_dir}")
+
+        if self.use_mlflow:
+            try:
+                import mlflow
+
+                experiment_name = self.config["monitoring"]["experiment_name"]
+                mlflow.set_experiment(experiment_name)
+                logger.info(f"MLflow experiment: {experiment_name}")
+            except (ImportError, AttributeError) as e:
+                logger.warning(
+                    f"MLflow not available or not properly initialized, disabling: {e}"
+                )
+                self.use_mlflow = False
+
+    def load_models(self) -> None:
+        """Load teacher and student models."""
+        # Load tokenizer
+        model_name = self.config["models"]["student"]["base_model"]
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+        # Add pad token if it doesn't exist
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        # Load student model - fix deprecated torch_dtype
+        dtype = (
+            torch.float16
+            if self.config["optimization"]["use_fp16"] and self.device.type == "cuda"
+            else torch.float32
+        )
+
+        self.student_model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            dtype=dtype,  # Use dtype instead of torch_dtype
+            device_map="auto" if self.device.type == "cuda" else None,
+            trust_remote_code=True,
+        )
+
+        # Apply LoRA to student model - fix target modules for DialoGPT
+        target_modules = self.config["lora"]["target_modules"]
+        # If using default transformer modules but this is DialoGPT, adjust
+        if target_modules == ["q_proj", "v_proj", "k_proj", "o_proj"]:
+            target_modules = ["c_attn", "c_proj", "c_fc"]  # DialoGPT modules
+            logger.info("Adjusted LoRA target modules for DialoGPT architecture")
+
+        lora_config = LoraConfig(
+            r=self.config["lora"]["r"],
+            lora_alpha=self.config["lora"]["lora_alpha"],
+            target_modules=target_modules,
+            lora_dropout=self.config["lora"]["lora_dropout"],
+            bias=self.config["lora"]["bias"],
+            task_type=TaskType.CAUSAL_LM,
+        )
+
+        self.student_model = get_peft_model(self.student_model, lora_config)
+        self.student_model.to(self.device)
+
+        # Setup teacher model
+        self.teacher_manager = TeacherModelManager(
+            self.config["models"]["teacher"], self.tokenizer
+        )
+
+        logger.info("Loaded student model with LoRA")
+        logger.info(
+            f"Trainable parameters: {self.student_model.num_parameters(only_trainable=True):,}"
+        )
+        logger.info("Loaded teacher model")
+
+    def load_datasets(self, agent_type: str) -> tuple:
+        """Load training and validation datasets for specific agent."""
+        data_dir = Path("data/processed")
+
+        train_path = data_dir / f"{agent_type}_train.jsonl"
+        val_path = data_dir / f"{agent_type}_validation.jsonl"
+
+        if not train_path.exists():
+            raise FileNotFoundError(f"Training data not found: {train_path}")
+        if not val_path.exists():
+            raise FileNotFoundError(f"Validation data not found: {val_path}")
+
+        max_length = self.config["data"]["max_sequence_length"]
+
+        train_dataset = ConversationDataset(train_path, self.tokenizer, max_length)
+        val_dataset = ConversationDataset(val_path, self.tokenizer, max_length)
+
+        logger.info(
+            f"Loaded datasets: {len(train_dataset)} train, {len(val_dataset)} val"
+        )
+
+        return train_dataset, val_dataset
+
+    def create_data_loaders(self, train_dataset, val_dataset) -> tuple:
+        """Create data loaders for training and validation."""
+        batch_size = self.config["training"]["batch_size"]
+        num_workers = self.config["optimization"]["dataloader_num_workers"]
+        pin_memory = self.config["optimization"]["pin_memory"]
+
+        train_loader = DataLoader(
+            train_dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            drop_last=True,
+        )
+
+        val_loader = DataLoader(
+            val_dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            drop_last=False,
+        )
+
+        return train_loader, val_loader
+
+    def setup_training(self, train_dataset_size: int) -> None:
+        """Set up optimizer, scheduler, and loss function."""
+        # Calculate training steps
+        batch_size = self.config["training"]["batch_size"]
+        gradient_accumulation_steps = self.config["training"][
+            "gradient_accumulation_steps"
+        ]
+        num_epochs = self.config["training"]["num_epochs"]
+
+        steps_per_epoch = train_dataset_size // (
+            batch_size * gradient_accumulation_steps
+        )
+        self.total_steps = steps_per_epoch * num_epochs
+
+        # Setup optimizer
+        self.optimizer = torch.optim.AdamW(
+            self.student_model.parameters(),
+            lr=self.config["training"]["learning_rate"],
+            weight_decay=0.01,
+        )
+
+        # Setup scheduler
+        self.scheduler = get_linear_schedule_with_warmup(
+            self.optimizer,
+            num_warmup_steps=self.config["training"]["warmup_steps"],
+            num_training_steps=self.total_steps,
+        )
+
+        # Setup loss function
+        self.distill_loss = DistillationLoss(
+            alpha=self.config["distillation"]["alpha"],
+            temperature=self.config["distillation"]["temperature"],
+        )
+
+        logger.info(f"Setup training: {self.total_steps} total steps")
+
+    def train_epoch(self, train_loader: DataLoader, epoch: int) -> Dict[str, float]:
+        """Train for one epoch."""
+        self.student_model.train()
+
+        total_loss = 0.0
+        total_task_loss = 0.0
+        total_distill_loss = 0.0
+        num_batches = 0
+
+        progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}", disable=False)
+
+        for batch_idx, batch in enumerate(progress_bar):
+            # Move batch to device
+            input_ids = batch["input_ids"].to(self.device)
+            attention_mask = batch["attention_mask"].to(self.device)
+            labels = batch["labels"].to(self.device)
+
+            # Get student outputs
+            student_outputs = self.student_model(
+                input_ids=input_ids, attention_mask=attention_mask
+            )
+            student_logits = student_outputs.logits
+
+            # Get teacher outputs
+            with torch.no_grad():
+                teacher_logits = self.teacher_manager.get_logits(
+                    input_ids, attention_mask
+                )
+
+            # Compute distillation loss
+            loss, loss_dict = self.distill_loss.compute_loss(
+                student_logits, teacher_logits, labels, attention_mask
+            )
+
+            # Backward pass with gradient accumulation
+            loss = loss / self.config["training"]["gradient_accumulation_steps"]
+            loss.backward()
+
+            # Update model
+            if (batch_idx + 1) % self.config["training"][
+                "gradient_accumulation_steps"
+            ] == 0:
+                torch.nn.utils.clip_grad_norm_(
+                    self.student_model.parameters(),
+                    self.config["training"]["max_grad_norm"],
+                )
+                self.optimizer.step()
+                self.scheduler.step()
+                self.optimizer.zero_grad()
+
+            # Track metrics
+            total_loss += loss_dict["total_loss"]
+            total_task_loss += loss_dict["task_loss"]
+            total_distill_loss += loss_dict["distillation_loss"]
+            num_batches += 1
+
+            # Update progress bar
+            progress_bar.set_postfix(
+                {
+                    "loss": f"{loss_dict['total_loss']:.4f}",
+                    "task": f"{loss_dict['task_loss']:.4f}",
+                    "distill": f"{loss_dict['distillation_loss']:.4f}",
+                }
+            )
+
+            # Log to tensorboard
+            if (
+                self.use_tensorboard
+                and batch_idx % self.config["training"]["logging_steps"] == 0
+            ):
+                step = epoch * len(train_loader) + batch_idx
+                self.tb_writer.add_scalar(
+                    "train/total_loss", loss_dict["total_loss"], step
+                )
+                self.tb_writer.add_scalar(
+                    "train/task_loss", loss_dict["task_loss"], step
+                )
+                self.tb_writer.add_scalar(
+                    "train/distillation_loss", loss_dict["distillation_loss"], step
+                )
+
+        # Calculate epoch averages
+        epoch_metrics = {
+            "avg_loss": total_loss / num_batches,
+            "avg_task_loss": total_task_loss / num_batches,
+            "avg_distill_loss": total_distill_loss / num_batches,
+        }
+
+        return epoch_metrics
+
+    def evaluate(self, val_loader: DataLoader) -> Dict[str, float]:
+        """Evaluate model on validation set."""
+        self.student_model.eval()
+
+        total_loss = 0.0
+        total_task_loss = 0.0
+        total_distill_loss = 0.0
+        num_batches = 0
+
+        with torch.no_grad():
+            for batch in tqdm(val_loader, desc="Evaluating"):
+                # Move batch to device
+                input_ids = batch["input_ids"].to(self.device)
+                attention_mask = batch["attention_mask"].to(self.device)
+                labels = batch["labels"].to(self.device)
+
+                # Get model outputs
+                student_outputs = self.student_model(
+                    input_ids=input_ids, attention_mask=attention_mask
+                )
+                student_logits = student_outputs.logits
+
+                # Get teacher outputs
+                teacher_logits = self.teacher_manager.get_logits(
+                    input_ids, attention_mask
+                )
+
+                # Compute loss
+                loss, loss_dict = self.distill_loss.compute_loss(
+                    student_logits, teacher_logits, labels, attention_mask
+                )
+
+                total_loss += loss_dict["total_loss"]
+                total_task_loss += loss_dict["task_loss"]
+                total_distill_loss += loss_dict["distillation_loss"]
+                num_batches += 1
+
+        val_metrics = {
+            "val_loss": total_loss / num_batches,
+            "val_task_loss": total_task_loss / num_batches,
+            "val_distill_loss": total_distill_loss / num_batches,
+        }
+
+        return val_metrics
+
+    def save_model(self, output_dir: str, agent_type: str, epoch: int) -> None:
+        """Save model checkpoint."""
+        output_path = Path(output_dir) / agent_type / f"epoch_{epoch}"
+        output_path.mkdir(parents=True, exist_ok=True)
+
+        # Save LoRA adapter
+        self.student_model.save_pretrained(output_path)
+
+        # Save tokenizer
+        self.tokenizer.save_pretrained(output_path)
+
+        # Save training config
+        config_path = output_path / "training_config.yaml"
+        with open(config_path, "w") as f:
+            yaml.dump(self.config, f)
+
+        logger.info(f"Saved model to: {output_path}")
+
+    def train_agent(self, agent_type: str) -> None:
+        """Train a specific agent with knowledge distillation."""
+        logger.info(f"Starting training for {agent_type} agent")
+
+        # Load models if not already loaded
+        if not hasattr(self, "student_model"):
+            self.load_models()
+
+        # Load datasets
+        train_dataset, val_dataset = self.load_datasets(agent_type)
+        train_loader, val_loader = self.create_data_loaders(train_dataset, val_dataset)
+
+        # Setup training components
+        self.setup_training(len(train_dataset))
+
+        # Start MLflow run
+        if self.use_mlflow:
+            mlflow.start_run(
+                run_name=f"{agent_type}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+            )
+            mlflow.log_params(
+                {
+                    "agent_type": agent_type,
+                    "model_name": self.config["models"]["student"]["base_model"],
+                    "lora_r": self.config["lora"]["r"],
+                    "lora_alpha": self.config["lora"]["lora_alpha"],
+                    "batch_size": self.config["training"]["batch_size"],
+                    "learning_rate": self.config["training"]["learning_rate"],
+                    "distillation_alpha": self.config["distillation"]["alpha"],
+                    "temperature": self.config["distillation"]["temperature"],
+                }
+            )
+
+        try:
+            # Training loop
+            best_val_loss = float("inf")
+            num_epochs = self.config["training"]["num_epochs"]
+
+            for epoch in range(num_epochs):
+                logger.info(f"Epoch {epoch+1}/{num_epochs}")
+
+                # Train
+                train_metrics = self.train_epoch(train_loader, epoch)
+                logger.info(
+                    f"Train - Loss: {train_metrics['avg_loss']:.4f}, "
+                    f"Task: {train_metrics['avg_task_loss']:.4f}, "
+                    f"Distill: {train_metrics['avg_distill_loss']:.4f}"
+                )
+
+                # Evaluate
+                val_metrics = self.evaluate(val_loader)
+                logger.info(
+                    f"Val - Loss: {val_metrics['val_loss']:.4f}, "
+                    f"Task: {val_metrics['val_task_loss']:.4f}, "
+                    f"Distill: {val_metrics['val_distill_loss']:.4f}"
+                )
+
+                # Log to MLflow
+                if self.use_mlflow:
+                    mlflow.log_metrics({**train_metrics, **val_metrics}, step=epoch)
+
+                # Log to TensorBoard
+                if self.use_tensorboard:
+                    for key, value in train_metrics.items():
+                        self.tb_writer.add_scalar(f"epoch/{key}", value, epoch)
+                    for key, value in val_metrics.items():
+                        self.tb_writer.add_scalar(f"epoch/{key}", value, epoch)
+
+                # Save checkpoint if best model
+                if val_metrics["val_loss"] < best_val_loss:
+                    best_val_loss = val_metrics["val_loss"]
+                    self.save_model(
+                        self.config["output"]["base_dir"], agent_type, epoch
+                    )
+                    logger.info(f"New best model saved (val_loss: {best_val_loss:.4f})")
+
+        finally:
+            if self.use_mlflow:
+                mlflow.end_run()
+
+        logger.info(f"Training completed for {agent_type} agent")
+
+
+class TeacherModelManager:
+    """Manages teacher model interactions (API or local)."""
+
+    def __init__(self, teacher_config: Dict, tokenizer):
+        self.config = teacher_config
+        self.tokenizer = tokenizer
+
+        if teacher_config["type"] == "api":
+            self.setup_api_teacher()
+        else:
+            self.setup_local_teacher()
+
+    def setup_api_teacher(self) -> None:
+        """Set up API-based teacher model."""
+        self.model_name = self.config["model_name"]
+        logger.info(f"Using API teacher model: {self.model_name}")
+
+        # This would integrate with OpenAI/Anthropic APIs
+        # For now, we'll use a placeholder that returns random logits
+        # In production, you'd implement actual API calls here
+
+    def setup_local_teacher(self) -> None:
+        """Set up local teacher model."""
+        model_path = self.config.get("local_model_path", "microsoft/DialoGPT-large")
+
+        self.teacher_model = AutoModelForCausalLM.from_pretrained(
+            model_path, torch_dtype=torch.float16, device_map="auto"
+        )
+        logger.info(f"Loaded local teacher model: {model_path}")
+
+    def get_logits(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
+    ) -> torch.Tensor:
+        """Get teacher model logits."""
+        if self.config["type"] == "api":
+            # Placeholder for API-based teacher
+            # In practice, you'd call the API and convert responses to logits
+            batch_size, seq_len = input_ids.shape
+            vocab_size = self.tokenizer.vocab_size
+            return torch.randn(batch_size, seq_len, vocab_size).to(input_ids.device)
+        else:
+            # Local teacher model
+            with torch.no_grad():
+                outputs = self.teacher_model(
+                    input_ids=input_ids, attention_mask=attention_mask
+                )
+                return outputs.logits
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Train MangoMAS agents with LoRA and knowledge distillation"
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        default="config/training/distillation.yaml",
+        help="Path to training configuration file",
+    )
+    parser.add_argument(
+        "--agent",
+        type=str,
+        choices=["infrastructure", "devsecops", "risk_assessment", "all"],
+        default="all",
+        help="Which agent to train",
+    )
+    parser.add_argument("--data", type=str, help="Path to training data file")
+
+    args = parser.parse_args()
+
+    # Initialize trainer
+    trainer = LoRADistillationTrainer(args.config)
+
+    # If data path is provided, update the trainer to use it
+    if args.data:
+        trainer.custom_data_path = args.data
+
+    # Train specified agent(s)
+    if args.agent == "all":
+        agents = ["infrastructure", "devsecops", "risk_assessment"]
+    else:
+        agents = [args.agent]
+
+    for agent_type in agents:
+        trainer.train_agent(agent_type)
+
+
+if __name__ == "__main__":
+    main()

src/multi_agent_training/modular_distillation_loss.py

@@ -0,0 +1,263 @@
+"""
+Refactored Knowledge Distillation Loss using modular architecture.
+
+This module implements a clean, testable loss function that follows the interface contracts
+and provides better separation of concerns.
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..core.base_components import BaseLossFunction
+from ..core.exceptions import TrainingError
+
+logger = logging.getLogger(__name__)
+
+
+class ModularDistillationLoss(BaseLossFunction):
+    """
+    Modular distillation loss function implementing clean interface contracts.
+    """
+
+    def __init__(
+        self,
+        alpha: float = 0.5,
+        temperature: float = 3.0,
+        task_loss_fn: Optional[nn.Module] = None,
+    ):
+        """
+        Initialize the modular distillation loss.
+
+        Args:
+            alpha: Balance between task loss and distillation loss (0.0-1.0)
+            temperature: Temperature for softmax in knowledge distillation
+            task_loss_fn: Custom task loss function (defaults to CrossEntropyLoss)
+        """
+        super().__init__({"alpha": alpha, "temperature": temperature})
+
+        if not 0.0 <= alpha <= 1.0:
+            raise ValueError(f"Alpha must be between 0 and 1, got {alpha}")
+        if temperature <= 0.0:
+            raise ValueError(f"Temperature must be positive, got {temperature}")
+
+        self.alpha = alpha
+        self.temperature = temperature
+        self.task_loss_fn = task_loss_fn or nn.CrossEntropyLoss(ignore_index=-100)
+        self.kl_loss_fn = nn.KLDivLoss(reduction="batchmean")
+
+        logger.info(
+            f"Initialized ModularDistillationLoss with alpha={alpha}, temperature={temperature}"
+        )
+
+    def compute(
+        self,
+        student_logits: torch.Tensor,
+        teacher_logits: torch.Tensor,
+        labels: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Compute the combined distillation loss.
+
+        Args:
+            student_logits: Logits from student model
+            teacher_logits: Logits from teacher model
+            labels: Target labels for task loss
+            attention_mask: Attention mask for valid tokens
+
+        Returns:
+            Combined loss tensor
+        """
+        try:
+            # Compute task loss (if labels provided)
+            task_loss = torch.tensor(0.0, device=student_logits.device)
+            if labels is not None:
+                task_loss = self._compute_task_loss(
+                    student_logits, labels, attention_mask
+                )
+
+            # Compute distillation loss
+            distill_loss = self._compute_distillation_loss(
+                student_logits, teacher_logits, attention_mask
+            )
+
+            # Combined loss
+            total_loss = self.alpha * task_loss + (1 - self.alpha) * distill_loss
+
+            # Track metrics
+            metrics = {
+                "task_loss": task_loss.item(),
+                "distillation_loss": distill_loss.item(),
+                "total_loss": total_loss.item(),
+                "alpha": self.alpha,
+                "temperature": self.temperature,
+            }
+            self._track_metrics(metrics)
+            self._track_loss(total_loss.item())
+
+            return total_loss
+
+        except Exception as e:
+            raise TrainingError(
+                f"Loss computation failed: {str(e)}", "LOSS_COMPUTATION_ERROR"
+            )
+
+    def _compute_task_loss(
+        self,
+        logits: torch.Tensor,
+        labels: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Compute standard cross-entropy task loss."""
+        # Reshape for cross entropy: [batch_size * seq_len, vocab_size]
+        shift_logits = logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+
+        # Flatten the tokens
+        flat_logits = shift_logits.view(-1, shift_logits.size(-1))
+        flat_labels = shift_labels.view(-1)
+
+        # Apply attention mask if provided
+        if attention_mask is not None:
+            shift_attention_mask = attention_mask[..., 1:].contiguous()
+            flat_attention_mask = shift_attention_mask.view(-1)
+
+            # Only compute loss on non-masked tokens
+            active_logits = flat_logits[flat_attention_mask == 1]
+            active_labels = flat_labels[flat_attention_mask == 1]
+
+            if active_logits.size(0) > 0:
+                return self.task_loss_fn(active_logits, active_labels)
+            else:
+                return torch.tensor(0.0, device=logits.device)
+
+        return self.task_loss_fn(flat_logits, flat_labels)
+
+    def _compute_distillation_loss(
+        self,
+        student_logits: torch.Tensor,
+        teacher_logits: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Compute knowledge distillation loss using KL divergence."""
+        # Apply temperature scaling
+        student_soft = F.log_softmax(student_logits / self.temperature, dim=-1)
+        teacher_soft = F.softmax(teacher_logits / self.temperature, dim=-1)
+
+        # Compute KL divergence
+        if attention_mask is not None:
+            # Mask out padding tokens
+            mask = attention_mask.unsqueeze(-1).expand_as(student_soft)
+            student_soft_masked = student_soft * mask
+            teacher_soft_masked = teacher_soft * mask
+
+            # Compute loss only on valid tokens
+            kl_loss = self.kl_loss_fn(student_soft_masked, teacher_soft_masked)
+        else:
+            kl_loss = self.kl_loss_fn(student_soft, teacher_soft)
+
+        # Scale by temperature squared (as per distillation literature)
+        return kl_loss * (self.temperature**2)
+
+    def update_alpha(self, new_alpha: float) -> None:
+        """Update the alpha parameter for dynamic loss weighting."""
+        if not 0.0 <= new_alpha <= 1.0:
+            raise ValueError(f"Alpha must be between 0 and 1, got {new_alpha}")
+
+        old_alpha = self.alpha
+        self.alpha = new_alpha
+        self.config["alpha"] = new_alpha
+
+        logger.info(f"Updated alpha from {old_alpha} to {new_alpha}")
+
+    def update_temperature(self, new_temperature: float) -> None:
+        """Update the temperature parameter for dynamic distillation."""
+        if new_temperature <= 0.0:
+            raise ValueError(f"Temperature must be positive, got {new_temperature}")
+
+        old_temperature = self.temperature
+        self.temperature = new_temperature
+        self.config["temperature"] = new_temperature
+
+        logger.info(f"Updated temperature from {old_temperature} to {new_temperature}")
+
+
+class AdaptiveDistillationLoss(ModularDistillationLoss):
+    """
+    Adaptive distillation loss that adjusts alpha based on training progress.
+    """
+
+    def __init__(
+        self,
+        alpha: float = 0.5,
+        temperature: float = 3.0,
+        adaptation_strategy: str = "linear_decay",
+        adaptation_config: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize adaptive distillation loss.
+
+        Args:
+            alpha: Initial alpha value
+            temperature: Temperature for distillation
+            adaptation_strategy: Strategy for adapting alpha ("linear_decay", "cosine_decay", "step_decay")
+            adaptation_config: Configuration for adaptation strategy
+        """
+        super().__init__(alpha, temperature)
+
+        self.initial_alpha = alpha
+        self.adaptation_strategy = adaptation_strategy
+        self.adaptation_config = adaptation_config or {}
+        self.step_count = 0
+
+        logger.info(
+            f"Initialized AdaptiveDistillationLoss with strategy: {adaptation_strategy}"
+        )
+
+    def compute(self, *args, **kwargs) -> torch.Tensor:
+        """Compute loss with adaptive alpha adjustment."""
+        # Update alpha based on training progress
+        self._update_alpha_adaptive()
+
+        # Increment step count
+        self.step_count += 1
+
+        return super().compute(*args, **kwargs)
+
+    def _update_alpha_adaptive(self) -> None:
+        """Update alpha based on the selected adaptation strategy."""
+        total_steps = self.adaptation_config.get("total_steps", 1000)
+
+        if self.adaptation_strategy == "linear_decay":
+            # Linearly decay alpha from initial value to 0
+            progress = min(self.step_count / total_steps, 1.0)
+            new_alpha = self.initial_alpha * (1.0 - progress)
+
+        elif self.adaptation_strategy == "cosine_decay":
+            # Cosine decay
+            import math
+
+            progress = min(self.step_count / total_steps, 1.0)
+            new_alpha = self.initial_alpha * (1 + math.cos(math.pi * progress)) / 2
+
+        elif self.adaptation_strategy == "step_decay":
+            # Step decay at specified intervals
+            decay_steps = self.adaptation_config.get("decay_steps", [500, 750])
+            decay_factor = self.adaptation_config.get("decay_factor", 0.5)
+
+            new_alpha = self.initial_alpha
+            for decay_step in decay_steps:
+                if self.step_count >= decay_step:
+                    new_alpha *= decay_factor
+        else:
+            # No adaptation
+            new_alpha = self.alpha
+
+        # Update if changed significantly
+        if abs(new_alpha - self.alpha) > 1e-6:
+            self.alpha = new_alpha

src/multi_agent_training/modular_trainer.py

@@ -0,0 +1,229 @@
+"""
+Refactored LoRA Knowledge Distillation Trainer using modular architecture.
+
+This module implements a clean, testable trainer that follows the interface contracts
+and provides better separation of concerns.
+"""
+
+import logging
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+from ..core.base_components import BaseTrainer
+from ..core.exceptions import TrainingError
+from ..core.interfaces import TrainingConfig
+
+logger = logging.getLogger(__name__)
+
+
+class ModularLoRATrainer(BaseTrainer):
+    """Modular LoRA trainer with clean separation of concerns."""
+
+    def __init__(
+        self,
+        model: nn.Module,
+        optimizer: torch.optim.Optimizer,
+        loss_function,
+        device: str = "cpu",
+        teacher_model: Optional[nn.Module] = None,
+    ):
+        """
+        Initialize the modular LoRA trainer.
+
+        Args:
+            model: Student model to train
+            optimizer: Optimizer for training
+            loss_function: Loss function implementing ILossFunction
+            device: Training device
+            teacher_model: Optional teacher model for distillation
+        """
+        super().__init__(model, optimizer, device)
+        self.loss_function = loss_function
+        self.teacher_model = teacher_model
+        if self.teacher_model:
+            self.teacher_model.to(self.device)
+            self.teacher_model.eval()
+
+        self.custom_loss_fn = None
+
+    def set_custom_loss_fn(self, loss_fn):
+        """Set custom loss function for specialized training."""
+        self.custom_loss_fn = loss_fn
+
+    def compute_distillation_loss(self, student_outputs, teacher_outputs, batch):
+        """Compute standard distillation loss."""
+        return self.loss_function.compute(
+            student_outputs.logits,
+            (
+                teacher_outputs.logits
+                if hasattr(teacher_outputs, "logits")
+                else teacher_outputs
+            ),
+            labels=batch.get("labels"),
+        )
+
+    def train(self, dataloader: DataLoader, config: TrainingConfig) -> Dict[str, Any]:
+        """
+        Train the model with the given configuration.
+
+        Args:
+            dataloader: Training data loader
+            config: Training configuration
+
+        Returns:
+            Training results and metrics
+        """
+        try:
+            self.model.train()
+            total_loss = 0.0
+            num_batches = 0
+            training_metrics = {}
+
+            for epoch in range(config.num_epochs):
+                epoch_loss = 0.0
+                epoch_batches = 0
+
+                for batch_idx, batch in enumerate(dataloader):
+                    # Move batch to device
+                    batch = self._move_batch_to_device(batch)
+
+                    # Forward pass
+                    self.optimizer.zero_grad()
+
+                    # Student model forward pass
+                    student_outputs = self.model(**batch)
+
+                    # Teacher model forward pass (if available)
+                    teacher_outputs = None
+                    if self.teacher_model:
+                        with torch.no_grad():
+                            teacher_outputs = self.teacher_model(**batch)
+
+                    # Compute loss
+                    if self.custom_loss_fn:
+                        loss = self.custom_loss_fn(
+                            student_outputs, teacher_outputs, batch
+                        )
+                    else:
+                        loss = self.loss_function.compute(
+                            (
+                                student_outputs.logits
+                                if hasattr(student_outputs, "logits")
+                                else student_outputs
+                            ),
+                            batch.get("labels", batch.get("input_ids")),
+                        )
+
+                    # Backward pass
+                    loss.backward()
+                    self.optimizer.step()
+
+                    # Track metrics
+                    epoch_loss += loss.item()
+                    epoch_batches += 1
+
+                    # Log training step
+                    if batch_idx % config.save_steps == 0:
+                        step_metrics = self.loss_function.get_metrics()
+                        self._log_training_step(
+                            epoch, batch_idx, loss.item(), step_metrics
+                        )
+
+                        logger.info(
+                            f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}"
+                        )
+
+                # End of epoch
+                avg_epoch_loss = (
+                    epoch_loss / epoch_batches if epoch_batches > 0 else 0.0
+                )
+                total_loss += epoch_loss
+                num_batches += epoch_batches
+
+                training_metrics[f"epoch_{epoch}_loss"] = avg_epoch_loss
+
+                # Save checkpoint
+                if epoch % config.save_steps == 0:
+                    checkpoint_path = (
+                        Path(config.output_dir) / f"checkpoint_epoch_{epoch}.pt"
+                    )
+                    self.save_checkpoint(checkpoint_path, epoch)
+
+            # Final results
+            avg_loss = total_loss / num_batches if num_batches > 0 else 0.0
+
+            results = {
+                "average_loss": avg_loss,
+                "total_epochs": config.num_epochs,
+                "total_batches": num_batches,
+                "training_metrics": training_metrics,
+                "loss_function_metrics": self.loss_function.get_metrics(),
+            }
+
+            logger.info(f"Training completed. Average loss: {avg_loss:.4f}")
+            return results
+
+        except Exception as e:
+            raise TrainingError(
+                f"Training failed: {str(e)}",
+                "TRAINING_FAILED",
+                {"epoch": getattr(self, "current_epoch", 0)},
+            )
+
+    def evaluate(self, dataloader: DataLoader) -> Dict[str, float]:
+        """
+        Evaluate the model on the given dataset.
+
+        Args:
+            dataloader: Evaluation data loader
+
+        Returns:
+            Evaluation metrics
+        """
+        try:
+            self.model.eval()
+            total_loss = 0.0
+            num_batches = 0
+
+            with torch.no_grad():
+                for batch in dataloader:
+                    batch = self._move_batch_to_device(batch)
+
+                    # Forward pass
+                    outputs = self.model(**batch)
+
+                    # Compute loss
+                    loss = self.loss_function.compute(
+                        outputs.logits if hasattr(outputs, "logits") else outputs,
+                        batch.get("labels", batch.get("input_ids")),
+                    )
+
+                    total_loss += loss.item()
+                    num_batches += 1
+
+            avg_loss = total_loss / num_batches if num_batches > 0 else 0.0
+
+            results = {"eval_loss": avg_loss, "eval_batches": num_batches}
+            results.update(self.loss_function.get_metrics())
+
+            logger.info(f"Evaluation completed. Average loss: {avg_loss:.4f}")
+            return results
+
+        except Exception as e:
+            raise TrainingError(f"Evaluation failed: {str(e)}", "EVALUATION_FAILED")
+
+    def _move_batch_to_device(
+        self, batch: Dict[str, torch.Tensor]
+    ) -> Dict[str, torch.Tensor]:
+        """Move batch tensors to the training device."""
+        device_batch = {}
+        for key, value in batch.items():
+            if isinstance(value, torch.Tensor):
+                device_batch[key] = value.to(self.device)
+            else:
+                device_batch[key] = value
+        return device_batch

src/multi_agent_training/multi_agent_trainer.py

@@ -0,0 +1,547 @@
+#!/usr/bin/env python3
+"""
+CPU-Optimized Multi-Agent Trainer
+
+This module provides comprehensive multi-agent training capabilities optimized for CPU execution,
+including LoRA fine-tuning, agent-specific conditioning, and integration with existing training infrastructure.
+"""
+
+import os
+import json
+import math
+import random
+import logging
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Any, Union
+from dataclasses import dataclass, field
+
+import torch
+import yaml
+from datasets import DatasetDict, Dataset
+from transformers import (
+    AutoModelForCausalLM, 
+    AutoTokenizer, 
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling
+)
+from trl import SFTTrainer
+from peft import LoraConfig, get_peft_model, TaskType
+from huggingface_hub import HfApi, create_repo
+
+from ..multi_agent_datasets.multi_agent_loader import MultiAgentDatasetLoader, MultiAgentDatasetConfig
+from ..multi_agent_tokenization.agent_tokenizer import AgentTokenManager, AgentTokenConfig, AgentTokenizer
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class MultiAgentTrainingConfig:
+    """Configuration for multi-agent training"""
+    # Model configuration
+    base_model: str = "microsoft/Phi-3.5-MoE-instruct"
+    model_cache_dir: Optional[str] = None
+    trust_remote_code: bool = True
+    
+    # Training configuration
+    output_dir: str = "./outputs"
+    max_steps: int = 50
+    num_train_epochs: int = 1
+    per_device_train_batch_size: int = 1
+    per_device_eval_batch_size: int = 1
+    gradient_accumulation_steps: int = 8
+    learning_rate: float = 2e-5
+    lr_scheduler_type: str = "cosine"
+    warmup_steps: int = 0
+    
+    # LoRA configuration
+    lora_r: int = 8
+    lora_alpha: int = 16
+    lora_dropout: float = 0.05
+    lora_target_modules: str = "all-linear"
+    lora_bias: str = "none"
+    
+    # CPU optimization
+    use_cpu: bool = True
+    bf16: bool = False
+    fp16: bool = False
+    gradient_checkpointing: bool = True
+    dataloader_num_workers: int = 0
+    remove_unused_columns: bool = False
+    
+    # Multi-agent specific
+    agent_prefix: str = "<|agent:"
+    agent_suffix: str = "|>"
+    balance_agents: bool = False
+    balance_cap: Optional[int] = None
+    
+    # Logging and monitoring
+    logging_steps: int = 5
+    save_steps: int = 50
+    eval_steps: int = 25
+    save_total_limit: int = 1
+    logging_dir: str = "./logs"
+    report_to: str = "none"
+    
+    # Hugging Face Hub
+    hub_repo_id: Optional[str] = None
+    push_to_hub: bool = False
+    hub_token: Optional[str] = None
+    
+    # Dataset configuration
+    dataset_config: Optional[MultiAgentDatasetConfig] = None
+
+class CPUOptimizedMultiAgentTrainer:
+    """
+    CPU-optimized multi-agent trainer with LoRA fine-tuning
+    """
+    
+    def __init__(self, config: MultiAgentTrainingConfig):
+        self.config = config
+        self.tokenizer: Optional[AutoTokenizer] = None
+        self.model: Optional[torch.nn.Module] = None
+        self.agent_manager: Optional[AgentTokenManager] = None
+        self.dataset_loader: Optional[MultiAgentDatasetLoader] = None
+        self.trainer: Optional[SFTTrainer] = None
+        self.agents: List[str] = []
+        self.dataset_stats: Dict[str, Any] = {}
+        
+        # Setup logging
+        self._setup_logging()
+        
+    def _setup_logging(self):
+        """Setup logging configuration"""
+        log_level = logging.INFO
+        logging.basicConfig(
+            level=log_level,
+            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+            handlers=[
+                logging.StreamHandler(),
+                logging.FileHandler(os.path.join(self.config.logging_dir, 'training.log'))
+            ]
+        )
+        
+        # Create logs directory
+        os.makedirs(self.config.logging_dir, exist_ok=True)
+        
+    def load_model_and_tokenizer(self) -> Tuple[AutoTokenizer, torch.nn.Module]:
+        """Load model and tokenizer optimized for CPU"""
+        logger.info(f"Loading model and tokenizer: {self.config.base_model}")
+        
+        # Load tokenizer
+        tokenizer_kwargs = {
+            "trust_remote_code": self.config.trust_remote_code,
+            "cache_dir": self.config.model_cache_dir
+        }
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.config.base_model,
+            **tokenizer_kwargs
+        )
+        
+        # Configure tokenizer for CPU training
+        self.tokenizer.model_max_length = 2048
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.unk_token or self.tokenizer.eos_token
+            self.tokenizer.pad_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.pad_token)
+        self.tokenizer.padding_side = "right"
+        
+        # Load model with CPU optimizations
+        model_kwargs = {
+            "trust_remote_code": self.config.trust_remote_code,
+            "torch_dtype": torch.float32,  # Use float32 for CPU
+            "device_map": "cpu",
+            "attn_implementation": "eager",  # Force CPU-compatible attention
+            "use_cache": False,  # Disable cache for training
+            "cache_dir": self.config.model_cache_dir
+        }
+        
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.config.base_model,
+            **model_kwargs
+        )
+        
+        logger.info(f"Model loaded with {self.model.num_parameters():,} parameters")
+        return self.tokenizer, self.model
+    
+    def setup_agent_tokens(self, agents: List[str]) -> AgentTokenManager:
+        """Setup agent token management"""
+        logger.info(f"Setting up agent tokens for {len(agents)} agents")
+        
+        agent_config = AgentTokenConfig(
+            agent_prefix=self.config.agent_prefix,
+            agent_suffix=self.config.agent_suffix,
+            resize_embeddings=True
+        )
+        
+        self.agent_manager = AgentTokenManager(agent_config)
+        
+        # Add agent tokens to tokenizer
+        self.tokenizer, agent_tokens = self.agent_manager.add_agent_tokens_to_tokenizer(
+            self.tokenizer, agents
+        )
+        
+        # Resize model embeddings
+        self.model = self.agent_manager.resize_model_embeddings(self.model, self.tokenizer)
+        
+        logger.info(f"Agent tokens setup complete. Tokens: {agent_tokens}")
+        return self.agent_manager
+    
+    def load_dataset(self, dataset_path: str) -> Tuple[DatasetDict, List[str], Dict[str, Any]]:
+        """Load and process multi-agent dataset"""
+        logger.info(f"Loading dataset from: {dataset_path}")
+        
+        # Create dataset configuration
+        if self.config.dataset_config is None:
+            dataset_config = MultiAgentDatasetConfig(
+                dataset_path=dataset_path,
+                agent_prefix=self.config.agent_prefix,
+                agent_suffix=self.config.agent_suffix,
+                balance_agents=self.config.balance_agents,
+                balance_cap=self.config.balance_cap
+            )
+        else:
+            dataset_config = self.config.dataset_config
+            dataset_config.dataset_path = dataset_path
+        
+        # Create dataset loader
+        self.dataset_loader = MultiAgentDatasetLoader(dataset_config)
+        
+        # Load and process dataset
+        dataset, agents, stats = self.dataset_loader.load_and_process(self.tokenizer)
+        
+        self.agents = agents
+        self.dataset_stats = stats
+        
+        logger.info(f"Dataset loaded: {len(agents)} agents, {stats['total_samples']} samples")
+        return dataset, agents, stats
+    
+    def create_lora_config(self) -> LoraConfig:
+        """Create LoRA configuration optimized for CPU"""
+        logger.info("Creating LoRA configuration")
+        
+        lora_config = LoraConfig(
+            r=self.config.lora_r,
+            lora_alpha=self.config.lora_alpha,
+            lora_dropout=self.config.lora_dropout,
+            bias=self.config.lora_bias,
+            task_type=TaskType.CAUSAL_LM,
+            target_modules=self.config.lora_target_modules
+        )
+        
+        logger.info(f"LoRA config: r={lora_config.r}, alpha={lora_config.lora_alpha}, dropout={lora_config.lora_dropout}")
+        return lora_config
+    
+    def create_training_arguments(self) -> TrainingArguments:
+        """Create training arguments optimized for CPU"""
+        logger.info("Creating training arguments")
+        
+        training_args = TrainingArguments(
+            output_dir=self.config.output_dir,
+            overwrite_output_dir=True,
+            num_train_epochs=self.config.num_train_epochs,
+            max_steps=self.config.max_steps,
+            per_device_train_batch_size=self.config.per_device_train_batch_size,
+            per_device_eval_batch_size=self.config.per_device_eval_batch_size,
+            gradient_accumulation_steps=self.config.gradient_accumulation_steps,
+            learning_rate=self.config.learning_rate,
+            lr_scheduler_type=self.config.lr_scheduler_type,
+            warmup_steps=self.config.warmup_steps,
+            
+            # CPU optimizations
+            bf16=self.config.bf16,
+            fp16=self.config.fp16,
+            gradient_checkpointing=self.config.gradient_checkpointing,
+            dataloader_num_workers=self.config.dataloader_num_workers,
+            remove_unused_columns=self.config.remove_unused_columns,
+            
+            # Logging and saving
+            logging_steps=self.config.logging_steps,
+            save_steps=self.config.save_steps,
+            eval_steps=self.config.eval_steps,
+            save_total_limit=self.config.save_total_limit,
+            logging_dir=self.config.logging_dir,
+            report_to=self.config.report_to,
+            
+            # Evaluation
+            evaluation_strategy="steps" if self.config.eval_steps > 0 else "no",
+            
+            # Optimization
+            optim="adamw_torch",
+            weight_decay=0.01,
+            max_grad_norm=1.0,
+            
+            # Hub integration
+            push_to_hub=self.config.push_to_hub,
+            hub_model_id=self.config.hub_repo_id,
+            hub_token=self.config.hub_token,
+        )
+        
+        logger.info(f"Training arguments created: {training_args.output_dir}")
+        return training_args
+    
+    def create_trainer(self, dataset: DatasetDict, lora_config: LoraConfig, training_args: TrainingArguments) -> SFTTrainer:
+        """Create SFT trainer for multi-agent training"""
+        logger.info("Creating SFT trainer")
+        
+        # Get training and evaluation datasets
+        train_dataset = dataset["train"]
+        eval_dataset = dataset.get("test", None)
+        
+        # Create trainer
+        self.trainer = SFTTrainer(
+            model=self.model,
+            args=training_args,
+            peft_config=lora_config,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            tokenizer=self.tokenizer,
+            max_seq_length=2048,
+            dataset_text_field="text",
+            packing=True,  # Enable packing for efficiency
+            data_collator=None,  # Use default
+        )
+        
+        logger.info("SFT trainer created successfully")
+        return self.trainer
+    
+    def train(self) -> Dict[str, Any]:
+        """Execute training process"""
+        logger.info("Starting training process")
+        
+        if self.trainer is None:
+            raise ValueError("Trainer not initialized. Call create_trainer() first.")
+        
+        # Start training
+        training_result = self.trainer.train()
+        
+        # Save model and tokenizer
+        self.save_model()
+        
+        # Save agent tokens
+        if self.agent_manager:
+            self.agent_manager.save_agent_tokens(self.config.output_dir)
+        
+        # Generate training report
+        report = self.generate_training_report(training_result)
+        
+        logger.info("Training completed successfully")
+        return report
+    
+    def save_model(self):
+        """Save trained model and tokenizer"""
+        logger.info(f"Saving model to {self.config.output_dir}")
+        
+        os.makedirs(self.config.output_dir, exist_ok=True)
+        
+        # Save model
+        self.trainer.model.save_pretrained(self.config.output_dir)
+        
+        # Save tokenizer
+        self.tokenizer.save_pretrained(self.config.output_dir)
+        
+        # Save training configuration
+        config_file = os.path.join(self.config.output_dir, "training_config.json")
+        with open(config_file, 'w') as f:
+            json.dump(self.config.__dict__, f, indent=2, default=str)
+        
+        logger.info("Model saved successfully")
+    
+    def generate_training_report(self, training_result: Any) -> Dict[str, Any]:
+        """Generate comprehensive training report"""
+        report = {
+            "training_config": self.config.__dict__,
+            "dataset_stats": self.dataset_stats,
+            "agents": self.agents,
+            "agent_tokens": self.agent_manager.get_agent_statistics() if self.agent_manager else {},
+            "training_metrics": {
+                "train_loss": getattr(training_result, 'train_loss', None),
+                "train_runtime": getattr(training_result, 'train_runtime', None),
+                "train_samples_per_second": getattr(training_result, 'train_samples_per_second', None),
+                "train_steps_per_second": getattr(training_result, 'train_steps_per_second', None),
+            },
+            "model_info": {
+                "base_model": self.config.base_model,
+                "num_parameters": self.model.num_parameters() if self.model else None,
+                "vocab_size": len(self.tokenizer) if self.tokenizer else None,
+            }
+        }
+        
+        # Save report
+        report_file = os.path.join(self.config.output_dir, "training_report.json")
+        with open(report_file, 'w') as f:
+            json.dump(report, f, indent=2, default=str)
+        
+        logger.info(f"Training report saved to {report_file}")
+        return report
+    
+    def push_to_hub(self, repo_id: Optional[str] = None, commit_message: str = "Multi-agent LoRA adapter"):
+        """Push trained model to Hugging Face Hub"""
+        if not self.config.push_to_hub:
+            logger.info("Push to hub disabled")
+            return
+        
+        repo_id = repo_id or self.config.hub_repo_id
+        if not repo_id:
+            raise ValueError("Repository ID not specified")
+        
+        if not self.config.hub_token:
+            raise ValueError("Hub token not provided")
+        
+        logger.info(f"Pushing model to Hub: {repo_id}")
+        
+        # Create repository
+        create_repo(repo_id, repo_type="model", exist_ok=True, token=self.config.hub_token)
+        
+        # Upload model
+        api = HfApi(token=self.config.hub_token)
+        api.upload_folder(
+            folder_path=self.config.output_dir,
+            repo_id=repo_id,
+            repo_type="model",
+            commit_message=commit_message,
+            allow_patterns=["*.json", "*.md", "*.bin", "*.yaml", "*.txt"]
+        )
+        
+        logger.info(f"Model pushed to https://huggingface.co/{repo_id}")
+    
+    def create_readme(self) -> str:
+        """Create README for the trained model"""
+        readme_content = f"""# Multi-Agent LoRA Adapter for {self.config.base_model}
+
+## Overview
+This is a LoRA (Low-Rank Adaptation) adapter trained for multi-agent scenarios using {self.config.base_model}.
+
+## Agent Conditioning Tokens
+This adapter expects agent-specific tokens to condition the model behavior:
+
+"""
+        
+        if self.agents:
+            for agent in self.agents:
+                token = f"{self.config.agent_prefix}{agent}{self.config.agent_suffix}"
+                readme_content += f"- `{token}` - {agent} agent\n"
+            
+            readme_content += f"""
+## Usage Example
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
+
+# Load base model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("{self.config.base_model}")
+model = AutoModelForCausalLM.from_pretrained("{self.config.base_model}")
+
+# Load LoRA adapter
+model = PeftModel.from_pretrained(model, "{self.config.hub_repo_id}")
+
+# Example usage
+prompt = "How do I implement a binary search algorithm?"
+agent_token = "{self.config.agent_prefix}SWE{self.config.agent_suffix}\\n"
+full_prompt = agent_token + prompt
+
+inputs = tokenizer(full_prompt, return_tensors="pt")
+outputs = model.generate(**inputs, max_new_tokens=100)
+response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+print(response)
+```
+
+## Training Configuration
+- **Base Model**: {self.config.base_model}
+- **LoRA Rank**: {self.config.lora_r}
+- **LoRA Alpha**: {self.config.lora_alpha}
+- **Learning Rate**: {self.config.learning_rate}
+- **Max Steps**: {self.config.max_steps}
+- **Batch Size**: {self.config.per_device_train_batch_size}
+
+## Dataset Statistics
+- **Total Samples**: {self.dataset_stats.get('total_samples', 'N/A')}
+- **Agents**: {', '.join(self.agents) if self.agents else 'N/A'}
+
+## License
+This model is released under the same license as the base model.
+"""
+        else:
+            readme_content += "No specific agents were configured for this adapter.\n"
+        
+        # Save README
+        readme_file = os.path.join(self.config.output_dir, "README.md")
+        with open(readme_file, 'w') as f:
+            f.write(readme_content)
+        
+        logger.info(f"README created: {readme_file}")
+        return readme_file
+
+class MultiAgentTrainingPipeline:
+    """
+    Complete pipeline for multi-agent training
+    """
+    
+    def __init__(self, config: MultiAgentTrainingConfig):
+        self.config = config
+        self.trainer = CPUOptimizedMultiAgentTrainer(config)
+    
+    def run_training(self, dataset_path: str) -> Dict[str, Any]:
+        """Run complete training pipeline"""
+        logger.info("Starting multi-agent training pipeline")
+        
+        try:
+            # Load model and tokenizer
+            self.trainer.load_model_and_tokenizer()
+            
+            # Load dataset
+            dataset, agents, stats = self.trainer.load_dataset(dataset_path)
+            
+            # Setup agent tokens
+            self.trainer.setup_agent_tokens(agents)
+            
+            # Create LoRA config
+            lora_config = self.trainer.create_lora_config()
+            
+            # Create training arguments
+            training_args = self.trainer.create_training_arguments()
+            
+            # Create trainer
+            self.trainer.create_trainer(dataset, lora_config, training_args)
+            
+            # Train model
+            training_result = self.trainer.train()
+            
+            # Create README
+            self.trainer.create_readme()
+            
+            # Push to hub if configured
+            if self.config.push_to_hub:
+                self.trainer.push_to_hub()
+            
+            logger.info("Training pipeline completed successfully")
+            return training_result
+            
+        except Exception as e:
+            logger.error(f"Training pipeline failed: {e}")
+            raise
+
+# Example usage and testing
+if __name__ == "__main__":
+    # Configure logging
+    logging.basicConfig(level=logging.INFO)
+    
+    # Example configuration
+    config = MultiAgentTrainingConfig(
+        base_model="microsoft/Phi-3.5-MoE-instruct",
+        output_dir="./outputs/multi_agent_test",
+        max_steps=10,  # Small number for testing
+        hub_repo_id="test/multi-agent-adapter",
+        push_to_hub=False  # Set to True for actual deployment
+    )
+    
+    # Create training pipeline
+    pipeline = MultiAgentTrainingPipeline(config)
+    
+    try:
+        # Run training (would need actual dataset path)
+        # result = pipeline.run_training("/path/to/dataset")
+        print("Multi-agent training pipeline ready")
+        
+    except Exception as e:
+        print(f"Error: {e}")

src/multi_agent_training/specialized/curiosity_module.py

@@ -0,0 +1,516 @@
+"""
+Curiosity Training Module for MangoMAS Local
+
+This module implements specialized training for curiosity and exploration capabilities,
+adapted from the AWS backup system for local training.
+"""
+
+import json
+import logging
+import os
+import random
+from typing import Any, Dict, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+
+from ..core_framework import SpecializedTrainingModule, TrainingModuleConfig
+
+logger = logging.getLogger(__name__)
+
+
+class CuriosityDataset(Dataset):
+    """Dataset for training curiosity and exploration capabilities."""
+
+    def __init__(self, data_path: str, tokenizer, max_length: int = 768):
+        """
+        Initialize the curiosity dataset.
+
+        Args:
+            data_path: Path to the curiosity data file
+            tokenizer: Tokenizer for text processing
+            max_length: Maximum sequence length
+        """
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.data = self._load_data(data_path)
+
+        logger.info(f"Loaded curiosity dataset with {len(self.data)} examples")
+
+    def _load_data(self, data_path: str) -> List[Dict]:
+        """Load curiosity training data."""
+        data = []
+        with open(data_path, "r", encoding="utf-8") as f:
+            for line in f:
+                try:
+                    item = json.loads(line.strip())
+                    # Validate required fields for curiosity data
+                    if (
+                        "scenario" in item
+                        and "curiosity_questions" in item
+                        and "exploration_directions" in item
+                    ):
+                        data.append(item)
+                except (json.JSONDecodeError, KeyError) as e:
+                    logger.warning(f"Skipping invalid curiosity data: {e}")
+        return data
+
+    def __len__(self) -> int:
+        """Return the number of examples in the dataset."""
+        return len(self.data)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        """Get a training example."""
+        item = self.data[idx]
+
+        # Format the curiosity example
+        scenario = item["scenario"]
+        curiosity_questions = item["curiosity_questions"]
+        exploration_directions = item["exploration_directions"]
+
+        # Format as text
+        text = f"Scenario: {scenario}\n\n"
+
+        text += "Curiosity Questions:\n"
+        for i, question in enumerate(curiosity_questions):
+            text += f"{i+1}. {question}\n"
+        text += "\n"
+
+        text += "Exploration Directions:\n"
+        for i, direction in enumerate(exploration_directions):
+            text += f"{i+1}. {direction}\n"
+
+        # Tokenize
+        encoding = self.tokenizer(
+            text,
+            truncation=True,
+            padding="max_length",
+            max_length=self.max_length,
+            return_tensors="pt",
+        )
+
+        return {
+            "input_ids": encoding["input_ids"].squeeze(),
+            "attention_mask": encoding["attention_mask"].squeeze(),
+            "labels": encoding["input_ids"].squeeze().clone(),
+            "scenario": scenario,
+            "curiosity_questions": curiosity_questions,
+            "exploration_directions": exploration_directions,
+        }
+
+
+class CuriosityEvaluator:
+    """Evaluator for curiosity and exploration capabilities."""
+
+    def __init__(self, tokenizer):
+        """
+        Initialize the curiosity evaluator.
+
+        Args:
+            tokenizer: Tokenizer for text processing
+        """
+        self.tokenizer = tokenizer
+        self.metrics = {
+            "question_diversity": 0.0,
+            "exploration_breadth": 0.0,
+            "uncertainty_identification": 0.0,
+            "assumption_challenging": 0.0,
+        }
+
+        # Question starters for evaluating diversity
+        self.question_starters = [
+            "what",
+            "how",
+            "why",
+            "when",
+            "where",
+            "who",
+            "which",
+            "could",
+            "would",
+            "is",
+            "are",
+            "do",
+            "does",
+            "have",
+            "has",
+        ]
+
+        # Exploration markers
+        self.exploration_markers = [
+            "alternative",
+            "perspective",
+            "consider",
+            "explore",
+            "investigate",
+            "possibility",
+            "approach",
+            "angle",
+            "viewpoint",
+            "scenario",
+        ]
+
+        # Uncertainty phrases
+        self.uncertainty_phrases = [
+            "unclear",
+            "unknown",
+            "uncertain",
+            "not sure",
+            "ambiguous",
+            "might be",
+            "could be",
+            "possibly",
+            "perhaps",
+            "may",
+        ]
+
+        # Assumption challenging phrases
+        self.assumption_phrases = [
+            "assuming",
+            "assumption",
+            "presuppose",
+            "presupposition",
+            "take for granted",
+            "implicit",
+            "unstated",
+            "underlying",
+        ]
+
+    def evaluate(self, model, eval_dataset: CuriosityDataset) -> Dict[str, float]:
+        """
+        Evaluate curiosity capabilities on the provided dataset.
+
+        Args:
+            model: The model to evaluate
+            eval_dataset: Dataset of curiosity examples
+
+        Returns:
+            Dictionary of evaluation metrics
+        """
+        model.eval()
+        device = next(model.parameters()).device
+
+        # Reset metrics
+        for key in self.metrics:
+            self.metrics[key] = 0.0
+
+        total_examples = min(
+            len(eval_dataset), 50
+        )  # Limit to 50 examples for efficiency
+
+        with torch.no_grad():
+            for idx in range(total_examples):
+                example = eval_dataset[idx]
+                scenario = example["scenario"]
+
+                # Generate questions for the scenario
+                prompt = f"Scenario: {scenario}\n\nGenerate curious questions to explore this further:"
+                input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(
+                    device
+                )
+
+                generated_ids = model.generate(
+                    input_ids,
+                    max_length=256,
+                    temperature=0.8,  # Slightly higher temperature for creativity
+                    num_return_sequences=1,
+                )
+
+                generated_text = self.tokenizer.decode(
+                    generated_ids[0], skip_special_tokens=True
+                )
+
+                # Evaluate curiosity quality
+                self._evaluate_curiosity(
+                    scenario=scenario,
+                    expected_questions=example["curiosity_questions"],
+                    expected_directions=example["exploration_directions"],
+                    generated_text=generated_text,
+                )
+
+        # Calculate averages
+        for key in self.metrics:
+            self.metrics[key] /= total_examples
+
+        return self.metrics
+
+    def _evaluate_curiosity(
+        self,
+        scenario: str,
+        expected_questions: List[str],
+        expected_directions: List[str],
+        generated_text: str,
+    ) -> None:
+        """
+        Evaluate curiosity quality for a specific example.
+
+        Args:
+            scenario: The scenario to explore
+            expected_questions: Expected curiosity questions
+            expected_directions: Expected exploration directions
+            generated_text: The text generated by the model
+        """
+        # Extract questions from generated text (simple approach)
+        generated_questions = [
+            line.strip()
+            for line in generated_text.split("\n")
+            if line.strip().endswith("?")
+        ]
+
+        if not generated_questions:
+            # Try to extract numbered questions
+            for line in generated_text.split("\n"):
+                if any(f"{i}." in line for i in range(1, 10)) and "?" in line:
+                    generated_questions.append(line.strip())
+
+        # 1. Question diversity - variety of question types
+        starter_counts = {starter: 0 for starter in self.question_starters}
+        for question in generated_questions:
+            for starter in self.question_starters:
+                if (
+                    question.lower().startswith(starter)
+                    or f" {starter} " in question.lower()
+                ):
+                    starter_counts[starter] += 1
+
+        unique_starters = sum(1 for count in starter_counts.values() if count > 0)
+        self.metrics["question_diversity"] += min(
+            1.0, unique_starters / 5
+        )  # Normalize to 5 unique types
+
+        # 2. Exploration breadth - check for exploration markers
+        exploration_marker_count = sum(
+            1 for marker in self.exploration_markers if marker in generated_text.lower()
+        )
+        self.metrics["exploration_breadth"] += min(1.0, exploration_marker_count / 3)
+
+        # 3. Uncertainty identification - check for uncertainty phrases
+        uncertainty_phrase_count = sum(
+            1 for phrase in self.uncertainty_phrases if phrase in generated_text.lower()
+        )
+        self.metrics["uncertainty_identification"] += min(
+            1.0, uncertainty_phrase_count / 2
+        )
+
+        # 4. Assumption challenging - check for phrases that challenge assumptions
+        assumption_phrase_count = sum(
+            1 for phrase in self.assumption_phrases if phrase in generated_text.lower()
+        )
+        self.metrics["assumption_challenging"] += min(1.0, assumption_phrase_count / 1)
+
+
+class CuriosityTrainingModule(SpecializedTrainingModule):
+    """Specialized training module for curiosity and exploration capabilities."""
+
+    def __init__(self, config: TrainingModuleConfig, tokenizer):
+        """
+        Initialize the curiosity training module.
+
+        Args:
+            config: Module configuration
+            tokenizer: Tokenizer for text processing
+        """
+        super().__init__(config, tokenizer)
+
+        # Initialize curiosity-specific components
+        self.data_path = config.data_path or "data/processed/curiosity_train.jsonl"
+        self.evaluator = CuriosityEvaluator(tokenizer)
+
+        # Curiosity-specific loss with higher temperature
+        self.curiosity_temp = config.module_config.get("temperature", 1.5)
+        self.curiosity_loss = nn.CrossEntropyLoss(ignore_index=-100)
+
+        # Training metrics
+        self.metrics = {
+            "curiosity_loss": 0.0,
+            "question_generation_score": 0.0,
+            "exploration_score": 0.0,
+        }
+
+        logger.info(
+            f"Initialized curiosity training module with temperature: {self.curiosity_temp}"
+        )
+
+    def prepare_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """
+        Prepare a batch of data for curiosity training.
+
+        Args:
+            batch: The input batch from the dataloader
+
+        Returns:
+            Processed batch ready for curiosity training
+        """
+        # Extract curiosity-specific elements if they exist
+        if all(
+            key in batch
+            for key in ["scenario", "curiosity_questions", "exploration_directions"]
+        ):
+            # This is already a curiosity-specific batch
+            return batch
+
+        # For general conversation batches, we could extract potential exploration scenarios
+        # This is a simplified placeholder implementation
+        return batch
+
+    def compute_loss(
+        self, student_outputs: Any, teacher_outputs: Any, batch: Dict[str, torch.Tensor]
+    ) -> torch.Tensor:
+        """
+        Compute the curiosity-specific loss.
+
+        Args:
+            student_outputs: Outputs from the student model
+            teacher_outputs: Outputs from the teacher model
+            batch: The processed input batch
+
+        Returns:
+            Curiosity-specific loss tensor
+        """
+        # Get logits from outputs
+        student_logits = (
+            student_outputs.logits
+            if hasattr(student_outputs, "logits")
+            else student_outputs
+        )
+        teacher_logits = (
+            teacher_outputs.logits
+            if hasattr(teacher_outputs, "logits")
+            else teacher_outputs
+        )
+
+        # Standard distillation loss calculation
+        student_logits = student_logits[:, :-1, :].contiguous()
+        teacher_logits = teacher_logits[:, :-1, :].contiguous()
+        target_ids = batch["labels"][:, 1:].contiguous()
+
+        # For curiosity, we use a higher temperature to encourage more diverse outputs
+        temperature = self.curiosity_temp
+        kl_loss = F.kl_div(
+            F.log_softmax(student_logits / temperature, dim=-1),
+            F.softmax(teacher_logits / temperature, dim=-1),
+            reduction="batchmean",
+        ) * (temperature**2)
+
+        # Cross-entropy loss against labels
+        ce_loss = self.curiosity_loss(
+            student_logits.view(-1, student_logits.size(-1)), target_ids.view(-1)
+        )
+
+        # Combined loss with curiosity focus
+        # We weight KL divergence higher to encourage exploration
+        loss = 0.4 * ce_loss + 0.6 * kl_loss
+
+        # Update metrics
+        self.metrics["curiosity_loss"] = loss.item()
+
+        return loss
+
+    def get_metrics(self) -> Dict[str, float]:
+        """
+        Get metrics specific to curiosity training.
+
+        Returns:
+            Dictionary of metric names and values
+        """
+        return self.metrics
+
+    def generate_synthetic_curiosity_data(
+        self, output_path: str, num_samples: int = 1000
+    ) -> None:
+        """
+        Generate synthetic curiosity training data.
+
+        Args:
+            output_path: Path to save the generated data
+            num_samples: Number of samples to generate
+        """
+        # This is a simplified implementation based on the AWS backup
+        # In a full implementation, this would be much more sophisticated
+
+        curiosity_templates = [
+            {
+                "scenario": "A company is developing a new voice assistant technology.",
+                "curiosity_questions": [
+                    "How might this technology affect people's privacy in their homes?",
+                    "What unexpected ways might users interact with this technology?",
+                    "How could this technology evolve over the next five years?",
+                    "What ethical considerations might arise from widespread adoption?",
+                    "How might this technology affect different demographic groups differently?",
+                ],
+                "exploration_directions": [
+                    "Consider alternative interaction models beyond voice commands",
+                    "Explore potential integration with other smart home systems",
+                    "Investigate privacy-preserving design approaches",
+                    "Consider accessibility implications for diverse user groups",
+                    "Examine potential unintended consequences of ambient listening",
+                ],
+            },
+            {
+                "scenario": "Scientists have discovered a new species of deep-sea organism that can survive extreme pressure and temperature.",
+                "curiosity_questions": [
+                    "What adaptations allow this organism to survive such extreme conditions?",
+                    "Could these adaptations be applied to human technology or medicine?",
+                    "What might this discovery tell us about the possibility of life on other planets?",
+                    "How might climate change affect deep-sea ecosystems and this organism?",
+                    "What other undiscovered species might exist in similar environments?",
+                ],
+                "exploration_directions": [
+                    "Examine evolutionary pathways for extreme environment adaptation",
+                    "Consider biomimicry applications in engineering and materials science",
+                    "Explore implications for astrobiology and extraterrestrial life",
+                    "Investigate ecological relationships in extreme environments",
+                    "Consider ethical dimensions of deep-sea exploration and bioprospecting",
+                ],
+            },
+            {
+                "scenario": "A small town is experiencing rapid population growth due to remote workers relocating from urban areas.",
+                "curiosity_questions": [
+                    "How might this demographic shift affect the town's culture and community?",
+                    "What infrastructure challenges might arise from rapid population growth?",
+                    "How could this trend impact local housing prices and affordability?",
+                    "What economic opportunities and challenges might emerge?",
+                    "How might long-term residents and newcomers develop different perspectives?",
+                ],
+                "exploration_directions": [
+                    "Investigate similar historical population shifts and their outcomes",
+                    "Consider varying perspectives from different stakeholder groups",
+                    "Explore potential policy approaches to manage growth sustainably",
+                    "Examine social integration mechanisms between established and new residents",
+                    "Consider environmental impacts of changing land use patterns",
+                ],
+            },
+        ]
+
+        # Generate variations
+        output_data = []
+        for _ in range(num_samples):
+            template = random.choice(curiosity_templates)
+
+            # Create a variation to avoid exact duplicates
+            variation = template.copy()
+
+            # Add metadata
+            variation["metadata"] = {
+                "generated": True,
+                "timestamp": (
+                    torch.cuda.get_device_name(0)
+                    if torch.cuda.is_available()
+                    else "CPU"
+                ),
+                "requires_exploration": True,
+            }
+
+            output_data.append(variation)
+
+        # Save to file
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        with open(output_path, "w", encoding="utf-8") as f:
+            for item in output_data:
+                f.write(json.dumps(item) + "\n")
+
+        logger.info(
+            f"Generated {len(output_data)} synthetic curiosity examples at {output_path}"
+        )

src/multi_agent_training/specialized/empathy_module.py

@@ -0,0 +1,502 @@
+"""
+Empathy Training Module for MangoMAS Local
+
+This module implements specialized training for empathy and emotional intelligence,
+adapted from the AWS backup system for local training.
+"""
+
+import json
+import logging
+import os
+import random
+from typing import Any, Dict, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+
+from ..core_framework import SpecializedTrainingModule, TrainingModuleConfig
+
+logger = logging.getLogger(__name__)
+
+
+class EmpathyDataset(Dataset):
+    """Dataset for training empathy and emotional intelligence capabilities."""
+
+    def __init__(self, data_path: str, tokenizer, max_length: int = 768):
+        """
+        Initialize the empathy dataset.
+
+        Args:
+            data_path: Path to the empathy data file
+            tokenizer: Tokenizer for text processing
+            max_length: Maximum sequence length
+        """
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.data = self._load_data(data_path)
+
+        logger.info(f"Loaded empathy dataset with {len(self.data)} examples")
+
+    def _load_data(self, data_path: str) -> List[Dict]:
+        """Load empathy training data."""
+        data = []
+        with open(data_path, "r", encoding="utf-8") as f:
+            for line in f:
+                try:
+                    item = json.loads(line.strip())
+                    # Validate required fields for empathy data
+                    if (
+                        "user_message" in item
+                        and "emotional_state" in item
+                        and "empathetic_response" in item
+                    ):
+                        data.append(item)
+                except (json.JSONDecodeError, KeyError) as e:
+                    logger.warning(f"Skipping invalid empathy data: {e}")
+        return data
+
+    def __len__(self) -> int:
+        """Return the number of examples in the dataset."""
+        return len(self.data)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        """Get a training example."""
+        item = self.data[idx]
+
+        # Format the empathy example
+        user_message = item["user_message"]
+        emotional_state = item["emotional_state"]
+        empathetic_response = item["empathetic_response"]
+
+        # Additional fields if available
+        emotional_cues = item.get("emotional_cues", [])
+        context = item.get("context", "")
+
+        # Construct the text with empathy markers
+        text = f"User: {user_message}\n\n"
+
+        # Include emotional analysis section for training
+        text += f"Emotional State: {emotional_state}\n"
+
+        if emotional_cues:
+            text += "Emotional Cues:\n"
+            for i, cue in enumerate(emotional_cues):
+                text += f"{i+1}. {cue}\n"
+            text += "\n"
+
+        if context:
+            text += f"Context: {context}\n\n"
+
+        text += f"Empathetic Response: {empathetic_response}"
+
+        # Tokenize
+        encoding = self.tokenizer(
+            text,
+            truncation=True,
+            padding="max_length",
+            max_length=self.max_length,
+            return_tensors="pt",
+        )
+
+        return {
+            "input_ids": encoding["input_ids"].squeeze(),
+            "attention_mask": encoding["attention_mask"].squeeze(),
+            "labels": encoding["input_ids"].squeeze().clone(),
+            "user_message": user_message,
+            "emotional_state": emotional_state,
+            "empathetic_response": empathetic_response,
+        }
+
+
+class EmpathyEvaluator:
+    """Evaluator for empathy and emotional intelligence capabilities."""
+
+    def __init__(self, tokenizer):
+        """
+        Initialize the empathy evaluator.
+
+        Args:
+            tokenizer: Tokenizer for text processing
+        """
+        self.tokenizer = tokenizer
+        self.metrics = {
+            "emotional_recognition": 0.0,
+            "empathetic_language": 0.0,
+            "supportive_tone": 0.0,
+            "personalization": 0.0,
+        }
+
+        # Empathetic language markers
+        self.empathetic_phrases = [
+            "understand",
+            "feel",
+            "appreciate",
+            "recognize",
+            "acknowledge",
+            "must be",
+            "sounds like",
+            "seems like",
+            "I hear you",
+            "that's difficult",
+            "that's challenging",
+            "I'm sorry",
+            "thank you for sharing",
+            "I can imagine",
+        ]
+
+        # Emotional state categories
+        self.emotional_states = {
+            "positive": [
+                "happy",
+                "excited",
+                "grateful",
+                "proud",
+                "hopeful",
+                "inspired",
+            ],
+            "negative": [
+                "sad",
+                "angry",
+                "frustrated",
+                "anxious",
+                "disappointed",
+                "overwhelmed",
+            ],
+            "neutral": [
+                "confused",
+                "uncertain",
+                "curious",
+                "surprised",
+                "contemplative",
+            ],
+        }
+
+    def evaluate(self, model, eval_dataset: EmpathyDataset) -> Dict[str, float]:
+        """
+        Evaluate empathy capabilities on the provided dataset.
+
+        Args:
+            model: The model to evaluate
+            eval_dataset: Dataset of empathy examples
+
+        Returns:
+            Dictionary of evaluation metrics
+        """
+        model.eval()
+        device = next(model.parameters()).device
+
+        # Reset metrics
+        for key in self.metrics:
+            self.metrics[key] = 0.0
+
+        total_examples = min(
+            len(eval_dataset), 50
+        )  # Limit to 50 examples for efficiency
+
+        with torch.no_grad():
+            for idx in range(total_examples):
+                example = eval_dataset[idx]
+                user_message = example["user_message"]
+                expected_emotional_state = example["emotional_state"]
+
+                # Generate response without providing emotional state
+                prompt = f"User: {user_message}\n\nProvide an empathetic response:"
+                input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(
+                    device
+                )
+
+                generated_ids = model.generate(
+                    input_ids, max_length=256, temperature=0.7, num_return_sequences=1
+                )
+
+                generated_text = self.tokenizer.decode(
+                    generated_ids[0], skip_special_tokens=True
+                )
+
+                # Evaluate empathy quality
+                self._evaluate_empathy(
+                    user_message=user_message,
+                    expected_emotional_state=expected_emotional_state,
+                    expected_response=example["empathetic_response"],
+                    generated_response=generated_text,
+                )
+
+        # Calculate averages
+        for key in self.metrics:
+            self.metrics[key] /= total_examples
+
+        return self.metrics
+
+    def _evaluate_empathy(
+        self,
+        user_message: str,
+        expected_emotional_state: str,
+        expected_response: str,
+        generated_response: str,
+    ) -> None:
+        """
+        Evaluate empathy quality for a specific example.
+
+        Args:
+            user_message: The user's message
+            expected_emotional_state: Expected identified emotional state
+            expected_response: Expected empathetic response
+            generated_response: The response generated by the model
+        """
+        # 1. Emotional recognition - check if response acknowledges correct emotion
+        emotional_category = None
+        for category, emotions in self.emotional_states.items():
+            if any(emotion in expected_emotional_state.lower() for emotion in emotions):
+                emotional_category = category
+                break
+
+        if emotional_category:
+            # Check if response contains words matching the emotional category
+            emotion_words = self.emotional_states[emotional_category]
+            emotion_recognition = any(
+                word in generated_response.lower() for word in emotion_words
+            )
+            self.metrics["emotional_recognition"] += 1.0 if emotion_recognition else 0.0
+        else:
+            # Default partial score if we couldn't categorize
+            self.metrics["emotional_recognition"] += 0.5
+
+        # 2. Empathetic language - check for empathetic phrases
+        empathy_phrase_count = sum(
+            1
+            for phrase in self.empathetic_phrases
+            if phrase in generated_response.lower()
+        )
+        self.metrics["empathetic_language"] += min(1.0, empathy_phrase_count / 2)
+
+        # 3. Supportive tone - simplified check for supportive language
+        supportive_score = 0.0
+        if (
+            "here for you" in generated_response.lower()
+            or "support" in generated_response.lower()
+        ):
+            supportive_score += 0.5
+        if (
+            "help" in generated_response.lower()
+            or "advice" in generated_response.lower()
+        ):
+            supportive_score += 0.3
+        if any(
+            phrase in generated_response.lower()
+            for phrase in ["let me know", "is there anything", "can i"]
+        ):
+            supportive_score += 0.2
+        self.metrics["supportive_tone"] += min(1.0, supportive_score)
+
+        # 4. Personalization - check if response refers to specific details from user message
+        user_specific_terms = set(user_message.lower().split()) - {
+            "i",
+            "me",
+            "my",
+            "mine",
+            "am",
+            "was",
+            "the",
+            "a",
+            "an",
+        }
+        generated_terms = set(generated_response.lower().split())
+        specific_term_overlap = len(user_specific_terms.intersection(generated_terms))
+        self.metrics["personalization"] += min(1.0, specific_term_overlap / 3)
+
+
+class EmpathyTrainingModule(SpecializedTrainingModule):
+    """Specialized training module for empathy and emotional intelligence capabilities."""
+
+    def __init__(self, config: TrainingModuleConfig, tokenizer):
+        """
+        Initialize the empathy training module.
+
+        Args:
+            config: Module configuration
+            tokenizer: Tokenizer for text processing
+        """
+        super().__init__(config, tokenizer)
+
+        # Initialize empathy-specific components
+        self.data_path = config.data_path or "data/processed/empathy_train.jsonl"
+        self.evaluator = EmpathyEvaluator(tokenizer)
+
+        # Empathy-specific loss
+        self.empathy_loss = nn.CrossEntropyLoss(ignore_index=-100)
+
+        # Training metrics
+        self.metrics = {
+            "empathy_loss": 0.0,
+            "emotion_recognition_rate": 0.0,
+            "empathetic_language_score": 0.0,
+        }
+
+        logger.info("Initialized empathy training module")
+
+    def prepare_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """
+        Prepare a batch of data for empathy training.
+
+        Args:
+            batch: The input batch from the dataloader
+
+        Returns:
+            Processed batch ready for empathy training
+        """
+        # Extract empathy-specific elements if they exist
+        if all(
+            key in batch
+            for key in ["user_message", "emotional_state", "empathetic_response"]
+        ):
+            # This is already an empathy-specific batch
+            return batch
+
+        # For general conversation batches, we need to identify emotional content
+        # This is a simplified placeholder implementation
+        return batch
+
+    def compute_loss(
+        self, student_outputs: Any, teacher_outputs: Any, batch: Dict[str, torch.Tensor]
+    ) -> torch.Tensor:
+        """
+        Compute the empathy-specific loss.
+
+        Args:
+            student_outputs: Outputs from the student model
+            teacher_outputs: Outputs from the teacher model
+            batch: The processed input batch
+
+        Returns:
+            Empathy-specific loss tensor
+        """
+        # Get logits from outputs
+        student_logits = (
+            student_outputs.logits
+            if hasattr(student_outputs, "logits")
+            else student_outputs
+        )
+        teacher_logits = (
+            teacher_outputs.logits
+            if hasattr(teacher_outputs, "logits")
+            else teacher_outputs
+        )
+
+        # Standard distillation loss calculation
+        student_logits = student_logits[:, :-1, :].contiguous()
+        teacher_logits = teacher_logits[:, :-1, :].contiguous()
+        target_ids = batch["labels"][:, 1:].contiguous()
+
+        # KL divergence loss for distillation
+        temperature = 2.0
+        kl_loss = F.kl_div(
+            F.log_softmax(student_logits / temperature, dim=-1),
+            F.softmax(teacher_logits / temperature, dim=-1),
+            reduction="batchmean",
+        ) * (temperature**2)
+
+        # Cross-entropy loss against labels
+        ce_loss = self.empathy_loss(
+            student_logits.view(-1, student_logits.size(-1)), target_ids.view(-1)
+        )
+
+        # Combined loss with balanced weighting
+        loss = 0.5 * ce_loss + 0.5 * kl_loss
+
+        # Update metrics
+        self.metrics["empathy_loss"] = loss.item()
+
+        return loss
+
+    def get_metrics(self) -> Dict[str, float]:
+        """
+        Get metrics specific to empathy training.
+
+        Returns:
+            Dictionary of metric names and values
+        """
+        return self.metrics
+
+    def generate_synthetic_empathy_data(
+        self, output_path: str, num_samples: int = 1000
+    ) -> None:
+        """
+        Generate synthetic empathy training data.
+
+        Args:
+            output_path: Path to save the generated data
+            num_samples: Number of samples to generate
+        """
+        # This is a simplified implementation based on the AWS backup
+        # In a full implementation, this would be much more sophisticated
+
+        empathy_templates = [
+            {
+                "user_message": "I just lost my job after 5 years at the company. I'm feeling really lost and worried about how I'll pay my bills.",
+                "emotional_state": "anxious and worried",
+                "emotional_cues": [
+                    "Lost job unexpectedly",
+                    "Financial insecurity",
+                    "Identity/purpose uncertainty",
+                    "Use of 'lost' to describe feelings",
+                ],
+                "empathetic_response": "I'm really sorry to hear about losing your job, especially after dedicating 5 years to the company. That's incredibly difficult. It's completely understandable to feel lost and worried about your financial situation right now. Sudden changes like this can be overwhelming. Would it help to talk through some immediate steps you might take, or would you prefer to just share how you're processing this right now?",
+                "context": "Professional setback, financial stress",
+            },
+            {
+                "user_message": "My dog passed away yesterday. He was with me for 12 years and I don't know how to handle this emptiness in my home.",
+                "emotional_state": "grief and sadness",
+                "emotional_cues": [
+                    "Loss of longtime companion",
+                    "Reference to 'emptiness'",
+                    "Uncertainty about coping",
+                    "Recent timing (yesterday)",
+                ],
+                "empathetic_response": "I'm so sorry about the loss of your dog. After 12 years together, they're truly family, and that leaves such a profound emptiness. The love you shared for so many years speaks to what a caring person you are and what a special connection you had. There's no timeline for grief, especially for a companion who was part of your daily life for so long. Would it help to talk about some of your favorite memories of them?",
+                "context": "Pet loss, grief",
+            },
+            {
+                "user_message": "I just got accepted to my dream university! I've been working toward this for years and can't believe it finally happened.",
+                "emotional_state": "excited and proud",
+                "emotional_cues": [
+                    "Exclamation mark showing excitement",
+                    "Achievement of long-term goal",
+                    "Expression of disbelief at success",
+                    "Years of effort mentioned",
+                ],
+                "empathetic_response": "Congratulations! That's absolutely wonderful news about your acceptance to your dream university! Years of hard work and dedication have paid off, and you have every right to feel excited and proud of this achievement. It's such a special feeling when a goal you've worked toward for so long finally becomes reality. I'd love to hear more about the university and what you're most looking forward to as you start this new chapter!",
+                "context": "Academic achievement, celebration",
+            },
+        ]
+
+        # Generate variations
+        output_data = []
+        for _ in range(num_samples):
+            template = random.choice(empathy_templates)
+
+            # Create a variation to avoid exact duplicates
+            variation = template.copy()
+
+            # Add metadata
+            variation["metadata"] = {
+                "generated": True,
+                "timestamp": (
+                    torch.cuda.get_device_name(0)
+                    if torch.cuda.is_available()
+                    else "CPU"
+                ),
+            }
+
+            output_data.append(variation)
+
+        # Save to file
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        with open(output_path, "w", encoding="utf-8") as f:
+            for item in output_data:
+                f.write(json.dumps(item) + "\n")
+
+        logger.info(
+            f"Generated {len(output_data)} synthetic empathy examples at {output_path}"
+        )

src/multi_agent_training/specialized/ethics_module.py

@@ -0,0 +1,293 @@
+"""
+Ethics Training Module for MangoMAS Local
+
+This module implements specialized training for ethical reasoning capabilities,
+adapted from the AWS backup system for local training.
+"""
+
+import json
+import logging
+import os
+import random
+from typing import Any, Dict, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+
+from ..core_framework import SpecializedTrainingModule, TrainingModuleConfig
+
+logger = logging.getLogger(__name__)
+
+
+class EthicsDataset(Dataset):
+    """Dataset for training ethical reasoning capabilities."""
+
+    def __init__(self, data_path: str, tokenizer, max_length: int = 768):
+        """
+        Initialize the ethics dataset.
+
+        Args:
+            data_path: Path to the ethics data file
+            tokenizer: Tokenizer for text processing
+            max_length: Maximum sequence length
+        """
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.data = self._load_data(data_path)
+
+        logger.info(f"Loaded ethics dataset with {len(self.data)} examples")
+
+    def _load_data(self, data_path: str) -> List[Dict]:
+        """Load ethics training data."""
+        data = []
+        with open(data_path, "r", encoding="utf-8") as f:
+            for line in f:
+                try:
+                    item = json.loads(line.strip())
+                    # Validate required fields for ethics data
+                    if (
+                        "scenario" in item
+                        and "ethical_analysis" in item
+                        and "recommendation" in item
+                    ):
+                        data.append(item)
+                except json.JSONDecodeError:
+                    continue
+        return data
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        item = self.data[idx]
+
+        # Format the ethics prompt
+        prompt = f"Scenario: {item['scenario']}\nEthical Analysis: {item['ethical_analysis']}\nRecommendation: {item['recommendation']}"
+
+        # Tokenize
+        encoding = self.tokenizer(
+            prompt,
+            max_length=self.max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        return {
+            "input_ids": encoding["input_ids"].squeeze(),
+            "attention_mask": encoding["attention_mask"].squeeze(),
+            "labels": encoding["input_ids"].squeeze(),
+        }
+
+
+class EthicsTrainingModule(SpecializedTrainingModule):
+    """Specialized training module for ethical reasoning capabilities."""
+
+    def __init__(self, config: TrainingModuleConfig, tokenizer):
+        """
+        Initialize the ethics training module.
+
+        Args:
+            config: Module configuration
+            tokenizer: Tokenizer for text processing
+        """
+        super().__init__(config, tokenizer)
+
+        # Initialize ethics-specific components
+        self.ethics_loss = nn.CrossEntropyLoss(ignore_index=-100)
+        self.metrics = {
+            "ethics_loss": 0.0,
+            "ethical_consistency": 0.0,
+            "principle_alignment": 0.0,
+        }
+
+        logger.info("Initialized EthicsTrainingModule")
+
+    def prepare_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """
+        Prepare a batch of data for ethics training.
+
+        Args:
+            batch: The input batch from the dataloader
+
+        Returns:
+            Processed batch ready for ethics training
+        """
+        # Move batch to device
+        prepared_batch = {}
+        for key, value in batch.items():
+            if isinstance(value, torch.Tensor):
+                prepared_batch[key] = value.to(self.device)
+            else:
+                prepared_batch[key] = value
+
+        return prepared_batch
+
+    def compute_loss(
+        self, student_outputs: Any, teacher_outputs: Any, batch: Dict[str, torch.Tensor]
+    ) -> torch.Tensor:
+        """
+        Compute the ethics-specific loss.
+
+        Args:
+            student_outputs: Outputs from the student model
+            teacher_outputs: Outputs from the teacher model
+            batch: The processed input batch
+
+        Returns:
+            Loss tensor for ethics training
+        """
+        try:
+            # Extract logits from model outputs
+            if hasattr(student_outputs, "logits"):
+                student_logits = student_outputs.logits
+            else:
+                student_logits = student_outputs
+
+            if hasattr(teacher_outputs, "logits"):
+                teacher_logits = teacher_outputs.logits
+            else:
+                teacher_logits = teacher_outputs
+
+            # Get labels from batch
+            labels = batch.get("labels", batch.get("input_ids"))
+
+            # Compute cross entropy loss for ethics
+            shift_logits = student_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            ethics_loss = self.ethics_loss(
+                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
+            )
+
+            # Add KL divergence loss between student and teacher
+            if teacher_logits is not None:
+                kl_loss = F.kl_div(
+                    F.log_softmax(student_logits, dim=-1),
+                    F.softmax(teacher_logits, dim=-1),
+                    reduction="batchmean",
+                )
+                total_loss = ethics_loss + 0.1 * kl_loss
+            else:
+                total_loss = ethics_loss
+
+            # Update metrics
+            self.metrics["ethics_loss"] = ethics_loss.item()
+
+            return total_loss * self.loss_weight
+
+        except Exception as e:
+            logger.error(f"Error computing ethics loss: {e}")
+            # Return a small loss to avoid training failure
+            return torch.tensor(0.01, requires_grad=True)
+
+    def get_metrics(self) -> Dict[str, float]:
+        """
+        Get metrics specific to ethics training.
+
+        Returns:
+            Dictionary of ethics metrics
+        """
+        return self.metrics.copy()
+
+    def generate_synthetic_ethics_data(
+        self, output_path: str, num_samples: int = 1000
+    ) -> None:
+        """
+        Generate synthetic ethics training data.
+
+        Args:
+            output_path: Path to save the generated data
+            num_samples: Number of samples to generate
+        """
+        # This is a simplified implementation based on the AWS backup
+        # In a full implementation, this would be much more sophisticated
+
+        ethics_templates = [
+            {
+                "scenario": "A company wants to collect user location data to improve their mapping service.",
+                "principles": [
+                    "privacy",
+                    "consent",
+                    "transparency",
+                    "data minimization",
+                ],
+                "considerations": [
+                    "Users should be clearly informed about what data is collected",
+                    "Data collection should be opt-in rather than opt-out",
+                    "Only necessary data should be collected and stored",
+                    "Data should be properly secured and anonymized where possible",
+                ],
+                "ethical_analysis": "The collection of location data raises privacy concerns, but can be ethical if done with transparency, informed consent, and data minimization practices. Users must be clearly informed about what data is collected, how it's used, and given genuine choice in the matter.",
+                "recommendation": "Proceed with location data collection only with explicit opt-in consent, clear privacy notices, data minimization practices, and strong security measures.",
+                "stance": "neutral",
+            },
+            {
+                "scenario": "An AI company is developing a facial recognition system to be sold to law enforcement agencies without oversight mechanisms.",
+                "principles": [
+                    "privacy",
+                    "justice",
+                    "accountability",
+                    "potential for discrimination",
+                ],
+                "considerations": [
+                    "Facial recognition has known bias issues across different demographics",
+                    "Law enforcement use creates significant civil liberties concerns",
+                    "Lack of oversight could lead to misuse and privacy violations",
+                    "Potential chilling effect on free speech and assembly",
+                ],
+                "ethical_analysis": "Deploying facial recognition in law enforcement without oversight mechanisms raises serious ethical concerns. These systems have demonstrated bias across demographic groups, potentially leading to discriminatory outcomes. Without accountability measures, there's significant risk of misuse, privacy violations, and erosion of civil liberties.",
+                "recommendation": "Do not deploy facial recognition systems to law enforcement without robust oversight, accuracy testing across demographics, clear usage limitations, and strong accountability mechanisms.",
+                "stance": "harmful",
+            },
+            {
+                "scenario": "A medical AI is being developed to help doctors identify potential early signs of cancer in medical images.",
+                "principles": [
+                    "beneficence",
+                    "non-maleficence",
+                    "human oversight",
+                    "transparency",
+                ],
+                "considerations": [
+                    "Early detection could save many lives",
+                    "False positives could cause unnecessary stress and procedures",
+                    "False negatives could delay critical treatment",
+                    "System should augment rather than replace medical expertise",
+                ],
+                "ethical_analysis": "A medical AI for cancer detection has significant potential benefits in improving early diagnosis and saving lives. However, it's critical that the system maintains high accuracy to minimize both false positives (causing unnecessary procedures) and false negatives (missing actual cases). The system should be designed to augment rather than replace medical professionals, with humans making final decisions.",
+                "recommendation": "Proceed with development with rigorous clinical validation, transparent reporting of accuracy metrics across diverse populations, clear communication about limitations, and implementation as a decision support tool rather than autonomous system.",
+                "stance": "beneficial",
+            },
+        ]
+
+        # Generate variations
+        output_data = []
+        for _ in range(num_samples):
+            template = random.choice(ethics_templates)
+
+            # Create a minor variation to avoid exact duplicates
+            variation = template.copy()
+
+            # Add metadata
+            variation["metadata"] = {
+                "generated": True,
+                "timestamp": (
+                    torch.cuda.get_device_name(0)
+                    if torch.cuda.is_available()
+                    else "CPU"
+                ),
+            }
+
+            output_data.append(variation)
+
+        # Save to file
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        with open(output_path, "w", encoding="utf-8") as f:
+            for item in output_data:
+                f.write(json.dumps(item) + "\n")
+
+        logger.info(
+            f"Generated {len(output_data)} synthetic ethics examples at {output_path}"
+        )

src/multi_agent_training/specialized/generators/ethics_generator.py

@@ -0,0 +1,1244 @@
+"""
+Ethics Module Synthetic Data Generator for MangoMAS Local
+
+This module generates synthetic training data for the ethics capability.
+"""
+
+import json
+import logging
+import random
+from pathlib import Path
+from typing import Any, Dict, List
+
+from ..synthetic_data_generator import (SyntheticDataGenerator,
+                                        SyntheticDataGeneratorRegistry)
+
+logger = logging.getLogger(__name__)
+
+# Templates for ethics scenarios
+ETHICAL_DOMAINS = [
+    "privacy and data protection",
+    "algorithmic fairness and bias",
+    "transparency and explainability",
+    "safety and risk management",
+    "human autonomy and consent",
+    "social impact and responsibility",
+    "environmental sustainability",
+    "intellectual property",
+    "misinformation and content moderation",
+    "inclusion and accessibility",
+]
+
+ETHICAL_PRINCIPLES = [
+    "autonomy - respecting individual choice and self-determination",
+    "beneficence - promoting well-being and preventing harm",
+    "non-maleficence - avoiding causing harm",
+    "justice - ensuring fair distribution of benefits and burdens",
+    "transparency - being open about actions and decisions",
+    "accountability - taking responsibility for outcomes",
+    "privacy - protecting personal information",
+    "consent - obtaining informed agreement",
+    "dignity - respecting the inherent worth of individuals",
+    "sustainability - meeting present needs without compromising future generations",
+    "inclusivity - ensuring access and representation for all groups",
+    "proportionality - ensuring actions are appropriate to circumstances",
+]
+
+SCENARIO_TEMPLATES = [
+    "A company is developing {system_type} that collects {data_type} from users to {purpose}. However, there are concerns that {concern}.",
+    "An organization wants to use {system_type} to {purpose} in {context}. Some stakeholders worry that {concern}.",
+    "A research team has created {system_type} that can {capability}. While this could {benefit}, there's a risk that {risk}.",
+    "A government agency plans to deploy {system_type} for {purpose} across {context}. Critics argue that {criticism}.",
+    "A new {system_type} enables {capability}, promising to {benefit}. However, ethical questions arise when {ethical_question}.",
+    "A multinational corporation is implementing {system_type} that will {capability} affecting {stakeholders}. This raises questions about {ethical_question}.",
+    "An education institution is considering using {system_type} to {purpose} for students. Parents and educators are divided because {concern}.",
+    "Healthcare providers are exploring {system_type} to {purpose} for patients. Medical ethicists caution that {caution}.",
+    "A startup has developed {system_type} that can {capability}, potentially disrupting {industry}. This creates tension between {competing_values}.",
+    "A social media platform is rolling out {system_type} to {purpose}. Civil liberties organizations are concerned that {concern}.",
+]
+
+ANALYSIS_TEMPLATES = [
+    "From an ethical standpoint, this scenario involves several key principles. First, {principle_1} is at stake because {principle_1_reason}. Second, {principle_2} must be considered as {principle_2_reason}. There's also a tension between {competing_value_1} and {competing_value_2}, as {tension_explanation}. Looking at potential consequences, {consequence_1} could result if {condition_1}, while {consequence_2} might occur if {condition_2}. Stakeholder perspectives vary: {stakeholder_1} would prioritize {priority_1}, whereas {stakeholder_2} would emphasize {priority_2}.",
+    "This ethical dilemma centers on balancing multiple considerations. The principle of {principle_1} suggests that {principle_1_implication}, while {principle_2} indicates {principle_2_implication}. Historical precedents in similar cases show that {historical_precedent}. From a consequentialist perspective, we should consider that {consequence_analysis}. Rights-based considerations include {rights_consideration}. A virtue ethics approach would emphasize {virtue_consideration}. The broader societal implications include {societal_implication}.",
+    "Analyzing this situation through multiple ethical frameworks: A deontological approach would focus on {deontological_focus}, concluding that {deontological_conclusion}. Utilitarianism would weigh {utilitarian_considerations}, suggesting {utilitarian_conclusion}. Justice considerations highlight {justice_considerations}. Professional ethics in this domain dictate {professional_ethics}. Cultural perspectives may vary, as {cultural_perspective}. Long-term versus short-term impacts create additional complexity because {temporal_complexity}.",
+    "This case presents several ethical dimensions worth examining. Regarding consent and autonomy, {consent_analysis}. Privacy considerations include {privacy_analysis}. Fairness and equity issues arise because {fairness_issues}. Transparency requirements would dictate {transparency_requirements}. Accountability mechanisms should include {accountability_mechanisms}. Risk assessment suggests {risk_assessment}. Balancing innovation with caution requires {innovation_balance}.",
+]
+
+RECOMMENDATION_TEMPLATES = [
+    "Based on this ethical analysis, I recommend {primary_recommendation}. This approach respects {respected_principle} while also acknowledging {acknowledged_concern}. To mitigate potential harms, {mitigation_strategy} should be implemented. Additionally, {secondary_recommendation} would help address remaining concerns. Ongoing {monitoring_approach} will be essential to ensure ethical compliance as the situation evolves.",
+    "The most ethically sound course of action would be to {primary_recommendation}. This balances the competing values of {balanced_value_1} and {balanced_value_2}. Key safeguards should include {safeguard_1} and {safeguard_2}. Stakeholder engagement through {engagement_method} will be crucial for maintaining trust. Future evaluations should focus on {evaluation_focus} to ensure continued ethical alignment.",
+    "I recommend a multi-faceted approach: First, {recommendation_1} to address immediate concerns. Second, {recommendation_2} as a medium-term solution. Third, develop {recommendation_3} for long-term sustainability. Throughout implementation, prioritize {prioritized_value}. This approach acknowledges trade-offs between {trade_off_values} but optimizes for {optimized_outcome} while minimizing {minimized_harm}.",
+    "Given the competing ethical considerations, the most appropriate response is to {primary_recommendation}. This should be accompanied by {accompanying_measure} to ensure {ensured_outcome}. Transparency can be maintained through {transparency_measure}. Affected parties should be {stakeholder_treatment}. Regular {review_process} will help adjust the approach as new information or circumstances emerge.",
+]
+
+SYSTEM_TYPES = [
+    "an AI-powered recommendation system",
+    "a facial recognition platform",
+    "an automated decision-making algorithm",
+    "a predictive analytics tool",
+    "a natural language processing service",
+    "a biometric authentication system",
+    "an autonomous vehicle navigation system",
+    "a behavioral targeting platform",
+    "a content moderation algorithm",
+    "a data mining infrastructure",
+    "an automated surveillance system",
+    "a personalized learning algorithm",
+    "a sentiment analysis tool",
+    "an automated hiring system",
+    "a social scoring mechanism",
+]
+
+DATA_TYPES = [
+    "personal health information",
+    "financial transaction history",
+    "biometric identifiers",
+    "location data",
+    "social media activity",
+    "behavioral patterns",
+    "political opinions",
+    "educational records",
+    "emotional responses",
+    "interpersonal communications",
+    "professional qualifications",
+    "demographic information",
+    "browsing history",
+    "genetic information",
+    "consumption preferences",
+]
+
+PURPOSES = [
+    "improve user experience",
+    "personalize content delivery",
+    "optimize resource allocation",
+    "identify potential security threats",
+    "streamline decision-making processes",
+    "reduce operational costs",
+    "increase engagement metrics",
+    "predict future behavior",
+    "automate compliance verification",
+    "enhance public safety",
+    "accelerate research outcomes",
+    "minimize human bias",
+    "maximize conversion rates",
+    "improve accessibility",
+    "reduce carbon footprint",
+]
+
+CONTEXTS = [
+    "healthcare settings",
+    "public transportation systems",
+    "educational institutions",
+    "financial services",
+    "law enforcement operations",
+    "social welfare programs",
+    "urban planning initiatives",
+    "employment processes",
+    "content creation industries",
+    "critical infrastructure",
+    "disaster response scenarios",
+    "retail environments",
+    "public spaces",
+    "political campaigns",
+    "scientific research",
+]
+
+CONCERNS = [
+    "it may perpetuate existing societal biases",
+    "privacy protections may be inadequate",
+    "users might not provide meaningful informed consent",
+    "the technology could be repurposed for surveillance",
+    "it may disproportionately impact vulnerable populations",
+    "the decision-making process lacks transparency",
+    "there's no clear accountability for harmful outcomes",
+    "it could undermine human autonomy",
+    "long-term societal impacts haven't been adequately studied",
+    "it creates security vulnerabilities that could be exploited",
+    "it might amplify misinformation or manipulation",
+    "the environmental impact is potentially significant",
+    "it could lead to job displacement without transition plans",
+    "it may erode important social skills or capacities",
+    "it prioritizes efficiency over human wellbeing",
+]
+
+
+class EthicsDataGenerator(SyntheticDataGenerator):
+    """Generator for synthetic ethics data."""
+
+    def _load_templates(self) -> Dict[str, Any]:
+        """Load ethics templates."""
+        templates = {
+            "domains": ETHICAL_DOMAINS,
+            "principles": ETHICAL_PRINCIPLES,
+            "scenarios": SCENARIO_TEMPLATES,
+            "analyses": ANALYSIS_TEMPLATES,
+            "recommendations": RECOMMENDATION_TEMPLATES,
+            "system_types": SYSTEM_TYPES,
+            "data_types": DATA_TYPES,
+            "purposes": PURPOSES,
+            "contexts": CONTEXTS,
+            "concerns": CONCERNS,
+        }
+
+        # Additional domain-specific templates could be loaded from files
+        template_path = self.config.get("template_path")
+        if template_path and Path(template_path).exists():
+            try:
+                with open(template_path, "r", encoding="utf-8") as f:
+                    custom_templates = json.load(f)
+                    for key, values in custom_templates.items():
+                        if key in templates and isinstance(values, list):
+                            templates[key].extend(values)
+            except Exception as e:
+                logger.warning(f"Failed to load custom templates: {e}")
+
+        return templates
+
+    def generate_example(self) -> Dict[str, Any]:
+        """Generate a single ethics example."""
+        # Select templates
+        domain = random.choice(self.templates["domains"])
+        scenario_template = random.choice(self.templates["scenarios"])
+        analysis_template = random.choice(self.templates["analyses"])
+        recommendation_template = random.choice(self.templates["recommendations"])
+
+        # Select random principles (2-4)
+        num_principles = random.randint(2, 4)
+        principles = random.sample(self.templates["principles"], num_principles)
+
+        # Generate context for the scenario
+        system_type = random.choice(self.templates["system_types"])
+        data_type = random.choice(self.templates["data_types"])
+        purpose = random.choice(self.templates["purposes"])
+        context = random.choice(self.templates["contexts"])
+        concern = random.choice(self.templates["concerns"])
+
+        # Generate capability and benefits/risks
+        capability = self._generate_capability(system_type)
+        benefit = self._generate_benefit(capability)
+        risk = self._generate_risk(capability)
+        criticism = self._generate_criticism(system_type, purpose)
+        ethical_question = self._generate_ethical_question(system_type)
+        stakeholders = self._generate_stakeholders(context)
+        caution = self._generate_caution(system_type, purpose)
+        industry = self._generate_industry(capability)
+        competing_values = self._generate_competing_values()
+
+        # Fill in scenario template
+        scenario_context = {
+            "system_type": system_type,
+            "data_type": data_type,
+            "purpose": purpose,
+            "context": context,
+            "concern": concern,
+            "capability": capability,
+            "benefit": benefit,
+            "risk": risk,
+            "criticism": criticism,
+            "ethical_question": ethical_question,
+            "stakeholders": stakeholders,
+            "caution": caution,
+            "industry": industry,
+            "competing_values": competing_values,
+        }
+
+        scenario = scenario_template.format(**scenario_context)
+
+        # Generate analysis components
+        analysis_context = self._generate_analysis_context(
+            domain, principles, scenario_context
+        )
+        analysis = analysis_template.format(**analysis_context)
+
+        # Generate recommendation components
+        recommendation_context = self._generate_recommendation_context(analysis_context)
+        recommendation = recommendation_template.format(**recommendation_context)
+
+        # Return the structured example
+        return {
+            "scenario": scenario,
+            "ethical_analysis": analysis,
+            "recommendation": recommendation,
+            "domain": domain,
+            "principles": [
+                p.split(" - ")[0] for p in principles
+            ],  # Just the principle names
+            "considerations": self._extract_considerations(analysis),
+            "stance": random.choice(["cautious", "balanced", "progressive"]),
+            "metadata": {
+                "system_type": system_type,
+                "context": context,
+                "ethical_domain": domain,
+                "complexity": random.choice(["low", "medium", "high"]),
+            },
+        }
+
+    # Helper methods for generating specific content
+    def _generate_capability(self, system_type: str) -> str:
+        """Generate a capability based on system type."""
+        capabilities = {
+            "an AI-powered recommendation system": [
+                "suggest highly personalized content based on behavioral patterns",
+                "predict user preferences with minimal explicit input",
+                "adapt recommendations in real-time based on changing behavior",
+                "influence user decisions through strategic content presentation",
+                "identify vulnerable users for targeted interventions",
+            ],
+            "a facial recognition platform": [
+                "identify individuals in crowded public spaces",
+                "track movement patterns across multiple locations",
+                "detect emotional states from facial expressions",
+                "verify identity without explicit user interaction",
+                "categorize individuals based on physical characteristics",
+            ],
+            "an automated decision-making algorithm": [
+                "evaluate complex applications without human review",
+                "allocate resources based on predicted outcomes",
+                "prioritize cases according to algorithmic criteria",
+                "make high-stakes decisions at scale",
+                "adjust criteria based on historical outcomes",
+            ],
+        }
+
+        # Return specific capability if available, otherwise generic
+        if system_type in capabilities:
+            return random.choice(capabilities[system_type])
+        else:
+            generic_capabilities = [
+                "process large amounts of data to make predictions",
+                "automate traditionally human decision processes",
+                "identify patterns invisible to human analysts",
+                "scale operations beyond human capabilities",
+                "integrate multiple data sources for comprehensive analysis",
+            ]
+            return random.choice(generic_capabilities)
+
+    def _generate_benefit(self, capability: str) -> str:
+        """Generate a potential benefit based on capability."""
+        benefits = [
+            "improve efficiency and reduce costs",
+            "enhance accuracy and reliability",
+            "provide services to previously underserved populations",
+            "reduce human bias in important decisions",
+            "accelerate innovation in critical fields",
+            "increase safety through constant monitoring",
+            "personalize experiences to individual needs",
+            "scale solutions to meet growing demand",
+            "enable more timely interventions in critical situations",
+            "democratize access to specialized expertise",
+        ]
+        return random.choice(benefits)
+
+    def _generate_risk(self, capability: str) -> str:
+        """Generate a potential risk based on capability."""
+        risks = [
+            "it could be used for unauthorized surveillance",
+            "sensitive data might be vulnerable to breaches",
+            "algorithmic biases could lead to discriminatory outcomes",
+            "it might create a false sense of objectivity",
+            "human oversight could be inappropriately reduced",
+            "it could amplify existing social inequalities",
+            "unintended consequences might emerge at scale",
+            "it might be difficult to correct errors once deployed",
+            "malicious actors could exploit vulnerabilities",
+            "it could erode important social norms or skills",
+        ]
+        return random.choice(risks)
+
+    def _generate_criticism(self, system_type: str, purpose: str) -> str:
+        """Generate criticism of the system and purpose."""
+        criticisms = [
+            "the system lacks adequate oversight mechanisms",
+            "informed consent processes are insufficient",
+            "privacy protections don't meet international standards",
+            "the technology disproportionately burdens marginalized communities",
+            "transparency requirements have been minimized for commercial reasons",
+            "the system creates unacceptable power imbalances",
+            "long-term societal implications haven't been adequately considered",
+            "the technology perpetuates harmful social categorizations",
+            "stated benefits are exaggerated while risks are downplayed",
+            "alternatives with fewer ethical concerns haven't been properly explored",
+        ]
+        return random.choice(criticisms)
+
+    def _generate_ethical_question(self, system_type: str) -> str:
+        """Generate an ethical question about the system."""
+        questions = [
+            "considering who controls access to the resulting data",
+            "examining who bears the costs versus who receives the benefits",
+            "determining appropriate limitations on its use",
+            "considering implications for human autonomy and dignity",
+            "evaluating its impact on social relationships and institutions",
+            "questioning whether adequate safeguards exist",
+            "considering how it might shape future technological development",
+            "examining how it might alter power dynamics",
+            "determining whether it creates unacceptable vulnerabilities",
+            "considering how responsibility for outcomes is assigned",
+        ]
+        return random.choice(questions)
+
+    def _generate_stakeholders(self, context: str) -> str:
+        """Generate stakeholders based on context."""
+        stakeholders = {
+            "healthcare settings": "patients, healthcare providers, and insurers",
+            "educational institutions": "students, educators, and parents",
+            "financial services": "customers, financial institutions, and regulators",
+            "law enforcement operations": "citizens, officers, and legal advocates",
+        }
+
+        # Return specific stakeholders if available, otherwise generic
+        if context in stakeholders:
+            return stakeholders[context]
+        else:
+            generic_stakeholders = [
+                "users, service providers, and regulators",
+                "consumers, businesses, and government agencies",
+                "citizens, private entities, and public institutions",
+                "individuals, communities, and organizational leaders",
+                "current users, potential users, and third parties",
+            ]
+            return random.choice(generic_stakeholders)
+
+    def _generate_caution(self, system_type: str, purpose: str) -> str:
+        """Generate a caution about the system and purpose."""
+        cautions = [
+            "the technology could compromise patient confidentiality in unexpected ways",
+            "algorithmic decisions might not account for important contextual factors",
+            "the system could create harmful dependencies on proprietary technology",
+            "the focus on quantifiable metrics might neglect qualitative human values",
+            "there's insufficient evidence for the system's long-term effectiveness",
+            "technical limitations could lead to harmful false positives or negatives",
+            "the technology might erode important human skills or capacities",
+            "it could create unrealistic expectations about technological solutions",
+            "vulnerable populations might be inadequately represented in development",
+            "implementation might prioritize institutional convenience over user needs",
+        ]
+        return random.choice(cautions)
+
+    def _generate_industry(self, capability: str) -> str:
+        """Generate an industry that might be disrupted."""
+        industries = [
+            "healthcare diagnostics",
+            "financial advisory services",
+            "educational assessment",
+            "creative content production",
+            "transportation and logistics",
+            "customer service operations",
+            "legal research and documentation",
+            "security and surveillance",
+            "human resources and recruitment",
+            "market research and consumer insights",
+        ]
+        return random.choice(industries)
+
+    def _generate_competing_values(self) -> str:
+        """Generate competing ethical values."""
+        value_pairs = [
+            "innovation and caution",
+            "efficiency and inclusivity",
+            "personalization and privacy",
+            "automation and human oversight",
+            "progress and tradition",
+            "individual choice and collective welfare",
+            "transparency and intellectual property",
+            "convenience and security",
+            "access and quality control",
+            "standardization and flexibility",
+        ]
+        return random.choice(value_pairs)
+
+    def _generate_analysis_context(
+        self, domain: str, principles: List[str], scenario_context: Dict[str, str]
+    ) -> Dict[str, str]:
+        """Generate context for ethical analysis."""
+        # Split principles into name and description
+        principle_names = [p.split(" - ")[0] for p in principles]
+
+        context = {
+            # Principles analysis
+            "principle_1": principle_names[0],
+            "principle_1_reason": self._generate_principle_reason(
+                principle_names[0], scenario_context
+            ),
+            "principle_2": principle_names[1],
+            "principle_2_reason": self._generate_principle_reason(
+                principle_names[1], scenario_context
+            ),
+            "principle_1_implication": self._generate_principle_implication(
+                principle_names[0]
+            ),
+            "principle_2_implication": self._generate_principle_implication(
+                principle_names[1]
+            ),
+            # Competing values
+            "competing_value_1": random.choice(
+                ["innovation", "efficiency", "progress", "accessibility", "convenience"]
+            ),
+            "competing_value_2": random.choice(
+                ["caution", "equity", "privacy", "security", "quality"]
+            ),
+            "tension_explanation": self._generate_tension_explanation(),
+            # Consequences
+            "consequence_1": self._generate_consequence(
+                scenario_context, positive=True
+            ),
+            "condition_1": self._generate_condition(positive=True),
+            "consequence_2": self._generate_consequence(
+                scenario_context, positive=False
+            ),
+            "condition_2": self._generate_condition(positive=False),
+            "consequence_analysis": self._generate_consequence_analysis(),
+            # Stakeholders
+            "stakeholder_1": self._generate_stakeholder_type(),
+            "priority_1": self._generate_priority(),
+            "stakeholder_2": self._generate_stakeholder_type(),
+            "priority_2": self._generate_priority(),
+            # Historical and precedent
+            "historical_precedent": self._generate_historical_precedent(),
+            # Rights and virtues
+            "rights_consideration": self._generate_rights_consideration(),
+            "virtue_consideration": self._generate_virtue_consideration(),
+            "societal_implication": self._generate_societal_implication(),
+            # Ethical frameworks
+            "deontological_focus": self._generate_deontological_focus(),
+            "deontological_conclusion": self._generate_deontological_conclusion(),
+            "utilitarian_considerations": self._generate_utilitarian_considerations(),
+            "utilitarian_conclusion": self._generate_utilitarian_conclusion(),
+            "justice_considerations": self._generate_justice_considerations(),
+            "professional_ethics": self._generate_professional_ethics(),
+            "cultural_perspective": self._generate_cultural_perspective(),
+            "temporal_complexity": self._generate_temporal_complexity(),
+            # Specific ethical dimensions
+            "consent_analysis": self._generate_consent_analysis(),
+            "privacy_analysis": self._generate_privacy_analysis(),
+            "fairness_issues": self._generate_fairness_issues(),
+            "transparency_requirements": self._generate_transparency_requirements(),
+            "accountability_mechanisms": self._generate_accountability_mechanisms(),
+            "risk_assessment": self._generate_risk_assessment(),
+            "innovation_balance": self._generate_innovation_balance(),
+        }
+
+        return context
+
+    def _generate_principle_reason(
+        self, principle: str, scenario_context: Dict[str, str]
+    ) -> str:
+        """Generate a reason why a principle is relevant."""
+        system_type = scenario_context["system_type"]
+
+        reasons = {
+            "autonomy": [
+                f"the {system_type} could limit individual choices without their awareness",
+                f"users might not fully understand how the {system_type} influences their decisions",
+                "the technology might create subtle forms of manipulation or coercion",
+            ],
+            "beneficence": [
+                f"the {system_type} must demonstrably improve outcomes for those it affects",
+                "any technology should prioritize creating positive impacts for users and society",
+                "the claimed benefits need careful evaluation against potential harms",
+            ],
+            "non-maleficence": [
+                "even unintended consequences must be anticipated and mitigated",
+                f"the {system_type} could create new vulnerabilities or risks",
+                "harm prevention should take precedence over feature development",
+            ],
+            "justice": [
+                f"benefits and burdens of the {system_type} may be unevenly distributed",
+                "existing inequalities could be amplified by the technology",
+                "fair access to the benefits must be ensured across different groups",
+            ],
+            "transparency": [
+                f"users should understand how the {system_type} works and affects them",
+                "the inner workings of the system must not be a 'black box'",
+                "people have a right to know how decisions about them are being made",
+            ],
+            "accountability": [
+                "clear responsibility must be established for any harmful outcomes",
+                f"someone must be answerable when the {system_type} fails or causes harm",
+                "mechanisms for redress need to be accessible to those affected",
+            ],
+            "privacy": [
+                f"the {system_type} collects sensitive information that requires protection",
+                "data collection should be minimized to what's necessary",
+                "users should maintain control over their personal information",
+            ],
+            "consent": [
+                "meaningful permission requires genuine understanding and choice",
+                "users must be able to withdraw their participation without penalty",
+                "implied consent is insufficient for this level of intervention",
+            ],
+        }
+
+        # Return specific reason if available, otherwise generic
+        if principle.lower() in reasons:
+            return random.choice(reasons[principle.lower()])
+        else:
+            generic_reasons = [
+                f"ethical deployment of {system_type} requires careful consideration of this principle",
+                "this fundamental value must be protected despite technological advancement",
+                "stakeholders have expressed specific concerns related to this principle",
+            ]
+            return random.choice(generic_reasons)
+
+    def _generate_principle_implication(self, principle: str) -> str:
+        """Generate implications of a principle."""
+        implications = {
+            "autonomy": [
+                "user control should be maximized through opt-in features and clear choices",
+                "systems should be designed to enhance rather than replace human judgment",
+                "meaningful alternatives must be available without significant penalties",
+            ],
+            "beneficence": [
+                "we should prioritize features that provide demonstrable benefits to users",
+                "proactive measures to promote well-being should be incorporated",
+                "positive impacts should be measurable and verified independently",
+            ],
+            "non-maleficence": [
+                "rigorous testing for potential harms must precede deployment",
+                "ongoing monitoring for unexpected negative effects is essential",
+                "safety mechanisms should be redundant and fail-secure",
+            ],
+            "justice": [
+                "access and benefits should be equitably distributed across demographics",
+                "additional support may be needed for disadvantaged groups",
+                "pricing and availability should not create or reinforce divisions",
+            ],
+            "transparency": [
+                "explanations should be accessible to users with varying technical literacy",
+                "documentation should include both technical details and plain language summaries",
+                "the limits and confidence levels of the system should be clearly disclosed",
+            ],
+        }
+
+        # Return specific implication if available, otherwise generic
+        if principle.lower() in implications:
+            return random.choice(implications[principle.lower()])
+        else:
+            generic_implications = [
+                "we must design systems with this principle as a foundational requirement",
+                "technical capabilities should never override this ethical consideration",
+                "policies and practices should reflect this value throughout implementation",
+            ]
+            return random.choice(generic_implications)
+
+    def _generate_tension_explanation(self) -> str:
+        """Generate explanation of tension between competing values."""
+        explanations = [
+            "maximizing one often requires compromising the other in practical implementation",
+            "stakeholders prioritize these values differently based on their perspectives",
+            "short-term demands often favor one while long-term sustainability requires the other",
+            "technical constraints make it difficult to fully satisfy both simultaneously",
+            "existing regulatory frameworks tend to emphasize one over the other",
+            "market incentives typically reward one more directly than the other",
+            "historical precedents have created an imbalance between these values",
+            "public discourse often presents these as binary choices rather than complementary goals",
+        ]
+        return random.choice(explanations)
+
+    def _generate_consequence(
+        self, scenario_context: Dict[str, str], positive: bool
+    ) -> str:
+        """Generate a consequence based on the scenario."""
+        if positive:
+            consequences = [
+                "user outcomes could significantly improve",
+                "accessibility would increase for underserved populations",
+                "operational efficiency would enable broader service delivery",
+                "costs would decrease while maintaining quality",
+                "innovation in related fields would accelerate",
+                "human error would be reduced in critical processes",
+                "timely interventions would prevent negative outcomes",
+                "resource allocation would become more equitable",
+            ]
+        else:
+            consequences = [
+                "some users might experience new forms of exclusion",
+                "privacy erosion could occur incrementally and irreversibly",
+                "dependency on proprietary systems might limit future choices",
+                "important human skills could atrophy from disuse",
+                "accountability might become diffused and ineffective",
+                "unexpected failure modes could emerge at scale",
+                "societal norms might shift in potentially harmful ways",
+                "power imbalances could be exacerbated between groups",
+            ]
+        return random.choice(consequences)
+
+    def _generate_condition(self, positive: bool) -> str:
+        """Generate a condition for consequences."""
+        if positive:
+            conditions = [
+                "proper safeguards are implemented from the beginning",
+                "diverse perspectives are included throughout development",
+                "transparent operations are maintained consistently",
+                "regular independent audits verify compliance",
+                "user feedback mechanisms are robust and responsive",
+                "adaptive governance evolves with the technology",
+                "benefits are explicitly shared across stakeholder groups",
+                "continuous education keeps users informed and empowered",
+            ]
+        else:
+            conditions = [
+                "commercial pressures override ethical considerations",
+                "deployment occurs without adequate testing",
+                "regulatory oversight fails to keep pace with innovation",
+                "user understanding lags behind implementation",
+                "technical complexity obscures important trade-offs",
+                "powerful interests dominate decision-making processes",
+                "short-term metrics drive optimization at the expense of long-term values",
+                "security vulnerabilities remain unaddressed",
+            ]
+        return random.choice(conditions)
+
+    def _generate_stakeholder_type(self) -> str:
+        """Generate a stakeholder type."""
+        stakeholders = [
+            "direct users",
+            "service providers",
+            "regulatory bodies",
+            "vulnerable populations",
+            "industry competitors",
+            "civil society organizations",
+            "technical developers",
+            "academic researchers",
+            "future generations",
+            "institutional purchasers",
+        ]
+        return random.choice(stakeholders)
+
+    def _generate_priority(self) -> str:
+        """Generate a stakeholder priority."""
+        priorities = [
+            "maximum functionality and convenience",
+            "data security and privacy protections",
+            "cost efficiency and scalability",
+            "accessibility and inclusivity",
+            "reliability and consistency",
+            "innovation and competitive advantage",
+            "compliance and risk mitigation",
+            "user control and customization",
+            "simplicity and ease of understanding",
+            "integration with existing systems",
+        ]
+        return random.choice(priorities)
+
+    def _generate_consequence_analysis(self) -> str:
+        """Generate analysis of consequences."""
+        analyses = [
+            "immediate benefits must be weighed against potential long-term harms",
+            "both intended and unintended consequences require thorough examination",
+            "the distribution of benefits and harms across different groups is asymmetric",
+            "second-order effects may be more significant than primary outcomes",
+            "probabilistic harms must be taken seriously even if uncertain",
+            "comparative analysis with alternatives shows important trade-offs",
+            "systemic impacts may differ significantly from individual-level effects",
+            "threshold effects could create tipping points with cascading consequences",
+        ]
+        return random.choice(analyses)
+
+    def _generate_historical_precedent(self) -> str:
+        """Generate historical precedent."""
+        precedents = [
+            "early adoption without adequate safeguards has repeatedly led to harmful outcomes",
+            "similar technologies have evolved to include stronger protections over time",
+            "initial concerns often materialize in unexpected ways during scaling",
+            "retroactive regulation has proven difficult once systems are widely deployed",
+            "stakeholder inclusion has consistently improved outcomes in similar cases",
+            "transparency requirements have increased following documented abuses",
+            "ethical frameworks have matured through iterative real-world applications",
+            "public trust has proven difficult to rebuild once compromised",
+        ]
+        return random.choice(precedents)
+
+    def _generate_rights_consideration(self) -> str:
+        """Generate rights-based considerations."""
+        considerations = [
+            "the right to privacy requires meaningful control over personal information",
+            "informational self-determination is fundamental to human dignity",
+            "equal treatment before automated systems is an emerging right",
+            "the right to explanation is increasingly recognized in algorithmic contexts",
+            "freedom from manipulation requires protection from subtle influences",
+            "cognitive liberty includes freedom from excessive technological nudging",
+            "meaningful human review should be available for significant decisions",
+            "rights to access and rectification are essential for data subjects",
+        ]
+        return random.choice(considerations)
+
+    def _generate_virtue_consideration(self) -> str:
+        """Generate virtue ethics considerations."""
+        considerations = [
+            "technologies should be designed to encourage human flourishing",
+            "systems should promote rather than undermine ethical character development",
+            "virtues like honesty and fairness should be embedded in system design",
+            "responsible innovation requires practical wisdom and moderation",
+            "compassion should guide decisions about vulnerable users",
+            "courage is needed to prioritize ethics over expediency",
+            "intellectual humility recognizes the limits of algorithmic solutions",
+            "prudence is essential when deploying technologies with uncertain impacts",
+        ]
+        return random.choice(considerations)
+
+    def _generate_societal_implication(self) -> str:
+        """Generate societal implications."""
+        implications = [
+            "social cohesion could be affected as technology mediates more interactions",
+            "civic discourse might be transformed in both positive and negative ways",
+            "power dynamics between institutions and individuals may shift significantly",
+            "collective decision-making processes could be enhanced or undermined",
+            "social norms regarding privacy and disclosure may evolve rapidly",
+            "intergenerational equity requires considering future impacts",
+            "cultural diversity in ethical standards must be respected",
+            "democratic oversight of technological development may face new challenges",
+        ]
+        return random.choice(implications)
+
+    def _generate_deontological_focus(self) -> str:
+        """Generate deontological focus."""
+        focuses = [
+            "whether the technology respects the inherent dignity of all persons",
+            "if users are treated as ends in themselves rather than merely as means",
+            "whether actions would be acceptable if universally applied",
+            "the intentions behind the system's design and deployment",
+            "whether the technology honors fundamental duties to users",
+            "if rights are being respected regardless of consequences",
+            "whether all stakeholders are treated with equal moral consideration",
+            "if the design process respects the moral autonomy of users",
+        ]
+        return random.choice(focuses)
+
+    def _generate_deontological_conclusion(self) -> str:
+        """Generate deontological conclusion."""
+        conclusions = [
+            "the current approach requires modification to fully respect human dignity",
+            "certain elements violate the categorical imperative and must be reconsidered",
+            "the design treats some users merely as means rather than as ends",
+            "fundamental duties to users are being upheld in principle but not in practice",
+            "rights-based considerations suggest significant changes are needed",
+            "the approach is permissible if informed consent is genuinely secured",
+            "intentions appear sound but implementation fails to honor core duties",
+            "moral consistency would require extending similar protections to all affected parties",
+        ]
+        return random.choice(conclusions)
+
+    def _generate_utilitarian_considerations(self) -> str:
+        """Generate utilitarian considerations."""
+        considerations = [
+            "the aggregate benefits and harms across all affected parties",
+            "whether alternative approaches might produce greater net welfare",
+            "how to quantify intangible impacts like privacy or autonomy",
+            "the distribution of utility across different stakeholder groups",
+            "long-term versus short-term utility calculations",
+            "the reliability of predictions about future consequences",
+            "preference satisfaction versus objective welfare measures",
+            "rule-utilitarian versus act-utilitarian perspectives",
+        ]
+        return random.choice(considerations)
+
+    def _generate_utilitarian_conclusion(self) -> str:
+        """Generate utilitarian conclusion."""
+        conclusions = [
+            "the likely benefits appear to outweigh potential harms if safeguards are implemented",
+            "alternative approaches could produce greater overall utility with less risk",
+            "a more targeted implementation would optimize the benefit-to-harm ratio",
+            "the uncertainty of long-term impacts suggests a more cautious approach",
+            "a phased rollout would allow for utility assessment before full deployment",
+            "potential harms to vulnerable minorities outweigh aggregate benefits",
+            "the utility calculation changes significantly when considering second-order effects",
+            "maximizing long-term utility requires stronger oversight mechanisms",
+        ]
+        return random.choice(conclusions)
+
+    def _generate_justice_considerations(self) -> str:
+        """Generate justice considerations."""
+        considerations = [
+            "disparate impacts across demographic groups require mitigation",
+            "access barriers could create new forms of inequality",
+            "procedural justice demands transparent decision criteria",
+            "distribution of benefits should prioritize the least advantaged",
+            "historical injustices might be perpetuated or amplified",
+            "fairness requires similar treatment in similar circumstances",
+            "capabilities and opportunities may be affected asymmetrically",
+            "intergenerational justice considerations for long-lived systems",
+        ]
+        return random.choice(considerations)
+
+    def _generate_professional_ethics(self) -> str:
+        """Generate professional ethics considerations."""
+        considerations = [
+            "practitioners must prioritize public welfare above other interests",
+            "professional codes of conduct require transparency about limitations",
+            "conflicts of interest must be disclosed and managed appropriately",
+            "the duty of care extends to foreseeable impacts of technical work",
+            "professional standards require ongoing competence as technology evolves",
+            "collaboration across disciplines is needed for ethical assessment",
+            "whistleblower protections should be available for ethical concerns",
+            "professional integrity requires refusing work that violates ethical principles",
+        ]
+        return random.choice(considerations)
+
+    def _generate_cultural_perspective(self) -> str:
+        """Generate cultural perspective considerations."""
+        considerations = [
+            "privacy expectations vary significantly across cultural contexts",
+            "concepts of fairness and justice have cultural dimensions",
+            "some communities may value collective benefits differently than individual rights",
+            "technological solutionism itself reflects particular cultural assumptions",
+            "diverse cultural perspectives should inform ethical frameworks",
+            "global deployment requires cultural sensitivity in implementation",
+            "different historical experiences shape risk perception across cultures",
+            "cultural values influence acceptable trade-offs and priorities",
+        ]
+        return random.choice(considerations)
+
+    def _generate_temporal_complexity(self) -> str:
+        """Generate temporal complexity considerations."""
+        considerations = [
+            "short-term efficiency gains may come at the cost of long-term resilience",
+            "path dependencies created today will shape options for future generations",
+            "adaptive capacity must be preserved for responding to emerging issues",
+            "technological lock-in can limit future ethical choices",
+            "values and norms evolve while deployed systems may remain static",
+            "reversibility becomes increasingly difficult as adoption spreads",
+            "learning curves for new technologies create temporary vulnerabilities",
+            "long-term maintenance of ethical standards requires ongoing vigilance",
+        ]
+        return random.choice(considerations)
+
+    def _generate_consent_analysis(self) -> str:
+        """Generate consent analysis."""
+        analyses = [
+            "truly informed consent requires genuine understanding of complex systems",
+            "default opt-in approaches often fail to secure meaningful consent",
+            "power imbalances can render formal consent procedures inadequate",
+            "consent should be dynamic and revocable as circumstances change",
+            "alternatives must be available without punitive consequences",
+            "consent interfaces often prioritize convenience over comprehension",
+            "special considerations apply for vulnerable or dependent users",
+            "presumed consent is inappropriate for novel applications with uncertain risks",
+        ]
+        return random.choice(analyses)
+
+    def _generate_privacy_analysis(self) -> str:
+        """Generate privacy analysis."""
+        analyses = [
+            "data minimization should be a foundational design principle",
+            "purpose limitation prevents function creep and unexpected uses",
+            "privacy by design requires proactive rather than reactive measures",
+            "differential privacy techniques could provide mathematical guarantees",
+            "data retention policies should align with necessity rather than utility",
+            "re-identification risks increase as data sets are combined",
+            "privacy impact assessments should precede new data uses",
+            "contextual integrity considers appropriate information flows",
+        ]
+        return random.choice(analyses)
+
+    def _generate_fairness_issues(self) -> str:
+        """Generate fairness issues."""
+        issues = [
+            "historical data often contains embedded biases that models reproduce",
+            "different fairness metrics may be mutually incompatible",
+            "seemingly neutral criteria can create disparate impacts",
+            "underrepresented groups often face higher error rates",
+            "feedback loops can amplify initial inequalities over time",
+            "proxy variables may reinstate bias even when protected attributes are excluded",
+            "fairness requires ongoing monitoring rather than one-time certification",
+            "different stakeholders may have legitimate but conflicting fairness definitions",
+        ]
+        return random.choice(issues)
+
+    def _generate_transparency_requirements(self) -> str:
+        """Generate transparency requirements."""
+        requirements = [
+            "explanations should be accessible to users with varying technical literacy",
+            "decision criteria must be documentable and justifiable",
+            "both technical mechanisms and governance processes require transparency",
+            "meaningful transparency includes understanding system limitations",
+            "audit trails should track key decisions and modifications",
+            "transparency may need to be tailored differently for various stakeholders",
+            "trade secrets should not prevent appropriate oversight",
+            "algorithmic impact assessments should be publicly available",
+        ]
+        return random.choice(requirements)
+
+    def _generate_accountability_mechanisms(self) -> str:
+        """Generate accountability mechanisms."""
+        mechanisms = [
+            "clear lines of responsibility for outcomes regardless of complexity",
+            "accessible processes for contesting decisions and seeking redress",
+            "regular independent audits of performance and impacts",
+            "penalties proportionate to potential harms for non-compliance",
+            "whistleblower protections for reporting ethical concerns",
+            "ongoing monitoring with stakeholder involvement",
+            "documentation requirements that enable external review",
+            "grievance mechanisms that include support for affected parties",
+        ]
+        return random.choice(mechanisms)
+
+    def _generate_risk_assessment(self) -> str:
+        """Generate risk assessment."""
+        assessments = [
+            "potential harms should be evaluated systematically before deployment",
+            "low-probability but high-impact risks warrant serious consideration",
+            "risk distribution across different populations requires explicit analysis",
+            "compounding risks from system interactions may be overlooked",
+            "unknown risks require adaptive governance and ongoing vigilance",
+            "reversibility and containment options limit potential damage",
+            "risk thresholds should be determined through inclusive deliberation",
+            "precautionary principles apply when risks are uncertain but potentially severe",
+        ]
+        return random.choice(assessments)
+
+    def _generate_innovation_balance(self) -> str:
+        """Generate innovation balance."""
+        balances = [
+            "ethical guardrails should channel rather than prevent innovation",
+            "staged deployment can balance progress with caution",
+            "innovation sandboxes allow testing under controlled conditions",
+            "principles-based approaches provide flexibility with accountability",
+            "participatory design incorporates diverse perspectives from the outset",
+            "value-sensitive design embeds ethical considerations in development",
+            "anticipatory governance prepares for emerging challenges",
+            "responsible innovation frameworks align technical and social progress",
+        ]
+        return random.choice(balances)
+
+    def _generate_recommendation_context(
+        self, analysis_context: Dict[str, str]
+    ) -> Dict[str, str]:
+        """Generate context for recommendation."""
+
+        # Primary recommendations based on analysis
+        primary_recommendations = [
+            "implement the system with enhanced transparency measures and opt-in consent",
+            "proceed with a limited pilot program with robust monitoring mechanisms",
+            "redesign key components to address identified ethical concerns",
+            "establish an independent oversight committee with diverse representation",
+            "develop clear usage guidelines with explicit ethical boundaries",
+            "adopt a phased implementation approach with checkpoints for review",
+            "incorporate additional safeguards for vulnerable populations",
+            "create a comprehensive data governance framework before proceeding",
+        ]
+
+        # Mitigation strategies
+        mitigation_strategies = [
+            "regular algorithmic audits by independent third parties",
+            "ongoing monitoring with clear thresholds for intervention",
+            "robust anonymization techniques for all personal data",
+            "user-friendly controls for privacy preferences",
+            "accessible explanations of system operation and limitations",
+            "formal impact assessments before major updates",
+            "channels for stakeholder feedback and reported concerns",
+            "sunset provisions requiring periodic reauthorization",
+        ]
+
+        # Secondary recommendations
+        secondary_recommendations = [
+            "develop educational resources to promote informed user decisions",
+            "establish an ethics advisory board with rotating membership",
+            "create a transparent incident response protocol",
+            "implement differential access levels based on necessity",
+            "require regular public reporting on system performance and impacts",
+            "develop industry standards for similar technologies",
+            "establish compensation mechanisms for demonstrated harms",
+            "invest in research addressing identified ethical challenges",
+        ]
+
+        # Generate the recommendation context
+        context = {
+            "primary_recommendation": random.choice(primary_recommendations),
+            "respected_principle": random.choice(
+                ["autonomy", "beneficence", "justice", "transparency", "privacy"]
+            ),
+            "acknowledged_concern": random.choice(
+                [
+                    "potential biases",
+                    "privacy risks",
+                    "accessibility barriers",
+                    "security vulnerabilities",
+                    "oversight limitations",
+                ]
+            ),
+            "mitigation_strategy": random.choice(mitigation_strategies),
+            "secondary_recommendation": random.choice(secondary_recommendations),
+            "monitoring_approach": random.choice(
+                [
+                    "impact assessment",
+                    "stakeholder feedback",
+                    "outcome metrics",
+                    "independent auditing",
+                    "automated monitoring",
+                ]
+            ),
+            # For the second template
+            "balanced_value_1": analysis_context["competing_value_1"],
+            "balanced_value_2": analysis_context["competing_value_2"],
+            "safeguard_1": random.choice(
+                [
+                    "data minimization protocols",
+                    "explicit user consent mechanisms",
+                    "algorithmic fairness checks",
+                    "human oversight for critical decisions",
+                    "robust security measures",
+                ]
+            ),
+            "safeguard_2": random.choice(
+                [
+                    "regular third-party audits",
+                    "transparent documentation",
+                    "accessible appeals process",
+                    "ongoing stakeholder consultation",
+                    "sunset provisions",
+                ]
+            ),
+            "engagement_method": random.choice(
+                [
+                    "public consultations",
+                    "representative advisory panels",
+                    "user feedback mechanisms",
+                    "community review boards",
+                    "multi-stakeholder governance",
+                ]
+            ),
+            "evaluation_focus": random.choice(
+                [
+                    "impact on vulnerable groups",
+                    "unintended consequences",
+                    "efficacy of safeguards",
+                    "changing risk landscape",
+                    "evolving social norms",
+                ]
+            ),
+            # For the third template
+            "recommendation_1": random.choice(
+                [
+                    "implement enhanced transparency measures",
+                    "strengthen consent procedures",
+                    "establish clear boundaries for use",
+                    "create oversight mechanisms",
+                    "develop explicit ethical guidelines",
+                ]
+            ),
+            "recommendation_2": random.choice(
+                [
+                    "regular independent audits",
+                    "ongoing stakeholder engagement",
+                    "robust monitoring protocols",
+                    "expanded user controls",
+                    "comprehensive documentation",
+                ]
+            ),
+            "recommendation_3": random.choice(
+                [
+                    "industry-wide ethical standards",
+                    "policy frameworks for governance",
+                    "research programs on ethical dimensions",
+                    "educational initiatives for users",
+                    "collaborative oversight mechanisms",
+                ]
+            ),
+            "prioritized_value": random.choice(
+                [
+                    "human dignity",
+                    "informed autonomy",
+                    "equitable access",
+                    "harm prevention",
+                    "accountability",
+                ]
+            ),
+            "trade_off_values": random.choice(
+                [
+                    "innovation and caution",
+                    "efficiency and inclusivity",
+                    "personalization and privacy",
+                    "automation and oversight",
+                    "scale and sensitivity",
+                ]
+            ),
+            "optimized_outcome": random.choice(
+                [
+                    "long-term trust",
+                    "sustainable innovation",
+                    "equitable benefits",
+                    "resilient systems",
+                    "ongoing improvement",
+                ]
+            ),
+            "minimized_harm": random.choice(
+                [
+                    "exclusion of vulnerable groups",
+                    "privacy violations",
+                    "erosion of autonomy",
+                    "unfair treatment",
+                    "unaccountable outcomes",
+                ]
+            ),
+            # For the fourth template
+            "accompanying_measure": random.choice(
+                [
+                    "robust oversight mechanisms",
+                    "comprehensive documentation",
+                    "ongoing impact assessment",
+                    "accessible control options",
+                    "proactive risk mitigation",
+                ]
+            ),
+            "ensured_outcome": random.choice(
+                [
+                    "equitable access",
+                    "meaningful consent",
+                    "appropriate use",
+                    "accountability for outcomes",
+                    "ongoing improvement",
+                ]
+            ),
+            "transparency_measure": random.choice(
+                [
+                    "public documentation",
+                    "algorithmic explainability",
+                    "regular disclosure reports",
+                    "accessible descriptions of processes",
+                    "notification of significant changes",
+                ]
+            ),
+            "stakeholder_treatment": random.choice(
+                [
+                    "actively consulted throughout implementation",
+                    "provided with meaningful opt-out options",
+                    "given accessible channels for feedback",
+                    "included in oversight processes",
+                    "informed of their rights and options",
+                ]
+            ),
+            "review_process": random.choice(
+                [
+                    "independent audits",
+                    "stakeholder consultations",
+                    "impact assessments",
+                    "performance evaluations",
+                    "ethical compliance reviews",
+                ]
+            ),
+        }
+
+        return context
+
+    def _extract_considerations(self, analysis: str) -> List[str]:
+        """Extract key considerations from the analysis."""
+        # This is a simplified version - in a real implementation,
+        # this could use NLP techniques to extract actual considerations
+        considerations = [
+            "Balancing innovation with ethical safeguards",
+            "Ensuring meaningful informed consent",
+            "Protecting privacy and data security",
+            "Preventing algorithmic bias and discrimination",
+            "Maintaining appropriate human oversight",
+            "Establishing clear lines of accountability",
+            "Providing transparency in operation and decisions",
+            "Considering impacts on vulnerable populations",
+        ]
+
+        # Randomly select 3-5 considerations
+        return random.sample(considerations, random.randint(3, 5))
+
+
+# Register the generator
+SyntheticDataGeneratorRegistry.register("ethics", EthicsDataGenerator)

src/multi_agent_training/specialized/generators/reasoning_generator.py

@@ -0,0 +1,581 @@
+"""
+Reasoning Module Synthetic Data Generator for MangoMAS Local
+
+This module generates synthetic training data for the reasoning capability.
+"""
+
+import json
+import logging
+import random
+from pathlib import Path
+from typing import Any, Dict, List
+
+from ..synthetic_data_generator import (SyntheticDataGenerator,
+                                        SyntheticDataGeneratorRegistry)
+
+logger = logging.getLogger(__name__)
+
+# Templates for reasoning scenarios
+REASONING_DOMAINS = [
+    "logical deduction",
+    "causal analysis",
+    "statistical reasoning",
+    "comparative analysis",
+    "analogical reasoning",
+    "counterfactual reasoning",
+    "inductive reasoning",
+    "abductive reasoning",
+    "conditional reasoning",
+    "syllogistic reasoning",
+]
+
+PREMISE_TEMPLATES = [
+    "In a study of {topic}, researchers observed that {observation}.",
+    "The data shows that {observation} when {condition}.",
+    "Historical records indicate that {observation} throughout {period}.",
+    "Given that {fact_1} and {fact_2}, we need to determine {question}.",
+    "In the context of {domain}, we observe that {observation}.",
+    "When analyzing {topic}, we can see that {observation} correlates with {factor}.",
+    "System logs reveal that {observation} occurs whenever {condition}.",
+    "The problem states that {observation} under conditions where {condition}.",
+    "In the experiment, {observation} was measured when {condition}.",
+    "The scenario presents a situation where {observation} after {event}.",
+]
+
+REASONING_TEMPLATES = [
+    "First, we need to identify the key variables: {variables}. Looking at the relationship between {var1} and {var2}, we can see that {relationship}. This suggests that {inference}. Furthermore, considering {var3}, we can deduce that {deduction}.",
+    "Let's analyze this step by step. If {fact_1}, then {consequence_1}. Given that {fact_2}, we can infer {consequence_2}. Combining these insights, {combined_inference}. Additionally, {extra_insight} further supports this reasoning.",
+    "Breaking this down systematically: 1) {step_1}, 2) {step_2}, 3) {step_3}, 4) {step_4}. The logical connection between steps 2 and 3 shows that {connection}, which leads us to {inference}.",
+    "The key insight here is that {key_insight}. This means that when {condition}, the result is {result}. We can verify this by examining {evidence}, which confirms {confirmation}.",
+    "To solve this, I'll use {method}. Starting with {starting_point}, I can determine that {determination}. This leads to {intermediate_conclusion}, and finally to {final_step}.",
+]
+
+CONCLUSION_TEMPLATES = [
+    "Therefore, we can conclude that {conclusion}.",
+    "Based on this analysis, the answer is {conclusion}.",
+    "The evidence strongly suggests that {conclusion}.",
+    "This reasoning leads to the conclusion that {conclusion}.",
+    "We can definitively state that {conclusion} based on the above analysis.",
+]
+
+# Topics and facts to fill in templates
+TOPICS = [
+    "climate patterns",
+    "market trends",
+    "neural network performance",
+    "population demographics",
+    "traffic flow optimization",
+    "disease spread",
+    "economic indicators",
+    "algorithm efficiency",
+    "material properties",
+    "social network dynamics",
+    "cognitive biases",
+    "language acquisition",
+    "genetic inheritance",
+    "planetary motion",
+    "quantum phenomena",
+]
+
+FACTS = [
+    "energy consumption increases with temperature",
+    "user engagement drops after 3 minutes",
+    "performance plateaus with more than 8 layers",
+    "error rates increase exponentially with load",
+    "response time correlates with system memory",
+    "conversion rates are highest on Tuesdays",
+    "signal strength decreases with distance squared",
+    "learning rate affects convergence time",
+    "failure probability doubles every 5 years",
+    "growth rate is proportional to nutrient concentration",
+]
+
+
+class ReasoningDataGenerator(SyntheticDataGenerator):
+    """Generator for synthetic reasoning data."""
+
+    def _load_templates(self) -> List[Dict[str, Any]]:
+        """Load reasoning templates."""
+        templates = {
+            "domains": REASONING_DOMAINS,
+            "premises": PREMISE_TEMPLATES,
+            "reasoning": REASONING_TEMPLATES,
+            "conclusions": CONCLUSION_TEMPLATES,
+            "topics": TOPICS,
+            "facts": FACTS,
+        }
+
+        # Additional domain-specific templates could be loaded from files
+        template_path = self.config.get("template_path")
+        if template_path and Path(template_path).exists():
+            try:
+                with open(template_path, "r", encoding="utf-8") as f:
+                    custom_templates = json.load(f)
+                    for key, values in custom_templates.items():
+                        if key in templates and isinstance(values, list):
+                            templates[key].extend(values)
+            except Exception as e:
+                logger.warning(f"Failed to load custom templates: {e}")
+
+        return templates
+
+    def generate_example(self) -> Dict[str, Any]:
+        """Generate a single reasoning example."""
+        # Select templates
+        domain = random.choice(self.templates["domains"])
+        premise_template = random.choice(self.templates["premises"])
+        reasoning_template = random.choice(self.templates["reasoning"])
+        conclusion_template = random.choice(self.templates["conclusions"])
+
+        # Generate content for templates
+        topic = random.choice(self.templates["topics"])
+        fact_1 = random.choice(self.templates["facts"])
+        fact_2 = random.choice(self.templates["facts"])
+        while fact_2 == fact_1:  # Ensure different facts
+            fact_2 = random.choice(self.templates["facts"])
+
+        # Generate dynamic content
+        observation = self._generate_observation(topic)
+        condition = self._generate_condition(topic)
+        variables = self._generate_variables(domain)
+        var1, var2, var3 = variables[:3]
+        relationship = self._generate_relationship(var1, var2)
+        inference = self._generate_inference(relationship)
+        deduction = self._generate_deduction(var3)
+
+        # Fill in templates
+        context = {
+            "domain": domain,
+            "topic": topic,
+            "observation": observation,
+            "condition": condition,
+            "fact_1": fact_1,
+            "fact_2": fact_2,
+            "variables": ", ".join(variables),
+            "var1": var1,
+            "var2": var2,
+            "var3": var3,
+            "relationship": relationship,
+            "inference": inference,
+            "deduction": deduction,
+            "question": self._generate_question(topic),
+            "period": self._generate_period(),
+            "factor": self._generate_factor(topic),
+            "event": self._generate_event(),
+            "consequence_1": self._generate_consequence(fact_1),
+            "consequence_2": self._generate_consequence(fact_2),
+            "combined_inference": self._generate_combined_inference(fact_1, fact_2),
+            "extra_insight": self._generate_extra_insight(topic),
+            "step_1": self._generate_step(1, domain),
+            "step_2": self._generate_step(2, domain),
+            "step_3": self._generate_step(3, domain),
+            "step_4": self._generate_step(4, domain),
+            "connection": self._generate_connection(),
+            "key_insight": self._generate_key_insight(domain),
+            "result": self._generate_result(domain),
+            "evidence": self._generate_evidence(domain),
+            "confirmation": self._generate_confirmation(domain),
+            "method": self._generate_method(domain),
+            "starting_point": self._generate_starting_point(domain),
+            "determination": self._generate_determination(domain),
+            "intermediate_conclusion": self._generate_intermediate_conclusion(domain),
+            "final_step": self._generate_final_step(domain),
+        }
+
+        # Format strings with context
+        premise = premise_template.format(**context)
+        reasoning = reasoning_template.format(**context)
+
+        # Generate conclusion based on the premise and reasoning
+        conclusion = self._generate_conclusion(premise, reasoning)
+        context["conclusion"] = conclusion
+        conclusion_text = conclusion_template.format(**context)
+
+        # Return the structured example
+        return {
+            "premise": premise,
+            "reasoning": reasoning,
+            "conclusion": conclusion_text,
+            "domain": domain,
+            "complexity": random.choice(["low", "medium", "high"]),
+            "metadata": {
+                "variables": variables,
+                "topic": topic,
+                "reasoning_type": domain,
+            },
+        }
+
+    # Helper methods to generate dynamic content
+    def _generate_observation(self, topic: str) -> str:
+        observations = [
+            "the rate of change increases over time",
+            "there is a strong correlation between input and output variables",
+            "performance degrades under specific conditions",
+            "the system exhibits unexpected behavior when stressed",
+            "outliers significantly impact the overall trend",
+            "recurring patterns emerge after sufficient iterations",
+            "the distribution follows a power law rather than normal distribution",
+            "feedback loops amplify small initial differences",
+            "thresholds exist beyond which behavior changes dramatically",
+            "cyclical patterns emerge with a period of varying length",
+        ]
+        return random.choice(observations)
+
+    def _generate_condition(self, topic: str) -> str:
+        conditions = [
+            "the system is under heavy load",
+            "external factors remain constant",
+            "all variables are optimized simultaneously",
+            "specific constraints are applied",
+            "the environment changes unexpectedly",
+            "resource limitations come into play",
+            "feedback mechanisms are activated",
+            "multiple agents interact simultaneously",
+            "time delays exceed a critical threshold",
+            "boundary conditions are enforced",
+        ]
+        return random.choice(conditions)
+
+    def _generate_variables(self, domain: str) -> List[str]:
+        variable_sets = {
+            "logical deduction": [
+                "premise validity",
+                "logical consistency",
+                "conclusion strength",
+                "assumption bias",
+            ],
+            "causal analysis": [
+                "cause magnitude",
+                "effect delay",
+                "confounding factors",
+                "intervention efficacy",
+            ],
+            "statistical reasoning": [
+                "sample size",
+                "confidence interval",
+                "p-value",
+                "effect size",
+                "statistical power",
+            ],
+            "comparative analysis": [
+                "baseline performance",
+                "improvement margin",
+                "relative efficiency",
+                "comparison fairness",
+            ],
+            "analogical reasoning": [
+                "source similarity",
+                "target applicability",
+                "mapping strength",
+                "inference validity",
+            ],
+        }
+
+        # Get domain-specific variables or use generic ones
+        variables = variable_sets.get(
+            domain,
+            [
+                "factor A",
+                "factor B",
+                "factor C",
+                "response variable",
+                "control variable",
+            ],
+        )
+
+        # Shuffle and return
+        random.shuffle(variables)
+        return variables
+
+    def _generate_relationship(self, var1: str, var2: str) -> str:
+        relationships = [
+            f"an increase in {var1} leads to a proportional increase in {var2}",
+            f"{var1} and {var2} have an inverse relationship",
+            f"changes in {var1} precede changes in {var2} by a consistent time interval",
+            f"{var1} influences {var2} only after exceeding a critical threshold",
+            f"the relationship between {var1} and {var2} is non-linear and follows a power law",
+            f"{var1} and {var2} are conditionally independent given certain conditions",
+            f"extreme values of {var1} have a disproportionate effect on {var2}",
+            f"the correlation between {var1} and {var2} changes direction over time",
+        ]
+        return random.choice(relationships)
+
+    def _generate_inference(self, relationship: str) -> str:
+        inferences = [
+            "we should focus our optimization efforts on the most sensitive parameters",
+            "the system will likely reach equilibrium after sufficient time",
+            "interventions should target root causes rather than symptoms",
+            "small changes can potentially lead to significant improvements",
+            "we need to account for interaction effects between variables",
+            "the observed behavior is likely part of a larger pattern",
+            "we should implement fail-safes for extreme conditions",
+            "multiple pathways may lead to the same outcome",
+        ]
+        return random.choice(inferences)
+
+    def _generate_deduction(self, var: str) -> str:
+        deductions = [
+            f"optimizing {var} alone will not solve the underlying problem",
+            f"changes in {var} represent a leading indicator for system performance",
+            f"the role of {var} has been previously underestimated",
+            f"{var} acts as a moderating variable in this context",
+            f"the impact of {var} follows a diminishing returns pattern",
+            f"{var} exhibits threshold effects that must be accounted for",
+            f"historical data on {var} supports this conclusion",
+            f"contrary to conventional wisdom, {var} is not the limiting factor",
+        ]
+        return random.choice(deductions)
+
+    def _generate_question(self, topic: str) -> str:
+        questions = [
+            "how to optimize performance under these conditions",
+            "whether the observed pattern will continue in the future",
+            "which factors contribute most significantly to the outcome",
+            "how to mitigate negative effects while preserving benefits",
+            "what intervention would produce the most efficient solution",
+            "how robust the system is to unexpected disturbances",
+            "whether the findings can be generalized to other contexts",
+            "how to distinguish correlation from causation in this case",
+        ]
+        return random.choice(questions)
+
+    def _generate_period(self) -> str:
+        periods = [
+            "the past decade",
+            "periods of economic volatility",
+            "the system's entire operational history",
+            "multiple successive iterations",
+            "both growth and decline phases",
+            "controlled experimental conditions",
+            "repeated stress-test cycles",
+            "varying environmental conditions",
+        ]
+        return random.choice(periods)
+
+    def _generate_factor(self, topic: str) -> str:
+        factors = [
+            "resource utilization",
+            "system complexity",
+            "external pressure",
+            "user engagement",
+            "adaptation rate",
+            "failure frequency",
+            "communication efficiency",
+            "innovation adoption",
+        ]
+        return random.choice(factors)
+
+    def _generate_event(self) -> str:
+        events = [
+            "system initialization",
+            "critical resource depletion",
+            "unexpected environmental change",
+            "crossing a performance threshold",
+            "implementing a major upgrade",
+            "integrating new components",
+            "encountering novel inputs",
+            "recovering from failure",
+        ]
+        return random.choice(events)
+
+    def _generate_consequence(self, fact: str) -> str:
+        return "the system will adapt by adjusting its parameters accordingly"
+
+    def _generate_combined_inference(self, fact1: str, fact2: str) -> str:
+        return "we can establish a clear causal relationship between the observed phenomena"
+
+    def _generate_extra_insight(self, topic: str) -> str:
+        insights = [
+            "temporal patterns reveal cyclical behavior",
+            "boundary conditions significantly affect outcomes",
+            "network effects amplify individual contributions",
+            "emergent properties cannot be predicted from components alone",
+            "system resilience depends on redundant pathways",
+            "optimization often involves trade-offs between competing goals",
+            "adaptation requires continuous feedback and adjustment",
+            "complex systems often exhibit counterintuitive behavior",
+        ]
+        return random.choice(insights)
+
+    def _generate_step(self, step_num: int, domain: str) -> str:
+        if step_num == 1:
+            steps = [
+                "Identify the key variables and their relationships",
+                "Establish the initial conditions and constraints",
+                "Define the problem space and boundaries",
+                "Gather relevant data and observations",
+                "Frame the question in precise terms",
+            ]
+        elif step_num == 2:
+            steps = [
+                "Analyze the patterns and correlations in the data",
+                "Apply appropriate analytical methods",
+                "Consider alternative explanations",
+                "Map the causal relationships between factors",
+                "Identify potential confounding variables",
+            ]
+        elif step_num == 3:
+            steps = [
+                "Evaluate the strength of evidence for each possibility",
+                "Synthesize insights from multiple perspectives",
+                "Test hypotheses against available data",
+                "Assess the logical consistency of arguments",
+                "Consider edge cases and exceptions",
+            ]
+        else:
+            steps = [
+                "Draw conclusions based on the strongest evidence",
+                "Formulate actionable recommendations",
+                "Identify remaining uncertainties",
+                "Propose methods to validate findings",
+                "Connect conclusions to the original question",
+            ]
+        return random.choice(steps)
+
+    def _generate_connection(self) -> str:
+        connections = [
+            "there's a causal relationship rather than mere correlation",
+            "feedback mechanisms create self-reinforcing patterns",
+            "threshold effects trigger qualitative changes in behavior",
+            "multiple factors interact in non-linear ways",
+            "temporal sequences reveal important dependencies",
+            "structural constraints limit possible outcomes",
+            "probabilistic influences accumulate deterministically",
+            "conditional dependencies reveal deeper patterns",
+        ]
+        return random.choice(connections)
+
+    def _generate_key_insight(self, domain: str) -> str:
+        insights = [
+            "optimizing for average cases often fails at the extremes",
+            "emergent properties cannot be reduced to component behaviors",
+            "apparent contradictions point to incomplete models",
+            "historical patterns constrain future possibilities",
+            "local optimizations can lead to global suboptimality",
+            "precision must be balanced with generalizability",
+            "second-order effects often dominate in the long run",
+            "robust systems prioritize adaptation over optimization",
+        ]
+        return random.choice(insights)
+
+    def _generate_result(self, domain: str) -> str:
+        results = [
+            "performance improves non-linearly",
+            "stability increases at the cost of responsiveness",
+            "resource utilization becomes more efficient",
+            "adaptability improves in novel situations",
+            "resilience to disturbances increases",
+            "error rates decrease systematically",
+            "learning accelerates with experience",
+            "coordination emerges without central control",
+        ]
+        return random.choice(results)
+
+    def _generate_evidence(self, domain: str) -> str:
+        evidence = [
+            "historical performance data",
+            "controlled experimental results",
+            "comparative case studies",
+            "simulation outcomes under varied conditions",
+            "natural experiments from system perturbations",
+            "user feedback and behavioral patterns",
+            "statistical analysis of large datasets",
+            "theoretical models with empirical validation",
+        ]
+        return random.choice(evidence)
+
+    def _generate_confirmation(self, domain: str) -> str:
+        confirmations = [
+            "the hypothesized mechanism actually operates as expected",
+            "predicted outcomes match observed results",
+            "alternative explanations can be ruled out",
+            "the pattern holds across different contexts",
+            "edge cases follow the same principles",
+            "the model successfully predicts future behavior",
+            "interventions produce expected effects",
+            "independent measures converge on the same conclusion",
+        ]
+        return random.choice(confirmations)
+
+    def _generate_method(self, domain: str) -> str:
+        methods = [
+            "systematic decomposition into components",
+            "counterfactual analysis",
+            "process tracing through causal chains",
+            "comparative analysis of similar cases",
+            "statistical inference from patterns",
+            "first principles reasoning",
+            "model-based simulation",
+            "abductive inference to the best explanation",
+        ]
+        return random.choice(methods)
+
+    def _generate_starting_point(self, domain: str) -> str:
+        starting_points = [
+            "the fundamental constraints of the system",
+            "established principles in this domain",
+            "patterns observed in similar situations",
+            "the minimal necessary conditions",
+            "key defining relationships",
+            "initial boundary conditions",
+            "critical assumptions that must hold",
+            "invariant properties across contexts",
+        ]
+        return random.choice(starting_points)
+
+    def _generate_determination(self, domain: str) -> str:
+        determinations = [
+            "certain factors exert disproportionate influence",
+            "system behavior follows predictable patterns under specific conditions",
+            "apparent anomalies actually confirm deeper principles",
+            "constraints channel possible outcomes in specific directions",
+            "dynamic equilibria emerge from competing forces",
+            "feedback loops stabilize or amplify depending on parameters",
+            "path dependencies limit future possibilities",
+            "critical thresholds separate qualitatively different regimes",
+        ]
+        return random.choice(determinations)
+
+    def _generate_intermediate_conclusion(self, domain: str) -> str:
+        conclusions = [
+            "we need to reconsider fundamental assumptions",
+            "the system exhibits emergent properties not predictable from components",
+            "apparent contradictions resolve at a higher level of analysis",
+            "complex interactions require a more nuanced approach",
+            "optimal solutions balance multiple competing objectives",
+            "robustness comes at the cost of peak performance",
+            "adaptability requires maintaining strategic flexibility",
+            "precision must be traded off against generalizability",
+        ]
+        return random.choice(conclusions)
+
+    def _generate_final_step(self, domain: str) -> str:
+        final_steps = [
+            "we can formulate a general principle that applies broadly",
+            "we can predict system behavior under novel conditions",
+            "we can design interventions that leverage key mechanisms",
+            "we can identify early warning signals for critical transitions",
+            "we can optimize for robust performance across scenarios",
+            "we can balance competing objectives through targeted trade-offs",
+            "we can establish boundaries of applicability for our conclusions",
+            "we can translate insights into actionable recommendations",
+        ]
+        return random.choice(final_steps)
+
+    def _generate_conclusion(self, premise: str, reasoning: str) -> str:
+        """Generate a coherent conclusion based on premise and reasoning."""
+        conclusions = [
+            "the observed patterns indicate a fundamental relationship between key variables",
+            "we should prioritize interventions that address root causes rather than symptoms",
+            "the system's behavior can be predicted with reasonable accuracy under specified conditions",
+            "optimizing for extreme cases provides more robust performance than optimizing for average cases",
+            "adaptation mechanisms are essential for maintaining performance in changing environments",
+            "complex interactions between components create emergent properties at the system level",
+            "resource allocation should follow a dynamic rather than static strategy",
+            "feedback loops must be carefully managed to prevent unintended consequences",
+        ]
+        return random.choice(conclusions)
+
+
+# Register the generator
+SyntheticDataGeneratorRegistry.register("reasoning", ReasoningDataGenerator)

src/multi_agent_training/specialized/memory_module.py

@@ -0,0 +1,317 @@
+"""
+Memory Training Module for MangoMAS Local
+
+This module implements specialized training for memory and context retention capabilities,
+adapted from the AWS backup system for local training.
+"""
+
+import json
+import logging
+import os
+import random
+from typing import Any, Dict, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+
+from ..core_framework import SpecializedTrainingModule, TrainingModuleConfig
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryDataset(Dataset):
+    """Dataset for training memory and context retention capabilities."""
+
+    def __init__(self, data_path: str, tokenizer, max_length: int = 1024):
+        """
+        Initialize the memory dataset.
+
+        Args:
+            data_path: Path to the memory training data file
+            tokenizer: Tokenizer for text processing
+            max_length: Maximum sequence length
+        """
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.data = self._load_data(data_path)
+
+        logger.info(f"Loaded memory dataset with {len(self.data)} examples")
+
+    def _load_data(self, data_path: str) -> List[Dict]:
+        """Load memory training data."""
+        data = []
+        with open(data_path, "r", encoding="utf-8") as f:
+            for line in f:
+                try:
+                    item = json.loads(line.strip())
+                    # Validate required fields
+                    if "conversation" in item and isinstance(
+                        item["conversation"], list
+                    ):
+                        data.append(item)
+                except json.JSONDecodeError:
+                    continue
+        return data
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        item = self.data[idx]
+
+        # Format the conversation for memory training
+        conversation = item["conversation"]
+        context = "\n".join(
+            [f"{turn['role']}: {turn['content']}" for turn in conversation[:-1]]
+        )
+        target = conversation[-1]["content"]
+
+        prompt = f"Context:\n{context}\nResponse: {target}"
+
+        # Tokenize
+        encoding = self.tokenizer(
+            prompt,
+            max_length=self.max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        return {
+            "input_ids": encoding["input_ids"].squeeze(),
+            "attention_mask": encoding["attention_mask"].squeeze(),
+            "labels": encoding["input_ids"].squeeze(),
+        }
+
+
+class MemoryTrainingModule(SpecializedTrainingModule):
+    """Specialized training module for memory and context retention capabilities."""
+
+    def __init__(self, config: TrainingModuleConfig, tokenizer):
+        """
+        Initialize the memory training module.
+
+        Args:
+            config: Module configuration
+            tokenizer: Tokenizer for text processing
+        """
+        super().__init__(config, tokenizer)
+
+        # Initialize memory-specific components
+        self.memory_loss = nn.CrossEntropyLoss(ignore_index=-100)
+        self.metrics = {
+            "memory_loss": 0.0,
+            "context_retention": 0.0,
+            "coherence_score": 0.0,
+        }
+
+        logger.info("Initialized MemoryTrainingModule")
+
+    def prepare_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """
+        Prepare a batch of data for memory training.
+
+        Args:
+            batch: The input batch from the dataloader
+
+        Returns:
+            Processed batch ready for memory training
+        """
+        # Move batch to device
+        prepared_batch = {}
+        for key, value in batch.items():
+            if isinstance(value, torch.Tensor):
+                prepared_batch[key] = value.to(self.device)
+            else:
+                prepared_batch[key] = value
+
+        return prepared_batch
+
+    def compute_loss(
+        self, student_outputs: Any, teacher_outputs: Any, batch: Dict[str, torch.Tensor]
+    ) -> torch.Tensor:
+        """
+        Compute the memory-specific loss.
+
+        Args:
+            student_outputs: Outputs from the student model
+            teacher_outputs: Outputs from the teacher model
+            batch: The processed input batch
+
+        Returns:
+            Loss tensor for memory training
+        """
+        try:
+            # Extract logits from model outputs
+            if hasattr(student_outputs, "logits"):
+                student_logits = student_outputs.logits
+            else:
+                student_logits = student_outputs
+
+            if hasattr(teacher_outputs, "logits"):
+                teacher_logits = teacher_outputs.logits
+            else:
+                teacher_logits = teacher_outputs
+
+            # Get labels from batch
+            labels = batch.get("labels", batch.get("input_ids"))
+
+            # Compute cross entropy loss for memory
+            shift_logits = student_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            memory_loss = self.memory_loss(
+                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
+            )
+
+            # Add KL divergence loss between student and teacher
+            if teacher_logits is not None:
+                kl_loss = F.kl_div(
+                    F.log_softmax(student_logits, dim=-1),
+                    F.softmax(teacher_logits, dim=-1),
+                    reduction="batchmean",
+                )
+                total_loss = memory_loss + 0.1 * kl_loss
+            else:
+                total_loss = memory_loss
+
+            # Update metrics
+            self.metrics["memory_loss"] = memory_loss.item()
+
+            return total_loss * self.loss_weight
+
+        except Exception as e:
+            logger.error(f"Error computing memory loss: {e}")
+            # Return a small loss to avoid training failure
+            return torch.tensor(0.01, requires_grad=True)
+
+    def get_metrics(self) -> Dict[str, float]:
+        """
+        Get metrics specific to memory training.
+
+        Returns:
+            Dictionary of memory metrics
+        """
+        return self.metrics.copy()
+
+    def generate_synthetic_memory_data(
+        self, output_path: str, num_samples: int = 1000
+    ) -> None:
+        """
+        Generate synthetic memory training data.
+
+        Args:
+            output_path: Path to save the generated data
+            num_samples: Number of samples to generate
+        """
+        # This is a simplified implementation based on the AWS backup
+        # In a full implementation, this would be much more sophisticated
+
+        conversation_templates = [
+            [
+                {
+                    "role": "user",
+                    "content": "Hi, my name is Alex and I'm interested in machine learning.",
+                },
+                {
+                    "role": "assistant",
+                    "content": "Hello Alex! I'd be happy to discuss machine learning with you. What aspects are you most interested in?",
+                },
+                {
+                    "role": "user",
+                    "content": "I'm particularly interested in natural language processing.",
+                },
+                {
+                    "role": "assistant",
+                    "content": "NLP is a fascinating field! It's used for tasks like translation, summarization, and question answering.",
+                },
+                {
+                    "role": "user",
+                    "content": "What do you think would be a good first project?",
+                },
+                {
+                    "role": "assistant",
+                    "content": "For a beginner in NLP, I'd recommend starting with a text classification project, like sentiment analysis.",
+                },
+            ],
+            [
+                {
+                    "role": "user",
+                    "content": "I'm planning a trip to Japan next spring.",
+                },
+                {
+                    "role": "assistant",
+                    "content": "That sounds exciting! Japan is beautiful in spring with cherry blossoms. What cities are you planning to visit?",
+                },
+                {
+                    "role": "user",
+                    "content": "I'm thinking Tokyo, Kyoto, and maybe Osaka.",
+                },
+                {
+                    "role": "assistant",
+                    "content": "Great choices! Tokyo has modern attractions, Kyoto has historical temples, and Osaka is known for amazing food.",
+                },
+                {
+                    "role": "user",
+                    "content": "What's the best way to travel between these cities?",
+                },
+                {
+                    "role": "assistant",
+                    "content": "The Shinkansen (bullet train) is the most efficient way to travel between these cities. It's fast, comfortable, and reliable.",
+                },
+            ],
+        ]
+
+        recall_templates = [
+            {
+                "recall_context": "what was my name again?",
+                "recall_target": "Your name is Alex, as you mentioned at the beginning of our conversation.",
+            },
+            {
+                "recall_context": "which cities did I say I wanted to visit?",
+                "recall_target": "You mentioned you're planning to visit Tokyo, Kyoto, and possibly Osaka during your trip to Japan.",
+            },
+        ]
+
+        # Generate variations
+        output_data = []
+        for _ in range(num_samples):
+            template_idx = random.randint(0, len(conversation_templates) - 1)
+            conversation = conversation_templates[template_idx].copy()
+
+            # Add a recall question if this is the right template
+            if template_idx < len(recall_templates):
+                recall_template = recall_templates[template_idx]
+
+                # Add a user question asking for recall
+                conversation.append(
+                    {"role": "user", "content": recall_template["recall_context"]}
+                )
+
+                # Create the full example with recall targets
+                example = {
+                    "conversation": conversation,
+                    "recall_context": recall_template["recall_context"],
+                    "recall_target": recall_template["recall_target"],
+                    "metadata": {"generated": True, "requires_memory": True},
+                }
+            else:
+                # Regular conversation without specific recall target
+                example = {
+                    "conversation": conversation,
+                    "metadata": {"generated": True, "requires_memory": False},
+                }
+
+            output_data.append(example)
+
+        # Save to file
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        with open(output_path, "w", encoding="utf-8") as f:
+            for item in output_data:
+                f.write(json.dumps(item) + "\n")
+
+        logger.info(
+            f"Generated {len(output_data)} synthetic memory examples at {output_path}"
+        )

src/multi_agent_training/specialized/reasoning_module.py

@@ -0,0 +1,404 @@
+"""
+Reasoning Training Module for MangoMAS Local
+
+This module implements specialized training for reasoning capabilities,
+adapted from the AWS backup system for local training.
+"""
+
+import json
+import logging
+import os
+import random
+import re
+from typing import Any, Dict, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+
+from ..core_framework import SpecializedTrainingModule, TrainingModuleConfig
+
+logger = logging.getLogger(__name__)
+
+
+class ReasoningDataset(Dataset):
+    """Dataset for training reasoning capabilities."""
+
+    def __init__(self, data_path: str, tokenizer, max_length: int = 512):
+        """
+        Initialize the reasoning dataset.
+
+        Args:
+            data_path: Path to the reasoning data file
+            tokenizer: Tokenizer for text processing
+            max_length: Maximum sequence length
+        """
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.data = self._load_data(data_path)
+
+        logger.info(f"Loaded reasoning dataset with {len(self.data)} examples")
+
+    def _load_data(self, data_path: str) -> List[Dict]:
+        """Load reasoning training data."""
+        data = []
+        with open(data_path, "r", encoding="utf-8") as f:
+            for line in f:
+                try:
+                    item = json.loads(line.strip())
+                    # Validate required fields
+                    if "question" in item and "reasoning" in item and "answer" in item:
+                        data.append(item)
+                except json.JSONDecodeError:
+                    continue
+        return data
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        item = self.data[idx]
+
+        # Format the reasoning prompt
+        prompt = f"Question: {item['question']}\nReasoning: {item['reasoning']}\nAnswer: {item['answer']}"
+
+        # Tokenize
+        encoding = self.tokenizer(
+            prompt,
+            max_length=self.max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        return {
+            "input_ids": encoding["input_ids"].squeeze(),
+            "attention_mask": encoding["attention_mask"].squeeze(),
+            "labels": encoding["input_ids"].squeeze(),
+        }
+
+
+class ReasoningEvaluator:
+    """Evaluator for reasoning capabilities."""
+
+    def __init__(self, tokenizer):
+        """
+        Initialize the reasoning evaluator.
+
+        Args:
+            tokenizer: Tokenizer for text processing
+        """
+        self.tokenizer = tokenizer
+        self.metrics = {
+            "logical_consistency": 0.0,
+            "premise_relevance": 0.0,
+            "conclusion_validity": 0.0,
+            "steps_coherence": 0.0,
+        }
+
+    def evaluate(self, model, eval_dataset: ReasoningDataset) -> Dict[str, float]:
+        """
+        Evaluate reasoning capabilities on the provided dataset.
+
+        Args:
+            model: The model to evaluate
+            eval_dataset: Dataset of reasoning examples
+
+        Returns:
+            Dictionary of evaluation metrics
+        """
+        model.eval()
+        device = next(model.parameters()).device
+
+        # Reset metrics
+        for key in self.metrics:
+            self.metrics[key] = 0.0
+
+        total_examples = min(
+            len(eval_dataset), 100
+        )  # Limit to 100 examples for efficiency
+
+        with torch.no_grad():
+            for idx in range(total_examples):
+                example = eval_dataset[idx]
+                premise = example["premise"]
+
+                # Generate reasoning and conclusion from premise
+                prompt = f"Premise: {premise}\nReasoning:"
+                input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(
+                    device
+                )
+
+                generated_ids = model.generate(
+                    input_ids, max_length=512, temperature=0.7, num_return_sequences=1
+                )
+
+                generated_text = self.tokenizer.decode(
+                    generated_ids[0], skip_special_tokens=True
+                )
+
+                # Extract reasoning and conclusion from generated text
+                try:
+                    generated_reasoning = re.search(
+                        r"Reasoning:(.*?)(?:Conclusion:|$)", generated_text, re.DOTALL
+                    )
+                    generated_conclusion = re.search(
+                        r"Conclusion:(.*?)$", generated_text, re.DOTALL
+                    )
+
+                    if generated_reasoning:
+                        gen_reasoning = generated_reasoning.group(1).strip()
+                    else:
+                        gen_reasoning = ""
+
+                    if generated_conclusion:
+                        gen_conclusion = generated_conclusion.group(1).strip()
+                    else:
+                        gen_conclusion = ""
+
+                    # Evaluate reasoning quality
+                    self._update_metrics(
+                        premise=premise,
+                        expected_reasoning=example["reasoning"],
+                        expected_conclusion=example["conclusion"],
+                        generated_reasoning=gen_reasoning,
+                        generated_conclusion=gen_conclusion,
+                    )
+                except Exception as e:
+                    logger.error(f"Error evaluating reasoning: {e}")
+
+        # Calculate averages
+        for key in self.metrics:
+            self.metrics[key] /= total_examples
+
+        return self.metrics
+
+    def _update_metrics(
+        self,
+        premise: str,
+        expected_reasoning: str,
+        expected_conclusion: str,
+        generated_reasoning: str,
+        generated_conclusion: str,
+    ) -> None:
+        """
+        Update reasoning metrics based on a single example.
+
+        Args:
+            premise: Input premise
+            expected_reasoning: Expected reasoning steps
+            expected_conclusion: Expected conclusion
+            generated_reasoning: Generated reasoning steps
+            generated_conclusion: Generated conclusion
+        """
+        # Very simplified evaluation - in a real system, this would use more sophisticated
+        # semantic similarity and logical consistency checking
+
+        # Logical consistency - check if reasoning follows from premise
+        self.metrics["logical_consistency"] += 0.5  # Simplified placeholder
+
+        # Premise relevance - check if reasoning references key terms from premise
+        premise_terms = set(premise.lower().split())
+        reasoning_terms = set(generated_reasoning.lower().split())
+        term_overlap = len(premise_terms.intersection(reasoning_terms)) / max(
+            len(premise_terms), 1
+        )
+        self.metrics["premise_relevance"] += term_overlap
+
+        # Conclusion validity - check if conclusion follows from reasoning
+        if generated_conclusion and "therefore" in generated_conclusion.lower():
+            self.metrics["conclusion_validity"] += 0.7  # Simplified placeholder
+        else:
+            self.metrics["conclusion_validity"] += 0.3
+
+        # Steps coherence - check for logical flow markers
+        flow_markers = [
+            "first",
+            "second",
+            "third",
+            "then",
+            "next",
+            "finally",
+            "because",
+            "thus",
+            "hence",
+        ]
+        marker_count = sum(
+            1 for marker in flow_markers if marker in generated_reasoning.lower()
+        )
+        self.metrics["steps_coherence"] += min(1.0, marker_count / 3)
+
+
+class ReasoningTrainingModule(SpecializedTrainingModule):
+    """Specialized training module for reasoning capabilities."""
+
+    def __init__(self, config: TrainingModuleConfig, tokenizer):
+        """
+        Initialize the reasoning training module.
+
+        Args:
+            config: Module configuration
+            tokenizer: Tokenizer for text processing
+        """
+        super().__init__(config, tokenizer)
+
+        # Initialize reasoning-specific components
+        self.reasoning_loss = nn.CrossEntropyLoss(ignore_index=-100)
+        self.metrics = {
+            "reasoning_loss": 0.0,
+            "reasoning_accuracy": 0.0,
+            "reasoning_perplexity": 0.0,
+        }
+
+        logger.info("Initialized ReasoningTrainingModule")
+
+    def prepare_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """
+        Prepare a batch of data for reasoning training.
+
+        Args:
+            batch: The input batch from the dataloader
+
+        Returns:
+            Processed batch ready for reasoning training
+        """
+        # Move batch to device
+        prepared_batch = {}
+        for key, value in batch.items():
+            if isinstance(value, torch.Tensor):
+                prepared_batch[key] = value.to(self.device)
+            else:
+                prepared_batch[key] = value
+
+        return prepared_batch
+
+    def compute_loss(
+        self, student_outputs: Any, teacher_outputs: Any, batch: Dict[str, torch.Tensor]
+    ) -> torch.Tensor:
+        """
+        Compute the reasoning-specific loss.
+
+        Args:
+            student_outputs: Outputs from the student model
+            teacher_outputs: Outputs from the teacher model
+            batch: The processed input batch
+
+        Returns:
+            Loss tensor for reasoning training
+        """
+        try:
+            # Extract logits from model outputs
+            if hasattr(student_outputs, "logits"):
+                student_logits = student_outputs.logits
+            else:
+                student_logits = student_outputs
+
+            if hasattr(teacher_outputs, "logits"):
+                teacher_logits = teacher_outputs.logits
+            else:
+                teacher_logits = teacher_outputs
+
+            # Get labels from batch
+            labels = batch.get("labels", batch.get("input_ids"))
+
+            # Compute cross entropy loss for reasoning
+            shift_logits = student_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            reasoning_loss = self.reasoning_loss(
+                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
+            )
+
+            # Add KL divergence loss between student and teacher
+            if teacher_logits is not None:
+                kl_loss = F.kl_div(
+                    F.log_softmax(student_logits, dim=-1),
+                    F.softmax(teacher_logits, dim=-1),
+                    reduction="batchmean",
+                )
+                total_loss = reasoning_loss + 0.1 * kl_loss
+            else:
+                total_loss = reasoning_loss
+
+            # Update metrics
+            self.metrics["reasoning_loss"] = reasoning_loss.item()
+
+            return total_loss * self.loss_weight
+
+        except Exception as e:
+            logger.error(f"Error computing reasoning loss: {e}")
+            # Return a small loss to avoid training failure
+            return torch.tensor(0.01, requires_grad=True)
+
+    def get_metrics(self) -> Dict[str, float]:
+        """
+        Get metrics specific to reasoning training.
+
+        Returns:
+            Dictionary of reasoning metrics
+        """
+        return self.metrics.copy()
+
+    def generate_synthetic_reasoning_data(
+        self, output_path: str, num_samples: int = 1000
+    ) -> None:
+        """
+        Generate synthetic reasoning data for training.
+
+        Args:
+            output_path: Path to save the generated data
+            num_samples: Number of samples to generate
+        """
+        # This is a simplified implementation based on the AWS backup's synthetic_generator
+        # In a full implementation, this would be much more sophisticated
+
+        templates = [
+            {
+                "premise": "If it rains, the ground gets wet. It is raining now.",
+                "reasoning": "Since it is raining, and rain makes the ground wet, we can conclude that the ground is getting wet.",
+                "conclusion": "Therefore, the ground is wet.",
+            },
+            {
+                "premise": "All mammals are warm-blooded. Whales are mammals.",
+                "reasoning": "Whales are classified as mammals. All mammals are warm-blooded animals. Therefore, as a mammal, a whale must be warm-blooded.",
+                "conclusion": "Therefore, whales are warm-blooded.",
+            },
+            {
+                "premise": "If you study hard, you will pass the exam. You studied hard.",
+                "reasoning": "The premise states a conditional relationship between studying hard and passing the exam. Since you studied hard, the condition is met.",
+                "conclusion": "Therefore, you will pass the exam.",
+            },
+        ]
+
+        # Generate variations of the templates
+        output_data = []
+        for _ in range(num_samples):
+            template = random.choice(templates)
+
+            # Create a variation (very simplified)
+            variation = {
+                "premise": template["premise"],
+                "reasoning": template["reasoning"],
+                "conclusion": template["conclusion"],
+                "metadata": {
+                    "generated": True,
+                    "timestamp": str(
+                        torch.cuda.get_device_name(0)
+                        if torch.cuda.is_available()
+                        else "CPU"
+                    ),
+                },
+            }
+
+            output_data.append(variation)
+
+        # Save to file
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        with open(output_path, "w", encoding="utf-8") as f:
+            for item in output_data:
+                f.write(json.dumps(item) + "\n")
+
+        logger.info(
+            f"Generated {len(output_data)} synthetic reasoning examples at {output_path}"
+        )

src/multi_agent_training/synthetic_data_generator.py

@@ -0,0 +1,162 @@
+"""
+Synthetic Data Generator for MangoMAS Local
+
+This module provides a framework for generating synthetic training data
+for different specialized capabilities, adaptable across all modules.
+"""
+
+import json
+import logging
+import os
+import random
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import yaml
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+
+class SyntheticDataGenerator(ABC):
+    """
+    Abstract base class for synthetic data generators.
+    Each specialized module can implement this interface for data generation.
+    """
+
+    def __init__(self, config: Dict[str, Any], output_dir: str = "data/processed"):
+        """
+        Initialize the synthetic data generator.
+
+        Args:
+            config: Configuration for the data generator
+            output_dir: Directory to save generated data
+        """
+        self.config = config
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        # Set number of examples to generate
+        self.num_examples = config.get("synthetic_examples", 1000)
+
+        # Template pool for generation
+        self.templates = self._load_templates()
+
+        logger.info(f"Initialized {self.__class__.__name__} with {self.num_examples} examples")
+
+    @abstractmethod
+    def _load_templates(self) -> List[Dict[str, Any]]:
+        """
+        Load templates for data generation.
+        Each implementation should define its own templates.
+
+        Returns:
+            List of template dictionaries
+        """
+        pass
+
+    @abstractmethod
+    def generate_example(self) -> Dict[str, Any]:
+        """
+        Generate a single synthetic training example.
+
+        Returns:
+            Dictionary with the generated example
+        """
+        pass
+
+    def generate_dataset(self, filename: str, num_examples: Optional[int] = None) -> str:
+        """
+        Generate a synthetic dataset and save to a JSONL file.
+
+        Args:
+            filename: Name of the output file
+            num_examples: Number of examples to generate (overrides config)
+
+        Returns:
+            Path to the generated dataset file
+        """
+        n = num_examples if num_examples is not None else self.num_examples
+        output_file = self.output_dir / filename
+
+        logger.info(f"Generating {n} synthetic examples for {self.__class__.__name__}")
+
+        with open(output_file, 'w', encoding='utf-8') as f:
+            for _ in tqdm(range(n), desc=f"Generating {filename}"):
+                example = self.generate_example()
+                f.write(json.dumps(example) + '\n')
+
+        logger.info(f"Generated dataset saved to {output_file}")
+        return str(output_file)
+
+    def augment_existing_dataset(self, input_file: str, output_file: Optional[str] = None,
+                                ratio: float = 0.5) -> str:
+        """
+        Augment an existing dataset with synthetic examples.
+
+        Args:
+            input_file: Path to the existing dataset
+            output_file: Path to save the augmented dataset (or None to overwrite)
+            ratio: Ratio of synthetic to original examples
+
+        Returns:
+            Path to the augmented dataset
+        """
+        if output_file is None:
+            output_file = input_file
+
+        # Load existing data
+        existing_data = []
+        try:
+            with open(input_file, 'r', encoding='utf-8') as f:
+                for line in f:
+                    existing_data.append(json.loads(line.strip()))
+        except (FileNotFoundError, json.JSONDecodeError) as e:
+            logger.warning(f"Error loading existing data: {e}")
+            existing_data = []
+
+        # Calculate number of synthetic examples to generate
+        n_existing = len(existing_data)
+        n_synthetic = int(n_existing * ratio)
+
+        # Generate synthetic examples
+        synthetic_data = [self.generate_example() for _ in tqdm(range(n_synthetic),
+                         desc=f"Generating augmentation data")]
+
+        # Combine datasets
+        combined_data = existing_data + synthetic_data
+        random.shuffle(combined_data)
+
+        # Save augmented dataset
+        with open(output_file, 'w', encoding='utf-8') as f:
+            for item in combined_data:
+                f.write(json.dumps(item) + '\n')
+
+        logger.info(f"Augmented dataset with {n_synthetic} synthetic examples, saved to {output_file}")
+        return output_file
+
+
+class SyntheticDataGeneratorRegistry:
+    """Registry for all synthetic data generators in the system."""
+
+    _generators = {}
+
+    @classmethod
+    def register(cls, module_type: str, generator_class):
+        """Register a generator class for a module type."""
+        cls._generators[module_type] = generator_class
+
+    @classmethod
+    def get_generator(cls, module_type: str, config: Dict[str, Any], output_dir: str) -> SyntheticDataGenerator:
+        """Get a generator instance for a module type."""
+        if module_type not in cls._generators:
+            raise ValueError(f"No generator registered for module type: {module_type}")
+
+        return cls._generators[module_type](config, output_dir)
+
+    @classmethod
+    def list_generators(cls) -> List[str]:
+        """List all registered generator types."""
+        return list(cls._generators.keys())
+"""