utils/hf_export.py

"""Export .gtkv token streams to Parquet / Hugging Face datasets.

The Parquet schema mirrors the container's hierarchical token layout —
one row per chunk, one column per token layer.  AI models can consume
the dataset directly without touching pixels.
"""
from pathlib import Path

from mediatok.container.gtkv import GtkvReader
from mediatok.entropy import entropy_decode

LAYER_COLUMNS = ["scene", "camera", "object", "motion", "texture", "residual"]


def gtkv_to_rows(gtkv_path: str) -> list[dict]:
    """Read a .gtkv file and return one dict per chunk."""
    reader = GtkvReader(gtkv_path)
    h = reader.header
    rows = []
    try:
        for ci in range(reader.num_chunks):
            block = reader.read_video_block(ci)
            tokens = entropy_decode(block.entropy_payload, h.entropy_codec_id,
                                    block.token_count)
            row = {
                "chunk_index": ci,
                "num_frames": h.chunk_size_frames,
                "width": h.width,
                "height": h.height,
                "fps": h.fps,
            }
            offset = 0
            for li, col in enumerate(LAYER_COLUMNS):
                n = block.layer_sizes[li] if li < len(block.layer_sizes) else 0
                row[f"{col}_tokens"] = tokens[offset:offset + n] if n > 0 else []
                offset += n
            rows.append(row)
    finally:
        reader.close()
    return rows


def gtkv_to_parquet(gtkv_path: str, out_path: str) -> str:
    """Write a .gtkv token stream to a Parquet file."""
    import pyarrow as pa
    import pyarrow.parquet as pq

    rows = gtkv_to_rows(gtkv_path)
    table = pa.Table.from_pylist(rows)
    pq.write_table(table, out_path)
    return out_path


def gtkv_to_hf_dataset(gtkv_path: str):
    """Return a Hugging Face Dataset of the token stream."""
    from datasets import Dataset
    return Dataset.from_list(gtkv_to_rows(gtkv_path))


def push_to_hub(gtkv_path: str, repo_id: str, token: str | None = None,
                private: bool = False):
    """Upload the token dataset to the Hugging Face Hub (Parquet-backed).

    Requires `huggingface_hub` credentials (HF_TOKEN env var or token arg).
    """
    ds = gtkv_to_hf_dataset(gtkv_path)
    return ds.push_to_hub(repo_id, token=token, private=private)