1#!/usr/bin/env python32"""WP44 — embed title + abstract of corpus works with a scientific sentence-embedding model, sharded for multi-node runs.34 --prepare : DuckDB exports data/processed/embeddings/input/shard_NNN.parquet (work_id, text) for eligible works with an5 abstract (or title only when --title-only-fallback), N shards of --shard-rows rows.6 --run : encode shards [--shard-from, --shard-to) with --model on MPS/CPU → data/processed/embeddings/<model_tag>/7 shard_NNN.npy (float16, L2-normalised) + shard_NNN.ids.npy. Resumable (existing shards skipped).8Provenance: data/processed/embeddings/<model_tag>/manifest.json (model, revision, max_seq_length, host, timings).9"""10import argparse, glob, json, os, socket, time, datetime as dt11import numpy as np121314def prepare(root, shard_rows, title_only_fallback):15 import duckdb16 con = duckdb.connect(os.path.join(root, "data/processed/wp44.duckdb"), read_only=True); con.execute("SET threads TO 16;")17 out = os.path.join(root, "data/processed/embeddings/input"); os.makedirs(out, exist_ok=True)18 fb = " OR TRUE" if title_only_fallback else ""19 con.execute(f"""CREATE OR REPLACE TEMP TABLE emb_in AS20 SELECT w.id AS work_id, w.title || '. ' || coalesce(substr(a.abstract, 1, 2000), '') AS text21 FROM corpus_works w LEFT JOIN oa_abstracts a ON a.work_id = w.id22 WHERE w.eligible AND (a.abstract IS NOT NULL{fb}) ORDER BY w.id""")23 n = con.execute("SELECT count(*) FROM emb_in").fetchone()[0]24 con.execute(f"COPY (SELECT *, (row_number() OVER (ORDER BY work_id) - 1) // {shard_rows} AS shard FROM emb_in) TO '{out}' (FORMAT PARQUET, PARTITION_BY (shard), OVERWRITE_OR_IGNORE)")25 print("prepared", n, "texts in", len(glob.glob(out + "/shard=*")), "shards")262728def run(root, model_name, shard_from, shard_to, batch, max_len):29 import torch, pyarrow.parquet as pq30 from sentence_transformers import SentenceTransformer31 tag = model_name.replace("/", "__"); out = os.path.join(root, "data/processed/embeddings", tag); os.makedirs(out, exist_ok=True)32 dev = "mps" if torch.backends.mps.is_available() else "cpu"33 model = SentenceTransformer(model_name, device=dev); model.max_seq_length = max_len34 shards = sorted(glob.glob(os.path.join(root, "data/processed/embeddings/input", "shard=*", "*.parquet")))35 man_p = os.path.join(out, f"manifest_{socket.gethostname()}.json"); man = json.load(open(man_p)) if os.path.exists(man_p) else {"model": model_name, "max_seq_length": max_len, "device": dev, "host": socket.gethostname(), "shards": {}}36 for p in shards:37 k = int(os.path.basename(os.path.dirname(p)).split("=")[1])38 if not (shard_from <= k < shard_to):39 continue40 o = os.path.join(out, f"shard_{k:04d}.npy")41 if os.path.exists(o):42 continue43 t = pq.read_table(p).to_pydict(); t0 = time.time()44 E = model.encode(t["text"], batch_size=batch, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False).astype(np.float16)45 np.save(o + ".part.npy", E); os.replace(o + ".part.npy", o); np.save(o.replace(".npy", ".ids.npy"), np.asarray(t["work_id"], dtype=np.int64))46 man["shards"][str(k)] = {"n": len(t["text"]), "seconds": round(time.time() - t0, 1), "when": dt.datetime.now(dt.UTC).isoformat()}47 json.dump(man, open(man_p, "w"), indent=1)48 print(f"shard {k}: {len(t['text'])} texts in {time.time()-t0:.0f}s ({len(t['text'])/(time.time()-t0):.0f}/s)", flush=True)495051def main():52 ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--prepare", action="store_true"); ap.add_argument("--run", action="store_true")53 ap.add_argument("--model", default="allenai/specter2_base"); ap.add_argument("--shard-rows", type=int, default=50000); ap.add_argument("--shard-from", type=int, default=0)54 ap.add_argument("--shard-to", type=int, default=10**6); ap.add_argument("--batch", type=int, default=128); ap.add_argument("--max-len", type=int, default=256); ap.add_argument("--title-only-fallback", action="store_true")55 a = ap.parse_args()56 if a.prepare:57 prepare(a.root, a.shard_rows, a.title_only_fallback)58 if a.run:59 run(a.root, a.model, a.shard_from, a.shard_to, a.batch, a.max_len)606162if __name__ == "__main__":63 main()64