scripts/hf_auth.py

"""Shared Hugging Face auth + dataset upload helpers."""

from __future__ import annotations

import os
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
TRAIN_FILE = ROOT / "data" / "finetune" / "study_coach_train.jsonl"
DATASET_README = ROOT / "data" / "finetune" / "README.md"

DATASET_REPO = os.environ.get("PMS_FT_DATASET_REPO", "GJB99/plane-mode-study-coach-ft-data")
OUTPUT_REPO = os.environ.get("PMS_FINETUNED_MODEL", "GuusBouwensNL/plane-mode-nemotron-4b-study-coach")
BASE_MODEL = os.environ.get("PMS_FT_BASE_MODEL", "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16")
UNSLOTH_MODEL = os.environ.get(
    "PMS_UNSLOTH_BASE_MODEL",
    "unsloth/NVIDIA-Nemotron-3-Nano-4B-BF16",
)


def resolve_hf_token() -> str | None:
    for key in (
        "HF_TOKEN",
        "HUGGING_FACE_HUB_TOKEN",
        "HUGGINGFACEHUB_API_TOKEN",
        "HF_WRITE_TOKEN",
        "HF_API_TOKEN",
    ):
        val = os.environ.get(key, "").strip()
        if val:
            return val
    try:
        from huggingface_hub import HfFolder

        stored = HfFolder.get_token()
        if stored:
            return stored
    except Exception:
        pass
    try:
        from huggingface_hub import get_token

        stored = get_token()
        if stored:
            return stored
    except Exception:
        pass
    return None


def resolve_hf_username(token: str | None) -> str:
    username = os.environ.get("HF_USERNAME") or os.environ.get("HF_USER")
    if username:
        return username
    if token:
        try:
            from huggingface_hub import HfApi

            return HfApi(token=token).whoami()["name"]
        except Exception:
            pass
    return "GJB99"


def push_dataset(token: str, dataset_repo: str = DATASET_REPO) -> None:
    from huggingface_hub import HfApi, upload_file

    if not TRAIN_FILE.exists():
        raise FileNotFoundError(f"Missing {TRAIN_FILE} — run scripts/build_ft_dataset.py")

    api = HfApi(token=token)
    api.create_repo(dataset_repo, repo_type="dataset", exist_ok=True)
    upload_file(
        path_or_fileobj=str(TRAIN_FILE),
        path_in_repo="train.jsonl",
        repo_id=dataset_repo,
        repo_type="dataset",
        token=token,
        commit_message="Study coach conversation SFT data",
    )
    if DATASET_README.exists():
        upload_file(
            path_or_fileobj=str(DATASET_README),
            path_in_repo="README.md",
            repo_id=dataset_repo,
            repo_type="dataset",
            token=token,
            commit_message="Dataset card",
        )
    print(f"Dataset: https://huggingface.co/datasets/{dataset_repo}")


def token_error_message() -> str:
    return (
        "No Hugging Face token found.\n"
        "  Set HF_TOKEN (write) in Cursor Cloud Agent secrets or:\n"
        "  export HF_TOKEN=hf_... && python scripts/run_unsloth_finetune.py\n"
        "  Or open notebooks/unsloth_nemotron4b_colab.ipynb in Google Colab."
    )