diff --git a/benchmarks/imf-runtime/modal-bench.py b/benchmarks/imf-runtime/modal-bench.py index af751dc..4a75952 100644 --- a/benchmarks/imf-runtime/modal-bench.py +++ b/benchmarks/imf-runtime/modal-bench.py @@ -69,6 +69,7 @@ def main( if not filename: import yaml - index = yaml.safe_load(open("models.yaml", encoding="utf-8")) + with open("models.yaml", encoding="utf-8") as fh: + index = yaml.safe_load(fh) filename = index["models"][model_id]["filename"] print(bench.remote(model_id, filename)) diff --git a/runtime/src/interscript_ml/model.py b/runtime/src/interscript_ml/model.py index 2491c39..a033bde 100644 --- a/runtime/src/interscript_ml/model.py +++ b/runtime/src/interscript_ml/model.py @@ -6,8 +6,8 @@ import numpy as np -from interscript_ml.tokens import EOS_ID, PAD_ID, decode, encode from interscript_ml.loader import load_manifest, verify_and_read +from interscript_ml.tokens import EOS_ID, PAD_ID, decode, encode class Model: @@ -45,7 +45,7 @@ def __init__(self, zip_path: Path | str): self._output_names = [o.name for o in self._decoder.get_outputs()] @classmethod - def load(cls, path_or_id: Path | str, index_url: str | None = None) -> "Model": + def load(cls, path_or_id: Path | str, index_url: str | None = None) -> Model: """Accepts a zip path OR a model id from models.yaml (dynamic fetch: download -> verify -> cache).""" candidate = str(path_or_id) diff --git a/runtime/tests/test_model.py b/runtime/tests/test_model.py index b72677d..36be14e 100644 --- a/runtime/tests/test_model.py +++ b/runtime/tests/test_model.py @@ -19,11 +19,10 @@ ort = pytest.importorskip("onnxruntime") onnx = pytest.importorskip("onnx") +import numpy as np # noqa: E402 from interscript_ml import Model, ModelFormatError, decode, encode # noqa: E402 from onnx import TensorProto, helper, numpy_helper # noqa: E402 -import numpy as np # noqa: E402 - def _graph(opset: int = 14) -> bytes: graph = helper.make_graph( diff --git a/runtime/tests/test_registry.py b/runtime/tests/test_registry.py index ad6dbbe..4dd61d8 100644 --- a/runtime/tests/test_registry.py +++ b/runtime/tests/test_registry.py @@ -3,17 +3,14 @@ from __future__ import annotations import hashlib -import zipfile +import os # noqa: E402 from pathlib import Path import pytest import yaml - from interscript_ml.registry import RegistryError, resolve from tests_helpers import build_tiny_zip -import os # noqa: E402 - def _index_file(tmp_path: Path, zip_path: Path, sha256: str | None = None) -> Path: index = { diff --git a/scripts/publish_model.py b/scripts/publish_model.py index 13eab53..2e40480 100644 --- a/scripts/publish_model.py +++ b/scripts/publish_model.py @@ -30,9 +30,10 @@ sys.path.insert(0, str(REPO_ROOT / "src")) sys.path.insert(0, str(REPO_ROOT / "scripts")) -from imf.validator import validate_zip # noqa: E402 from split_release import split as split_zip # noqa: E402 +from imf.validator import validate_zip # noqa: E402 + # GitHub hard-caps release assets at 2,147,483,648 bytes; split well below. SPLIT_THRESHOLD = 2_000_000_000 DEFAULT_REPO = "interscript/interscript-ml" diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index 7ba823d..e76ccd8 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -429,8 +429,6 @@ def val_loss() -> float: ck.mkdir(exist_ok=True) torch.save(student.state_dict(), ck / "student.pt") torch.save(optimizer.state_dict(), ck / "optim.pt") - if mtp_head is not None: - torch.save(mtp_head.state_dict(), ck / "mtp_head.pt") CHECKPOINTS.commit() vl = val_loss() diff --git a/src/gpu/modal_export.py b/src/gpu/modal_export.py index 241c9c4..3300053 100644 --- a/src/gpu/modal_export.py +++ b/src/gpu/modal_export.py @@ -356,7 +356,9 @@ def stage(event: str) -> None: timeout=5 * 3600, volumes={**CHECKPOINT_VOLUMES, **DATASET_VOLUMES, "/outputs": MODELS_VOLUME}, ) -def margin_model(model_id: str, precisions: list[str], limit: int = 0, dump_positions: bool = False) -> dict[str, str]: +def margin_model( + model_id: str, precisions: list[str], limit: int = 0, dump_positions: bool = False +) -> dict[str, str]: """Margin analysis alone over already-exported zips — read-only for the zips (diagnostic JSON only); validates published artifacts without touching their metadata.""" @@ -414,7 +416,9 @@ def parity(model: str, precisions: str = "fp32,fp16,int8", limit: int = 0) -> No @app.local_entrypoint() -def margins(model: str, precisions: str = "fp32,fp16,int8", limit: int = 0, dump_positions: bool = False) -> None: +def margins( + model: str, precisions: str = "fp32,fp16,int8", limit: int = 0, dump_positions: bool = False +) -> None: reports = margin_model.remote(model, precisions.split(","), limit, dump_positions) for precision, status in reports.items(): print(f"{model} [{precision}] {status}") diff --git a/src/gpu/mtp.py b/src/gpu/mtp.py index 53ae66b..574a8fc 100644 --- a/src/gpu/mtp.py +++ b/src/gpu/mtp.py @@ -11,8 +11,8 @@ from __future__ import annotations import torch -from torch import nn import torch.nn.functional as F +from torch import nn class MTPHead(nn.Module): diff --git a/src/imf/parity.py b/src/imf/parity.py index c9b98e5..6be0894 100644 --- a/src/imf/parity.py +++ b/src/imf/parity.py @@ -253,7 +253,9 @@ def _onnx_forced_logits(enc_sess, dec_sess, source: str, target_ids: list[int]): return dict(zip(out_names, out, strict=True))["logits"][0] -def run_margin_analysis(model, zip_path, pairs, max_len: int = 256, dump_positions: Path | str | None = None) -> MarginReport: +def run_margin_analysis( + model, zip_path, pairs, max_len: int = 256, dump_positions: Path | str | None = None +) -> MarginReport: """pairs: iterable of (source_text, gold_target) — the same probe set the CER parity gate uses. Teacher-forces both sides and measures the argmax flip rate, reference top1−top2 margin quantiles, and KL @@ -294,7 +296,7 @@ def run_margin_analysis(model, zip_path, pairs, max_len: int = 256, dump_positio dump = Path(dump_positions) dump.parent.mkdir(parents=True, exist_ok=True) with dump.open("w", encoding="utf-8") as out: - for pair_idx, (m, f) in enumerate(zip(margin_chunks, flip_chunks)): + for pair_idx, (m, f) in enumerate(zip(margin_chunks, flip_chunks, strict=True)): out.write(_json.dumps( {"pair": pair_idx, "tokens": int(f.size), "flip_positions": [int(x) for x in np.nonzero(f)[0]], diff --git a/tests/test_imf_parity.py b/tests/test_imf_parity.py index 56a57a6..c36dde3 100644 --- a/tests/test_imf_parity.py +++ b/tests/test_imf_parity.py @@ -168,7 +168,9 @@ def test_margin_report_json_roundtrip(gated_zip: Path, fixture_model, tmp_path: assert data["flip_rate"] == report.flip_rate -def test_margin_analysis_dumps_per_pair_positions(gated_zip: Path, fixture_model, tmp_path: Path) -> None: +def test_margin_analysis_dumps_per_pair_positions( + gated_zip: Path, fixture_model, tmp_path: Path +) -> None: """TODO.training-work/05: the flip bootstrap needs per-pair token and flip counts, not just aggregates.""" dump = tmp_path / "positions.jsonl"