""" 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}")