utils.py

"""
Utility functions for GRACE-SVD.
"""

from pathlib import Path
from typing import Dict, Any, Union, Tuple

import torch


def load_state_dict(
    path: Union[str, Path],
    map_location: str = 'cpu',
    weights_only: bool = True
) -> Dict[str, torch.Tensor]:
    """Load a state dictionary from a .pt file."""
    path = Path(path)
    
    if not path.exists():
        raise FileNotFoundError(f"File not found: {path}")
    
    try:
        checkpoint = torch.load(
            path,
            map_location=map_location,
            weights_only=weights_only
        )
        
        if isinstance(checkpoint, torch.Tensor):
            return {'weights': checkpoint}
        elif isinstance(checkpoint, dict):
            if 'state_dict' in checkpoint:
                return checkpoint['state_dict']
            elif 'model' in checkpoint:
                return checkpoint['model']
            elif 'model_state_dict' in checkpoint:
                return checkpoint['model_state_dict']
            else:
                return checkpoint
        else:
            raise RuntimeError(f"Unexpected checkpoint format: {type(checkpoint)}")
            
    except Exception as e:
        raise RuntimeError(f"Error loading state dict from {path}: {str(e)}")


def save_state_dict(
    state_dict: Dict[str, torch.Tensor],
    path: Union[str, Path],
    create_dirs: bool = True
) -> Path:
    """Save a state dictionary to a .pt file."""
    path = Path(path)
    
    if create_dirs:
        path.parent.mkdir(parents=True, exist_ok=True)
    
    try:
        torch.save(state_dict, path)
        return path
    except Exception as e:
        raise OSError(f"Error saving state dict to {path}: {str(e)}")


def get_state_dict_info(state_dict: Dict[str, torch.Tensor]) -> Dict[str, Any]:
    """Get information about a state dictionary."""
    num_tensors = len(state_dict)
    total_params = sum(p.numel() for p in state_dict.values() if isinstance(p, torch.Tensor))
    total_size_mb = sum(p.numel() * p.element_size() for p in state_dict.values() if isinstance(p, torch.Tensor)) / (1024 * 1024)
    
    node_pairs = []
    decomposed_layers = set()
    for name in state_dict.keys():
        if "nodes" in name and name.endswith(".0"):
            node1_name = name[:-1] + "1"
            if node1_name in state_dict:
                node_pairs.append((name, node1_name))
                if "nodes." in name:
                    layer_prefix = name.split("nodes.")[0]
                    decomposed_layers.add(layer_prefix.rstrip("."))
    
    return {
        'num_tensors': num_tensors,
        'total_params': total_params,
        'total_size_mb': total_size_mb,
        'num_node_pairs': len(node_pairs),
        'node_pairs': node_pairs,
        'num_decomposed_layers': len(decomposed_layers),
        'decomposed_layers': sorted(decomposed_layers),
    }


def print_state_dict_info(state_dict: Dict[str, torch.Tensor]) -> None:
    """Print formatted information about a state dictionary."""
    info = get_state_dict_info(state_dict)
    
    print("=" * 60)
    print("State Dictionary Information")
    print("=" * 60)
    print(f"Number of tensors:       {info['num_tensors']}")
    print(f"Total parameters:        {info['total_params']:,}")
    print(f"Total size:              {info['total_size_mb']:.2f} MB")
    print(f"Decomposed layers:       {info['num_decomposed_layers']}")
    print("=" * 60)


def validate_state_dict_for_equalization(
    state_dict: Dict[str, torch.Tensor]
) -> Tuple[bool, str]:
    """Validate that a state dictionary is suitable for equalization."""
    if not state_dict:
        return False, "State dictionary is empty"
    
    node_pairs = []
    for name in state_dict.keys():
        if "nodes" in name and name.endswith(".0"):
            node1_name = name[:-1] + "1"
            if node1_name in state_dict:
                node_pairs.append((name, node1_name))
    
    if not node_pairs:
        return False, "No node pairs found in state dictionary (looking for nodes.X.0 and nodes.X.1)"
    
    for node0, node1 in node_pairs:
        if not isinstance(state_dict[node0], torch.Tensor):
            return False, f"{node0} is not a tensor"
        if not isinstance(state_dict[node1], torch.Tensor):
            return False, f"{node1} is not a tensor"
    
    return True, f"Valid state dictionary with {len(node_pairs)} node pairs"


def extract_layer_info(layer_name: str) -> Dict[str, Any]:
    """Extract layer information from a layer name."""
    info = {
        'layer_number': -1,
        'layer_type': 'Other',
        'node_index': -1,
        'base_name': layer_name
    }
    
    if 'layers.' in layer_name:
        try:
            parts = layer_name.split('layers.')
            if len(parts) > 1:
                info['layer_number'] = int(parts[1].split('.')[0])
        except (IndexError, ValueError):
            pass
    
    if 'self_attn' in layer_name:
        if 'q_proj' in layer_name:
            info['layer_type'] = 'Attention_Q'
        elif 'k_proj' in layer_name:
            info['layer_type'] = 'Attention_K'
        elif 'v_proj' in layer_name:
            info['layer_type'] = 'Attention_V'
        elif 'o_proj' in layer_name:
            info['layer_type'] = 'Attention_O'
    elif 'mlp' in layer_name:
        if 'gate_proj' in layer_name:
            info['layer_type'] = 'MLP_Gate'
        elif 'up_proj' in layer_name:
            info['layer_type'] = 'MLP_Up'
        elif 'down_proj' in layer_name:
            info['layer_type'] = 'MLP_Down'
    
    if 'nodes.' in layer_name:
        try:
            parts = layer_name.split('nodes.')
            if len(parts) > 1:
                node_str = parts[1].split('.')[0] if '.' in parts[1] else parts[1]
                info['node_index'] = int(node_str)
        except (IndexError, ValueError):
            pass
    
    return info