26 lines
689 B
Python
26 lines
689 B
Python
|
"""
|
||
|
|
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()
|