""" AgentFile Model Merger - Main Script Uses HuggingFace Transformers for model merging Supports GGUF, SafeTensors, and HuggingFace Hub models """ import sys import os import torch from typing import List, Dict, Optional import argparse import json import logging # Add src to path sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src')) from model_merger import ( ModelMerger, MergedModelConfig, ExpertConfig, MergeStrategy, create_merged_model ) from resource_manager import ( IntelligentResourceManager, ResourceBudget, DynamicExpertPool, QualityAwareRouter ) # Setup logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) class AgentFileModelMerger: """Main class for merging AI models using HuggingFace""" def __init__(self): self.merger = None self.resource_manager = None self.expert_pool = None self.router = None def merge_models( self, expert_paths: List[str], expert_names: List[str], output_path: str, merge_strategy: str = "adaptive_fusion", memory_budget: float = 8.0, max_experts: int = 4, quality_threshold: float = 0.8, load_in_4bit: bool = False, push_to_hub: bool = False, hub_model_id: Optional[str] = None ) -> None: """Merge multiple models into one""" print("=" * 60) print(" AgentFile Model Merger - Beyond Normal MoE") print("=" * 60) print() # Validate inputs if len(expert_paths) != len(expert_names): raise ValueError("Number of paths and names must match") if len(expert_paths) < 2: raise ValueError("Need at least 2 models to merge") print(f"[*] Merging {len(expert_paths)} models:") for i, (path, name) in enumerate(zip(expert_paths, expert_names)): print(f" {i+1}. {name} ({path})") print() # 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=max_experts, quality_threshold=quality_threshold, memory_budget=memory_budget, output_path=output_path, push_to_hub=push_to_hub, hub_model_id=hub_model_id ) # Initialize resource manager print("[*] Initializing intelligent resource manager...") budget = ResourceBudget( memory_gb=memory_budget, max_experts=max_experts, quality_threshold=quality_threshold ) self.resource_manager = IntelligentResourceManager(budget) # Initialize expert pool print("[*] Initializing dynamic expert pool...") self.expert_pool = DynamicExpertPool(max_experts=max_experts) # Initialize quality-aware router print("[*] Initializing quality-aware router...") self.router = QualityAwareRouter( num_experts=len(experts), quality_threshold=quality_threshold ) # Create merger print("[*] Creating deep merger...") self.merger = ModelMerger(config) # Load experts print("\n[*] Loading expert models...") for expert in experts: self.merger.load_expert(expert) # Merge models print("\n[*] Merging models...") merged_model = self.merger.merge_models() # Get tokenizer (use first tokenizer) tokenizer = self.merger.tokenizers[0] # Save merged model print("\n[*] Saving merged model...") self.merger.save_merged_model(merged_model, tokenizer, output_path) # Push to hub if requested if push_to_hub and hub_model_id: print(f"\n[*] Pushing model to HuggingFace Hub: {hub_model_id}") self.merger.push_to_hub(merged_model, tokenizer, hub_model_id) # Start resource monitoring print("\n[*] Starting resource monitoring...") self.resource_manager.start_monitoring() print("\n" + "=" * 60) print(" Merge Complete!") print("=" * 60) print(f"\n Output: {output_path}") print(f" Strategy: {merge_strategy}") print(f" Experts: {len(experts)}") if push_to_hub: print(f" Hub: {hub_model_id}") print() # Cleanup self.merger.cleanup() def analyze_models(self, model_paths: List[str]) -> Dict: """Analyze models before merging""" from transformers import AutoConfig analysis = { 'models': [], 'total_parameters': 0, 'compatible': True } for path in model_paths: try: config = AutoConfig.from_pretrained(path, trust_remote_code=True) # Get model info model_info = { 'path': path, 'hidden_size': config.hidden_size, 'num_layers': config.num_hidden_layers, 'num_heads': config.num_attention_heads, 'vocab_size': config.vocab_size, 'model_type': getattr(config, 'model_type', 'unknown') } # Try to estimate parameters try: # This is an approximation - actual params depend on architecture params = config.hidden_size * config.num_hidden_layers * 12 * config.hidden_size model_info['parameters'] = params analysis['total_parameters'] += params except: model_info['parameters'] = 0 analysis['models'].append(model_info) except Exception as e: logger.warning(f"Could not analyze {path}: {e}") analysis['compatible'] = False # Check compatibility if len(analysis['models']) > 1: hidden_sizes = [m['hidden_size'] for m in analysis['models'] if m.get('hidden_size')] if hidden_sizes and len(set(hidden_sizes)) > 1: print("[!] Warning: Models have different hidden sizes") print(" This may affect merge quality") return analysis def get_recommendations(self, analysis: Dict) -> Dict: """Get merge recommendations based on analysis""" recommendations = { 'strategy': 'adaptive_fusion', 'max_experts': 4, 'quality_threshold': 0.8, 'memory_budget': 8.0 } # Adjust based on model sizes total_params = analysis.get('total_parameters', 0) if total_params > 10e9: # > 10B parameters recommendations['memory_budget'] = 16.0 recommendations['max_experts'] = 2 elif total_params > 5e9: # > 5B parameters recommendations['memory_budget'] = 12.0 recommendations['max_experts'] = 3 else: recommendations['memory_budget'] = 8.0 recommendations['max_experts'] = 4 # Check if models are similar if len(analysis['models']) > 1: hidden_sizes = [m['hidden_size'] for m in analysis['models'] if m.get('hidden_size')] if hidden_sizes and max(hidden_sizes) / min(hidden_sizes) > 1.5: recommendations['strategy'] = 'deep_merge' print("[*] Using deep_merge strategy due to model differences") return recommendations def interactive_merge(self): """Interactive merge mode""" print("=" * 60) print(" AgentFile Model Merger - Interactive Mode") print("=" * 60) print() # Get model paths model_paths = [] model_names = [] print("Enter model paths (empty line to finish):") while True: path = input(" Model path: ").strip() if not path: break model_paths.append(path) name = input(" Model name: ").strip() if not name: name = os.path.basename(path) model_names.append(name) print() if len(model_paths) < 2: print("[-] Need at least 2 models to merge") return # Analyze models print("\n[*] Analyzing models...") analysis = self.analyze_models(model_paths) # Get recommendations recommendations = self.get_recommendations(analysis) print("\n[*] Analysis Results:") for i, model in enumerate(analysis['models'], 1): print(f" {i}. {model['path']}") print(f" Hidden size: {model.get('hidden_size', 'N/A')}") print(f" Layers: {model.get('num_layers', 'N/A')}") print(f" Model type: {model.get('model_type', 'N/A')}") print() print("[*] Recommendations:") print(f" Strategy: {recommendations['strategy']}") print(f" Max experts: {recommendations['max_experts']}") print(f" Memory budget: {recommendations['memory_budget']} GB") print() # Get merge settings print("Configure merge settings (press Enter for defaults):") strategy = input(f" Strategy [{recommendations['strategy']}]: ").strip() if not strategy: strategy = recommendations['strategy'] max_experts = input(f" Max experts [{recommendations['max_experts']}]: ").strip() if not max_experts: max_experts = recommendations['max_experts'] else: max_experts = int(max_experts) memory_budget = input(f" Memory budget (GB) [{recommendations['memory_budget']}]: ").strip() if not memory_budget: memory_budget = recommendations['memory_budget'] else: memory_budget = float(memory_budget) output_path = input(" Output path [models/merged_model]: ").strip() if not output_path: output_path = "models/merged_model" load_4bit = input(" Load in 4-bit for memory efficiency? (y/N): ").strip().lower() == 'y' push_hub = input(" Push to HuggingFace Hub? (y/N): ").strip().lower() == 'y' hub_model_id = None if push_hub: hub_model_id = input(" HuggingFace model ID: ").strip() if not hub_model_id: push_hub = False # Perform merge print("\n" + "=" * 60) print(" Starting Merge...") print("=" * 60) self.merge_models( expert_paths=model_paths, expert_names=model_names, output_path=output_path, merge_strategy=strategy, memory_budget=memory_budget, max_experts=max_experts, load_in_4bit=load_4bit, push_to_hub=push_hub, hub_model_id=hub_model_id ) print("\n[+] Merge completed successfully!") print(f" Output saved to: {output_path}") def main(): parser = argparse.ArgumentParser( description="AgentFile Model Merger - Beyond Normal MoE" ) subparsers = parser.add_subparsers(dest='command', help='Command to run') # Merge command merge_parser = subparsers.add_parser('merge', help='Merge multiple models') merge_parser.add_argument('--models', nargs='+', required=True, help='Model paths') merge_parser.add_argument('--names', nargs='+', help='Model names') merge_parser.add_argument('--output', default='models/merged_model', help='Output path') merge_parser.add_argument('--strategy', default='adaptive_fusion', choices=['ties', 'dare', 'deep_merge', 'adaptive_fusion', 'neural_synthesis', 'model_soup'], help='Merge strategy') merge_parser.add_argument('--memory-budget', type=float, default=8.0, help='Memory budget in GB') merge_parser.add_argument('--max-experts', type=int, default=4, help='Maximum experts per token') merge_parser.add_argument('--quality-threshold', type=float, default=0.8, help='Quality threshold') merge_parser.add_argument('--load-in-4bit', action='store_true', help='Load models in 4-bit quantization') merge_parser.add_argument('--push-to-hub', action='store_true', help='Push merged model to HuggingFace Hub') merge_parser.add_argument('--hub-model-id', type=str, help='HuggingFace Hub model ID') # Analyze command analyze_parser = subparsers.add_parser('analyze', help='Analyze models') analyze_parser.add_argument('--models', nargs='+', required=True, help='Model paths') # Interactive command interactive_parser = subparsers.add_parser('interactive', help='Interactive merge mode') args = parser.parse_args() merger = AgentFileModelMerger() if args.command == 'merge': # Generate names if not provided if args.names is None: args.names = [os.path.basename(p) for p in args.models] merger.merge_models( expert_paths=args.models, expert_names=args.names, output_path=args.output, merge_strategy=args.strategy, memory_budget=args.memory_budget, max_experts=args.max_experts, quality_threshold=args.quality_threshold, load_in_4bit=args.load_in_4bit, push_to_hub=args.push_to_hub, hub_model_id=args.hub_model_id ) elif args.command == 'analyze': analysis = merger.analyze_models(args.models) print("\nAnalysis Results:") print(json.dumps(analysis, indent=2)) recommendations = merger.get_recommendations(analysis) print("\nRecommendations:") print(json.dumps(recommendations, indent=2)) elif args.command == 'interactive': merger.interactive_merge() else: parser.print_help() if __name__ == "__main__": main()