1#!/usr/bin/env python32"""Benchmark candidate scientific embedding models on Apple Silicon (MPS): throughput and a quick sanity check.3Models: allenai/specter2_base (SPECTER2 base, no adapter), malteos/scincl, BAAI/bge-base-en-v1.5, sentence-transformers/all-MiniLM-L6-v2."""4import sys, time, json, torch, duckdb5from sentence_transformers import SentenceTransformer6R = sys.argv[1]; N = int(sys.argv[2]) if len(sys.argv) > 2 else 20007con = duckdb.connect(f"{R}/data/processed/wp44.duckdb", read_only=True)8rows = con.execute(f"SELECT w.title, a.abstract FROM oa_works w JOIN oa_abstracts a ON a.work_id=w.id WHERE w.publication_year=2024 LIMIT {N}").fetchall()9texts = [f"{t}. {ab[:1500]}" for t, ab in rows]10dev = "mps" if torch.backends.mps.is_available() else "cpu"; print("device", dev, "texts", len(texts), flush=True)11res = {}12for m in ["allenai/specter2_base", "malteos/scincl", "BAAI/bge-base-en-v1.5", "sentence-transformers/all-MiniLM-L6-v2"]:13 try:14 model = SentenceTransformer(m, device=dev); model.max_seq_length = 25615 _ = model.encode(texts[:64], batch_size=64)16 t0 = time.time(); E = model.encode(texts, batch_size=128, convert_to_numpy=True, normalize_embeddings=True); dt_ = time.time() - t017 res[m] = {"dim": int(E.shape[1]), "texts_per_s": round(len(texts) / dt_, 1), "seconds": round(dt_, 1)}18 print(m, res[m], flush=True)19 except Exception as e:20 res[m] = {"error": str(e)[:300]}; print(m, "ERROR", str(e)[:200], flush=True)21json.dump(res, open(f"{R}/data/metadata/embedding_benchmark.json", "w"), indent=1)22