equalizer.py

"""
Core Equalizer module for GRACE-SVD.
"""

from pathlib import Path
from typing import Dict, Union

import torch

from strategies import (
    EqualizationStrategy,
    BaseStrategy,
    get_strategy,
)
from utils import (
    load_state_dict,
    save_state_dict,
    print_state_dict_info,
    validate_state_dict_for_equalization,
)


class CoreEqualizer:
    """
    Main class for equalizing cores in SVD decomposed models.
    
    Available Strategies:
        - WHITENING_DIVISION_FACTOR: Divide whitening matrix by a given factor
        - MAX_MIN_RANGE: Max-min range based equalization
        - WASSERSTEIN_DISTANCE: Minimize Wasserstein distance using kurtosis-based range ratio
    """
    
    def __init__(
        self,
        strategy: Union[EqualizationStrategy, str] = EqualizationStrategy.WASSERSTEIN_DISTANCE
    ):
        """Initialize the CoreEqualizer."""
        self._strategy = self._parse_strategy(strategy)
        self._strategy_instance: BaseStrategy = get_strategy(self._strategy)
    
    def _parse_strategy(
        self,
        strategy: Union[EqualizationStrategy, str]
    ) -> EqualizationStrategy:
        """Parse strategy from enum or string."""
        if isinstance(strategy, EqualizationStrategy):
            return strategy
        
        strategy_map = {
            'whitening_division_factor': EqualizationStrategy.WHITENING_DIVISION_FACTOR,
            'max_min_range': EqualizationStrategy.MAX_MIN_RANGE,
            'wasserstein_distance': EqualizationStrategy.WASSERSTEIN_DISTANCE,
        }
        
        strategy_lower = strategy.lower().strip()
        if strategy_lower not in strategy_map:
            raise ValueError(
                f"Unknown strategy: '{strategy}'. "
                f"Available strategies: {list(strategy_map.keys())}"
            )
        
        return strategy_map[strategy_lower]
    
    @property
    def strategy(self) -> EqualizationStrategy:
        """Get the current equalization strategy."""
        return self._strategy
    
    @property
    def strategy_name(self) -> str:
        """Get the name of the current strategy."""
        return self._strategy_instance.name
    
    @property
    def strategy_description(self) -> str:
        """Get the description of the current strategy."""
        return self._strategy_instance.description
    
    def set_strategy(self, strategy: Union[EqualizationStrategy, str]) -> None:
        """Set the equalization strategy."""
        self._strategy = self._parse_strategy(strategy)
        self._strategy_instance = get_strategy(self._strategy)
    
    def equalize(
        self,
        state_dict: Dict[str, torch.Tensor],
        validate: bool = True,
        verbose: bool = True,
        **kwargs
    ) -> Dict[str, torch.Tensor]:
        """
        Equalize cores in the given state dictionary.
        
        Args:
            state_dict: Compressed model state dictionary containing weights.
            validate: Whether to validate the state dict before equalization.
            verbose: Whether to print progress information.
            **kwargs: Strategy-specific parameters.
        
        Returns:
            Equalized state dictionary.
        """
        if validate:
            is_valid, msg = validate_state_dict_for_equalization(state_dict)
            if not is_valid:
                raise ValueError(f"Invalid state dictionary: {msg}")
            if verbose:
                print(f"Validation passed: {msg}")
        
        if verbose:
            print(f"\nApplying '{self.strategy_name}' equalization strategy...")
            print_state_dict_info(state_dict)
        
        equalized_state_dict = self._strategy_instance.equalize(state_dict, **kwargs)
        
        if verbose:
            print("\nEqualization complete!")
            print_state_dict_info(equalized_state_dict)
        
        return equalized_state_dict
    
    def equalize_file(
        self,
        input_path: Union[str, Path],
        output_path: Union[str, Path],
        validate: bool = True,
        verbose: bool = True,
        **kwargs
    ) -> Dict[str, torch.Tensor]:
        """
        Load, equalize, and save a state dictionary.
        
        Args:
            input_path: Path to the input .pt file.
            output_path: Path to save the equalized .pt file.
            validate: Whether to validate the state dict before equalization.
            verbose: Whether to print progress information.
            **kwargs: Strategy-specific parameters.
        
        Returns:
            Equalized state dictionary.
        """
        input_path = Path(input_path)
        output_path = Path(output_path)
        
        if verbose:
            print(f"Loading state dict from: {input_path}")
        
        state_dict = load_state_dict(input_path)
        
        if verbose:
            print(f"Loaded {len(state_dict)} tensors")
        
        equalized_state_dict = self.equalize(
            state_dict,
            validate=validate,
            verbose=verbose,
            **kwargs
        )
        
        if verbose:
            print(f"\nSaving equalized state dict to: {output_path}")
        
        save_path = save_state_dict(equalized_state_dict, output_path)
        
        if verbose:
            print(f"Successfully saved to: {save_path}")
        
        return equalized_state_dict
    
    @staticmethod
    def list_strategies() -> Dict[str, str]:
        """List all available equalization strategies."""
        return {
            'whitening_division_factor': (
                "Divide whitening matrix by a given factor. "
                "Requires pre-computed scaling factor with which the whitening matrix has to be divided."
            ),
            'max_min_range': (
                "Max-min range based equalization. "
                "Computes scaling factor from the weight value range."
            ),
            'wasserstein_distance': (
                "Minimize Wasserstein distance using kurtosis-based range ratio. "
                "Adapts percentile range based on weight distribution kurtosis. "
                "Parameters: px (percentile limit), beta (kurtosis scaling factor)."
            ),
        }


def equalize_cores(
    input_path: Union[str, Path],
    output_path: Union[str, Path],
    strategy: Union[EqualizationStrategy, str] = "wasserstein_distance",
    **kwargs
) -> Dict[str, torch.Tensor]:
    """Convenience function for quick equalization."""
    equalizer = CoreEqualizer(strategy=strategy)
    return equalizer.equalize_file(input_path, output_path, **kwargs)