"""
Stage 0: Download model weights if they don't already exist.
"""

import os
from huggingface_hub import snapshot_download

MODEL_ID = os.environ.get("MODEL_ID", "HuggingFaceTB/SmolLM3-3B")
MODEL_DIR = os.environ.get("MODEL_DIR", "/workspace/models/SmolLM3-3B")

def main():
    if os.path.exists(os.path.join(MODEL_DIR, "config.json")):
        print(f"[stage0] Weights already exist at {MODEL_DIR}, skipping download.")
        return

    print(f"[stage0] Downloading {MODEL_ID} → {MODEL_DIR} ...")
    snapshot_download(
        repo_id=MODEL_ID,
        local_dir=MODEL_DIR,
    )
    print(f"[stage0] Done. Weights saved to {MODEL_DIR}")


if __name__ == "__main__":
    main()