Spaces:
Running
Running
File size: 2,723 Bytes
086a913 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """
Pre-compute query embeddings and cache them to disk.
The ef_sweep script (and any future queries-only experiment) reads from this
cache to skip the embedding step. The cache key includes the embedding model
name, so changing EMBEDDING_MODEL invalidates this cache automatically and
forces a recompute.
Usage:
uv run python scripts/cache_query_vectors.py
uv run python scripts/cache_query_vectors.py --recompute
"""
from __future__ import annotations
import argparse
import os
import pickle
import sys
import time
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src"))
import numpy as np
from sentence_transformers import SentenceTransformer
from codesearch.config import EMBEDDING_MODEL
from codesearch.data import load_codesearch
CACHE_DIR = ".cache"
def cache_path(model_name: str) -> str:
safe = model_name.replace("/", "_")
return os.path.join(CACHE_DIR, f"query_vectors_{safe}.pkl")
def main() -> None:
parser = argparse.ArgumentParser(
description="Cache query embeddings so the ef sweep can skip the encode step."
)
parser.add_argument(
"--recompute",
action="store_true",
help="Re-embed even if the cache already exists.",
)
args = parser.parse_args()
path = cache_path(EMBEDDING_MODEL)
if os.path.exists(path) and not args.recompute:
size_mb = os.path.getsize(path) / 1024 / 1024
print(f"[skip] Cache already exists at {path} ({size_mb:.1f} MB).")
print(" Pass --recompute to rebuild.")
return
os.makedirs(CACHE_DIR, exist_ok=True)
print("[1/3] Loading eval queries (test split only)...")
_, queries = load_codesearch(n=-1, queries_only=True)
print(f"[2/3] Loading embedding model: {EMBEDDING_MODEL}")
model = SentenceTransformer(EMBEDDING_MODEL)
print(f"[3/3] Embedding {len(queries):,} queries...")
t0 = time.perf_counter()
vectors = model.encode(
[q["query"] for q in queries],
normalize_embeddings=True,
show_progress_bar=True,
batch_size=256,
)
elapsed = time.perf_counter() - t0
print(f" Done in {elapsed:.1f}s ({elapsed / len(queries) * 1000:.2f} ms/query)")
print(f"Writing cache to {path}...")
with open(path, "wb") as f:
pickle.dump(
{
"model": EMBEDDING_MODEL,
"queries": queries,
"vectors": np.asarray(vectors, dtype=np.float32),
},
f,
protocol=pickle.HIGHEST_PROTOCOL,
)
size_mb = os.path.getsize(path) / 1024 / 1024
print(f"Cached {len(queries):,} query vectors ({size_mb:.1f} MB).")
if __name__ == "__main__":
main()
|