scripts/embed_column.py

# /// script
# requires-python = ">=3.11"
# dependencies = ["saturate[hf]>=0.1.2"]
#
# [tool.hf-jobs]
# flavor = "a10g-small"
# image = "ghcr.io/huggingface/text-embeddings-inference:86-latest"
# secrets = ["HF_TOKEN"]
# timeout = "4h"
# ///
"""embed one text column of davanstrien/prelinger-moments with bge-m3 on TEI.

Adapted from the `hf pipe embed --dump` driver. Differences from the ejected
driver, all deliberate:
  * --column picks caption | scene | events_text; output repo is per-column
  * rows are keyed by the dataset's own `id` column (ids="id"), not row index —
    the merge back into prelinger-moments joins on it
  * no limit (full 23,148 rows)
  * per-row ids ride along in the batch so `parse` can emit them alongside the
    vectors; a batch id alone cannot be unzipped back to rows
  * TRUNCATE raised to 6000 chars: longest scene/caption is ~5.1k chars
    (~1.3k tokens), far inside bge-m3's 8192 — the driver's 1200 would have
    silently cut a third of the corpus

Run (one job per column):
  hf jobs uv run --flavor a10g-small \
      --image ghcr.io/huggingface/text-embeddings-inference:86-latest \
      -s HF_TOKEN embed_column.py -- --column caption
"""

from __future__ import annotations

import argparse

from saturate import Auto, Engine, dataset_rows, pump
from saturate.source import shard_select

MODEL = "BAAI/bge-m3"
SOURCE = "davanstrien/prelinger-moments"
BATCH = 32
TRUNCATE = 6000  # chars; --auto-truncate is the token-level backstop

TEI_CMD = ["text-embeddings-router", "--model-id", MODEL, "--port", "8080",
           "--max-batch-tokens", "32768", "--max-client-batch-size", "128",
           "--auto-truncate"]


def batched(rows, column: str):
    # caller-side micro-batching (rows-only boundary); empty texts skipped,
    # tail batch flushes. Ids travel with the texts: parse needs them to write
    # one output row per source row.
    buf = []
    for rid, row in rows:
        text = (row.get(column) or "").strip()
        if not text:
            continue
        buf.append((rid, text[:TRUNCATE]))
        if len(buf) == BATCH:
            yield f"b-{buf[0][0]}", {"ids": [i for i, _ in buf],
                                     "texts": [t for _, t in buf]}
            buf = []
    if buf:
        yield f"b-{buf[0][0]}", {"ids": [i for i, _ in buf],
                                 "texts": [t for _, t in buf]}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--column", required=True,
                    choices=["caption", "scene", "events_text"])
    ap.add_argument("--rank", type=int, default=0)
    ap.add_argument("--world", type=int, default=1)
    args = ap.parse_args()

    repo = f"davanstrien/prelinger-moments-emb-{args.column.replace('_', '-')}"
    output = f"hf://datasets/{repo}/data"

    from huggingface_hub import HfApi

    api = HfApi()
    api.create_repo(repo, repo_type="dataset", private=True, exist_ok=True)
    if not api.file_exists(repo, "README.md", repo_type="dataset"):
        card = ("---\nconfigs:\n- config_name: data\n  data_files:\n"
                "  - split: train\n    path: data/part-*.parquet\n  default: true\n"
                "- config_name: manifest\n  data_files:\n"
                "  - split: train\n    path: data/_manifest/ids-*.parquet\n---\n"
                f"# {MODEL} embeddings of `{SOURCE}` column `{args.column}`\n\n"
                "Intermediate saturate output. Batched rows: one row per request,\n"
                "`ids`/`texts`/`embeddings` are parallel lists.\n")
        api.upload_file(path_or_fileobj=card.encode(), path_in_repo="README.md",
                        repo_id=repo, repo_type="dataset")

    def to_request(row: dict) -> dict:
        return {"model": MODEL, "input": row["texts"]}

    def parse(row: dict, body: dict) -> dict:
        data = body["data"]
        if len(data) != len(row["texts"]):  # schema-invalid: never mark as success
            raise ValueError(f"expected {len(row['texts'])} embeddings, got {len(data)}")
        # Reorder by the OpenAI-style `index` when TEI supplies it; fall back to
        # response order otherwise.
        if all("index" in d for d in data):
            vecs = [None] * len(data)
            for d in data:
                vecs[d["index"]] = d["embedding"]
            if any(v is None for v in vecs):
                raise ValueError("response indices did not cover the request")
        else:
            vecs = [d["embedding"] for d in data]
        return {"ids": row["ids"],
                "embeddings": vecs,
                "n_texts": len(data),
                "prompt_tokens": (body.get("usage") or {}).get("prompt_tokens")}

    rows = dataset_rows(SOURCE, config="default", split="train",
                        columns=[args.column], ids="id")
    rows = shard_select(batched(rows, args.column), rank=args.rank, world=args.world)
    with Engine(MODEL, cmd=TEI_CMD, port=8080,
                ready_route="/embeddings",
                ready_payload={"model": MODEL, "input": ["ready?"]},
                ready_accept=lambda r: r.status_code == 200) as endpoint:
        stats = pump(rows, to_request, parse, endpoint, output,
                     route="/embeddings",
                     window=Auto(target_waiting=8, initial=4, max_limit=64),
                     shard=(args.rank, args.world), flush_every=50)
    print(f"EMBED column={args.column} rank={args.rank} " + stats.to_json(), flush=True)


if __name__ == "__main__":
    main()