File size: 4,012 Bytes
29d1f68 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """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
|