diff --git a/src/api/inference.py b/src/api/inference.py index d82492a..1db7a26 100644 --- a/src/api/inference.py +++ b/src/api/inference.py @@ -26,7 +26,10 @@ .env({"IMAGE_REV": "6"}) .add_local_dir(str(Path(__file__).resolve().parent.parent), "/root/interscript-ml", copy=True) # deterministic mounts for the index-driven resolver + the index itself - .add_local_file(str(Path(__file__).resolve().parent / "model_resolution.py"), "/root/model_resolution.py") + .add_local_file( + str(Path(__file__).resolve().parent / "model_resolution.py"), + "/root/model_resolution.py", + ) .add_local_file(str(_REPO_ROOT / "models.yaml"), "/root/models.yaml") ) diff --git a/src/api/model_resolution.py b/src/api/model_resolution.py index 629faeb..1807487 100644 --- a/src/api/model_resolution.py +++ b/src/api/model_resolution.py @@ -13,9 +13,10 @@ from __future__ import annotations -import yaml from pathlib import Path +import yaml + def load_index(path: str | Path = "models.yaml") -> dict: with open(path, encoding="utf-8") as fh: diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index 1094499..19f6749 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -1617,7 +1617,7 @@ def probe_pkm_gates(spec_id: str = "ara-diac-small-pkm") -> dict: raise RuntimeError(f"no step checkpoints under {out_root}") _ensure_src_path() - from gpu.pkm import ProductKeyMemory, inject_pkm + from gpu.pkm import inject_pkm student = AutoModelForSeq2SeqLM.from_pretrained(spec["student_init"]) inject_pkm(student, **spec["pkm"]) diff --git a/src/gpu/muon.py b/src/gpu/muon.py index 4431dfc..b27d836 100644 --- a/src/gpu/muon.py +++ b/src/gpu/muon.py @@ -41,16 +41,27 @@ def __init__(self, params, lr: float = 0.01, momentum: float = 0.95, weight_decay: float = 0.0) -> None: super().__init__( list(params), - dict(lr=lr, momentum=momentum, nesterov=nesterov, - ns_steps=ns_steps, weight_decay=weight_decay, adamw=False), + { + "lr": lr, + "momentum": momentum, + "nesterov": nesterov, + "ns_steps": ns_steps, + "weight_decay": weight_decay, + "adamw": False, + }, ) def add_adamw_group(self, params, lr: float = 1e-4, betas=(0.9, 0.999), weight_decay: float = 0.0) -> None: """Embedding-like parameters: standard AdamW math, shared scheduler (the cosine scales every group's lr).""" - self.add_param_group(dict(params=list(params), lr=lr, betas=tuple(betas), - weight_decay=weight_decay, adamw=True)) + self.add_param_group({ + "params": list(params), + "lr": lr, + "betas": tuple(betas), + "weight_decay": weight_decay, + "adamw": True, + }) @torch.no_grad() def step(self, closure=None): # noqa: ARG002 diff --git a/src/imf/parity.py b/src/imf/parity.py index 62ff899..ef33ed3 100644 --- a/src/imf/parity.py +++ b/src/imf/parity.py @@ -277,9 +277,8 @@ def run_margin_analysis(model, zip_path, pairs, max_len: int = 256) -> MarginRep def write_margin_report(report: MarginReport, out_path: Path | str) -> Path: """Emit the margin analysis as JSON next to a release zip (diagnostic artifact; the release gate remains the CER parity block).""" - from dataclasses import asdict - import json + from dataclasses import asdict out_path = Path(out_path) out_path.parent.mkdir(parents=True, exist_ok=True) diff --git a/tests/test_decode_guard.py b/tests/test_decode_guard.py index 679d85e..282b7b3 100644 --- a/tests/test_decode_guard.py +++ b/tests/test_decode_guard.py @@ -22,24 +22,24 @@ def __init__(self): self.step = 0 def get_outputs(self): - class O: + class Out: def __init__(self, name): self.name = name - return [O("logits"), O("present_k"), O("present_v")] + return [Out("logits"), Out("present_k"), Out("present_v")] def get_inputs(self): - class I: + class In: def __init__(self, name, shape, typ="tensor(float)"): self.name = name self.shape = shape self.type = typ return [ - I("input_ids", [1, "seq"]), - I("encoder_hidden_states", [1, "seq", 4]), - I("past_k", [1, 4, "past", 8]), - I("past_v", [1, 4, "past", 8]), + In("input_ids", [1, "seq"]), + In("encoder_hidden_states", [1, "seq", 4]), + In("past_k", [1, 4, "past", 8]), + In("past_v", [1, 4, "past", 8]), ] def run(self, _, feeds): diff --git a/tests/test_imf_parity.py b/tests/test_imf_parity.py index 60c5f99..3510cef 100644 --- a/tests/test_imf_parity.py +++ b/tests/test_imf_parity.py @@ -23,7 +23,6 @@ make_fixture_checkpoint, ) from imf.parity import ( # noqa: E402 - MarginReport, ParityReport, _margin_stats, run_margin_analysis,