uv-scripts/embeddings

1# /// script2# requires-python = ">=3.10"3# dependencies = [4#     "datasets",5#     "vllm",6#     "huggingface-hub",7# ]8# ///9"""10High-throughput embedding generation with vLLM pooling mode — the "scale" variant of11generate-embeddings.py, for large *decoder* embedding models (e.g. Qwen3-Embedding). On12Qwen3-Embedding-0.6B this was ~2x the sentence-transformers throughput on the same GPU.13 14Prefer the plain sentence-transformers `generate-embeddings.py` unless you specifically need15vLLM throughput: this variant has a heavier cold-start and two footguns handled below16(the embedding-mode kwarg drifted across vLLM versions; vLLM does not auto-truncate).17 18Runs on the BARE uv image (vLLM ships the CUDA toolkit + flashinfer as wheels).19 20    hf jobs uv run --flavor l4x1 -s HF_TOKEN generate-embeddings-vllm.py \\21        stanfordnlp/imdb your-name/imdb-embeddings --column text --model Qwen/Qwen3-Embedding-0.6B --private22"""23import argparse24import logging25import os26import time27os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")28os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")29logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")30log = logging.getLogger("generate-embeddings-vllm")31 32 33def build_llm(LLM, model, max_model_len, gpu_mem_util):34    """vLLM's embedding-mode selector drifted: modern uses runner='pooling', old used35    task='embed'. A wrong kwarg raises TypeError at init (cheap) → fall through."""36    base = dict(enforce_eager=True, max_model_len=max_model_len, gpu_memory_utilization=gpu_mem_util)37    for label, extra in [("runner", {"runner": "pooling"}), ("task", {"task": "embed"}), ("auto", {})]:38        try:39            llm = LLM(model=model, **base, **extra)40            log.info(f"engine init via '{label}'")41            return llm42        except TypeError as te:43            log.warning(f"ctor '{label}' rejected: {te}")44    raise RuntimeError("no vLLM constructor form accepted")45 46 47def main():48    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)49    ap.add_argument("input_dataset")50    ap.add_argument("output_dataset")51    ap.add_argument("--model", default="Qwen/Qwen3-Embedding-0.6B")52    ap.add_argument("--column", default="text")53    ap.add_argument("--output-column", default="embeddings")54    ap.add_argument("--split", default="train")55    ap.add_argument("--max-samples", type=int, default=None)56    ap.add_argument("--config", default=None, help="dataset config name (e.g. wikipedia needs one)")57    ap.add_argument("--max-model-len", type=int, default=512)58    ap.add_argument("--gpu-mem-util", type=float, default=0.85)59    ap.add_argument("--private", action="store_true")60    args = ap.parse_args()61 62    import torch63    from datasets import load_dataset64    from huggingface_hub import DatasetCard, login65    from vllm import LLM66    if not torch.cuda.is_available():67        raise SystemExit("No CUDA GPU available — vLLM needs one. Run with a GPU flavor, e.g. "68                         "`hf jobs uv run --flavor l4x1 ...` (or use generate-embeddings.py on CPU).")69    if os.environ.get("HF_TOKEN"):70        login(token=os.environ["HF_TOKEN"])71 72    ds = (load_dataset(args.input_dataset, args.config, split=args.split) if args.config73          else load_dataset(args.input_dataset, split=args.split))74    if args.output_column in ds.column_names:75        raise SystemExit(f"Output column {args.output_column!r} already exists — pick another.")76    if args.max_samples:77        ds = ds.select(range(min(args.max_samples, len(ds))))78    texts = [t if isinstance(t, str) and t.strip() else " " for t in ds[args.column]]79    n = len(texts)80 81    llm = build_llm(LLM, args.model, args.max_model_len, args.gpu_mem_util)82    embed_fn = getattr(llm, "embed", None) or getattr(llm, "encode")83 84    # vLLM raises on inputs > max_model_len (no silent truncation) — pre-truncate at the tokenizer.85    # Tokenize each text once (not twice) — this pass is CPU-bound on large datasets.86    tk = llm.get_tokenizer()87    cap = max(8, args.max_model_len - 16)88    def _truncate(t):89        ids = tk.encode(t)90        return tk.decode(ids[:cap]) if len(ids) > cap else t91    texts = [_truncate(t) for t in texts]92 93    t0 = time.perf_counter()94    outs = embed_fn(texts)95    log.info(f"embedded {n} rows in {time.perf_counter()-t0:.1f}s")96 97    def vec(o):98        e = o.outputs99        e = getattr(e, "embedding", None) or getattr(e, "data", e)100        return list(e)101    ds = ds.add_column(args.output_column, [vec(o) for o in outs])102    dim = len(ds[0][args.output_column])103 104    card = DatasetCard(105        f"# {args.output_dataset}\n\nEmbeddings of `{args.input_dataset}` column `{args.column}` "106        f"with [`{args.model}`](https://huggingface.co/{args.model}) (dim {dim}, vLLM pooling).\n\n"107        f"Produced on Hugging Face Jobs with `uv-scripts/embeddings/generate-embeddings-vllm.py`.\n")108    # Retry the push with an XET-disable fallback — a transient failure would lose the paid run.109    max_retries = 3110    for attempt in range(1, max_retries + 1):111        try:112            if attempt > 1:113                log.warning("Disabling XET (fallback to HTTP upload)")114                os.environ["HF_HUB_DISABLE_XET"] = "1"115            ds.push_to_hub(args.output_dataset, private=args.private)116            break117        except Exception as e:118            log.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")119            if attempt < max_retries:120                delay = 30 * (2 ** (attempt - 1))121                log.info(f"Retrying in {delay}s...")122                time.sleep(delay)123            else:124                log.error("All upload attempts failed. Results are lost.")125                raise SystemExit(1)126    try:127        card.push_to_hub(args.output_dataset, repo_type="dataset")128    except Exception as e:129        log.warning(f"card push skipped: {e}")130    log.info(f"✅ https://huggingface.co/datasets/{args.output_dataset}")131 132 133if __name__ == "__main__":134    main()135