File size: 1,131 Bytes
1d9bd9b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | """Re-embed the corpus chunks with BGE-small on the GPU to regenerate escr_vectors.npy locally
(faster than transferring the 1.9GB float32 matrix over Thor's slow link). Embeds in FILE ORDER so
the vector index matches escr_chunks.jsonl line order, exactly as serve.py expects. Documents are
embedded PLAIN (no query instruction prefix — that's query-side only)."""
import json, time, os
import numpy as np
from sentence_transformers import SentenceTransformer
DATA = os.environ.get("THEMIS_DATA", ".")
DEV = os.environ.get("THEMIS_DEVICE", "cuda")
texts = []
with open(os.path.join(DATA, "escr_chunks.jsonl"), encoding="utf-8") as f:
for l in f:
texts.append(json.loads(l)["text"])
print(f"{len(texts)} chunks; embedding on {DEV} ...", flush=True)
st = SentenceTransformer("BAAI/bge-small-en-v1.5", device=DEV)
t0 = time.time()
M = st.encode(texts, batch_size=512, normalize_embeddings=True, convert_to_numpy=True,
show_progress_bar=True).astype(np.float32)
np.save(os.path.join(DATA, "escr_vectors.npy"), M)
print(f"done {M.shape} {M.dtype} in {time.time()-t0:.0f}s -> escr_vectors.npy", flush=True)
|