scripts/merge_embeddings.py

# /// script
# requires-python = ">=3.11"
# dependencies = ["pyarrow", "numpy", "huggingface_hub", "fsspec"]
# ///
"""Merge the three per-column saturate outputs into one `embeddings` config.

Input : hf://datasets/davanstrien/prelinger-moments-emb-{caption,scene,events-text}/data
        (batched rows: parallel `ids` / `embeddings` lists, one row per request)
Output: davanstrien/prelinger-moments :: embeddings/train-00000-of-00001.parquet
        id · emb_caption · emb_scene · emb_events, fp16, sorted by id,
        row_group_size=2000.

Validation is the point of the fp32 stage: a NULL *inside* a vector throws in
`list_cosine_similarity` and `IS NOT NULL` does not catch it, so every vector is
checked dense before it is ever written.
"""

from __future__ import annotations

import argparse
import io

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq

SOURCE = "davanstrien/prelinger-moments"
DIM = 1024
COLUMNS = {"caption": "emb_caption", "scene": "emb_scene", "events_text": "emb_events"}


def repo_for(column: str) -> str:
    return f"davanstrien/prelinger-moments-emb-{column.replace('_', '-')}"


def load_column(fs, column: str) -> tuple[list[str], np.ndarray, dict]:
    """Explode one intermediate repo into (ids, vectors[n,1024] float32, report)."""
    root = f"datasets/{repo_for(column)}/data"
    paths = sorted(p for p in fs.glob(f"{root}/part-*.parquet"))
    if not paths:
        raise SystemExit(f"no parts under {root}")
    all_ids: list[str] = []
    chunks: list[np.ndarray] = []
    errors = 0
    for p in paths:
        with fs.open(p, "rb") as f:
            t = pq.read_table(io.BytesIO(f.read()))
        cols = set(t.column_names)
        if "error" in cols:
            err = t.column("error")
            errors += len(err) - err.null_count
            t = t.filter(pa.compute.is_null(err))
        if t.num_rows == 0:
            continue
        ids_col = t.column("ids").combine_chunks()
        emb_col = t.column("embeddings").combine_chunks()
        if isinstance(ids_col, pa.ChunkedArray):
            ids_col = ids_col.chunk(0)
        if isinstance(emb_col, pa.ChunkedArray):
            emb_col = emb_col.chunk(0)
        flat_ids = ids_col.flatten()
        per_text = emb_col.flatten()          # one list<double> per text
        if flat_ids.null_count:
            raise SystemExit(f"{p}: NULL id inside an ids list")
        lens = pa.compute.list_value_length(per_text)
        if len(lens) and (pa.compute.min(lens).as_py() != DIM
                          or pa.compute.max(lens).as_py() != DIM):
            raise SystemExit(f"{p}: vector length != {DIM} "
                             f"(min {pa.compute.min(lens)}, max {pa.compute.max(lens)})")
        values = per_text.flatten()
        if values.null_count:
            raise SystemExit(f"{p}: NULL element inside a vector "
                             f"({values.null_count} of {len(values)}) — would poison "
                             "list_cosine_similarity and survive IS NOT NULL")
        if len(flat_ids) * DIM != len(values):
            raise SystemExit(f"{p}: {len(flat_ids)} ids vs {len(values) / DIM} vectors")
        all_ids.extend(flat_ids.to_pylist())
        chunks.append(np.asarray(values, dtype=np.float64)
                      .reshape(-1, DIM).astype(np.float32))
    vecs = np.concatenate(chunks) if chunks else np.zeros((0, DIM), np.float32)
    if not np.isfinite(vecs).all():
        raise SystemExit(f"{column}: non-finite values in {(~np.isfinite(vecs)).sum()} slots")
    norms = np.linalg.norm(vecs, axis=1)
    report = {"parts": len(paths), "error_rows": errors, "rows": len(all_ids),
              "unique": len(set(all_ids)),
              "norm_min": float(norms.min()) if len(norms) else 0.0,
              "norm_max": float(norms.max()) if len(norms) else 0.0}
    # duplicate ids: keep the first, but say so loudly
    if report["unique"] != report["rows"]:
        seen, keep = set(), []
        for i, rid in enumerate(all_ids):
            if rid not in seen:
                seen.add(rid)
                keep.append(i)
        all_ids = [all_ids[i] for i in keep]
        vecs = vecs[keep]
    return all_ids, vecs, report


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="embeddings.parquet")
    ap.add_argument("--upload", action="store_true")
    args = ap.parse_args()

    from huggingface_hub import HfApi, HfFileSystem, hf_hub_download

    fs = HfFileSystem()
    api = HfApi()

    # the id spine comes from the published default config, not from the embeddings
    src = pq.read_table(
        hf_hub_download(SOURCE, "data/train-00000-of-00001.parquet",
                        repo_type="dataset"), columns=["id"])
    spine = sorted(src.column("id").to_pylist())
    if len(spine) != len(set(spine)):
        raise SystemExit("source id column is not unique")
    index = {rid: i for i, rid in enumerate(spine)}
    n = len(spine)
    print(f"source spine: {n} unique ids")

    arrays = {"id": pa.array(spine, type=pa.string())}
    for column, out_name in COLUMNS.items():
        ids, vecs, rep = load_column(fs, column)
        missing = [i for i in ids if i not in index]
        if missing:
            raise SystemExit(f"{column}: {len(missing)} ids not in the source spine, "
                             f"e.g. {missing[:3]}")
        mat = np.zeros((n, DIM), dtype=np.float16)
        mask = np.zeros(n, dtype=bool)
        rows = np.fromiter((index[i] for i in ids), dtype=np.int64, count=len(ids))
        mat[rows] = vecs.astype(np.float16)
        mask[rows] = True
        cov = int(mask.sum())
        print(f"{column:12s} -> {out_name:12s} {rep} covered={cov}/{n} "
              f"({n - cov} zero-filled)")
        # Rows with no source text keep the zero vector `mat` was allocated with,
        # NOT a NULL. Two measured reasons (see diag_null.py / diag_zero.py):
        # a NULL fixed_size_list cannot round-trip parquet at all in pyarrow
        # ("Expected all lists to be of size=1024"), and in the variable-size
        # form DuckDB reports `emb IS NULL` as FALSE while still refusing the
        # value — so the NULL is both unguardable and fatal mid-query. A zero
        # vector scores exactly -1.0, sorts last, and never reaches a top-k.
        flat = pa.array(mat.reshape(-1), type=pa.float16())
        arrays[out_name] = pa.FixedSizeListArray.from_arrays(flat, DIM)

    table = pa.table(arrays)
    pq.write_table(table, args.out, compression="zstd", row_group_size=2000)

    md = pq.read_metadata(args.out)
    groups = [md.row_group(i).num_rows for i in range(md.num_row_groups)]
    print(f"\nwrote {args.out}: {md.num_rows} rows, {md.num_row_groups} row groups "
          f"{groups[:3]}...{groups[-1:]}, {md.serialized_size / 1e6:.1f} MB metadata")
    import os
    print(f"file size: {os.path.getsize(args.out) / 1e6:.1f} MB")
    assert set(groups[:-1]) <= {2000}, f"unexpected row-group sizes: {set(groups)}"
    schema = pq.read_schema(args.out)
    for name in COLUMNS.values():
        assert schema.field(name).type == pa.list_(pa.float16(), DIM), schema.field(name)

    # Full re-read: the check that catches a file DuckDB will half-accept but
    # pyarrow (and therefore `datasets`, and therefore the viewer) cannot open.
    back = pq.read_table(args.out)
    assert back.num_rows == md.num_rows, (back.num_rows, md.num_rows)
    for name in COLUMNS.values():
        col = back.column(name).combine_chunks()
        assert col.null_count == 0, f"{name}: {col.null_count} NULL vectors survived"
        assert col.flatten().null_count == 0, f"{name}: NULL element inside a vector"
    print("re-read OK: no NULL vectors, no NULL elements")

    if not args.upload:
        print("\n(dry run — pass --upload to push)")
        return

    api.upload_file(path_or_fileobj=args.out,
                    path_in_repo="embeddings/train-00000-of-00001.parquet",
                    repo_id=SOURCE, repo_type="dataset",
                    commit_message="Add embeddings config (bge-m3, fp16, 3 columns)")
    print("uploaded parquet")
    update_readme(api)


def update_readme(api) -> None:
    """Append the `embeddings` config to the card's frontmatter, touching nothing else."""
    import re

    from huggingface_hub import hf_hub_download

    text = open(hf_hub_download(SOURCE, "README.md", repo_type="dataset")).read()
    m = re.match(r"^---\n(.*?)\n---\n?(.*)$", text, re.S)
    if not m:
        raise SystemExit("README has no frontmatter block — refusing to guess")
    front, body = m.group(1), m.group(2)
    if re.search(r"^- config_name: embeddings$", front, re.M):
        print("embeddings config already present — README unchanged")
        return
    block = ("- config_name: embeddings\n"
             "  data_files:\n"
             "  - split: train\n"
             "    path: embeddings/train-*.parquet\n")
    # append to the existing configs list: insert after its last item, i.e. before
    # the next top-level key or at the end of the block
    lines = front.split("\n")
    try:
        start = next(i for i, ln in enumerate(lines) if ln.strip() == "configs:")
    except StopIteration:
        raise SystemExit("no `configs:` key in frontmatter — refusing to guess")
    end = len(lines)
    for i in range(start + 1, len(lines)):
        if lines[i] and not lines[i].startswith((" ", "-", "\t")):
            end = i
            break
    new_front = "\n".join(lines[:end] + block.rstrip("\n").split("\n") + lines[end:])
    api.upload_file(path_or_fileobj=f"---\n{new_front}\n---\n{body}".encode(),
                    path_in_repo="README.md", repo_id=SOURCE, repo_type="dataset",
                    commit_message="Declare embeddings config")
    print("README frontmatter updated")


if __name__ == "__main__":
    main()