scripts/push_dataset_to_hub.py

#!/usr/bin/env python3
"""
Envia o DatasetDict local (pasta gerada por save_to_disk) para o Hugging Face Hub.

Garante splits e colunas compatíveis com `load_dataset(repo_id)` em train.py.

Uso:
  cd ml/configs/huggingface_training_config
  export HF_TOKEN=...
  python scripts/push_dataset_to_hub.py
  python scripts/push_dataset_to_hub.py /caminho/absoluto/para/eda_training_dataset
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

from datasets import load_from_disk
from huggingface_hub import create_repo, whoami


def _repo_root() -> Path:
    return Path(__file__).resolve().parents[4]


def _ensure_hf_token_from_dotenv() -> None:
    if os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN"):
        return
    env_path = _repo_root() / ".env"
    if not env_path.is_file():
        return
    for raw in env_path.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line or line.startswith("#"):
            continue
        for prefix in ("HF_TOKEN=", "HUGGING_FACE_HUB_TOKEN="):
            if line.startswith(prefix):
                val = line[len(prefix) :].strip().strip('"').strip("'")
                if val:
                    os.environ["HF_TOKEN"] = val
                return


def _resolve_token() -> str | None:
    return (
        os.environ.get("HF_TOKEN")
        or os.environ.get("HUGGING_FACE_HUB_TOKEN")
        or None
    )


def main() -> None:
    _ensure_hf_token_from_dotenv()
    token = _resolve_token()
    if not token:
        print(
            "Defina HF_TOKEN ou HUGGING_FACE_HUB_TOKEN (token com escrita no Hub). "
            "Crie em https://huggingface.co/settings/tokens ou adicione ao .env na raiz do repo.",
            file=sys.stderr,
        )
        if os.environ.get("HF_INFERENCE_ENDPOINT_TOKEN"):
            print(
                "Nota: HF_INFERENCE_ENDPOINT_TOKEN nao substitui HF_TOKEN para upload ao Hub.",
                file=sys.stderr,
            )
        sys.exit(1)

    default_path = _repo_root() / "ml" / "datasets" / "eda_training_dataset"
    disk_path = Path(sys.argv[1]).resolve() if len(sys.argv) > 1 else default_path

    if not disk_path.is_dir():
        print(f"Pasta não encontrada: {disk_path}", file=sys.stderr)
        sys.exit(1)

    private = os.environ.get("HF_DATASET_PRIVATE", "true").lower() in ("1", "true", "yes")

    try:
        identity = whoami(token=token)
        hub_user = identity.get("name") or "?"
        print(f"Autenticado no Hub como: {hub_user}")
    except Exception as exc:
        print(
            "Falha na autenticacao (token invalido ou expirado). "
            "Gere um novo token em https://huggingface.co/settings/tokens com permissao Write.\n"
            f"Detalhe: {exc}",
            file=sys.stderr,
        )
        sys.exit(1)

    if os.environ.get("DATASET_REPO"):
        repo_id = os.environ["DATASET_REPO"].strip()
    else:
        repo_id = f"{hub_user}/eda-training-dataset"
        print(
            "DATASET_REPO nao definido: a usar o teu espaco de utilizador "
            f"(evita 401 se nao tiveres acesso a org beAnalytic).\n"
            f"  -> {repo_id}\n"
            "Para a org beAnalytic (com permissao):\n"
            "  export DATASET_REPO=beAnalytic/eda-training-dataset"
        )

    try:
        create_repo(
            repo_id=repo_id,
            repo_type="dataset",
            token=token,
            private=private,
            exist_ok=True,
        )
    except Exception as exc:
        print(
            f"Nao foi possivel criar ou aceder ao dataset `{repo_id}`: {exc}\n"
            "Se o repositorio e de uma organizacao (ex.: beAnalytic), a conta precisa de "
            "ser membro com permissao de escrita, ou use um repo no teu utilizador:\n"
            "  export DATASET_REPO=O_TEU_USER/eda-training-dataset",
            file=sys.stderr,
        )
        sys.exit(1)

    print(f"Carregando: {disk_path}")
    ds = load_from_disk(str(disk_path))
    print(ds)
    print(f"Enviando para {repo_id} (private={private})...")
    ds.push_to_hub(repo_id, private=private, token=token)
    print(f"Concluído: https://huggingface.co/datasets/{repo_id}")
    print()
    print("No treino (local ou Space), usa o mesmo repositorio de dataset:")
    print(f"  export DATASET_REPO={repo_id}")


if __name__ == "__main__":
    main()