tracklight_server/hf/push.py

# tracklight_server/hf/push.py

from datasets import Dataset
from huggingface_hub import HfApi, HfFolder
from ..db import duckdb
import os

def push_to_hub(repo_id: str, hf_token: str):
    """
    Pushes the local DuckDB data to a Hugging Face Dataset repo.
    """
    # Authenticate with Hugging Face
    HfFolder.save_token(hf_token)

    # Get all data from DuckDB
    with duckdb.get_connection() as con:
        df = con.execute(f"SELECT * FROM {duckdb.TABLE_NAME}").fetchdf()

    if df.empty:
        print("No data to push.")
        return

    # Create a Hugging Face Dataset
    dataset = Dataset.from_pandas(df)

    # Push the dataset to the Hub
    try:
        dataset.push_to_hub(repo_id=repo_id, private=True)
        print(f"Successfully pushed data to {repo_id}")
    except Exception as e:
        print(f"Failed to push data to {repo_id}: {e}")