scripts/train_lora.py

from __future__ import annotations

import argparse
import json
from hashlib import sha256
from pathlib import Path
from typing import Any, cast
import sys

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from app_kit.lora_training import (
    DEFAULT_TRAINING_EPOCHS,
    DEFAULT_TRAINING_LR,
    DEFAULT_TRAINING_SEED,
    DEFAULT_RANK,
    build_training_artifact_from_dataset_file,
    write_adapter_files,
)
from app_kit.modal_lora_training import app as modal_app, train_adapter_from_dataset_json


def _resolve(repo_root: Path, value: str) -> Path:
    path = Path(value).expanduser()
    if not path.is_absolute():
        path = (repo_root / path).resolve()
    return path


def _build_local(
    dataset_path: Path,
    output_path: Path,
    manifest_path: Path,
    *,
    source_dataset_label: str,
    rank: int,
    epochs: int,
    learning_rate: float,
    seed: int,
) -> dict[str, object]:
    return build_training_artifact_from_dataset_file(
        dataset_path,
        output_path,
        manifest_path=manifest_path,
        source_dataset_label=source_dataset_label,
        training_backend='local',
        training_command='python scripts/train_lora.py --backend local',
        rank=rank,
        epochs=epochs,
        learning_rate=learning_rate,
        seed=seed,
    )


def _build_modal(
    dataset_path: Path,
    output_path: Path,
    manifest_path: Path,
    *,
    source_dataset_label: str,
    rank: int,
    epochs: int,
    learning_rate: float,
    seed: int,
) -> dict[str, object]:
    if modal_app is None:
        raise RuntimeError('modal is not installed in this environment; install the modal package or use --backend local')

    dataset_json = dataset_path.read_text(encoding='utf-8')
    dataset_sha256 = sha256(dataset_json.encode('utf-8')).hexdigest()
    modal_runtime_app = cast(Any, modal_app)
    train_fn = cast(Any, train_adapter_from_dataset_json)
    with modal_runtime_app.run():
        artifact = train_fn.remote(
            dataset_json,
            source_dataset_label=source_dataset_label,
            source_dataset_sha256=dataset_sha256,
            training_command='python scripts/train_lora.py --backend modal',
            rank=rank,
            epochs=epochs,
            learning_rate=learning_rate,
            seed=seed,
        )
    return write_adapter_files(artifact, output_path, manifest_path=manifest_path)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description='Train the P2 voice journal adapter locally or on Modal.')
    parser.add_argument(
        '--backend',
        choices=('local', 'modal'),
        default='local',
        help='Training backend to use. Local builds the artifact in-place; modal launches the GPU job.',
    )
    parser.add_argument(
        '--dataset',
        default='data/well_tuned/p2_voice_journal/training_examples.json',
        help='Path to the synthetic training examples JSON (relative to repo root by default).',
    )
    parser.add_argument(
        '--output',
        default='models/lora/p2_voice_journal_adapter.json',
        help='Where to write the adapter artifact JSON (relative to repo root by default).',
    )
    parser.add_argument(
        '--manifest',
        default=None,
        help='Optional path for the manifest JSON (defaults to <output>.manifest.json).',
    )
    parser.add_argument('--rank', type=int, default=DEFAULT_RANK, help='Adapter rank to train.')
    parser.add_argument('--epochs', type=int, default=DEFAULT_TRAINING_EPOCHS, help='Training epochs.')
    parser.add_argument('--learning-rate', type=float, default=DEFAULT_TRAINING_LR, help='Training learning rate.')
    parser.add_argument('--seed', type=int, default=DEFAULT_TRAINING_SEED, help='Random seed.')
    args = parser.parse_args(argv)

    dataset_path = _resolve(REPO_ROOT, args.dataset)
    output_path = _resolve(REPO_ROOT, args.output)
    manifest_path = _resolve(REPO_ROOT, args.manifest) if args.manifest else output_path.with_suffix('.manifest.json')
    source_dataset_label = args.dataset

    if args.backend == 'modal':
        artifact = _build_modal(
            dataset_path,
            output_path,
            manifest_path,
            source_dataset_label=source_dataset_label,
            rank=args.rank,
            epochs=args.epochs,
            learning_rate=args.learning_rate,
            seed=args.seed,
        )
    else:
        artifact = _build_local(
            dataset_path,
            output_path,
            manifest_path,
            source_dataset_label=source_dataset_label,
            rank=args.rank,
            epochs=args.epochs,
            learning_rate=args.learning_rate,
            seed=args.seed,
        )

    summary = {
        'backend': args.backend,
        'output': str(output_path),
        'manifest': artifact['manifest_path'],
        'artifact_sha256': artifact['artifact_sha256'],
        'manifest_sha256': artifact['manifest_sha256'],
        'adapter_name': artifact['adapter_name'],
        'adapter_type': artifact['adapter_type'],
        'training_examples': artifact['training_examples'],
        'base_model_id': artifact['base_model_id'],
        'training_backend': artifact['training_backend'],
        'training_command': artifact['training_command'],
        'source_dataset_sha256': artifact['source_dataset_sha256'],
    }
    print(json.dumps(summary, indent=2, ensure_ascii=False, sort_keys=True))
    return 0


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