scripts/upload_dataset.py

#!/usr/bin/env python3
"""
Script para fazer upload do dataset de treinamento para o HuggingFace Hub.

Autenticação: defina HF_TOKEN ou use credenciais do `huggingface-cli login`.
"""

import os
import sys
from pathlib import Path

from huggingface_hub import HfApi, login

HF_USER = os.environ.get("HF_USER", "beAnalytic")
DATASET_REPO = os.environ.get("DATASET_REPO", f"{HF_USER}/eda-training-dataset")


def print_step(message: str, status: str = "info"):
    """Print formatado para diferentes tipos de mensagens"""
    colors = {
        "info": "\033[0;34m",
        "success": "\033[0;32m",
        "warning": "\033[1;33m",
        "error": "\033[0;31m",
    }
    symbols = {
        "info": "i",
        "success": "ok",
        "warning": "!",
        "error": "x",
    }
    reset = "\033[0m"
    print(f"{colors.get(status, colors['info'])}[{symbols.get(status, 'i')}] {message}{reset}")


def main():
    if len(sys.argv) < 2:
        print_step("Uso: python upload_dataset.py <caminho_para_dataset>", "error")
        print()
        print("Exemplo:")
        print("  python upload_dataset.py ../eda_training_dataset.jsonl")
        print("  python upload_dataset.py ../dataset_folder/")
        print()
        print("Requer HF_TOKEN no ambiente ou `huggingface-cli login`.")
        sys.exit(1)

    dataset_path = Path(sys.argv[1])

    if not dataset_path.exists():
        print_step(f"Erro: Caminho nao encontrado: {dataset_path}", "error")
        print(
            "Use o caminho real do ficheiro .jsonl ou pasta do dataset.",
            "Ex.: exporte o JSONL pelo pipeline EDA do projeto ou aponte para o ficheiro local.",
            sep="\n",
            file=sys.stderr,
        )
        sys.exit(1)

    hf_token = os.environ.get("HF_TOKEN")
    if not hf_token:
        print_step(
            "HF_TOKEN nao definido. Exporte HF_TOKEN ou execute: huggingface-cli login",
            "error",
        )
        sys.exit(1)

    print_step("Upload de Dataset para HuggingFace", "info")
    print()

    print_step("Autenticando no HuggingFace...", "info")
    try:
        login(token=hf_token, add_to_git_credential=False)
        print_step("Autenticado com sucesso", "success")
    except Exception as e:
        print_step(f"Erro ao autenticar: {e}", "error")
        sys.exit(1)

    api = HfApi(token=hf_token)

    print()
    print_step(f"Fazendo upload para: {DATASET_REPO}", "info")
    print_step(f"Origem: {dataset_path.absolute()}", "info")
    print()

    try:
        if dataset_path.is_file():
            print_step("Upload de arquivo unico...", "info")
            api.upload_file(
                path_or_fileobj=str(dataset_path),
                path_in_repo=dataset_path.name,
                repo_id=DATASET_REPO,
                repo_type="dataset",
            )
        else:
            print_step("Upload de pasta...", "info")
            api.upload_folder(
                folder_path=str(dataset_path),
                repo_id=DATASET_REPO,
                repo_type="dataset",
            )

        print()
        print_step("Upload concluido com sucesso.", "success")
        print()
        print(f"Dataset: https://huggingface.co/datasets/{DATASET_REPO}")

    except Exception as e:
        print_step(f"Erro durante upload: {e}", "error")
        sys.exit(1)


if __name__ == "__main__":
    main()