scripts/train_lora_sft.py

#!/usr/bin/env python3
"""LoRA fine-tune Qwen/Qwen2.5-Coder-0.5B-Instruct on Mythos SFT messages."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from mythos_lora_core import train_lora


def project_root() -> Path:
    return Path(__file__).resolve().parent.parent


def parse_args() -> argparse.Namespace:
    root = project_root()
    parser = argparse.ArgumentParser(description="Train Mythos-Coder LoRA adapter.")
    parser.add_argument(
        "--model_name",
        default="Qwen/Qwen2.5-Coder-0.5B-Instruct",
        help="Hugging Face base model ID",
    )
    parser.add_argument(
        "--train_file",
        default=str(root / "data" / "train" / "mythos_sft_messages_clean.jsonl"),
        help="SFT messages JSONL path",
    )
    parser.add_argument(
        "--output_dir",
        default=str(root / "models" / "mythos-coder-qwen-0.5b-lora"),
        help="Directory to save LoRA adapter",
    )
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--grad_accum", type=int, default=8)
    parser.add_argument("--learning_rate", type=float, default=1e-4)
    parser.add_argument("--max_seq_length", type=int, default=2048)
    return parser.parse_args()


def resolve_path(path_str: str) -> Path:
    path = Path(path_str)
    if not path.is_absolute():
        path = project_root() / path
    return path


def main() -> int:
    args = parse_args()
    train_path = resolve_path(args.train_file)
    output_dir = resolve_path(args.output_dir)

    if not train_path.exists():
        print(f"Error: training file not found: {train_path}", file=sys.stderr)
        return 1

    row_count = sum(1 for line in train_path.open(encoding="utf-8") if line.strip())
    print("=== Mythos-Coder LoRA Training ===")
    print(f"Training examples:  {row_count}")

    result = train_lora(
        train_path,
        output_dir,
        model_name=args.model_name,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.learning_rate,
        max_length=args.max_seq_length,
    )

    print("=== Training Summary ===")
    print(f"Status:             {result['status']}")
    print(f"CUDA:               {result['cuda']}")
    print(f"Train loss:         {result['train_loss']}")
    print(f"Runtime (seconds):  {result['runtime_seconds']:.1f}")
    print(f"Adapter saved to:   {result['output_dir']}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())