predictive-maintenance-pipeline/scripts/02_register_data.py

#!/usr/bin/env python3
"""Register dataset to Hugging Face Hub"""
import os
import logging
from huggingface_hub import HfApi, create_repo
from huggingface_hub.utils import RepositoryNotFoundError

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

HF_TOKEN = os.getenv("HF_TOKEN")
HF_USERNAME = os.getenv("HF_USERNAME", "SharleyK")
DATASET_NAME = os.getenv("DATASET_NAME", "PredictiveMaintenance")
repo_id = f"{HF_USERNAME}/{DATASET_NAME}"

if not HF_TOKEN:
    raise ValueError("HF_TOKEN not set!")

logger.info(f"Registering dataset: {repo_id}")

api = HfApi(token=HF_TOKEN)

try:
    api.repo_info(repo_id=repo_id, repo_type="dataset")
    logger.info(f"✓ Repository exists: {repo_id}")
except RepositoryNotFoundError:
    create_repo(repo_id=repo_id, repo_type="dataset", token=HF_TOKEN, private=False)
    logger.info(f"✓ Created repository: {repo_id}")

if os.path.exists("data/engine_data.csv"):
    api.upload_file(
        path_or_fileobj="data/engine_data.csv",
        path_in_repo="engine_data.csv",
        repo_id=repo_id,
        repo_type="dataset",
        token=HF_TOKEN
    )
    logger.info("✓ Uploaded engine_data.csv")

logger.info("✓ Data registration completed!")