Spaces:
Runtime error
Runtime error
File size: 4,358 Bytes
1cf3825 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | """Rebuild colab_bundle.zip with ALL transitive deps the v9b trainers need.
Previous miss (2026-06-02): train_v9b_stage1_jepa.py imports `_atomic_save`
from train_segmentation_v7.py, and v9b_model.py imports
`synthetic_brain_sdf_template` from research/geometric_prior.py. Neither
was in the first bundle, so Stage 1 crashed on Colab with
`ModuleNotFoundError: No module named 'src.train_segmentation_v7'`.
This script is the source of truth for what goes in the bundle.
"""
from __future__ import annotations
import time
import zipfile
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / 'colab_bundle.zip'
# Top-level colab_bundle/ files
BUNDLE_FILES = [
'colab_bundle/README_COLAB.md',
'colab_bundle/__init__.py',
'colab_bundle/requirements_colab.txt',
'colab_bundle/v9b_colab_train.ipynb',
]
# src/ files unzipped into /content/neurolens/src/ on the Colab VM. After
# the 2026-06-02 refactor, v9b stages import `atomic_save` from the new
# `src/checkpoint_utils.py` (no further local imports), so the v5+v7
# trainer chain is no longer needed in this bundle — only the actual v9b
# code + its research modules.
SRC_FILES = [
'src/__init__.py',
'src/utils.py',
'src/checkpoint_utils.py',
'src/train_v9b_stage1_jepa.py',
'src/train_v9b_stage2.py',
'src/train_v9b_andi_ddpm.py', # NEW (June 2026): proper ANDi
# DDPM training with pyramidal noise.
'src/train_v9c_stage1.py', # NEW (June 2026): v9c JEPA predictor
# on frozen DINOv2 backbone.
'src/v9b_inference.py',
'src/research/__init__.py',
'src/research/jepa.py',
'src/research/jepa_conformal.py',
'src/research/latent_diffusion_decoder.py',
'src/research/sdf_geometric_tower.py',
'src/research/symmetry_geometry.py', # NEW: deterministic symmetry
# geometry score, replaces SDF.
'src/research/pyramidal_noise.py', # NEW: ANDi pyramidal noise gen.
'src/research/andi_inference.py', # NEW: ANDi inference aggregation.
'src/research/two_tower_anomaly.py',
'src/research/geometric_prior.py',
'src/research/mesh_extraction.py',
'src/research/mni152_registration.py',
'src/research/v9b_model.py',
'src/research/v9b_advisory.py', # end-to-end advisory wrapper.
'src/research/v9c_dinov2_jepa.py', # NEW (June 2026): v9c model
# (frozen DINOv2 + JEPA predictor).
]
def main():
t0 = time.perf_counter()
if OUT.exists():
OUT.unlink()
missing: list[str] = []
with zipfile.ZipFile(OUT, 'w', zipfile.ZIP_DEFLATED, compresslevel=6) as zf:
for rel in BUNDLE_FILES + SRC_FILES:
p = ROOT / rel
if not p.exists():
missing.append(rel)
continue
# Inside the zip: BUNDLE_FILES go to the top level (strip
# "colab_bundle/" prefix), SRC_FILES keep their src/ path so
# the notebook can unzip and `python src/train_v9b_*` works.
if rel.startswith('colab_bundle/'):
arc = rel[len('colab_bundle/'):]
else:
arc = rel
zf.write(p, arcname=arc)
if missing:
print('WARNING: missing source files (not bundled):')
for m in missing:
print(f' - {m}')
size_kb = OUT.stat().st_size / 1024
print(f'\n[done] colab_bundle.zip = {size_kb:.1f} KB '
f'({len(BUNDLE_FILES) + len(SRC_FILES) - len(missing)} files) '
f'in {time.perf_counter()-t0:.1f}s')
# Verify checkpoint_utils + geometric_prior are in, and that the v7
# trainer (now removed) is NOT — confirms the bundle slimming worked.
with zipfile.ZipFile(OUT) as zf:
names = set(zf.namelist())
for need in ('src/checkpoint_utils.py', 'src/research/geometric_prior.py'):
print(f' {"[OK] " if need in names else "[FAIL]"} {need}')
for gone in ('src/train_segmentation_v5.py', 'src/train_segmentation_v7.py'):
print(f' {"[OK (removed)]" if gone not in names else "[STILL IN BUNDLE]"} {gone}')
if __name__ == '__main__':
main()
|