Files
Building/scripts_4/adapter.py
CrbnsCat10n 79f3de05f7 modified: .DS_Store
modified:   checkpoints_final/task1_final/task1_final_model.pkl
	new file:   checkpoints_final/task1_final/task1_tmd_adapter.pkl
	new file:   scripts_4/README.md
	new file:   scripts_4/__init__.py
	new file:   scripts_4/__pycache__/adapter.cpython-314.pyc
	new file:   scripts_4/adapter.py
	new file:   scripts_4/predict_tmd_single.py
	new file:   scripts_4/train_tmd_adapter.py
2026-05-06 20:57:43 +08:00

102 lines
3.9 KiB
Python

from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import pickle
import numpy as np
def _safe_float(value: float) -> float:
return float(np.asarray(value, dtype=np.float64).reshape(-1)[0])
@dataclass
class FrequencyScaleAdapter:
frequencies_hz: np.ndarray
scale_values: np.ndarray
shrinkage: float = 0.2
extrapolation_decay_hz: float = 0.25
min_scale: float = 0.4
max_scale: float = 2.5
def __post_init__(self) -> None:
freqs = np.asarray(self.frequencies_hz, dtype=np.float64).reshape(-1)
scales = np.asarray(self.scale_values, dtype=np.float64).reshape(-1)
if freqs.size == 0:
raise ValueError("frequencies_hz cannot be empty.")
if freqs.size != scales.size:
raise ValueError("frequencies_hz and scale_values must have the same length.")
order = np.argsort(freqs)
freqs = freqs[order]
scales = np.clip(scales[order], self.min_scale, self.max_scale)
# Pull scales toward 1.0 to reduce overfitting with tiny TMD samples.
scales = 1.0 + (scales - 1.0) * (1.0 - float(np.clip(self.shrinkage, 0.0, 0.95)))
self.frequencies_hz = freqs
self.scale_values = scales
def _inside_range_weight(self, frequency_hz: float) -> float:
left = self.frequencies_hz[0]
right = self.frequencies_hz[-1]
f = _safe_float(frequency_hz)
if left <= f <= right:
return 1.0
distance = min(abs(f - left), abs(f - right))
decay = max(_safe_float(self.extrapolation_decay_hz), 1e-6)
return float(np.exp(-distance / decay))
def scale_at(self, frequency_hz: float) -> float:
f = _safe_float(frequency_hz)
boundary_interp = float(np.interp(f, self.frequencies_hz, self.scale_values))
in_range_weight = self._inside_range_weight(f)
scale = 1.0 + in_range_weight * (boundary_interp - 1.0)
return float(np.clip(scale, self.min_scale, self.max_scale))
def predict(self, base_rms: float, frequency_hz: float) -> float:
return max(_safe_float(base_rms), 0.0) * self.scale_at(frequency_hz)
def to_payload(self) -> dict:
return {
"frequencies_hz": self.frequencies_hz.tolist(),
"scale_values": self.scale_values.tolist(),
# scale_values are already post-shrinkage values after __post_init__.
"shrinkage": 0.0,
"extrapolation_decay_hz": float(self.extrapolation_decay_hz),
"min_scale": float(self.min_scale),
"max_scale": float(self.max_scale),
"already_shrunk": True,
}
@classmethod
def from_payload(cls, payload: dict) -> "FrequencyScaleAdapter":
shrinkage = float(payload.get("shrinkage", 0.2))
if bool(payload.get("already_shrunk", False)):
shrinkage = 0.0
return cls(
frequencies_hz=np.asarray(payload["frequencies_hz"], dtype=np.float64),
scale_values=np.asarray(payload["scale_values"], dtype=np.float64),
shrinkage=shrinkage,
extrapolation_decay_hz=float(payload.get("extrapolation_decay_hz", 0.25)),
min_scale=float(payload.get("min_scale", 0.4)),
max_scale=float(payload.get("max_scale", 2.5)),
)
def save_adapter(adapter: FrequencyScaleAdapter, output_path: Path, extra: dict | None = None) -> None:
output_path.parent.mkdir(parents=True, exist_ok=True)
payload = {
"adapter_type": "frequency_scale_interp_v1",
"adapter": adapter.to_payload(),
"extra": extra or {},
}
with output_path.open("wb") as handle:
pickle.dump(payload, handle)
def load_adapter(adapter_path: Path) -> tuple[FrequencyScaleAdapter, dict]:
with adapter_path.open("rb") as handle:
payload = pickle.load(handle)
adapter = FrequencyScaleAdapter.from_payload(payload["adapter"])
extra = payload.get("extra", {})
return adapter, extra