Blame for scripts/train.py

"""Entraîne M.O.U.H.  Lancement :

  python scripts/train.py --preset mouh-tiny --data ./corpus --max_steps 200
  torchrun --nproc_per_node=8 scripts/train.py --preset mouh-1.5b --data <dataset_hub> --parallel fsdp
"""
import argparse
import dataclasses
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from mouh.training import TrainArgs, train  # noqa: E402


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    for f in dataclasses.fields(TrainArgs):
        t = f.type if isinstance(f.type, type) else str
        if f.name == "betas":
            ap.add_argument("--betas", type=float, nargs=2, default=list(f.default))
        elif f.type in (bool, "bool"):
            ap.add_argument(f"--{f.name}", action="store_true", default=f.default)
        else:
            ap.add_argument(f"--{f.name}", type={"int": int, "float": float}.get(str(f.type), str), default=f.default)
    ns = vars(ap.parse_args())
    ns["betas"] = tuple(ns["betas"])
    print(train(TrainArgs(**ns)))


if __name__ == "__main__":
    main()