"""Small CPU-only text retrieval example for the LeCompute EmbeddingGemma 2 guide. Index short, already prepared passages once; encode only the query at search time. The corpus is educational, not a retrieval benchmark or a production RAG system. """ import argparse import hashlib import importlib.metadata import json import os from pathlib import Path import platform import sys import time MODEL_ID = "google/embeddinggemma-2" REVISION = "f9972ac22e696a105c482202646979e33e4b5a1e" FORMATTING = "Document prompt over title + newline + text; SearchQuery for queries" def load_model(dimensions, threads, offline): os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1") if offline: os.environ["HF_HUB_OFFLINE"] = "1" import torch from sentence_transformers import SentenceTransformer torch.set_num_threads(threads) started = time.perf_counter() model = SentenceTransformer( MODEL_ID, revision=REVISION, device="cpu", local_files_only=offline, config_kwargs={"vision_config": None, "audio_config": None}, model_kwargs={"dtype": torch.float32}, truncate_dim=dimensions, ) return model, time.perf_counter() - started def describe(model, dimensions, threads, load_seconds): return { "model": MODEL_ID, "revision": REVISION, "dimensions": dimensions, "formatting": FORMATTING, "device": "cpu", "dtype": "float32", "threads": threads, "python": platform.python_version(), "packages": { name: importlib.metadata.version(name) for name in ("torch", "torchvision", "transformers", "sentence-transformers", "numpy", "Pillow") }, "loaded_parameters": sum(p.numel() for p in model.parameters()), "model_load_seconds": load_seconds, } def build_index(args): import numpy as np source = args.documents.read_bytes() documents = json.loads(source) if not isinstance(documents, list) or not documents: raise ValueError("Expected a non-empty JSON list of short passages") for doc in documents: if not isinstance(doc, dict) or any( not isinstance(doc.get(key), str) or not doc[key].strip() for key in ("id", "title", "text") ): raise ValueError("Every passage needs non-empty id, title and text strings") if len({doc["id"] for doc in documents}) != len(documents): raise ValueError("Passage ids must be unique") model, load_seconds = load_model(args.dimensions, args.threads, args.offline) passages = [doc["title"] + "\n" + doc["text"] for doc in documents] # Refuse silent truncation. This demo accepts short, pre-segmented text only. prefix = model.prompts["Document"] for doc, passage in zip(documents, passages): token_count = len(model.tokenizer.encode(prefix + passage)) if token_count > model.max_seq_length: raise ValueError(f"Passage {doc['id']} exceeds the model context; split it first") started = time.perf_counter() vectors = model.encode( passages, prompt_name="Document", normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False, ).astype(np.float32) encode_seconds = time.perf_counter() - started if not np.isfinite(vectors).all(): raise ValueError("Non-finite document embeddings") manifest = describe(model, args.dimensions, args.threads, load_seconds) manifest.update( source_sha256=hashlib.sha256(source).hexdigest(), documents=documents, indexing_seconds=encode_seconds, raw_vector_bytes=vectors.nbytes, ) args.index.parent.mkdir(parents=True, exist_ok=True) # A file handle prevents NumPy from silently changing the requested filename. with args.index.open("wb") as handle: np.savez_compressed(handle, vectors=vectors, manifest=json.dumps(manifest)) print(json.dumps({k: v for k, v in manifest.items() if k != "documents"}, indent=2)) print(f"Saved {len(documents)} passages to {args.index}") def search_index(args): import numpy as np with np.load(args.index, allow_pickle=False) as archive: vectors = archive["vectors"].copy() manifest = json.loads(str(archive["manifest"].item())) if (manifest["model"], manifest["revision"], manifest["formatting"]) != ( MODEL_ID, REVISION, FORMATTING ): raise ValueError("Index and query encoder differ; rebuild the index") if vectors.shape != (len(manifest["documents"]), manifest["dimensions"]): raise ValueError("Index dimensions or passage count do not match its manifest") if not np.isfinite(vectors).all(): raise ValueError("Non-finite index vectors") model, load_seconds = load_model( manifest["dimensions"], args.threads, args.offline ) query_text = model.prompts["SearchQuery"] + args.query if len(model.tokenizer.encode(query_text)) > model.max_seq_length: raise ValueError("Query exceeds the model context") started = time.perf_counter() query = model.encode( [args.query], prompt_name="SearchQuery", normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False, )[0].astype(np.float32) query_seconds = time.perf_counter() - started if not np.isfinite(query).all(): raise ValueError("Non-finite query embedding") started = time.perf_counter() scores = vectors @ query best = np.argsort(-scores, kind="stable")[:args.top_k] ranking_seconds = time.perf_counter() - started result = describe(model, manifest["dimensions"], args.threads, load_seconds) result.update( query=args.query, query_encode_seconds=query_seconds, ranking_seconds=ranking_seconds, results=[dict(manifest["documents"][i], score=float(scores[i])) for i in best], ) print(json.dumps(result, ensure_ascii=False, indent=2)) def main(): parser = argparse.ArgumentParser(description=__doc__) commands = parser.add_subparsers(dest="command", required=True) index = commands.add_parser("index") index.add_argument("--documents", required=True, type=Path) index.add_argument("--dimensions", type=int, choices=(128, 256, 512, 768), default=256) search = commands.add_parser("search") search.add_argument("--query", required=True) search.add_argument("--top-k", type=int, default=3) for command in (index, search): command.add_argument("--index", type=Path, default=Path("index.npz")) command.add_argument("--threads", type=int, default=4) command.add_argument("--offline", action="store_true") args = parser.parse_args() if args.threads < 1 or (args.command == "search" and args.top_k < 1): parser.error("threads and top-k must be positive") try: (build_index if args.command == "index" else search_index)(args) except (ValueError, OSError, KeyError) as exc: parser.exit(1, f"Error: {exc}\n") if __name__ == "__main__": sys.stdout.reconfigure(encoding="utf-8") sys.stderr.reconfigure(encoding="utf-8") main()