From 32f5de92e753dfbbc89eacbc81641b1b1697781b Mon Sep 17 00:00:00 2001 From: Daniel van Strien Date: Tue, 10 Nov 2026 17:23:20 +0000 Subject: [PATCH 1/4] embeddings: generate-embeddings-saturate.py, a vLLM server recipe for large runs The script starts a vLLM server on the pinned vllm/vllm-openai:v0.30.0 image and sends it one text per request with saturate. The output has one row per input row and a run can be resumed. On one A10G with 20,000 IMDb reviews, embedding step only, one run each: Qwen3-Embedding-0.6B 187 rows/s against 96 for generate-embeddings.py, and bge-m3 276 rows/s against 71. Texts are cut at the model's own maximum, which the server reports, unless --max-seq-len is set. Prompts follow generate-embeddings.py: --prompt, then the model's registered document prompt, then the known-family table. generate-embeddings-vllm.py stays, and its docstring names this recipe as its successor. Co-Authored-By: Claude Fable 5.2 Claude-Session: https://claude.ai/code/session_016isSPpNH4MzMcKDBpSwokh --- embeddings/README.md | 27 +- embeddings/generate-embeddings-saturate.py | 321 +++++++++++++++++++++ embeddings/generate-embeddings-vllm.py | 2 + 3 files changed, 348 insertions(+), 2 deletions(-) create mode 100644 embeddings/generate-embeddings-saturate.py diff --git a/embeddings/README.md b/embeddings/README.md index 2526f8b..8f74866 100644 --- a/embeddings/README.md +++ b/embeddings/README.md @@ -13,13 +13,14 @@ Generate embeddings for a Hugging Face dataset — text or images — with one c cloud GPU, no infra. The output lands back on the Hub as a new dataset (or, with the Lance variant, as a **searchable vector index you can query over `hf://` without downloading**). -There is one simple default and two variants; they are separate single-file scripts because +There is one simple default and a few variants; they are separate single-file scripts because their dependencies (sentence-transformers vs vLLM vs Lance) are too different to share one env. | Script | Use it for | Engine | |---|---|---| | `generate-embeddings.py` | The default. Text or images. Simple, fast, runs anywhere. | sentence-transformers | -| `generate-embeddings-vllm.py` | Max throughput on large *decoder* embedding models (Qwen3-Embedding). | vLLM pooling | +| `generate-embeddings-saturate.py` | **Large text runs.** A vLLM server in the Job, driven by [saturate](https://github.com/davanstrien/saturate). One row out per row in, resumable. | vLLM server + saturate | +| `generate-embeddings-vllm.py` | The older vLLM path (offline engine, unpinned `vllm`). Use `generate-embeddings-saturate.py` instead. | vLLM pooling | | `embed-to-lance.py` | Get a **searchable vector index as a Hub dataset** (the "vector DB" path). | sentence-transformers + Lance | | `launch-embedding-fleet.py` | **Fan one run out across N parallel Jobs** (+ `consolidate-shards.py`). | run_uv_job + Buckets | @@ -39,6 +40,28 @@ hf jobs uv run --flavor l4x1 -s HF_TOKEN \ Always try `--max-samples 100 --private` first. +## Large runs: `generate-embeddings-saturate.py` + +```bash +hf jobs uv run \ + https://huggingface.co/datasets/uv-scripts/embeddings/raw/main/generate-embeddings-saturate.py \ + stanfordnlp/imdb your-name/imdb-embeddings --column text --max-samples 100 +``` + +The script's header sets the `vllm/vllm-openai:v0.30.0` image, `a10g-small` and the `HF_TOKEN` +secret, so the command needs no job flags. Measured on one A10G with 20,000 IMDb reviews of up +to 300 words, one run each, embedding step only: + +| Model | `generate-embeddings.py` | `generate-embeddings-saturate.py` | +|---|---|---| +| Qwen/Qwen3-Embedding-0.6B | 96 rows/s | 187 rows/s | +| BAAI/bge-m3 | 71 rows/s | 276 rows/s | + +The server is ready 100 to 120 seconds after it starts, so `generate-embeddings.py` is quicker +for a few hundred rows. Texts are cut at the model's own maximum length unless you pass +`--max-seq-len`. Prompts follow the same rule as below: `--prompt`, then the model's registered +document prompt, then the known-family table. + ## Fan-out: N Jobs in parallel One L4 does ~900 rows/s with `all-MiniLM-L6-v2`; a fleet of 8 does ~7k. `launch-embedding-fleet.py` diff --git a/embeddings/generate-embeddings-saturate.py b/embeddings/generate-embeddings-saturate.py new file mode 100644 index 0000000..6b52c14 --- /dev/null +++ b/embeddings/generate-embeddings-saturate.py @@ -0,0 +1,321 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "saturate[hf]==0.3.1", +# "huggingface-hub>=1.12", +# ] +# +# [tool.hf-jobs] +# image = "vllm/vllm-openai:v0.30.0" +# flavor = "a10g-small" +# timeout = "2h" +# secrets = ["HF_TOKEN"] +# /// +""" +Generate embeddings for a text column of a Hugging Face dataset with a vLLM server, and write +them to a new dataset on the Hub. + +This is the fast variant for large runs. The script starts a vLLM server on the Job's GPU and +sends it one text per request with `saturate`, which sets the number of requests in flight. +vLLM groups the requests into batches itself. For a first try, for image embeddings, or for a +model vLLM does not serve, use generate-embeddings.py (sentence-transformers) instead. + + hf jobs uv run \\ + https://huggingface.co/datasets/uv-scripts/embeddings/raw/main/generate-embeddings-saturate.py \\ + stanfordnlp/imdb your-username/imdb-embeddings --column text --max-samples 100 + +Measured on one A10G with 20,000 IMDb reviews of up to 300 words (one run each): +Qwen/Qwen3-Embedding-0.6B 187 rows/s, against 96 rows/s for generate-embeddings.py. +BAAI/bge-m3 276 rows/s, against 71 rows/s for generate-embeddings.py at full precision. +The server was ready 100 to 120 seconds after it started, so a run of a few hundred rows is +quicker with generate-embeddings.py. + +On Jobs, the [tool.hf-jobs] header sets the vllm/vllm-openai:v0.30.0 image, the a10g-small +flavor, a 2h timeout and the HF_TOKEN secret (`hf` CLI 1.32+). --flavor and --timeout before +the script override it. + +PROMPTS: many retrieval models need a different prefix for documents and for queries. The +script embeds documents by default and picks the prefix in this order: --prompt, then the +"document" prompt the model registers in config_sentence_transformers.json, then a small table +of known families (e5, nomic), then no prefix. --query-mode uses the query side. The prefix is +logged and written to the dataset card. + +LONG TEXTS: each text is cut at the model's own maximum length, which the server reports +(for example 8,192 tokens for BAAI/bge-m3). Long documents cost more to embed. Pass +--max-seq-len 512 for a shorter cut and a predictable cost; a value above the model's maximum +is lowered to it. generate-embeddings.py cuts at 512 tokens by default. The cut is logged and +written to the dataset card. + +OUTPUT: one row per input row, with `id`, the text column, the embedding column, `model`, +`prompt_tokens` and `error`. A row that failed has a message in `error` and no embedding. +Run the same command again to resume: rows already written are skipped. Add --retry-errors to +try the failed rows again. +""" +import argparse +import json +import os +import re +import sys +import time +import urllib.request + +SCRIPT_URL = "https://huggingface.co/datasets/uv-scripts/embeddings/raw/main/generate-embeddings-saturate.py" + + +def known_family_prefixes(model_id): + """Best-effort (query_prefix, doc_prefix) for common families whose convention is + documented in the model card but NOT registered in config_sentence_transformers.json. + Returns None if unknown. Overridable with --prompt / --no-auto-prompt. + + Verified 2026-07-03 on HF Jobs: of e5 / nomic / bge-en / bge-m3 / Qwen3-Embedding, only + Qwen3-Embedding registers real ST prompts; the rest ship none and rely on manual prefixes. + """ + m = model_id.lower() + # Instruction-style embedders (e5-*-instruct, gte-Qwen, ...): prefer the model's REGISTERED + # prompt or an explicit --prompt; don't guess a literal prefix. + if "instruct" in m: + return None + if "nomic-embed-text" in m: + return ("search_query: ", "search_document: ") + if "bge-m3" in m: # bge-m3 uses no prompts + return ("", "") + # e5 family (e5-base/large/small, multilingual-e5-*), boundaried so e.g. "table5" or a + # model with "e5" mid-word can't silently pick up "query:/passage:" prefixes. + if re.search(r"(^|[/_-])e5([_-]|$)", m): + return ("query: ", "passage: ") + if "bge" in m and "-en" in m: # English bge retrieval: query instruction, docs raw + return ("Represent this sentence for searching relevant passages: ", "") + return None + + +def jobs_flavor() -> str: + """Return the HF Jobs hardware flavor this run is on, or "" when it cannot be told. + + HF Jobs sets ACCELERATOR to "cpu" or "gpu", not the flavor name, so the flavor is looked up + in the public Jobs hardware list: first by GPU model and count, then by the CPU_CORES and + MEMORY values that Jobs sets in the container (to tell a10g-small from a10g-large, or + cpu-basic from cpu-upgrade). Any failure (no network, older huggingface_hub) returns "". + """ + hardware = os.environ.get("ACCELERATOR") or "" + if "-" in hardware or any(character.isdigit() for character in hardware): + return hardware # already a flavor name + try: + from huggingface_hub import HfApi + + gpu_name, gpu_count = "", 0 + try: + import torch + + if torch.cuda.is_available(): + gpu_name = torch.cuda.get_device_name(0) + gpu_count = torch.cuda.device_count() + except ImportError: + # Torch-free drivers (e.g. the -server recipes): ask nvidia-smi instead. + import subprocess + + try: + listed = subprocess.run( + ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + capture_output=True, text=True, timeout=10, check=True, + ).stdout + names = [line.strip() for line in listed.splitlines() if line.strip()] + if names: + gpu_name, gpu_count = names[0], len(names) + except (OSError, subprocess.SubprocessError): + pass + if gpu_count == 0 and hardware == "gpu": + return "" # a GPU run whose GPU cannot be identified: do not guess a CPU flavor + + # Whole-word match, so "L4" does not match "NVIDIA L40S" ("NVIDIA A100-SXM4-80GB" -> A100). + gpu_words = gpu_name.upper().replace("-", " ").split() + candidates = [] + for flavor in HfApi(token=False).list_jobs_hardware(): + accelerator = flavor.accelerator + if gpu_count == 0 and accelerator is None: + candidates.append(flavor) + elif ( + gpu_count > 0 + and accelerator is not None + and accelerator.model.upper() in gpu_words + and int(accelerator.quantity) == gpu_count + ): + candidates.append(flavor) + if not candidates: + return "" + + cores = float(os.environ.get("CPU_CORES") or 0) + memory_gb = float((os.environ.get("MEMORY") or "0").upper().rstrip("GIB ") or 0) + + def distance(flavor) -> float: + # Relative gap to the listed vCPUs and RAM; Jobs reports slightly different numbers. + listed_cores = float(flavor.cpu.split()[0]) + listed_memory = float(flavor.ram.split()[0]) + return abs(listed_cores - cores) / listed_cores + abs(listed_memory - memory_gb) / listed_memory + + return min(candidates, key=distance).name + except Exception: + return "" + + +def registered_prompt(model_id, side): + """Return the prompt the model registers for "document" or "query", or None when it has none.""" + from huggingface_hub import hf_hub_download + from huggingface_hub.errors import HfHubHTTPError + + try: + path = hf_hub_download(model_id, "config_sentence_transformers.json") + except (HfHubHTTPError, OSError): + return None # the model has no such file + with open(path) as file: + prompts = json.load(file).get("prompts") or {} + return prompts.get(side) or None + + +def choose_prefix(args): + """Return (prefix, where it came from) for the text side this run embeds.""" + side = "query" if args.query_mode else "document" + if args.prompt is not None: + return args.prompt, "--prompt" + registered = registered_prompt(args.model, side) + if registered: + return registered, "registered by the model" + family = known_family_prefixes(args.model) + if family: + query_prefix, document_prefix = family + return (query_prefix if args.query_mode else document_prefix), "known model family" + return "", "none" + + +def model_max_tokens(endpoint, model_id): + """Return the most tokens the server accepts for one text of this model.""" + with urllib.request.urlopen(f"{endpoint}/models", timeout=60) as response: + served = json.load(response)["data"] + for entry in served: + if entry["id"] == model_id: + return entry["max_model_len"] + sys.exit(f"The server does not list {model_id}. It serves: {[entry['id'] for entry in served]}") + + +def write_card(args, prefix, prefix_source, dimension, max_tokens, stats): + """Write the dataset card, with the canonical uv-scripts reproduction section.""" + from huggingface_hub import HfApi + + on_jobs = os.environ.get("JOB_ID") is not None + flavor = jobs_flavor() + origin = ( + "Produced on [Hugging Face Jobs](https://huggingface.co/docs/huggingface_hub/guides/jobs)" + + (f" (`{flavor}`)" if flavor else "") + ) if on_jobs else "Generated" + jobs_tag = "\n- hf-jobs" if on_jobs else "" + side = "query" if args.query_mode else "document" + prefix_line = f"`{prefix}` ({prefix_source})" if prefix else "(none)" + cut_option = f" --max-seq-len {args.max_seq_len}" if args.max_seq_len else "" + card = ( + f"---\ntags:\n- embeddings\n- uv-script\n- generated{jobs_tag}\n---\n\n" + f"# {args.output_dataset}\n\n" + f"Embeddings of [`{args.input_dataset}`](https://huggingface.co/datasets/{args.input_dataset}) " + f"column `{args.column}`, made with a vLLM server.\n\n" + f"- Model: [`{args.model}`](https://huggingface.co/{args.model}) (dim {dimension})\n" + f"- Column: `{args.output_column}`, normalized\n" + f"- Prompt prepended ({side} side): {prefix_line}\n" + f"- Texts cut at {max_tokens} tokens\n" + f"- Rows: {stats.rows_processed} embedded, {stats.rows_failed} failed (see the `error` column)\n\n" + f"## Reproduction\n\n" + f"{origin} with the [`generate-embeddings-saturate.py`]({SCRIPT_URL}) recipe " + f"from [uv-scripts](https://huggingface.co/uv-scripts). Run it yourself:\n\n" + f"```bash\nhf jobs uv run {SCRIPT_URL} \\\n" + f" {args.input_dataset} --column {args.column} --model {args.model}{cut_option}\n```\n" + ) + HfApi().upload_file( + path_or_fileobj=card.encode(), path_in_repo="README.md", + repo_id=args.output_dataset, repo_type="dataset", + ) + + +def main(): + parser = argparse.ArgumentParser(description="Embed a text column with a vLLM server via saturate") + parser.add_argument("input_dataset", help="Input dataset ID on the Hugging Face Hub") + parser.add_argument("output_dataset", help="Output dataset ID (created if missing)") + parser.add_argument("--model", default="Qwen/Qwen3-Embedding-0.6B", help="An embedding model that vLLM serves") + parser.add_argument("--column", default="text", help="Text column to embed") + parser.add_argument("--output-column", default="embeddings") + parser.add_argument("--config", default=None, help="Dataset config name") + parser.add_argument("--split", default="train") + parser.add_argument("--id-column", default=None, help="Column to use as the row id (default: the row's position)") + parser.add_argument("--max-samples", type=int, default=None, help="Embed only the first N rows (for testing)") + parser.add_argument("--max-seq-len", type=int, default=None, + help="Cut each text at this many tokens (default: the model's own maximum)") + parser.add_argument("--prompt", default=None, + help="Prefix to put before every text, such as 'passage: '. Pass '' for no prefix.") + parser.add_argument("--query-mode", action="store_true", help="Embed the texts as queries, not documents") + parser.add_argument("--private", action="store_true", help="Make the output dataset private") + parser.add_argument("--retry-errors", action="store_true", help="Try again the rows whose only record is an error") + args = parser.parse_args() + + import pyarrow as pa + from huggingface_hub import HfApi + from saturate import Engine, dataset_rows, pump + + prefix, prefix_source = choose_prefix(args) + print(f"Prompt prefix: {prefix!r} ({prefix_source})", file=sys.stderr) + + HfApi().create_repo(args.output_dataset, repo_type="dataset", private=args.private, exist_ok=True) + rows = dataset_rows( + args.input_dataset, config=args.config, split=args.split, + columns=[args.column], ids=args.id_column or "index", limit=args.max_samples, + ) + + cut = {"max_tokens": None} # set once the server reports the model's maximum + + def to_request(row): + return { + "model": args.model, + "input": prefix + (row[args.column] or ""), + "truncate_prompt_tokens": cut["max_tokens"], + } + + # Declared, so a part holding only error rows has the same columns as the rest. + schema = pa.schema([ + ("id", pa.string()), ("error", pa.string()), (args.column, pa.string()), + (args.output_column, pa.list_(pa.float32())), ("model", pa.string()), ("prompt_tokens", pa.int64()), + ]) # fmt: skip + seen = {"dimension": None} + + def parse(row, body): + vector = body["data"][0]["embedding"] + seen["dimension"] = len(vector) + return { + args.column: row[args.column], + args.output_column: vector, + "model": args.model, + "prompt_tokens": (body.get("usage") or {}).get("prompt_tokens"), + } + + # The readiness check sends one embedding request, so the server is known to embed before the run starts. + engine = Engine( + args.model, engine="vllm", extra_args=["--runner", "pooling"], + ready_route="/embeddings", ready_payload={"model": args.model, "input": "ready"}, + ) + with engine as endpoint: + model_maximum = model_max_tokens(endpoint, args.model) + cut["max_tokens"] = min(args.max_seq_len or model_maximum, model_maximum) + print(f"Texts are cut at {cut['max_tokens']} tokens (the model's maximum is {model_maximum})", file=sys.stderr) + started = time.monotonic() + stats = pump( + rows, to_request, parse, endpoint, f"hf://datasets/{args.output_dataset}/data", + route="/embeddings", retry_errors=args.retry_errors, schema=schema, + ) + elapsed = time.monotonic() - started + + write_card(args, prefix, prefix_source, seen["dimension"], cut["max_tokens"], stats) + print( + f"https://huggingface.co/datasets/{args.output_dataset} " + f"({stats.rows_processed} ok, {stats.rows_failed} error rows, " + f"{stats.rows_processed / max(elapsed, 1e-6):.0f} rows/s, dim {seen['dimension']})", + file=sys.stderr, + ) + print("GENERATE_EMBEDDINGS_SATURATE " + stats.to_json(), flush=True) + + +if __name__ == "__main__": + main() diff --git a/embeddings/generate-embeddings-vllm.py b/embeddings/generate-embeddings-vllm.py index 4e4025a..878163b 100644 --- a/embeddings/generate-embeddings-vllm.py +++ b/embeddings/generate-embeddings-vllm.py @@ -7,6 +7,8 @@ # ] # /// """ +Superseded by generate-embeddings-saturate.py, which runs a vLLM server on a pinned image. + High-throughput embedding generation with vLLM pooling mode — the "scale" variant of generate-embeddings.py, for large *decoder* embedding models (e.g. Qwen3-Embedding). On Qwen3-Embedding-0.6B this was ~2x the sentence-transformers throughput on the same GPU. From 385edf15a2d7be7c4a4f00436de19f903bed55dc Mon Sep 17 00:00:00 2001 From: Daniel van Strien Date: Tue, 10 Nov 2026 17:26:28 +0000 Subject: [PATCH 2/4] maintenance: smoke records for generate-embeddings-saturate.py Four 100-row runs: the default model on a10g-small and l4x1, an e5 model with a cut above its maximum, and bge-m3 with --max-seq-len 128. Co-Authored-By: Claude Fable 5.2 Claude-Session: https://claude.ai/code/session_016isSPpNH4MzMcKDBpSwokh --- maintenance/STATUS.md | 3 ++- maintenance/verifications.jsonl | 4 ++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/maintenance/STATUS.md b/maintenance/STATUS.md index 5f607f5..f61ae55 100644 --- a/maintenance/STATUS.md +++ b/maintenance/STATUS.md @@ -3,7 +3,7 @@ Generated by `python3 tools/recipe-status.py --write`. Do not edit by hand. Tiers and rules: [MAINTAINING.md](../MAINTAINING.md). -Tiered: 1 of 50 recipes with a header. +Tiered: 1 of 51 recipes with a header. | Recipe | Tier | Image | Packages | Last smoke test | Result | |---|---|---|---|---|---| @@ -14,6 +14,7 @@ Tiered: 1 of 50 recipes with a header. | `classification/train-classifier.py` | — | default | latest | never | | | `classification/train-gliner2.py` | — | default | latest | never | | | `classification/train-kev.py` | — | default | latest | never | | +| `embeddings/generate-embeddings-saturate.py` | — | `vllm/vllm-openai:v0.30.0` | latest | 2026-11-10 | [pass](https://huggingface.co/jobs/davanstrien/6af35312ab34849891ac6fcf) | | `ocr/deepseek-ocr-vllm.py` | — | `vllm/vllm-openai:v0.29.0` | latest | never | | | `ocr/deepseek-ocr2-vllm.py` | — | `vllm/vllm-openai:v0.29.0` | latest | never | | | `ocr/dots-mocr.py` | — | `vllm/vllm-openai:v0.30.0` | latest | never | | diff --git a/maintenance/verifications.jsonl b/maintenance/verifications.jsonl index 53a590f..5546abf 100644 --- a/maintenance/verifications.jsonl +++ b/maintenance/verifications.jsonl @@ -4,3 +4,7 @@ {"date": "2026-11-06", "script": "ocr/ovis-ocr2.py", "commit": "6e37e09", "image": "vllm/vllm-openai:v0.29.0", "exclude_newer": "", "model_revision": "main", "flavor": "a10g-small", "input": "davanstrien/lighton-ocr3-speed-test (private; 3 shuffled pages)", "rows": 3, "outcome": "pass", "job_url": "https://huggingface.co/jobs/davanstrien/6aedd70cc78863f9c6346083", "notes": "datasets<5.1 pin (#179). 3/3 rows, no error rows, text read; row 1 is a blank page (empty for every model)"} {"date": "2026-11-06", "script": "ocr/unlimited-ocr-vllm.py", "commit": "be299b7", "image": "vllm/vllm-openai:unlimited-ocr", "exclude_newer": "", "model_revision": "07dea832e22aefee32ad281d4b80551282e1c168", "flavor": "a10g-small", "input": "davanstrien/lighton-ocr3-speed-test (private; 3 shuffled pages)", "rows": 3, "outcome": "pass", "job_url": "https://huggingface.co/jobs/davanstrien/6aede17e0e0338d1bb817171", "notes": "datasets<5.1 pin and model revision pin (#179); main fails since 2026-10-14 (remote code needs addict, matplotlib). 3/3 rows, no error rows, text read."} {"date": "2026-11-06", "script": "ocr/lighton-ocr2.py", "commit": "5ec2132", "image": "vllm/vllm-openai@sha256:953d3a06d5e64ab582985cd7401289d3abf2a2c14ef2158e9a84313daeec77d7", "exclude_newer": "2026-10-01T00:00:00Z", "model_revision": "be7c41313c30e940fe5571de59e5fabc779e00c8", "flavor": "a10g-small", "input": "davanstrien/lighton-ocr3-speed-test (private; 3 shuffled pages)", "rows": 3, "outcome": "pass", "job_url": "https://huggingface.co/jobs/davanstrien/6aede22ac78863f9c6347831", "notes": "First frozen run: digest, cutoff and revision applied (log shows the revision and vLLM v0.22.1). Same output lengths as the unfrozen run on the same pages. 3/3 rows, no error rows, text read; row 1 is a blank page (empty for every model)"} +{"date": "2026-11-10", "script": "embeddings/generate-embeddings-saturate.py", "commit": "32f5de9", "image": "vllm/vllm-openai:v0.30.0", "exclude_newer": "", "model_revision": "main", "flavor": "l4x1", "input": "stanfordnlp/imdb (first 100 train rows)", "rows": 100, "outcome": "pass", "job_url": "https://huggingface.co/jobs/davanstrien/6af3504bab34849891ac68ea", "notes": "--flavor l4x1, default model Qwen3-Embedding-0.6B. Earlier version with a fixed 512-token cut. 100/100 rows, 0 error rows, unit-length vectors, 100 distinct, 100 unique ids; card names the flavor, the prompt and the cut."} +{"date": "2026-11-10", "script": "embeddings/generate-embeddings-saturate.py", "commit": "32f5de9", "image": "vllm/vllm-openai:v0.30.0", "exclude_newer": "", "model_revision": "main", "flavor": "a10g-small", "input": "stanfordnlp/imdb (first 100 train rows)", "rows": 100, "outcome": "pass", "job_url": "https://huggingface.co/jobs/davanstrien/6af351559737d9a9e09119f2", "notes": "Header only, default model Qwen3-Embedding-0.6B, cut at the model maximum 32768 (longest text 1164 tokens). Ran the script before the card's reproduce line gained --max-seq-len; no other change. 100/100 rows, 0 error rows, unit-length vectors, 100 distinct, 100 unique ids; card names the flavor, the prompt and the cut."} +{"date": "2026-11-10", "script": "embeddings/generate-embeddings-saturate.py", "commit": "32f5de9", "image": "vllm/vllm-openai:v0.30.0", "exclude_newer": "", "model_revision": "main", "flavor": "a10g-small", "input": "stanfordnlp/imdb (first 100 train rows)", "rows": 100, "outcome": "pass", "job_url": "https://huggingface.co/jobs/davanstrien/6af351589737d9a9e09119fe", "notes": "--model intfloat/multilingual-e5-small --max-seq-len 100000: lowered to the model maximum 512, prefix `passage: `, 384-d. Ran the script before the card's reproduce line gained --max-seq-len; no other change. 100/100 rows, 0 error rows, unit-length vectors, 100 distinct, 100 unique ids; card names the flavor, the prompt and the cut."} +{"date": "2026-11-10", "script": "embeddings/generate-embeddings-saturate.py", "commit": "32f5de9", "image": "vllm/vllm-openai:v0.30.0", "exclude_newer": "", "model_revision": "main", "flavor": "a10g-small", "input": "stanfordnlp/imdb (first 100 train rows)", "rows": 100, "outcome": "pass", "job_url": "https://huggingface.co/jobs/davanstrien/6af35312ab34849891ac6fcf", "notes": "--model BAAI/bge-m3 --max-seq-len 128: longest text 128 tokens, no prefix, 1024-d. 100/100 rows, 0 error rows, unit-length vectors, 100 distinct, 100 unique ids; card names the flavor, the prompt and the cut."} From 16457ba55d49eb41f6d7e88cf790db781d1d85d7 Mon Sep 17 00:00:00 2001 From: Daniel van Strien Date: Tue, 10 Nov 2026 17:33:47 +0000 Subject: [PATCH 3/4] embeddings: fix four review findings in generate-embeddings-saturate.py - A resumed run that processed no rows wrote a card with "dim None" and "0 embedded". The dimension now comes from one request to the server, and the card gives the rows of the last run and the rows written before. - --id-column index and --id-column content were read as saturate id strategies, not as column names. They are now passed as a callable, as in lighton-ocr3-saturate.py. - The card's reproduce command now carries the options that change the result: config, split, id column, sample limit, text cut, prompt, query mode and output column. - The docstring says that a --retry-errors run keeps the old error record, and how to read one record per id. Co-Authored-By: Claude Fable 5.2 Claude-Session: https://claude.ai/code/session_016isSPpNH4MzMcKDBpSwokh --- embeddings/generate-embeddings-saturate.py | 72 ++++++++++++++++++---- 1 file changed, 59 insertions(+), 13 deletions(-) diff --git a/embeddings/generate-embeddings-saturate.py b/embeddings/generate-embeddings-saturate.py index 6b52c14..00ce994 100644 --- a/embeddings/generate-embeddings-saturate.py +++ b/embeddings/generate-embeddings-saturate.py @@ -46,15 +46,19 @@ is lowered to it. generate-embeddings.py cuts at 512 tokens by default. The cut is logged and written to the dataset card. -OUTPUT: one row per input row, with `id`, the text column, the embedding column, `model`, +OUTPUT: one record per input row, with `id`, the text column, the embedding column, `model`, `prompt_tokens` and `error`. A row that failed has a message in `error` and no embedding. Run the same command again to resume: rows already written are skipped. Add --retry-errors to -try the failed rows again. +try the failed rows again. A retry adds a second record for the row and keeps the old error +record, so after a --retry-errors run drop the records that have an `error`, or read the +output with `saturate.read_output`, which keeps one record per id (a success over an earlier +error). """ import argparse import json import os import re +import shlex import sys import time import urllib.request @@ -196,6 +200,39 @@ def model_max_tokens(endpoint, model_id): sys.exit(f"The server does not list {model_id}. It serves: {[entry['id'] for entry in served]}") +def embedding_dimension(endpoint, model_id): + """Return the length of the model's vectors, from one request to the server.""" + request = urllib.request.Request( + f"{endpoint}/embeddings", + data=json.dumps({"model": model_id, "input": "dimension check"}).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=600) as response: + return len(json.load(response)["data"][0]["embedding"]) + + +def reproduce_options(args): + """Return the options that give this dataset again, as a shell-quoted string.""" + options = ["--column", args.column, "--model", args.model] + if args.output_column != "embeddings": + options += ["--output-column", args.output_column] + if args.config: + options += ["--config", args.config] + if args.split != "train": + options += ["--split", args.split] + if args.id_column: + options += ["--id-column", args.id_column] + if args.max_samples: + options += ["--max-samples", str(args.max_samples)] + if args.max_seq_len: + options += ["--max-seq-len", str(args.max_seq_len)] + if args.prompt is not None: + options += ["--prompt", args.prompt] + if args.query_mode: + options += ["--query-mode"] + return " ".join(shlex.quote(option) for option in options) + + def write_card(args, prefix, prefix_source, dimension, max_tokens, stats): """Write the dataset card, with the canonical uv-scripts reproduction section.""" from huggingface_hub import HfApi @@ -209,7 +246,8 @@ def write_card(args, prefix, prefix_source, dimension, max_tokens, stats): jobs_tag = "\n- hf-jobs" if on_jobs else "" side = "query" if args.query_mode else "document" prefix_line = f"`{prefix}` ({prefix_source})" if prefix else "(none)" - cut_option = f" --max-seq-len {args.max_seq_len}" if args.max_seq_len else "" + # A resumed run processes only the rows that were missing, so the card gives both counts. + earlier = f", {stats.rows_done_prior} already written by an earlier run" if stats.rows_done_prior else "" card = ( f"---\ntags:\n- embeddings\n- uv-script\n- generated{jobs_tag}\n---\n\n" f"# {args.output_dataset}\n\n" @@ -219,12 +257,13 @@ def write_card(args, prefix, prefix_source, dimension, max_tokens, stats): f"- Column: `{args.output_column}`, normalized\n" f"- Prompt prepended ({side} side): {prefix_line}\n" f"- Texts cut at {max_tokens} tokens\n" - f"- Rows: {stats.rows_processed} embedded, {stats.rows_failed} failed (see the `error` column)\n\n" + f"- Rows in the last run: {stats.rows_processed} embedded, {stats.rows_failed} failed{earlier}. " + f"A failed row has a message in `error`.\n\n" f"## Reproduction\n\n" f"{origin} with the [`generate-embeddings-saturate.py`]({SCRIPT_URL}) recipe " f"from [uv-scripts](https://huggingface.co/uv-scripts). Run it yourself:\n\n" f"```bash\nhf jobs uv run {SCRIPT_URL} \\\n" - f" {args.input_dataset} --column {args.column} --model {args.model}{cut_option}\n```\n" + f" {args.input_dataset} {reproduce_options(args)}\n```\n" ) HfApi().upload_file( path_or_fileobj=card.encode(), path_in_repo="README.md", @@ -259,10 +298,20 @@ def main(): prefix, prefix_source = choose_prefix(args) print(f"Prompt prefix: {prefix!r} ({prefix_source})", file=sys.stderr) + if args.id_column is None: + ids = "index" + elif args.id_column in ("index", "content"): + # saturate reads these two names as id strategies, so a column with one of them + # needs a callable. It returns the raw value, so saturate still refuses a null id. + def ids(example): + return example[args.id_column] + else: + ids = args.id_column + HfApi().create_repo(args.output_dataset, repo_type="dataset", private=args.private, exist_ok=True) rows = dataset_rows( args.input_dataset, config=args.config, split=args.split, - columns=[args.column], ids=args.id_column or "index", limit=args.max_samples, + columns=[args.column], ids=ids, limit=args.max_samples, ) cut = {"max_tokens": None} # set once the server reports the model's maximum @@ -279,14 +328,10 @@ def to_request(row): ("id", pa.string()), ("error", pa.string()), (args.column, pa.string()), (args.output_column, pa.list_(pa.float32())), ("model", pa.string()), ("prompt_tokens", pa.int64()), ]) # fmt: skip - seen = {"dimension": None} - def parse(row, body): - vector = body["data"][0]["embedding"] - seen["dimension"] = len(vector) return { args.column: row[args.column], - args.output_column: vector, + args.output_column: body["data"][0]["embedding"], "model": args.model, "prompt_tokens": (body.get("usage") or {}).get("prompt_tokens"), } @@ -300,6 +345,7 @@ def parse(row, body): model_maximum = model_max_tokens(endpoint, args.model) cut["max_tokens"] = min(args.max_seq_len or model_maximum, model_maximum) print(f"Texts are cut at {cut['max_tokens']} tokens (the model's maximum is {model_maximum})", file=sys.stderr) + dimension = embedding_dimension(endpoint, args.model) started = time.monotonic() stats = pump( rows, to_request, parse, endpoint, f"hf://datasets/{args.output_dataset}/data", @@ -307,11 +353,11 @@ def parse(row, body): ) elapsed = time.monotonic() - started - write_card(args, prefix, prefix_source, seen["dimension"], cut["max_tokens"], stats) + write_card(args, prefix, prefix_source, dimension, cut["max_tokens"], stats) print( f"https://huggingface.co/datasets/{args.output_dataset} " f"({stats.rows_processed} ok, {stats.rows_failed} error rows, " - f"{stats.rows_processed / max(elapsed, 1e-6):.0f} rows/s, dim {seen['dimension']})", + f"{stats.rows_processed / max(elapsed, 1e-6):.0f} rows/s, dim {dimension})", file=sys.stderr, ) print("GENERATE_EMBEDDINGS_SATURATE " + stats.to_json(), flush=True) From dc55f084be6f3e312e650831eae4c23dbd5a040b Mon Sep 17 00:00:00 2001 From: Daniel van Strien Date: Tue, 10 Nov 2026 17:41:58 +0000 Subject: [PATCH 4/4] maintenance: smoke records for the review fixes in generate-embeddings-saturate.py A fresh run and a resumed run of the committed file on the same output. Co-Authored-By: Claude Fable 5.2 Claude-Session: https://claude.ai/code/session_016isSPpNH4MzMcKDBpSwokh --- maintenance/STATUS.md | 2 +- maintenance/verifications.jsonl | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/maintenance/STATUS.md b/maintenance/STATUS.md index f61ae55..52f35aa 100644 --- a/maintenance/STATUS.md +++ b/maintenance/STATUS.md @@ -14,7 +14,7 @@ Tiered: 1 of 51 recipes with a header. | `classification/train-classifier.py` | — | default | latest | never | | | `classification/train-gliner2.py` | — | default | latest | never | | | `classification/train-kev.py` | — | default | latest | never | | -| `embeddings/generate-embeddings-saturate.py` | — | `vllm/vllm-openai:v0.30.0` | latest | 2026-11-10 | [pass](https://huggingface.co/jobs/davanstrien/6af35312ab34849891ac6fcf) | +| `embeddings/generate-embeddings-saturate.py` | — | `vllm/vllm-openai:v0.30.0` | latest | 2026-11-10 | [pass](https://huggingface.co/jobs/davanstrien/6af356829737d9a9e0912648) | | `ocr/deepseek-ocr-vllm.py` | — | `vllm/vllm-openai:v0.29.0` | latest | never | | | `ocr/deepseek-ocr2-vllm.py` | — | `vllm/vllm-openai:v0.29.0` | latest | never | | | `ocr/dots-mocr.py` | — | `vllm/vllm-openai:v0.30.0` | latest | never | | diff --git a/maintenance/verifications.jsonl b/maintenance/verifications.jsonl index 5546abf..673663c 100644 --- a/maintenance/verifications.jsonl +++ b/maintenance/verifications.jsonl @@ -8,3 +8,5 @@ {"date": "2026-11-10", "script": "embeddings/generate-embeddings-saturate.py", "commit": "32f5de9", "image": "vllm/vllm-openai:v0.30.0", "exclude_newer": "", "model_revision": "main", "flavor": "a10g-small", "input": "stanfordnlp/imdb (first 100 train rows)", "rows": 100, "outcome": "pass", "job_url": "https://huggingface.co/jobs/davanstrien/6af351559737d9a9e09119f2", "notes": "Header only, default model Qwen3-Embedding-0.6B, cut at the model maximum 32768 (longest text 1164 tokens). Ran the script before the card's reproduce line gained --max-seq-len; no other change. 100/100 rows, 0 error rows, unit-length vectors, 100 distinct, 100 unique ids; card names the flavor, the prompt and the cut."} {"date": "2026-11-10", "script": "embeddings/generate-embeddings-saturate.py", "commit": "32f5de9", "image": "vllm/vllm-openai:v0.30.0", "exclude_newer": "", "model_revision": "main", "flavor": "a10g-small", "input": "stanfordnlp/imdb (first 100 train rows)", "rows": 100, "outcome": "pass", "job_url": "https://huggingface.co/jobs/davanstrien/6af351589737d9a9e09119fe", "notes": "--model intfloat/multilingual-e5-small --max-seq-len 100000: lowered to the model maximum 512, prefix `passage: `, 384-d. Ran the script before the card's reproduce line gained --max-seq-len; no other change. 100/100 rows, 0 error rows, unit-length vectors, 100 distinct, 100 unique ids; card names the flavor, the prompt and the cut."} {"date": "2026-11-10", "script": "embeddings/generate-embeddings-saturate.py", "commit": "32f5de9", "image": "vllm/vllm-openai:v0.30.0", "exclude_newer": "", "model_revision": "main", "flavor": "a10g-small", "input": "stanfordnlp/imdb (first 100 train rows)", "rows": 100, "outcome": "pass", "job_url": "https://huggingface.co/jobs/davanstrien/6af35312ab34849891ac6fcf", "notes": "--model BAAI/bge-m3 --max-seq-len 128: longest text 128 tokens, no prefix, 1024-d. 100/100 rows, 0 error rows, unit-length vectors, 100 distinct, 100 unique ids; card names the flavor, the prompt and the cut."} +{"date": "2026-11-10", "script": "embeddings/generate-embeddings-saturate.py", "commit": "16457ba", "image": "vllm/vllm-openai:v0.30.0", "exclude_newer": "", "model_revision": "main", "flavor": "a10g-small", "input": "stanfordnlp/imdb (first 100 train rows)", "rows": 100, "outcome": "pass", "job_url": "https://huggingface.co/jobs/davanstrien/6af35597ab34849891ac74f2", "notes": "After the review fixes. Header only, default model Qwen3-Embedding-0.6B. 100/100 rows, 0 error rows, unit-length vectors, 100 unique ids; card has dim 1024 and the reproduce line carries --max-samples 100."} +{"date": "2026-11-10", "script": "embeddings/generate-embeddings-saturate.py", "commit": "16457ba", "image": "vllm/vllm-openai:v0.30.0", "exclude_newer": "", "model_revision": "main", "flavor": "a10g-small", "input": "stanfordnlp/imdb (first 100 train rows)", "rows": 100, "outcome": "pass", "job_url": "https://huggingface.co/jobs/davanstrien/6af356829737d9a9e0912648", "notes": "Resume: the same command again on the same output. 0 rows processed, output still 100 rows with 100 unique ids; card keeps dim 1024 and says 100 rows were written by an earlier run."}