uv-scripts/embeddings

1# /// script2# requires-python = ">=3.10"3# dependencies = [4#     "datasets",5#     "sentence-transformers>=5.0.0",6#     "torch",7#     "numpy",8#     "einops",9#     "pyarrow",10#     "pylance",11#     "huggingface-hub",12# ]13# ///14"""15Embed a Hugging Face dataset and push it back as a Lance vector index — a Hub dataset that16IS a searchable vector database. Anyone you share it with can vector-search it over `hf://`17without downloading it:18 19    import lance20    ds = lance.dataset("hf://datasets/your-name/my-vecdb/vecdb.lance")   # opens fast, no download21    hits = ds.to_table(nearest={"column": "vector", "q": query_vector, "k": 5})22 23Best for share-and-search over a corpus; for high-QPS serving, pull the dataset local first.24 25PROMPTS: documents are embedded with the model's known DOCUMENT convention (e5 → "passage: ",26nomic → "search_document: "; bge-en/bge-m3 → none). At SEARCH time, embed your query with the27matching QUERY prefix (printed at the end of the run) or retrieval quality silently drops.28Override the document prefix with --prompt '<prefix>' (or --prompt '' for none).29 30    hf jobs uv run --flavor l4x1 -s HF_TOKEN embed-to-lance.py \\31        stanfordnlp/imdb your-name/imdb-vecdb --column text --model BAAI/bge-base-en-v1.5 --private32"""33import argparse34import logging35import os36import re37import shutil38import sys39import time40import numpy as np41import pyarrow as pa42 43logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")44log = logging.getLogger("embed-to-lance")45 46 47def known_convention(model_id):48    """(query_prefix, doc_prefix) for common families (documented in model cards, not registered49    in sentence-transformers config). Same table as generate-embeddings.py; None = unknown."""50    m = model_id.lower()51    if "instruct" in m:52        return None53    if "nomic-embed-text" in m:54        return ("search_query: ", "search_document: ")55    if "bge-m3" in m:56        return ("", "")57    if re.search(r"(^|[/_-])e5([_-]|$)", m):58        return ("query: ", "passage: ")59    if "bge" in m and "-en" in m:60        return ("Represent this sentence for searching relevant passages: ", "")61    return None62 63 64def main():65    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)66    ap.add_argument("input_dataset")67    ap.add_argument("output_repo")68    ap.add_argument("--column", default="text")69    ap.add_argument("--config", default=None, help="dataset config name (e.g. wikipedia needs one)")70    ap.add_argument("--split", default="train")71    ap.add_argument("--model", default="BAAI/bge-base-en-v1.5")72    ap.add_argument("--max-samples", type=int, default=None)73    ap.add_argument("--batch-size", type=int, default=64)74    ap.add_argument("--max-seq-len", type=int, default=512)75    ap.add_argument("--prompt", default=None,76                    help="Document prefix to prepend (default: auto from the known-family table; "77                         "pass '' to force none)")78    ap.add_argument("--private", action="store_true")79    args = ap.parse_args()80 81    import torch82    import lance83    from datasets import load_dataset84    from huggingface_hub import HfApi, login85    from sentence_transformers import SentenceTransformer86 87    if os.environ.get("HF_TOKEN"):88        login(token=os.environ["HF_TOKEN"])89 90    t_all = time.perf_counter()91    ds = load_dataset(args.input_dataset, args.config, split=args.split) if args.config \92        else load_dataset(args.input_dataset, split=args.split)93    if args.max_samples:94        ds = ds.select(range(min(args.max_samples, len(ds))))95    texts = [t if isinstance(t, str) and t.strip() else " " for t in ds[args.column]]96    n = len(texts)97 98    t_load = time.perf_counter()99 100    device = "cuda" if torch.cuda.is_available() else "cpu"101    model = SentenceTransformer(args.model, device=device, trust_remote_code=True)102    if getattr(model, "max_seq_length", None):103        model.max_seq_length = min(model.max_seq_length, args.max_seq_len)104    dim = model.get_sentence_embedding_dimension()105 106    # Document-side prompt: explicit --prompt wins (incl. '' for none), else the known-family107    # table; else None → encode_document() natively selects any REGISTERED document prompt108    # (and routes Router models by task).109    registered = {k: v for k, v in (getattr(model, "prompts", {}) or {}).items() if v}110    kc = known_convention(args.model)111    doc_prompt = args.prompt if args.prompt is not None else (kc[1] if kc else None)112    query_prompt = kc[0] if kc else registered.get("query", "")113    log.info(f"document prompt: {doc_prompt!r}" if doc_prompt114             else ("document prompt: native (registered)" if registered.get("document")115                   else "document prompt: (none)"))116 117    t0 = time.perf_counter()118    encode_kwargs = {"prompt": doc_prompt} if doc_prompt is not None else {}119    emb = model.encode_document(texts, batch_size=args.batch_size, show_progress_bar=True,120                                convert_to_numpy=True, normalize_embeddings=True,121                                **encode_kwargs).astype(np.float32)122    log.info(f"embedded {n} rows in {time.perf_counter()-t0:.1f}s, dim={dim}")123 124    tbl = pa.table({125        "id": pa.array(range(n), pa.int64()),126        "text": pa.array([t[:2000] for t in texts]),127        "vector": pa.FixedSizeListArray.from_arrays(pa.array(emb.reshape(-1), pa.float32()), dim),128    })129    local = "vecdb.lance"130    if os.path.exists(local):131        shutil.rmtree(local)132    lds = lance.write_dataset(tbl, local, mode="overwrite")133    try:134        parts = max(1, min(256, int(np.sqrt(n))))135        lds.create_index("vector", index_type="IVF_PQ", num_partitions=parts,136                         num_sub_vectors=max(1, dim // 16))137        log.info(f"built IVF_PQ index (partitions={parts})")138    except Exception as e:139        log.warning(f"index build skipped ({repr(e)[:120]}); flat search still works over hf://")140 141    # Retry the upload with an XET-disable fallback — a transient failure here would lose the142    # whole (paid) embedding run.143    api = HfApi()144    api.create_repo(args.output_repo, repo_type="dataset", private=args.private, exist_ok=True)145    max_retries = 3146    for attempt in range(1, max_retries + 1):147        try:148            if attempt > 1:149                log.warning("Disabling XET (fallback to HTTP upload)")150                os.environ["HF_HUB_DISABLE_XET"] = "1"151            api.upload_folder(folder_path=local, path_in_repo="vecdb.lance",152                              repo_id=args.output_repo, repo_type="dataset")153            break154        except Exception as e:155            log.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")156            if attempt < max_retries:157                delay = 30 * (2 ** (attempt - 1))158                log.info(f"Retrying in {delay}s...")159                time.sleep(delay)160            else:161                log.error("All upload attempts failed. Results are lost.")162                sys.exit(1)163    total_s = time.perf_counter() - t_all164    import json as _json165    log.info("ROUNDTRIP " + _json.dumps({166        "input": args.input_dataset, "n": n, "dim": dim, "model": args.model,167        "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu",168        "batch_size": args.batch_size, "load_s": round(t_load - t_all, 1),169        "total_roundtrip_s": round(total_s, 1), "rows_per_s_end_to_end": round(n / total_s, 1),170        "hf_path": f"hf://datasets/{args.output_repo}/vecdb.lance"}))171    log.info(f"✅ {n} rows → searchable vector DB in {total_s/60:.1f} min "172             f"(load→embed→index→push). hf://datasets/{args.output_repo}/vecdb.lance")173    if query_prompt or registered.get("query"):174        log.info("⚠️ At search time, embed queries with the QUERY convention — mismatched prompts "175                 "degrade retrieval. Easiest: model.encode_query([your_query])"176                 + (f", or explicitly: model.encode([{query_prompt!r} + your_query])" if query_prompt else "."))177 178 179if __name__ == "__main__":180    main()181