Spaces:
Sleeping
Sleeping
| """Thin glue the stage notebooks call so each one is ~5 cells. | |
| ``init_stage(profile)`` resolves the run profile, mounts Drive (Colab), and builds | |
| the shared ``ArtifactPaths`` — returning a ``StageContext`` whose ``run_step`` runs | |
| a pipeline CLI with ``PYTHONPATH`` set and ``restore``/``save`` move artifacts | |
| to/from Drive. On Colab the first notebook cell does a minimal locate-or-clone | |
| before importing this module (it can't import the package until ``scribe`` is | |
| on ``sys.path``). | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import subprocess | |
| import sys | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from gec import env | |
| from gec.paths import ArtifactPaths | |
| from gec.profiles import RunProfile, get_profile | |
| class StageContext: | |
| profile: RunProfile | |
| paths: ArtifactPaths | |
| backup: Path | None | |
| in_colab: bool | |
| dataset: str = "tensorxt/ViMedCSS" | |
| def run_step(self, args: list[str], env_extra: dict | None = None) -> None: | |
| """Run a ``scribe/training/scripts/*`` CLI with PYTHONPATH set, raising on failure.""" | |
| run_env = dict(os.environ) | |
| run_env["PYTHONPATH"] = os.pathsep.join(("scribe/training", "scribe")) | |
| run_env["PYTHONIOENCODING"] = "utf-8" | |
| if env_extra: | |
| run_env.update(env_extra) | |
| printable = " ".join(str(a) for a in args) | |
| print(">>>", printable, flush=True) | |
| proc = subprocess.run([sys.executable, *map(str, args)], env=run_env) | |
| if proc.returncode != 0: | |
| raise RuntimeError(f"step failed ({proc.returncode}): {printable}") | |
| def restore(self, rel_paths: list[str]) -> None: | |
| env.restore_artifacts(self.backup, rel_paths) | |
| def restore_optional(self, rel_paths: list[str]) -> None: | |
| """Copy artifacts from Drive *if present*, without failing when they aren't. | |
| Used for inputs that may legitimately be absent (no synthetic pairs, no | |
| labeled export) so a teammate continuing a run on a fresh Colab still pulls | |
| whatever exists on Drive. | |
| """ | |
| if self.backup is None: | |
| return # local: files already on disk (or intentionally absent) | |
| import shutil | |
| for rel in rel_paths: | |
| dst = Path(rel) | |
| src = self.backup / dst.name | |
| if src.exists(): | |
| dst.parent.mkdir(parents=True, exist_ok=True) | |
| shutil.copy(src, dst) | |
| print("restored", dst, "from", src) | |
| else: | |
| print("(optional, not on Drive):", dst) | |
| def save(self, rel_paths: list[str]) -> None: | |
| env.save_artifacts(self.backup, rel_paths) | |
| def durable(self, path: str | Path) -> str: | |
| """Output path that survives a Colab runtime recycle. | |
| Long, resumable stages (ASR pairs, TTS) write here so ``--resume`` can pick | |
| up after a disconnect: on Colab that's the Drive backup, where every flushed | |
| row already persists; locally it's the normal ``artifacts/`` path. Uses the | |
| same basename as ``save``/``restore`` so a downstream stage's ``restore`` | |
| finds it unchanged. | |
| """ | |
| p = Path(path) | |
| if self.backup is None: | |
| return str(p) | |
| dst = self.backup / p.name | |
| dst.parent.mkdir(parents=True, exist_ok=True) | |
| return str(dst) | |
| def init_stage(profile: str = "smoke", dataset: str = "tensorxt/ViMedCSS") -> StageContext: | |
| """Resolve profile + Drive backup + artifact paths for a stage notebook.""" | |
| prof = get_profile(profile) | |
| in_colab = env.in_colab() | |
| backup = env.setup_backup(in_colab) | |
| suffix = "" if prof.name == "full" else f"_{prof.name}" | |
| adapters_root = (backup / "gec_lora" / "qwen3") if backup else None | |
| paths = ArtifactPaths(root=Path("artifacts"), suffix=suffix, adapters_root=adapters_root) | |
| print(f"profile={prof.name} | n_best={prof.n_best} | seeds={prof.seeds} | adapters={paths.adapters}") | |
| return StageContext(profile=prof, paths=paths, backup=backup, in_colab=in_colab, dataset=dataset) | |