train/push_models.py

"""Push the trained Yui models to PRIVATE repos under HF_ORG.

Uploads (idempotent — safe to re-run, each upload is a no-op if unchanged):
  - SFT-1 LoRA adapter (data/checkpoints_backup/brain1-ha_actions-sft1/)
      -> {HF_ORG}/yui-brain1-sft1   (model repo)
  - wake word (voice/models/hey_yui.{tflite,json})
      -> {HF_ORG}/yui-wakeword-hey-yui   (model repo)

Auth uses HF_WRITE (a write token); it is never printed. Prints repo URLs and
the uploaded file list (fetched back from the Hub) so the upload is confirmed.

Run: uv run --group train python train/push_models.py
"""

import os
from pathlib import Path

from dotenv import load_dotenv
from huggingface_hub import HfApi

ROOT = Path(__file__).resolve().parent.parent
ADAPTER_DIR = ROOT / "data" / "checkpoints_backup" / "brain1-ha_actions-sft1"
WAKEWORD_FILES = [
    ROOT / "voice" / "models" / "hey_yui.tflite",
    ROOT / "voice" / "models" / "hey_yui.json",
]

BRAIN_REPO = "yui-brain1-sft1"
WAKEWORD_REPO = "yui-wakeword-hey-yui"


def _push_folder(api: HfApi, repo_id: str, folder: Path) -> None:
    api.create_repo(repo_id, repo_type="model", private=True, exist_ok=True)
    api.upload_folder(repo_id=repo_id, repo_type="model", folder_path=str(folder))


def _push_files(api: HfApi, repo_id: str, files: list[Path]) -> None:
    api.create_repo(repo_id, repo_type="model", private=True, exist_ok=True)
    for f in files:
        api.upload_file(
            path_or_fileobj=str(f),
            path_in_repo=f.name,
            repo_id=repo_id,
            repo_type="model",
        )


def main() -> None:
    load_dotenv(ROOT / ".env")
    org = os.environ["HF_ORG"]
    token = os.environ["HF_WRITE"]  # write token; never logged

    if not ADAPTER_DIR.is_dir():
        raise SystemExit(f"adapter dir missing: {ADAPTER_DIR}")
    for f in WAKEWORD_FILES:
        if not f.is_file():
            raise SystemExit(f"wakeword file missing: {f}")

    api = HfApi(token=token)

    brain_id = f"{org}/{BRAIN_REPO}"
    wake_id = f"{org}/{WAKEWORD_REPO}"

    print(f"[push] uploading SFT-1 adapter -> {brain_id}")
    _push_folder(api, brain_id, ADAPTER_DIR)

    print(f"[push] uploading wake word -> {wake_id}")
    _push_files(api, wake_id, WAKEWORD_FILES)

    for repo_id in (brain_id, wake_id):
        files = sorted(api.list_repo_files(repo_id, repo_type="model"))
        print(f"\n[push] {repo_id}")
        print(f"       https://huggingface.co/{repo_id}")
        for name in files:
            print(f"       - {name}")


if __name__ == "__main__":
    main()