scripts/publish_prelinger.py

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "saturate[hf]",
#     "huggingface-hub",
#     "pyarrow",
#     "pyyaml",
# ]
# ///
"""Join the Prelinger caption output into the wide chunk-level dataset.

Sources (all in the public bucket `davanstrien/prelinger-films`):
  captions/            saturate parquet output — first fleets (holds the
  captions2/           duplicate rows from the mis-sharded run); top-up fleet.
                       Both are read and unioned: one row per chunk carrying
                       identifier / chunk_start / chunk_end / scene / caption /
                       events / token counts. Read via `read_output`, which
                       applies the error-IS-NULL-wins rule and drops error-only
                       ids, so partial (still-running) output is safe to read.
  chunkmeta/{id}.json  recorded segment boundaries — the ONLY place the chunk
                       FILENAME lives, which is what thumb_url needs.
  meta/{id}.json       per-film IA sidecar — title, date, licenseurl.

Sidecars are fetched lazily for just the identifiers present in the caption
rows, so `--limit` is cheap.

Emits one row per chunk with the data-layer schema (the dataset IS the API),
writes the parquet directly (row groups of 2000 — see ROW_GROUP_SIZE) and
uploads it to a PRIVATE datasets repo alongside a README declaring the configs.

Run the FULL join on a Job with the bucket mounted, not from a laptop:
`read_output` opens every part- file serially, and the caption output is ~5.7k
small parquets, so remote reads cost ~20 min of pure round-trips.

    hf jobs uv run publish_prelinger.py --flavor cpu-upgrade \
      --volume hf://buckets/davanstrien/prelinger-films:/pf \
      --secrets HF_TOKEN=$(hf auth token) \
      -- --captions /pf/captions --captions /pf/captions2 --sidecars /pf

`--max-parts N` caps how many part- files are read per caption dir (they are
staged locally first, so `read_output` sees a small dir). Smoke tests only —
the full publish must read every part.

If the bucket moves to the biglam org, the two namespaces are independent on
purpose: --sidecars is where data is READ from, --bucket-namespace is what goes
into the emitted video_url/thumb_url. So you can publish biglam-facing URLs
while still reading from wherever the bytes currently live:

    --bucket-namespace biglam/prelinger-films

Re-pushing is safe now that the parquet is written here rather than by
`push_to_hub`: the README's stale `dataset_info` block (which is what caused
the CastError on a schema-changing re-push) is dropped on every write, and the
card body is preserved.
"""

import argparse
import json
import math
import re
import sys
import tempfile
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import pyarrow as pa

BUCKET = "davanstrien/prelinger-films"
DEFAULT_SIDECARS = f"hf://buckets/{BUCKET}"
DEFAULT_CAPTIONS = [f"{DEFAULT_SIDECARS}/captions", f"{DEFAULT_SIDECARS}/captions2"]
DEFAULT_REPO = "davanstrien/prelinger-moments"

# Non-negotiable (DuckDB partial-read research, 2026-07-30): a single row group
# makes even `SELECT * LIMIT 5` pull the whole file, and one fat group flirts
# with the Hub's ~300 MB regeneration ceiling once embeddings land.
ROW_GROUP_SIZE = 2000
PARQUET_PATH = "data/train-00000-of-00001.parquet"

# `id` is first and unique: identifier alone is not (one film, many chunks).
SCHEMA = pa.schema([
    ("id", pa.string()),
    ("identifier", pa.string()),
    ("title", pa.string()),
    ("date", pa.string()),
    ("licenseurl", pa.string()),
    ("ia_url", pa.string()),
    ("video_url", pa.string()),
    ("thumb_url", pa.string()),
    ("chunk_start", pa.float64()),
    ("chunk_end", pa.float64()),
    ("scene", pa.string()),
    ("caption", pa.string()),
    ("events", pa.string()),
    ("events_text", pa.string()),
    ("prompt_tokens", pa.int64()),
    ("completion_tokens", pa.int64()),
])

# Declared configs. The embeddings config (id + the three fp16 vector columns,
# its own dir) is additive — append a second dict here and upload under
# `embeddings/`; the Hub mis-resolves an undeclared dir, so it must be listed.
CONFIGS = [
    {"config_name": "default",
     "default": True,
     "data_files": [{"split": "train", "path": "data/train-*.parquet"}]},
]

PLACEHOLDER_BODY = """# Prelinger moments

Chunk-level captions over the Prelinger public-domain film mirror. One row per
video chunk; `video_url` and `thumb_url` point at the public bucket.

*Card pending — this body is a placeholder and is preserved across re-publishes.*
"""


def url(namespace: str, *segments: str) -> str:
    from urllib.parse import quote
    base = f"https://huggingface.co/buckets/{namespace}/resolve"
    return "/".join([base] + [quote(s, safe="") for s in segments])


def stage_parts(src: str, n: int, dest_root: Path, workers: int) -> str:
    """Copy the first `n` part- files of a saturate output dir locally.

    Smoke-test affordance only. `read_output` materializes every part before it
    yields, so `--limit` cannot stop a laptop-side read early; capping the parts
    is the only way to keep a live-bucket smoke test to seconds. Sorted the same
    way `read_output` sorts, so the subset is deterministic.
    """
    import fsspec
    fs, root = fsspec.url_to_fs(src)
    parts = sorted(fs.glob(f"{root}/part-*.parquet"))[:n]
    dest = dest_root / re.sub(r"\W+", "_", src.strip("/"))
    dest.mkdir(parents=True, exist_ok=True)

    def one(p):
        (dest / Path(p).name).write_bytes(fs.cat_file(p))
    with ThreadPoolExecutor(workers) as ex:
        list(ex.map(one, parts))
    print(f"[stage] {len(parts)} parts from {src} -> {dest}", flush=True)
    return str(dest)


def read_captions(dirs: list[str], max_parts: int, workers: int) -> list[dict]:
    """Winner-per-id across the union of saturate output dirs.

    `read_output` already resolves within a dir (error-IS-NULL wins, error-only
    ids dropped) and yields only successful records, so a plain first-writer-
    wins union across dirs is the whole rule: an id seen in `captions/` is never
    clobbered by the same id turning up again in `captions2/`.
    """
    from saturate import read_output
    best: dict[str, dict] = {}
    with tempfile.TemporaryDirectory() as tmp:
        for d in dirs:
            src = stage_parts(d, max_parts, Path(tmp), workers) if max_parts else d
            print(f"[read] {src}", flush=True)
            before = len(best)
            for id_, rec in read_output(src):
                if id_ not in best:
                    best[id_] = rec
            print(f"[read] {d}: +{len(best) - before} new ids "
                  f"({len(best)} total)", flush=True)
    return list(best.values())


def fetch_sidecars(root: str, prefix: str, idents: list[str],
                   workers: int) -> dict:
    """{identifier: parsed json or None} for one sidecar prefix under `root`.

    `root` may be an hf:// URI or a mounted path, so a Job run can read the
    sidecars off the mount instead of paying ~3.7k network round-trips.

    Retries: at 1.8k concurrent small remote reads a couple of percent time out,
    and a swallowed timeout would silently drop a whole film from the dataset.
    """
    import fsspec
    fs, base = fsspec.url_to_fs(root)

    def one(ident: str):
        path = f"{base}/{prefix}/{ident}.json"
        for attempt in range(4):
            try:
                with fs.open(path, "rb") as f:
                    return ident, json.loads(f.read())
            except FileNotFoundError:
                return ident, None
            except Exception:
                time.sleep(1.5 * (attempt + 1))
        return ident, None
    with ThreadPoolExecutor(workers) as ex:
        return dict(ex.map(one, idents))


def write_parquet(records: list[dict], path: Path, row_group_size: int) -> None:
    import pyarrow.parquet as pq
    table = pa.Table.from_pylist(records, schema=SCHEMA)
    pq.write_table(table, path, row_group_size=row_group_size, compression="snappy")

    md = pq.ParquetFile(path).metadata
    groups = [md.row_group(i).num_rows for i in range(md.num_row_groups)]
    expected = math.ceil(md.num_rows / row_group_size)
    if md.num_row_groups != expected or any(g > row_group_size for g in groups):
        sys.exit(f"[abort] row groups {groups} do not match "
                 f"row_group_size={row_group_size} over {md.num_rows} rows")
    print(f"[parquet] {path.name}: {md.num_rows} rows, "
          f"{md.num_row_groups} row group(s) {groups}, "
          f"{path.stat().st_size / 1e6:.1f} MB", flush=True)


def render_readme(existing: str | None) -> str:
    """New frontmatter (configs declared), existing card body preserved."""
    import yaml
    front, body = {}, PLACEHOLDER_BODY
    if existing and existing.startswith("---"):
        _, fm, rest = existing.split("---", 2)
        front = yaml.safe_load(fm) or {}
        body = rest.lstrip("\n")
    elif existing:
        body = existing
    # push_to_hub's dataset_info pins the OLD column set and is what turns a
    # schema-changing re-push into a CastError. The configs block replaces it.
    front.pop("dataset_info", None)
    front["configs"] = CONFIGS
    fm = yaml.dump(front, sort_keys=False, default_flow_style=False)
    return f"---\n{fm}---\n\n{body}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--captions", action="append", default=None,
                    help="saturate output dir; repeat the flag or pass a "
                         "comma-separated list. Winners are unioned across "
                         f"dirs, first dir wins (default: {DEFAULT_CAPTIONS})")
    ap.add_argument("--repo", default=DEFAULT_REPO)
    ap.add_argument("--bucket-namespace", default=BUCKET,
                    help="namespace/bucket that video_url and thumb_url point "
                         f"at (default {BUCKET}). Set to biglam/prelinger-films "
                         "once the bucket moves — this only rewrites the "
                         "emitted URLs, see --sidecars for where data is READ.")
    ap.add_argument("--sidecars", default=DEFAULT_SIDECARS,
                    help="root holding chunkmeta/ and meta/; hf:// URI or a "
                         f"mount path like /pf (default {DEFAULT_SIDECARS})")
    ap.add_argument("--limit", type=int, default=0,
                    help="cap emitted rows (smoke test); 0 = all")
    ap.add_argument("--max-parts", type=int, default=0,
                    help="read only the first N part- files per caption dir "
                         "(smoke test; 0 = all). Full publish must use 0.")
    ap.add_argument("--row-group-size", type=int, default=ROW_GROUP_SIZE)
    ap.add_argument("--workers", type=int, default=16)
    ap.add_argument("--dry-run", action="store_true",
                    help="build rows and print a preview, do not push")
    ap.add_argument("--allow-missing", action="store_true",
                    help="proceed when a film's sidecar can't be read "
                         "(default: refuse, so a flaky read can't silently "
                         "drop films from the published dataset)")
    args = ap.parse_args()

    dirs = []
    for spec in (args.captions or DEFAULT_CAPTIONS):
        dirs.extend(d.strip() for d in spec.split(",") if d.strip())

    caption_rows = read_captions(dirs, args.max_parts, args.workers)
    print(f"[read] {len(caption_rows)} caption rows across {len(dirs)} dir(s)",
          flush=True)
    if not caption_rows:
        sys.exit("no caption rows read")
    caption_rows.sort(key=lambda r: (r["identifier"], float(r["chunk_start"])))
    if args.limit:
        caption_rows = caption_rows[:args.limit]

    idents = sorted({r["identifier"] for r in caption_rows})
    print(f"[sidecars] {len(idents)} identifiers from {args.sidecars}",
          flush=True)
    cmeta = fetch_sidecars(args.sidecars, "chunkmeta", idents, args.workers)
    fmeta = fetch_sidecars(args.sidecars, "meta", idents, args.workers)
    missing = [i for i in idents if not cmeta.get(i) or not fmeta.get(i)]
    if missing:
        msg = (f"{len(missing)} identifiers missing a sidecar: {missing[:5]}")
        if not args.allow_missing:
            sys.exit(f"[abort] {msg}\n  re-run (reads are retried) or pass "
                     f"--allow-missing to publish without them")
        print(f"[warn] {msg}", flush=True)

    # (identifier, start) -> chunk filename. Both sides come from the same
    # chunkmeta floats, so rounding to 3dp matches exactly.
    chunk_file = {}
    for ident, meta in cmeta.items():
        for c in (meta or {}).get("chunks", []):
            chunk_file[(ident, round(float(c["start"]), 3))] = c["chunk"]

    ns = args.bucket_namespace
    records, skipped = [], {"no_chunk": 0, "no_meta": 0}
    for r in caption_rows:
        ident = r["identifier"]
        film = fmeta.get(ident)
        if not film:
            skipped["no_meta"] += 1
            continue
        start = round(float(r["chunk_start"]), 3)
        cf = chunk_file.get((ident, start))
        if not cf:
            skipped["no_chunk"] += 1
            continue
        stem = cf.rsplit(".", 1)[0]
        events = r.get("events") or "[]"
        chunk_start = float(r["chunk_start"])
        records.append({
            # Matches `identifier || '@' || chunk_start` evaluated in DuckDB,
            # so an agent can reconstruct the key from the other two columns.
            "id": f"{ident}@{chunk_start}",
            "identifier": ident,
            "title": film.get("title"),
            "date": film.get("date"),
            "licenseurl": film.get("licenseurl"),
            "ia_url": f"https://archive.org/details/{ident}",
            "video_url": url(ns, "films", f"{ident}.mp4"),
            "thumb_url": url(ns, "thumbs", ident, f"{stem}.jpg"),
            "chunk_start": chunk_start,
            "chunk_end": float(r["chunk_end"]),
            "scene": r.get("scene"),
            "caption": r.get("caption"),
            "events": events,
            # Plain-text join of event descriptions: the embed target for
            # emb_events (raw JSON would spend tokens on brackets/keys) and
            # the nicer ILIKE surface for agents.
            "events_text": "; ".join(
                e.get("text", "") for e in json.loads(events)),
            "prompt_tokens": r.get("prompt_tokens"),
            "completion_tokens": r.get("completion_tokens"),
        })

    records.sort(key=lambda x: (x["identifier"], x["chunk_start"]))
    # The caption output holds up to ~7 records per chunk (shard restarts
    # re-attempted rows whose manifest entry had not flushed). read_output
    # collapses those by id; assert it, because a duplicated chunk would
    # silently skew any downstream count or embedding run.
    ids = [x["id"] for x in records]
    if len(set(ids)) != len(ids):
        sys.exit(f"[abort] {len(ids) - len(set(ids))} duplicate ids survived "
                 f"the union")
    n_ev = sum(len(json.loads(x["events"])) for x in records)
    print(f"[build] {len(records)} rows, {len(set(x['identifier'] for x in records))} "
          f"films, {n_ev} events, skipped={skipped}", flush=True)
    if not records:
        sys.exit("no joinable rows")
    print(json.dumps({k: (v[:160] if isinstance(v, str) else v)
                      for k, v in records[0].items()}, indent=1), flush=True)

    with tempfile.TemporaryDirectory() as tmp:
        local = Path(tmp) / Path(PARQUET_PATH).name
        write_parquet(records, local, args.row_group_size)

        if args.dry_run:
            print("[dry-run] not pushing", flush=True)
            return

        from huggingface_hub import HfApi
        from huggingface_hub.errors import EntryNotFoundError
        api = HfApi()
        api.create_repo(args.repo, repo_type="dataset", private=True,
                        exist_ok=True)
        try:
            existing = api.hf_hub_download(args.repo, "README.md",
                                           repo_type="dataset")
            existing = Path(existing).read_text()
        except (EntryNotFoundError, FileNotFoundError):
            existing = None
        readme = Path(tmp) / "README.md"
        readme.write_text(render_readme(existing))

        print(f"[push] {args.repo} (private)", flush=True)
        api.upload_file(path_or_fileobj=local, path_in_repo=PARQUET_PATH,
                        repo_id=args.repo, repo_type="dataset",
                        commit_message=f"publish {len(records)} chunk rows")
        api.upload_file(path_or_fileobj=readme, path_in_repo="README.md",
                        repo_id=args.repo, repo_type="dataset",
                        commit_message="declare default config")
    print("[PUBLISH DONE]", flush=True)


if __name__ == "__main__":
    main()