style: ruff clean (unused import, line length, zip strict)

Author: Ronald Tse

src/gpu/modal_golden_matrix.py

@@ -39,19 +39,31 @@
 
 # index id -> volume zip path + test pairs (volume-relative)
 SOURCES = {
-    "khm-latn-1.0": ("imf/khm-latn/khm-latn-1.0-fp32.zip", "/secryst-datasets/khmer-translit/test.jsonl"),
+    "khm-latn-1.0":
+        ("imf/khm-latn/khm-latn-1.0-fp32.zip",
+         "/secryst-datasets/khmer-translit/test.jsonl"),
     "urd-g2p-1.0": ("imf/urd-g2p/urd-g2p-1.0-fp32.zip", "/ud-g2p/urdu-g2p/test.jsonl"),
     "urd-diac-1.0": ("imf/urd-diac/urd-diac-1.0-fp32.zip", "/ud-diacrit/urdu-diacrit/test.jsonl"),
-    "tha-g2p-base-1.0": ("imf/tha-g2p-base/tha-g2p-base-1.0-fp32.zip", "/secryst-datasets/thai-ipa/test.jsonl"),
-    "tha-g2p-small-1.0": ("imf/tha-g2p-small/tha-g2p-small-1.0-int8.zip", "/secryst-datasets/thai-ipa/test.jsonl"),
+    "tha-g2p-base-1.0":
+        ("imf/tha-g2p-base/tha-g2p-base-1.0-fp32.zip",
+         "/secryst-datasets/thai-ipa/test.jsonl"),
+    "tha-g2p-small-1.0":
+        ("imf/tha-g2p-small/tha-g2p-small-1.0-int8.zip",
+         "/secryst-datasets/thai-ipa/test.jsonl"),
     "fas-g2p-1.0": ("imf/fas-g2p/fas-g2p-1.0-fp32.zip", "/persian/persian-g2p/test.jsonl"),
     "heb-diac-1.0": ("imf/heb-diac/heb-diac-1.0-fp32.zip", "nakdimon/test-imf.jsonl"),
     "heb-diac-1.1": ("imf/heb-diac/heb.zip", "nakdimon/test-imf.jsonl"),
-    "heb-diac-small-1.0": ("imf/heb-diac-small/heb-diac-small-1.0-fp32.zip", "nakdimon/test-imf.jsonl"),
+    "heb-diac-small-1.0":
+        ("imf/heb-diac-small/heb-diac-small-1.0-fp32.zip",
+         "nakdimon/test-imf.jsonl"),
     "ara-diac-1.0": ("imf/ara-diac/ara-diac-1.0-fp32.zip", "arabic-sadeed-imf/test.jsonl"),
     "ara-diac-2.0-int8": ("imf/ara-diac2/ara-diac-2.0-int8.zip", "arabic-sadeed-imf/test.jsonl"),
-    "ara-diac-small-2.1-int8": ("imf/ara-diac-small-21/ara-diac-small-2.1-int8.zip", "arabic-sadeed-imf/test.jsonl"),
-    "ara-diac-layerdrop-1.0-int4": ("imf/ara-diac-layerdrop/ara-diac-layerdrop-1.0-int4.zip", "arabic-sadeed-imf/test.jsonl"),
+    "ara-diac-small-2.1-int8":
+        ("imf/ara-diac-small-21/ara-diac-small-2.1-int8.zip",
+         "arabic-sadeed-imf/test.jsonl"),
+    "ara-diac-layerdrop-1.0-int4":
+        ("imf/ara-diac-layerdrop/ara-diac-layerdrop-1.0-int4.zip",
+         "arabic-sadeed-imf/test.jsonl"),
 }
 
 

src/gpu/modal_teacher_sadeed.py

@@ -53,7 +53,6 @@ def teacher_preds() -> dict:
     import sys
 
     import pyarrow.parquet as pq
-    import torch
     from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
     # Modal copies the entry file to /root/<name>.py while the repo

tests/test_muon_headwise.py

@@ -43,7 +43,7 @@ def test_headwise_equals_per_slice_vanilla() -> None:
     hw.step()
 
     vanilla = Muon(slices, lr=0.01, momentum=0.95)
-    for p, g in zip(slices, slice_grads):
+    for p, g in zip(slices, slice_grads, strict=True):
         p.grad = g.clone()
     vanilla.step()
 
@@ -114,7 +114,8 @@ def test_split_parameters_unchanged_when_headwise_unused() -> None:
     group kind, it does not touch the default split."""
     named = [
         ("shared.weight", torch.zeros(3, 4, requires_grad=True)),
-        ("encoder.block.0.layer.1.DenseReluDense.wi_0.weight", torch.zeros(4, 3, requires_grad=True)),
+        ("encoder.block.0.layer.1.DenseReluDense.wi_0.weight",
+         torch.zeros(4, 3, requires_grad=True)),
         ("encoder.block.0.layer.0.layer_norm.weight", torch.zeros(3, requires_grad=True)),
     ]
     muon, adamw = split_parameters(named)