"""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)