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 prepare_delta(root, shard_rows, model_name, start_shard=100):29 """Shards for eligible works with an abstract that have no embedding yet (ids read from the existing shard_*.ids.npy)."""30 import duckdb, glob as _g31 tag = model_name.replace("/", "__"); done = np.concatenate([np.load(f) for f in _g.glob(os.path.join(root, "data/processed/embeddings", tag, "shard_*.ids.npy"))]) if _g.glob(os.path.join(root, "data/processed/embeddings", tag, "shard_*.ids.npy")) else np.array([], dtype=np.int64)32 con = duckdb.connect(os.path.join(root, "data/processed/wp44.duckdb"), read_only=True); con.execute("SET threads TO 8;")33 con.register("done_ids", __import__("pandas").DataFrame({"work_id": done}))34 out = os.path.join(root, "data/processed/embeddings/input"); os.makedirs(out, exist_ok=True)35 con.execute("""CREATE OR REPLACE TEMP TABLE emb_in AS36 SELECT w.id AS work_id, w.title || '. ' || coalesce(substr(a.abstract, 1, 2000), '') AS text37 FROM corpus_works w JOIN oa_abstracts a ON a.work_id = w.id LEFT JOIN done_ids d ON d.work_id = w.id38 WHERE w.eligible AND d.work_id IS NULL ORDER BY w.id""")39 n = con.execute("SELECT count(*) FROM emb_in").fetchone()[0]40 con.execute(f"COPY (SELECT *, {start_shard} + (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)")41 print("delta prepared", n, "texts from shard", start_shard, "to", start_shard + (n - 1) // shard_rows)424344def run(root, model_name, shard_from, shard_to, batch, max_len):45 import torch, pyarrow.parquet as pq46 from sentence_transformers import SentenceTransformer47 tag = model_name.replace("/", "__"); out = os.path.join(root, "data/processed/embeddings", tag); os.makedirs(out, exist_ok=True)48 dev = "mps" if torch.backends.mps.is_available() else "cpu"49 model = SentenceTransformer(model_name, device=dev); model.max_seq_length = max_len50 shards = sorted(glob.glob(os.path.join(root, "data/processed/embeddings/input", "shard=*", "*.parquet")))51 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": {}}52 for p in shards:53 k = int(os.path.basename(os.path.dirname(p)).split("=")[1])54 if not (shard_from <= k < shard_to):55 continue56 o = os.path.join(out, f"shard_{k:04d}.npy")57 if os.path.exists(o):58 continue59 t = pq.read_table(p).to_pydict(); t0 = time.time()60 E = model.encode(t["text"], batch_size=batch, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False).astype(np.float16)61 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))62 man["shards"][str(k)] = {"n": len(t["text"]), "seconds": round(time.time() - t0, 1), "when": dt.datetime.now(dt.UTC).isoformat()}63 json.dump(man, open(man_p, "w"), indent=1)64 print(f"shard {k}: {len(t['text'])} texts in {time.time()-t0:.0f}s ({len(t['text'])/(time.time()-t0):.0f}/s)", flush=True)656667def main():68 ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--prepare", action="store_true"); ap.add_argument("--run", action="store_true")69 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)70 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")71 ap.add_argument("--prepare-delta", action="store_true", help="shards (numbered from 100) for eligible works not yet embedded")72 a = ap.parse_args()73 if a.prepare_delta:74 prepare_delta(a.root, a.shard_rows, a.model)75 if a.prepare:76 prepare(a.root, a.shard_rows, a.title_only_fallback)77 if a.run:78 run(a.root, a.model, a.shard_from, a.shard_to, a.batch, a.max_len)798081if __name__ == "__main__":82 main()83