config.py

"""Configuration and profile loading helpers for HeadRouter."""

import json
import math
from pathlib import Path
from typing import Dict, Tuple, Union

import torch


PathLike = Union[str, Path]
_CONFIG_ROOT = Path(__file__).resolve().parents[1] / "configs" / "qwen2_5_omni_3b"
DEFAULT_PROFILE_DIR = _CONFIG_ROOT / "layer1"
DEFAULT_ROUTING_CONFIG = _CONFIG_ROOT / "routing.json"
PROFILE_NAMES = ("semantic", "uniform", "acoustic")


def _load_weight_array(path: Path) -> torch.Tensor:
    with path.open("r", encoding="utf-8") as handle:
        data = json.load(handle)
    weights = data.get("weights")
    if weights is None:
        layer_keys = [key for key in data if key.startswith("layer") and key.endswith("_weights")]
        if len(layer_keys) != 1:
            raise ValueError(f"{path} must contain 'weights' or one layer weights array")
        weights = data[layer_keys[0]]

    tensor = torch.tensor(weights, dtype=torch.float32)
    if tensor.ndim != 1 or tensor.numel() == 0:
        raise ValueError(f"{path} must contain a non-empty one-dimensional weights array")
    if not torch.isfinite(tensor).all() or (tensor < 0).any() or tensor.sum() <= 0:
        raise ValueError(f"{path} contains invalid head weights")
    return tensor / tensor.sum()


def load_profiles(profile_dir: PathLike = DEFAULT_PROFILE_DIR) -> Dict[str, torch.Tensor]:
    """Load semantic, uniform, and acoustic head-weight profiles."""

    directory = Path(profile_dir).expanduser().resolve()
    profiles = {name: _load_weight_array(directory / f"{name}.json") for name in PROFILE_NAMES}
    head_counts = {weights.numel() for weights in profiles.values()}
    if len(head_counts) != 1:
        raise ValueError("all HeadRouter profiles must contain the same number of heads")
    return profiles


def load_routing_config(
    path: PathLike = DEFAULT_ROUTING_CONFIG,
) -> Tuple[Dict[str, float], float]:
    """Load and validate Gaussian routing centers and bandwidth."""

    config_path = Path(path).expanduser().resolve()
    with config_path.open("r", encoding="utf-8") as handle:
        data = json.load(handle)

    centers = {name: float(data["centers"][name]) for name in PROFILE_NAMES}
    sigma = float(data["sigma"])
    if not all(math.isfinite(value) for value in centers.values()):
        raise ValueError("routing centers must be finite")
    if not math.isfinite(sigma) or sigma <= 0:
        raise ValueError("routing sigma must be a positive finite value")
    return centers, sigma


def _model_config(model):
    thinker = getattr(model, "thinker", None)
    thinker_model = getattr(thinker, "model", None)
    if thinker_model is not None and hasattr(thinker_model, "config"):
        return thinker_model.config
    if hasattr(model, "config"):
        return model.config
    raise TypeError("model must expose either model.config or model.thinker.model.config")


def configure_headrouter(
    model,
    pruning_ratio: float = 0.6,
    pruning_layer: int = 2,
    profile_dir: PathLike = DEFAULT_PROFILE_DIR,
    routing_config: PathLike = DEFAULT_ROUTING_CONFIG,
) -> None:
    """Enable HeadRouter on the bundled Qwen2.5-Omni implementation."""

    if not math.isfinite(pruning_ratio) or not 0.0 <= pruning_ratio < 1.0:
        raise ValueError("pruning_ratio must satisfy 0 <= pruning_ratio < 1")
    if pruning_layer < 1:
        raise ValueError("pruning_layer must be at least 1")

    profile_path = Path(profile_dir).expanduser().resolve()
    routing_path = Path(routing_config).expanduser().resolve()
    profiles = load_profiles(profile_path)
    centers, sigma = load_routing_config(routing_path)

    config = _model_config(model)
    config.headrouter = True
    config.headrouter_dynamic_routing = True
    config.headrouter_profile_dir = str(profile_path)
    config.headrouter_profile_centers = centers
    config.headrouter_routing_sigma = sigma
    config.headrouter_num_heads = next(iter(profiles.values())).numel()

    # Compatibility flags used by the original experiment integration.
    config.hawk = True
    config.hawk_lite = False
    config.hawk_paper_mode = True
    config.hawk_paper_use_rope = False
    config.hawk_soft_rope = False
    config.hawk_head_weights_path = str(profile_path / "uniform.json")
    config.hawk_head_weight_method = "ablation"
    config.hawk_thermal_smoothing = False
    config.sss_alpha = 0.0
    config.sss_energy_threshold = 0.9
    config.audio_layer_idx = int(pruning_layer)
    config.audio_prune_ratio = float(pruning_ratio)
    config.random = False
    config.frame = False