merge.py

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