fix(lint): ruff clean on main; drop MTP save lines from logit-mode distill

The F821 was a real latent bug: MTP checkpoint lines from the #114
wiring were pasted into the logit-mode distiller save block where
mtp_head never exists — any logit-mode run would NameError at its
first step checkpoint. The lines belong only in distill_sequence.
The rest: import sorting, unused zipfile, quoted annotation, E501
signature wraps, zip strict=, context manager for models.yaml.

Author: Ronald Tse

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))

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)

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(

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 = {

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"

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()

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}")

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):

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]],

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"