Regenerate knapsack_cache.zip for instanovo 1.2.2 format

knapsack_cache.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:50ed79f898e338ddb996e927bb199dec3243fd4d4106f637ab0eed8a432013b6
-size 61446529
+oid sha256:ba45f03fb19b29c7c3a8f47aa7535bbd1f8f1407c4aa0bf2e3799453a63eb437
+size 62161027

scripts/regenerate_knapsack.py

@@ -0,0 +1,78 @@
+"""Regenerate knapsack_cache/ and knapsack_cache.zip against the currently
+installed instanovo version.
+
+When InstaNovo bumps its Knapsack on-disk format (e.g. v1.1 -> v1.2.2 added a
+field to parameters.pkl, giving the "not enough values to unpack" error on
+load), the bundled cache becomes unreadable and the runtime falls back to
+multi-minute regeneration on every container cold-start. This script does
+that regeneration once on a workstation so the committed zip stays in sync.
+
+Run with:
+    .venv/bin/python scripts/regenerate_knapsack.py
+
+Expect a few minutes on CPU. Output:
+    ./knapsack_cache/{parameters.pkl,masses.npy,chart.npy}
+    ./knapsack_cache.zip   (the three files above, repacked for HF Spaces)
+"""
+
+from __future__ import annotations
+
+import shutil
+import time
+import zipfile
+from pathlib import Path
+
+from instanovo.constants import MASS_SCALE, MAX_MASS
+from instanovo.inference import Knapsack
+from instanovo.transformer.model import InstaNovo
+
+TRANSFORMER_MODEL_ID = "instanovo-v1.2.0"
+KNAPSACK_DIR = Path(__file__).resolve().parent.parent / "knapsack_cache"
+KNAPSACK_ZIP = KNAPSACK_DIR.with_suffix(".zip")
+
+
+def main() -> None:
+    print(f"Loading {TRANSFORMER_MODEL_ID} for residue set...")
+    model, _ = InstaNovo.from_pretrained(TRANSFORMER_MODEL_ID)
+    residue_set = model.residue_set
+
+    residue_masses = dict(residue_set.residue_masses)
+    excluded = set(residue_set.special_tokens) | {
+        k for k, v in residue_masses.items() if v <= 0
+    }
+    for res in excluded:
+        residue_masses.pop(res, None)
+    if not residue_masses:
+        raise SystemExit("No valid residues with positive mass; cannot regenerate.")
+    print(f"Excluding {sorted(excluded)} from knapsack; {len(residue_masses)} residues remain.")
+
+    if KNAPSACK_DIR.exists():
+        print(f"Clearing existing {KNAPSACK_DIR}...")
+        shutil.rmtree(KNAPSACK_DIR)
+    KNAPSACK_DIR.parent.mkdir(parents=True, exist_ok=True)
+    # Knapsack.save() creates KNAPSACK_DIR itself and refuses if it already exists.
+
+    start = time.time()
+    print("Generating knapsack (this can take several minutes)...")
+    knapsack = Knapsack.construct_knapsack(
+        residue_masses=residue_masses,
+        residue_indices=residue_set.residue_to_index,
+        max_mass=MAX_MASS,
+        mass_scale=MASS_SCALE,
+    )
+    knapsack.save(str(KNAPSACK_DIR))
+    print(f"Saved knapsack to {KNAPSACK_DIR} ({time.time() - start:.1f}s).")
+
+    print(f"Packing {KNAPSACK_ZIP}...")
+    if KNAPSACK_ZIP.exists():
+        KNAPSACK_ZIP.unlink()
+    with zipfile.ZipFile(KNAPSACK_ZIP, "w", zipfile.ZIP_DEFLATED) as zf:
+        for f in sorted(KNAPSACK_DIR.iterdir()):
+            zf.write(f, arcname=f.name)
+    print(f"Wrote {KNAPSACK_ZIP} ({KNAPSACK_ZIP.stat().st_size / 1e6:.1f} MB).")
+
+    print("Done. Commit the new zip with git (LFS-tracked) and push to the Space.")
+
+
+if __name__ == "__main__":
+    main()