spb/zyquo-mlx Public MIT

1#2#  zyquo_train.py3#  Zyquo MLX4#5#  Author: Simon-Pierre Boucher6#  Mail: contact@spboucher.ai7#8#  Training driver for the Zyquo MLX Swift app.9#10#  Wraps mlx_lm.lora.run() with a custom TrainingCallback and emits a stable11#  JSON-lines protocol on stdout (one JSON object per line, event-typed).12#  We NEVER rely on mlx-lm's own stdout format (it changed between 0.31.3 and13#  main — docs/TRAINING-RESEARCH.md §5). Pinned against mlx-lm==0.31.3.14#15#  Usage: zyquo_train.py --config <run-config.yaml>16#17#  Events: {"event":"start", ...} {"event":"train", ...} {"event":"val", ...}18#          {"event":"save", ...} {"event":"done"} {"event":"error", ...}1920import argparse21import json22import sys23import time24import types252627def emit(obj):28    sys.stdout.write(json.dumps(obj) + "\n")29    sys.stdout.flush()303132class ZyquoCallback:33    """Receives the stable mlx-lm TrainingCallback dict payloads34    (docs/TRAINING-RESEARCH.md §5.2) and re-emits them as JSON lines."""3536    def on_train_loss_report(self, info):37        emit({"event": "train", **info, "ts": time.time()})3839    def on_val_loss_report(self, info):40        emit({"event": "val", **info, "ts": time.time()})414243def main():44    parser = argparse.ArgumentParser()45    parser.add_argument("--config", required=True, help="YAML run config (mlx-lm schema)")46    args = parser.parse_args()4748    try:49        import numpy as np50        import yaml51        from mlx_lm import lora52        from mlx_lm.tuner.datasets import load_dataset53        from mlx_lm.utils import load5455        with open(args.config) as f:56            config = yaml.safe_load(f)5758        # Build the args namespace exactly like mlx_lm.lora's CLI does:59        # defaults first, then config overrides.60        run_args = dict(lora.CONFIG_DEFAULTS)61        run_args.update(config)62        ns = types.SimpleNamespace(**run_args)6364        emit({65            "event": "start",66            "model": ns.model,67            "fine_tune_type": ns.fine_tune_type,68            "iters": ns.iters,69            "batch_size": ns.batch_size,70            "learning_rate": ns.learning_rate,71            "adapter_path": ns.adapter_path,72        })7374        # NOTE: we deliberately do NOT call lora.run() — in mlx-lm 0.31.3 it75        # overwrites the training_callback argument with76        # get_reporting_callbacks(args.report_to) (None here), silently77        # discarding ours. Replicate run()'s exact flow instead.78        np.random.seed(ns.seed)79        emit({"event": "stage", "stage": "loading_model"})80        model, tokenizer = load(ns.model, tokenizer_config={"trust_remote_code": True})81        emit({"event": "stage", "stage": "loading_datasets"})82        train_set, valid_set, _test_set = load_dataset(ns, tokenizer)83        emit({"event": "stage", "stage": "training"})84        lora.train_model(ns, model, train_set, valid_set, ZyquoCallback())85        emit({"event": "done"})86    except KeyboardInterrupt:87        emit({"event": "error", "message": "cancelled"})88        sys.exit(130)89    except Exception as exc:  # noqa: BLE001 - single funnel to the app90        emit({"event": "error", "message": str(exc), "type": type(exc).__name__})91        sys.exit(1)929394if __name__ == "__main__":95    main()96