"""5.11 在 Mac 上用 MLX 给小模型训 LoRA。

把 data/pairs_train.jsonl、pairs_valid.jsonl 转成 mlx-lm 要的对话格式，然后调用 mlx_lm.lora 训练。
训练好的 LoRA 权重（adapter）存在 adapters/ 下，只有几十 MB，底座模型不变。

用法：
    python train_mlx.py                              # 默认 Qwen/Qwen3-1.7B，450 步
    python train_mlx.py --model Qwen/Qwen3-0.6B    # 内存小的机器
    python train_mlx.py --data my_data --iters 400   # 换成你自己的数据
"""
import argparse
import json
import subprocess
import sys
from pathlib import Path

from common import HERE, read_pairs, to_chat


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', default='Qwen/Qwen3-1.7B')
    ap.add_argument('--data', help='放 pairs_train.jsonl / pairs_valid.jsonl 的目录，默认 data/')
    ap.add_argument('--iters', type=int, default=450, help='训练步数；每步 batch-size 条，默认约等于把训练集过 2 遍')
    ap.add_argument('--batch-size', type=int, default=2)
    ap.add_argument('--learning-rate', type=float, default=1e-4)
    ap.add_argument('--adapter', default=str(HERE / 'adapters' / 'mlx'))
    args = ap.parse_args()

    mlx_dir = HERE / 'data_mlx'
    mlx_dir.mkdir(exist_ok=True)
    for split, name in [('train', 'train'), ('valid', 'valid')]:
        rows = read_pairs(split, args.data)
        (mlx_dir / f'{name}.jsonl').write_text(''.join(json.dumps(to_chat(r), ensure_ascii=False) + '\n' for r in rows),
                                               encoding='utf-8')
        print(f'{split}: {len(rows)} 对')

    cmd = [sys.executable, '-m', 'mlx_lm', 'lora', '--model', args.model, '--train', '--data', str(mlx_dir),
           '--iters', str(args.iters), '--batch-size', str(args.batch_size), '--learning-rate', str(args.learning_rate),
           '--num-layers', '-1',            # 每一层都加 LoRA
           '--mask-prompt',                 # 只在“改过的版本”上算损失，不学着复述输入
           '--grad-checkpoint',             # 用时间换内存
           '--steps-per-eval', '50', '--save-every', '50', '--adapter-path', args.adapter]
    print(' '.join(cmd[1:]))
    subprocess.run(cmd, check=True)
    print(f'\nLoRA 已存到 {Path(args.adapter).relative_to(HERE)}/，下一步：python evaluate.py --backend mlx')


if __name__ == '__main__':
    main()