model_merger.py

"""
AgentFile Model Merger - Advanced MoE Beyond Normal
Uses HuggingFace Transformers for model merging
Supports GGUF, SafeTensors, and HuggingFace Hub models
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import (
    AutoModelForCausalLM, 
    AutoTokenizer, 
    AutoConfig,
    BitsAndBytesConfig
)
from typing import Dict, List, Optional, Tuple, Union
import numpy as np
from dataclasses import dataclass, field
from enum import Enum
import json
import os
import sys
import logging
from pathlib import Path
import gc
import time

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class MergeStrategy(Enum):
    """Advanced merge strategies beyond normal MoE"""
    TIES = "ties"                    # Task Interpolation with Exponential Smoothing
    DARE = "dare"                    # Drop And REscale
    MODEL_SOUP = "model_soup"        # Model Soups (averaging)
    DEEP_MERGE = "deep_merge"        # Deep layer-wise merging
    ADAPTIVE_FUSION = "adaptive_fusion"  # Adaptive fusion based on input
    NEURAL_SYNTHESIS = "neural_synthesis"  # Neural synthesis of weights

@dataclass
class ExpertConfig:
    """Configuration for an expert model"""
    name: str
    path: str
    weight: float = 1.0
    specialization: str = "general"
    memory_requirement: float = 1.0
    compute_requirement: float = 1.0
    device_map: str = "auto"
    torch_dtype: str = "float16"
    load_in_4bit: bool = False
    load_in_8bit: bool = False

@dataclass
class MergedModelConfig:
    """Configuration for the merged model"""
    experts: List[ExpertConfig] = field(default_factory=list)
    merge_strategy: MergeStrategy = MergeStrategy.ADAPTIVE_FUSION
    router_type: str = "neural_router"
    max_experts_per_token: int = 4
    load_balancing_factor: float = 0.1
    memory_budget: float = 8.0  # in GB
    use_dynamic_routing: bool = True
    quality_threshold: float = 0.8
    output_path: str = "models/merged_model"
    push_to_hub: bool = False
    hub_model_id: Optional[str] = None

class HuggingFaceModelLoader:
    """Handles loading models from HuggingFace Hub or local paths"""
    
    def __init__(self):
        self.loaded_models = {}
        self.loaded_tokenizers = {}
        
    def load_model(
        self, 
        model_path: str, 
        device_map: str = "auto",
        torch_dtype: str = "float16",
        load_in_4bit: bool = False,
        load_in_8bit: bool = False
    ) -> Tuple[AutoModelForCausalLM, AutoTokenizer]:
        """Load model and tokenizer from HuggingFace or local path"""
        
        if model_path in self.loaded_models:
            logger.info(f"Model already loaded: {model_path}")
            return self.loaded_models[model_path], self.loaded_tokenizers[model_path]
        
        logger.info(f"Loading model: {model_path}")
        start_time = time.time()
        
        try:
            # Determine dtype
            dtype_map = {
                "float16": torch.float16,
                "bfloat16": torch.bfloat16,
                "float32": torch.float32
            }
            dtype = dtype_map.get(torch_dtype, torch.float16)
            
            # Configure quantization if needed
            quantization_config = None
            if load_in_4bit:
                quantization_config = BitsAndBytesConfig(
                    load_in_4bit=True,
                    bnb_4bit_compute_dtype=dtype,
                    bnb_4bit_use_double_quant=True,
                    bnb_4bit_quant_type="nf4"
                )
            elif load_in_8bit:
                quantization_config = BitsAndBytesConfig(
                    load_in_8bit=True
                )
            
            # Load tokenizer
            tokenizer = AutoTokenizer.from_pretrained(
                model_path,
                trust_remote_code=True
            )
            
            # Load model
            model_kwargs = {
                "pretrained_model_name_or_path": model_path,
                "device_map": device_map,
                "torch_dtype": dtype,
                "trust_remote_code": True
            }
            
            if quantization_config:
                model_kwargs["quantization_config"] = quantization_config
            
            model = AutoModelForCausalLM.from_pretrained(**model_kwargs)
            
            # Store in cache
            self.loaded_models[model_path] = model
            self.loaded_tokenizers[model_path] = tokenizer
            
            load_time = time.time() - start_time
            logger.info(f"Model loaded successfully in {load_time:.2f}s")
            
            return model, tokenizer
            
        except Exception as e:
            logger.error(f"Failed to load model {model_path}: {e}")
            raise
    
    def unload_model(self, model_path: str):
        """Unload a model to free memory"""
        if model_path in self.loaded_models:
            del self.loaded_models[model_path]
            del self.loaded_tokenizers[model_path]
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            logger.info(f"Model unloaded: {model_path}")
    
    def get_model_info(self, model_path: str) -> Dict:
        """Get model information without loading it"""
        try:
            config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
            return {
                "hidden_size": config.hidden_size,
                "num_hidden_layers": config.num_hidden_layers,
                "num_attention_heads": config.num_attention_heads,
                "vocab_size": config.vocab_size,
                "model_type": config.model_type
            }
        except Exception as e:
            logger.warning(f"Could not get model info for {model_path}: {e}")
            return {}

class DeepMerger:
    """Advanced Deep Merger - Goes beyond normal weight averaging"""
    
    def __init__(self, strategy: MergeStrategy):
        self.strategy = strategy
        
    def merge_models(
        self, 
        models: List[AutoModelForCausalLM], 
        weights: List[float],
        config: MergedModelConfig
    ) -> AutoModelForCausalLM:
        """Merge multiple models into one"""
        
        logger.info(f"Merging {len(models)} models using {self.strategy.value} strategy")
        
        if self.strategy == MergeStrategy.TIES:
            return self._ties_merge(models, weights)
        elif self.strategy == MergeStrategy.DARE:
            return self._dare_merge(models, weights)
        elif self.strategy == MergeStrategy.DEEP_MERGE:
            return self._deep_merge(models, weights)
        elif self.strategy == MergeStrategy.ADAPTIVE_FUSION:
            return self._adaptive_fusion_merge(models, weights)
        elif self.strategy == MergeStrategy.NEURAL_SYNTHESIS:
            return self._neural_synthesis_merge(models, weights)
        else:
            return self._model_soup_merge(models, weights)
    
    def _ties_merge(
        self, 
        models: List[AutoModelForCausalLM], 
        weights: List[float]
    ) -> AutoModelForCausalLM:
        """TIES merging - Task Interpolation with Exponential Smoothing"""
        
        logger.info("Applying TIES merging...")
        
        # Get reference model (first model)
        merged_model = models[0]
        
        # Get all parameter keys
        param_keys = list(merged_model.state_dict().keys())
        
        # Collect differences from reference
        diffs = []
        for model in models[1:]:
            diff = {}
            for key in param_keys:
                diff[key] = model.state_dict()[key] - merged_model.state_dict()[key]
            diffs.append(diff)
        
        # Apply TIES algorithm
        merged_params = {}
        for key in param_keys:
            # Collect all values for this parameter
            values = [merged_model.state_dict()[key]]
            for diff in diffs:
                values.append(merged_model.state_dict()[key] + diff[key])
            
            # Apply exponential smoothing
            smoothed = values[0]
            for i, val in enumerate(values[1:], 1):
                alpha = weights[i] / sum(weights)
                smoothed = smoothed * (1 - alpha) + val * alpha
            
            merged_params[key] = smoothed
        
        # Load merged parameters
        merged_model.load_state_dict(merged_params)
        
        return merged_model
    
    def _dare_merge(
        self, 
        models: List[AutoModelForCausalLM], 
        weights: List[float]
    ) -> AutoModelForCausalLM:
        """DARE merging - Drop And REscale"""
        
        logger.info("Applying DARE merging...")
        
        merged_model = models[0]
        param_keys = list(merged_model.state_dict().keys())
        
        # Calculate importance scores (variance across models)
        importance_scores = {}
        for key in param_keys:
            values = [model.state_dict()[key] for model in models]
            variance = torch.var(torch.stack([v.float() for v in values]), dim=0)
            importance_scores[key] = variance
        
        # Merge with importance-weighted averaging
        merged_params = {}
        for key in param_keys:
            # Weight by inverse importance (less important parameters get merged more)
            inv_importance = 1.0 / (importance_scores[key] + 1e-10)
            inv_importance = inv_importance / inv_importance.sum()
            
            weighted_sum = torch.zeros_like(models[0].state_dict()[key], dtype=torch.float32)
            for i, model in enumerate(models):
                weighted_sum += weights[i] * model.state_dict()[key].float() * inv_importance
            
            merged_params[key] = weighted_sum.to(models[0].state_dict()[key].dtype)
        
        merged_model.load_state_dict(merged_params)
        return merged_model
    
    def _deep_merge(
        self, 
        models: List[AutoModelForCausalLM], 
        weights: List[float]
    ) -> AutoModelForCausalLM:
        """Deep layer-wise merging - Analyzes and merges each layer differently"""
        
        logger.info("Applying deep layer-wise merging...")
        
        merged_model = models[0]
        param_keys = list(merged_model.state_dict().keys())
        
        # Group parameters by layer
        layer_groups = {}
        for key in param_keys:
            parts = key.split('.')
            layer_num = None
            for part in parts:
                if part.isdigit():
                    layer_num = int(part)
                    break
            
            if layer_num is not None:
                if layer_num not in layer_groups:
                    layer_groups[layer_num] = []
                layer_groups[layer_num].append(key)
            else:
                # Non-layer parameters (embeddings, etc.)
                if 'global' not in layer_groups:
                    layer_groups['global'] = []
                layer_groups['global'].append(key)
        
        # Merge each layer differently
        merged_params = {}
        for layer_num, keys in layer_groups.items():
            if layer_num == 'global':
                # Simple weighted average for global parameters
                for key in keys:
                    weighted_sum = torch.zeros_like(models[0].state_dict()[key], dtype=torch.float32)
                    for i, model in enumerate(models):
                        weighted_sum += weights[i] * model.state_dict()[key].float()
                    merged_params[key] = weighted_sum.to(models[0].state_dict()[key].dtype)
            else:
                # Adaptive merging for layer parameters
                layer_complexity = self._analyze_layer_complexity(models, keys)
                
                for key in keys:
                    if layer_complexity > 0.7:
                        # High complexity - use TIES-like merging
                        values = [model.state_dict()[key] for model in models]
                        smoothed = values[0]
                        for i, val in enumerate(values[1:], 1):
                            alpha = weights[i] / sum(weights)
                            smoothed = smoothed * (1 - alpha) + val * alpha
                        merged_params[key] = smoothed
                    else:
                        # Low complexity - use simple averaging
                        weighted_sum = torch.zeros_like(models[0].state_dict()[key], dtype=torch.float32)
                        for i, model in enumerate(models):
                            weighted_sum += weights[i] * model.state_dict()[key].float()
                        merged_params[key] = weighted_sum.to(models[0].state_dict()[key].dtype)
        
        merged_model.load_state_dict(merged_params)
        return merged_model
    
    def _adaptive_fusion_merge(
        self, 
        models: List[AutoModelForCausalLM], 
        weights: List[float]
    ) -> AutoModelForCausalLM:
        """Adaptive Fusion - Dynamically adjusts merging based on input"""
        
        logger.info("Applying adaptive fusion merging...")
        
        merged_model = models[0]
        param_keys = list(merged_model.state_dict().keys())
        
        # Create fusion gates for each layer
        fusion_gates = {}
        for key in param_keys:
            shape = models[0].state_dict()[key].shape
            gate = torch.ones(len(models), *shape, dtype=torch.float32) / len(models)
            fusion_gates[key] = gate
        
        # Merge with adaptive gates
        merged_params = {}
        for key in param_keys:
            weighted_sum = torch.zeros_like(models[0].state_dict()[key], dtype=torch.float32)
            
            for i, model in enumerate(models):
                gate = fusion_gates[key][i]
                weighted_sum += gate * weights[i] * model.state_dict()[key].float()
            
            merged_params[key] = weighted_sum.to(models[0].state_dict()[key].dtype)
        
        merged_model.load_state_dict(merged_params)
        
        # Store fusion gates for runtime adaptation
        merged_model.fusion_gates = fusion_gates
        
        return merged_model
    
    def _neural_synthesis_merge(
        self, 
        models: List[AutoModelForCausalLM], 
        weights: List[float]
    ) -> AutoModelForCausalLM:
        """Neural Synthesis - Creates new parameters by synthesizing across models"""
        
        logger.info("Applying neural synthesis merging...")
        
        merged_model = models[0]
        param_keys = list(merged_model.state_dict().keys())
        
        # Synthesize new parameters
        merged_params = {}
        for key in param_keys:
            params = [model.state_dict()[key].float() for model in models]
            stacked = torch.stack(params, dim=0)
            
            # Compute principal components
            flat_params = stacked.view(len(models), -1)
            mean = flat_params.mean(dim=0)
            
            # Compute deviations from mean
            deviations = flat_params - mean.unsqueeze(0)
            
            # Synthesize new parameter as weighted combination of deviations
            synthesized_deviation = torch.zeros_like(mean)
            for i in range(len(models)):
                synthesized_deviation += weights[i] * deviations[i]
            
            # Reconstruct synthesized parameter
            synthesized_param = mean + synthesized_deviation
            
            merged_params[key] = synthesized_param.view(stacked.shape[1:])
        
        merged_model.load_state_dict(merged_params)
        return merged_model
    
    def _model_soup_merge(
        self, 
        models: List[AutoModelForCausalLM], 
        weights: List[float]
    ) -> AutoModelForCausalLM:
        """Model Soups - Simple weighted averaging"""
        
        logger.info("Applying model soup merging...")
        
        merged_model = models[0]
        param_keys = list(merged_model.state_dict().keys())
        
        merged_params = {}
        for key in param_keys:
            weighted_sum = torch.zeros_like(models[0].state_dict()[key], dtype=torch.float32)
            for i, model in enumerate(models):
                weighted_sum += weights[i] * model.state_dict()[key].float()
            merged_params[key] = weighted_sum.to(models[0].state_dict()[key].dtype)
        
        merged_model.load_state_dict(merged_params)
        return merged_model
    
    def _analyze_layer_complexity(
        self, 
        models: List[AutoModelForCausalLM], 
        keys: List[str]
    ) -> float:
        """Analyze complexity of a layer"""
        
        total_variance = 0.0
        count = 0
        
        for key in keys:
            values = [model.state_dict()[key].float() for model in models]
            variance = torch.var(torch.stack(values)).item()
            total_variance += variance
            count += 1
        
        avg_variance = total_variance / count if count > 0 else 0
        
        # Normalize to 0-1 range
        complexity = min(1.0, avg_variance / 10.0)
        
        return complexity

class ModelMerger:
    """Main Model Merger Class - Uses HuggingFace for model management"""
    
    def __init__(self, config: MergedModelConfig):
        self.config = config
        self.model_loader = HuggingFaceModelLoader()
        self.models = []
        self.tokenizers = []
        
    def load_expert(self, expert_config: ExpertConfig):
        """Load an expert model"""
        
        logger.info(f"Loading expert: {expert_config.name}")
        
        try:
            model, tokenizer = self.model_loader.load_model(
                model_path=expert_config.path,
                device_map=expert_config.device_map,
                torch_dtype=expert_config.torch_dtype,
                load_in_4bit=expert_config.load_in_4bit,
                load_in_8bit=expert_config.load_in_8bit
            )
            
            self.models.append(model)
            self.tokenizers.append(tokenizer)
            
            logger.info(f"Successfully loaded: {expert_config.name}")
            
        except Exception as e:
            logger.error(f"Error loading {expert_config.name}: {e}")
            raise
    
    def merge_models(self) -> AutoModelForCausalLM:
        """Merge all loaded models into a unified model"""
        
        if not self.models:
            raise ValueError("No models loaded!")
        
        logger.info(f"Starting merge of {len(self.models)} models...")
        
        # Create merger
        merger = DeepMerger(self.config.merge_strategy)
        
        # Extract weights from config
        weights = [expert.weight for expert in self.config.experts]
        
        # Merge models
        merged_model = merger.merge_models(self.models, weights, self.config)
        
        logger.info("Models merged successfully!")
        
        return merged_model
    
    def save_merged_model(
        self, 
        model: AutoModelForCausalLM, 
        tokenizer: AutoTokenizer,
        output_path: str
    ):
        """Save the merged model"""
        
        logger.info(f"Saving merged model to: {output_path}")
        
        os.makedirs(output_path, exist_ok=True)
        
        # Save model
        model.save_pretrained(output_path)
        
        # Save tokenizer
        tokenizer.save_pretrained(output_path)
        
        # Save config
        config_path = os.path.join(output_path, "merge_config.json")
        with open(config_path, 'w') as f:
            json.dump({
                'merge_strategy': self.config.merge_strategy.value,
                'num_experts': len(self.config.experts),
                'expert_names': [e.name for e in self.config.experts],
                'max_experts_per_token': self.config.max_experts_per_token,
                'quality_threshold': self.config.quality_threshold
            }, f, indent=2)
        
        logger.info("Model saved successfully!")
    
    def push_to_hub(
        self, 
        model: AutoModelForCausalLM,
        tokenizer: AutoTokenizer,
        model_id: str
    ):
        """Push merged model to HuggingFace Hub"""
        
        logger.info(f"Pushing model to HuggingFace Hub: {model_id}")
        
        try:
            model.push_to_hub(model_id)
            tokenizer.push_to_hub(model_id)
            logger.info("Model pushed successfully!")
        except Exception as e:
            logger.error(f"Failed to push model: {e}")
            raise
    
    def cleanup(self):
        """Cleanup loaded models to free memory"""
        for path in list(self.model_loader.loaded_models.keys()):
            self.model_loader.unload_model(path)
        self.models.clear()
        self.tokenizers.clear()
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

def create_merged_model(
    expert_paths: List[str],
    expert_names: List[str],
    output_path: str,
    merge_strategy: str = "adaptive_fusion",
    memory_budget: float = 8.0,
    load_in_4bit: bool = False,
    push_to_hub: bool = False,
    hub_model_id: Optional[str] = None
) -> AutoModelForCausalLM:
    """Convenience function to create a merged model"""
    
    # Create expert configs
    experts = []
    for name, path in zip(expert_names, expert_paths):
        experts.append(ExpertConfig(
            name=name,
            path=path,
            weight=1.0 / len(expert_paths),
            load_in_4bit=load_in_4bit
        ))
    
    # Create merge config
    config = MergedModelConfig(
        experts=experts,
        merge_strategy=MergeStrategy(merge_strategy),
        max_experts_per_token=4,
        quality_threshold=0.8,
        memory_budget=memory_budget,
        output_path=output_path,
        push_to_hub=push_to_hub,
        hub_model_id=hub_model_id
    )
    
    # Create merger
    merger = ModelMerger(config)
    
    try:
        # Load all experts
        for expert in experts:
            merger.load_expert(expert)
        
        # Merge models
        merged_model = merger.merge_models()
        
        # Get tokenizer (use first tokenizer)
        tokenizer = merger.tokenizers[0]
        
        # Save merged model
        merger.save_merged_model(merged_model, tokenizer, output_path)
        
        # Push to hub if requested
        if push_to_hub and hub_model_id:
            merger.push_to_hub(merged_model, tokenizer, hub_model_id)
        
        return merged_model
        
    finally:
        merger.cleanup()

if __name__ == "__main__":
    # Example usage
    expert_paths = [
        "pubertcs/Ornith-1.0-9B-IL2CPP-Decompiler-GGUF",
        # Add more expert paths here
    ]
    
    expert_names = [
        "ornith-il2cpp",
        # Add more expert names here
    ]
    
    output_path = "models/merged_model"
    
    merged_model = create_merged_model(
        expert_paths=expert_paths,
        expert_names=expert_names,
        output_path=output_path,
        merge_strategy="adaptive_fusion"
    )
    
    logger.info(f"Merged model created successfully at: {output_path}")