weather-llm-initial / scripts /train_pretrain.py
NagacharanVemula
Initial code-first model repo publish.
78dea75
Raw
History Blame Contribute Delete
1.18 kB
#!/usr/bin/env python3
"""Run causal LM pretraining (config-driven); optional resume from a checkpoint."""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from weather_llm.config_loader import repo_root_from
from weather_llm.training.pretrain import run_pretrain
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--pretrain_config", type=Path, default=None)
ap.add_argument("--model_config", type=Path, default=None)
ap.add_argument(
"--resume_from_checkpoint",
type=Path,
default=None,
help="Trainer checkpoint-* dir under output_dir (has trainer_state.json); overrides YAML",
)
args = ap.parse_args()
root = repo_root_from(ROOT)
pc = args.pretrain_config or (root / "configs/pretrain_prototype.yaml")
mc = args.model_config or (root / "configs/model_prototype.yaml")
run_pretrain(
pc.resolve(),
mc.resolve(),
resume_from_checkpoint=args.resume_from_checkpoint,
)
return 0
if __name__ == "__main__":
raise SystemExit(main())