#!/usr/bin/env python3
"""Benchmark candidate scientific embedding models on Apple Silicon (MPS): throughput and a quick sanity check.
Models: allenai/specter2_base (SPECTER2 base, no adapter), malteos/scincl, BAAI/bge-base-en-v1.5, sentence-transformers/all-MiniLM-L6-v2."""
import sys, time, json, torch, duckdb
from sentence_transformers import SentenceTransformer
R = sys.argv[1]; N = int(sys.argv[2]) if len(sys.argv) > 2 else 2000
con = duckdb.connect(f"{R}/data/processed/wp44.duckdb", read_only=True)
rows = 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()
texts = [f"{t}. {ab[:1500]}" for t, ab in rows]
dev = "mps" if torch.backends.mps.is_available() else "cpu"; print("device", dev, "texts", len(texts), flush=True)
res = {}
for m in ["allenai/specter2_base", "malteos/scincl", "BAAI/bge-base-en-v1.5", "sentence-transformers/all-MiniLM-L6-v2"]:
    try:
        model = SentenceTransformer(m, device=dev); model.max_seq_length = 256
        _ = model.encode(texts[:64], batch_size=64)
        t0 = time.time(); E = model.encode(texts, batch_size=128, convert_to_numpy=True, normalize_embeddings=True); dt_ = time.time() - t0
        res[m] = {"dim": int(E.shape[1]), "texts_per_s": round(len(texts) / dt_, 1), "seconds": round(dt_, 1)}
        print(m, res[m], flush=True)
    except Exception as e:
        res[m] = {"error": str(e)[:300]}; print(m, "ERROR", str(e)[:200], flush=True)
json.dump(res, open(f"{R}/data/metadata/embedding_benchmark.json", "w"), indent=1)