router_loader.py

"""Minimal loaders for the public MoS/DFlash request-level router artifacts."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import torch
from safetensors.torch import load_file
from torch import nn


class OfflineV2Head(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, dropout: float):
        super().__init__()
        self.n = nn.LayerNorm(input_dim)
        self.f1 = nn.Linear(input_dim, hidden_dim)
        self.dropout = nn.Dropout(dropout)
        self.f2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        hidden = torch.nn.functional.gelu(self.f1(self.n(features)))
        return self.f2(self.dropout(hidden))


class SidecarHead(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, dropout: float):
        super().__init__()
        self.router_head = nn.Sequential(
            nn.LayerNorm(input_dim),
            nn.Linear(input_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, output_dim),
        )

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        return self.router_head(features)


def _read_config(path: str | Path) -> dict[str, Any]:
    value = json.loads(Path(path).read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise ValueError("router config must be a JSON object")
    return value


def load_offline_v2(
    weights_path: str | Path,
    config_path: str | Path,
    *,
    device: str | torch.device = "cpu",
) -> tuple[OfflineV2Head, torch.Tensor, torch.Tensor, dict[str, Any]]:
    config = _read_config(config_path)
    tensors = load_file(str(weights_path), device=str(device))
    expected = {
        "n.weight",
        "n.bias",
        "f1.weight",
        "f1.bias",
        "f2.weight",
        "f2.bias",
        "feature_mean",
        "feature_std",
    }
    if set(tensors) != expected:
        raise ValueError(f"unexpected offline-v2 tensor keys: {sorted(tensors)}")
    head_config = config["head"]
    head = OfflineV2Head(
        input_dim=int(head_config["input_dim"]),
        hidden_dim=int(head_config["hidden_dim"]),
        output_dim=int(head_config["output_dim"]),
        dropout=float(head_config["dropout"]),
    )
    head.load_state_dict(
        {key: value for key, value in tensors.items() if key not in {"feature_mean", "feature_std"}},
        strict=True,
    )
    head.to(device).eval()
    return head, tensors["feature_mean"], tensors["feature_std"], config


def prepare_offline_features(
    features: torch.Tensor,
    feature_mean: torch.Tensor,
    feature_std: torch.Tensor,
    *,
    dead_std_threshold: float = 2e-6,
    clamp_abs: float = 10.0,
) -> torch.Tensor:
    safe_std = feature_std.clamp_min(1e-6)
    normalized = (features - feature_mean) / safe_std
    normalized = normalized.masked_fill(feature_std <= dead_std_threshold, 0)
    return normalized.clamp(-clamp_abs, clamp_abs)


def load_sidecar(
    weights_path: str | Path,
    config_path: str | Path,
    *,
    device: str | torch.device = "cpu",
) -> tuple[SidecarHead, dict[str, Any]]:
    config = _read_config(config_path)
    tensors = load_file(str(weights_path), device=str(device))
    expected = {
        "router_head.0.weight",
        "router_head.0.bias",
        "router_head.1.weight",
        "router_head.1.bias",
        "router_head.4.weight",
        "router_head.4.bias",
    }
    if set(tensors) != expected:
        raise ValueError(f"unexpected sidecar tensor keys: {sorted(tensors)}")
    head_config = config["head"]
    head = SidecarHead(
        input_dim=int(head_config["input_dim"]),
        hidden_dim=int(head_config["hidden_dim"]),
        output_dim=int(head_config["output_dim"]),
        dropout=float(head_config["dropout"]),
    )
    head.load_state_dict(tensors, strict=True)
    head.to(device).eval()
    return head, config