Art 0.1.0

Built from art-dev 78fb10b.

art/screen/screen.py

@@ -0,0 +1,148 @@
+"""`bingham screen`: embed, score against each profile, re-rank, highlight.
+
+For each profile, the batch is every stored item not yet screened with the
+current profile version, embedder, and re-ranker. Each paper in the batch gets
+
+  cosine  max similarity to the profile description and its exemplars (bge-large)
+  bm25    lexical score against the profile keywords and description
+
+and the union of the top candidates by each (the larger of screen.top_k and
+screen.top_frac of the batch), plus any paper by a tracked author, is scored
+by the cross-encoder. A paper is highlighted when its re-rank score exceeds
+the profile's threshold. Scores are stored for every paper, nominated or not,
+so the ablations can be run from the database.
+"""
+
+import math
+import time
+
+import numpy as np
+
+from art import env, store
+from art.feeds import fetch
+from art.profile import profiles as prof
+from art.screen import bm25, embed, rerank
+
+
+def load_items(conn, where: str = "1", params: tuple = ()) -> list[dict]:
+    """Stored items with text to screen (see store.SCREENABLE)."""
+    rows = conn.execute(f"SELECT * FROM items WHERE {store.SCREENABLE} AND {where}", params)
+    return [store.item_dict(r) for r in rows]
+
+
+def exemplar_texts(conn, profile: prof.Profile) -> dict[str, str]:
+    """{ref: "title\\n\\nabstract"} for the profile's exemplars, fetching unknown ones once."""
+    out = {}
+    for ex in profile.exemplars:
+        row = conn.execute("SELECT * FROM exemplars WHERE ref = ?", (ex.ref,)).fetchone()
+        if ex.abstract:   # pasted in profiles.yaml; always wins
+            title = row["title"] if row else ""
+            conn.execute("INSERT OR REPLACE INTO exemplars VALUES (?, ?, ?, ?, ?)",
+                         (ex.ref, title, ex.abstract.strip(), "profiles.yaml", store.now()))
+        elif row is None:
+            try:
+                it = fetch.get_crossref_work(ex.doi) if ex.doi else fetch.get_arxiv_record(ex.arxiv)
+            except Exception as exc:
+                print(f"  exemplar {ex.ref}: fetch failed ({exc}); skipped this run")
+                continue
+            if it is None:
+                print(f"  exemplar {ex.ref}: not found; skipped")
+                continue
+            conn.execute("INSERT INTO exemplars VALUES (?, ?, ?, ?, ?)",
+                         (ex.ref, it["title"], it["abstract"],
+                          "crossref" if ex.doi else "arxiv-oai", store.now()))
+            if not it["abstract"]:
+                print(f"  exemplar {ex.ref}: no abstract available; using title only")
+        conn.commit()
+        row = conn.execute("SELECT * FROM exemplars WHERE ref = ?", (ex.ref,)).fetchone()
+        if row:
+            out[ex.ref] = f"{row['title']}\n\n{row['abstract'] or ''}".strip()
+    return out
+
+
+def profile_vectors(conn, profile: prof.Profile) -> tuple[list[str], np.ndarray]:
+    """Labels ("description" or exemplar refs) and unit vectors for one profile."""
+    texts = {f"profile:{profile.id}:{profile.version}": profile.description.strip()}
+    for ref, text in exemplar_texts(conn, profile).items():
+        texts[f"exemplar:{ref}"] = text
+    vecs = embed.embed(conn, texts)
+    labels = ["description" if k.startswith("profile:") else k.removeprefix("exemplar:")
+              for k in texts]
+    return labels, np.stack([vecs[k] for k in texts])
+
+
+def score_batch(conn, profile, items, vecs, index, tracked, cfg) -> list[dict]:
+    """Cosine, BM25, nomination, and re-rank for one profile over a batch of items."""
+    labels, pvecs = profile_vectors(conn, profile)
+    sims = np.stack([vecs[it["id"]] for it in items]) @ pvecs.T
+    cos, best = sims.max(axis=1), sims.argmax(axis=1)
+    all_bm25 = index.scores(bm25.profile_query(profile))
+    bm = np.array([all_bm25[index.pos[it["id"]]] for it in items])
+
+    k = max(cfg["top_k"], math.ceil(cfg["top_frac"] * len(items)))
+    nominated: dict[int, list[str]] = {}
+    for how, s in (("cosine", cos), ("bm25", bm)):
+        for i in np.argsort(-s)[:k]:
+            if how == "bm25" and s[i] <= 0:
+                continue
+            nominated.setdefault(int(i), []).append(how)
+    for i, it in enumerate(items):
+        if prof.tracked_matches(it["authors"], tracked):
+            nominated.setdefault(i, []).append("author")
+
+    idx = sorted(nominated)
+    scores = dict(zip(idx, rerank.score(profile.description.strip(),
+                                        [embed.paper_text(items[i]) for i in idx])))
+    return [{"item": it, "cosine": float(cos[i]), "cosine_best": labels[best[i]],
+             "bm25": float(bm[i]), "nominated_by": ",".join(nominated[i]) if i in nominated else None,
+             "rerank": scores.get(i)} for i, it in enumerate(items)]
+
+
+def run(conn, run_id: int) -> tuple[dict, dict]:
+    """Screen everything not yet screened. Returns (stats, stage seconds)."""
+    cfg = env.config()["screen"]
+    thresholds = cfg["thresholds"]
+    data = prof.load()
+    em, rm = embed.model_name(), rerank.model_name()
+    stages = {"embed": 0.0, "screen": 0.0, "rerank": 0.0}
+    stats = {"items": 0, "embedded": 0, "screened_pairs": 0, "reranked_pairs": 0, "highlighted": 0}
+
+    items = load_items(conn)
+    if not items:
+        return stats, stages
+    t = time.perf_counter()
+    vecs = embed.embed(conn, {it["id"]: embed.paper_text(it) for it in items})
+    stages["embed"] = time.perf_counter() - t
+    stats["items"] = len(items)
+
+    t = time.perf_counter()
+    index = bm25.Index([it["id"] for it in items], [bm25.paper_text(it) for it in items])
+    stages["screen"] += time.perf_counter() - t
+
+    for p in data.profiles:
+        done = {r[0] for r in conn.execute(
+            "SELECT item_id FROM screens WHERE profile_id = ? AND profile_version = ? "
+            "AND embed_model = ? AND rerank_model = ?", (p.id, p.version, em, rm))}
+        batch = [it for it in items if it["id"] not in done]
+        if not batch:
+            continue
+        threshold = thresholds.get(p.id, thresholds["default"])
+        t = time.perf_counter()
+        rows = score_batch(conn, p, batch, vecs, index, data.tracked_authors, cfg)
+        elapsed = time.perf_counter() - t
+        n_rr = sum(r["rerank"] is not None for r in rows)
+        n_hi = 0
+        for r in rows:
+            hi = r["rerank"] is not None and r["rerank"] >= threshold
+            n_hi += hi
+            conn.execute("INSERT OR REPLACE INTO screens VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+                         (r["item"]["id"], p.id, p.version, em, rm, r["cosine"], r["cosine_best"],
+                          r["bm25"], r["nominated_by"], r["rerank"], threshold, int(hi), run_id))
+        conn.commit()
+        stages["rerank"] += elapsed
+        stats["screened_pairs"] += len(rows)
+        stats["reranked_pairs"] += n_rr
+        stats["highlighted"] += n_hi
+        print(f"  {p.id:24s} batch {len(batch):5d}  re-ranked {n_rr:4d}  highlighted {n_hi:3d}"
+              f"  (threshold {threshold}, {elapsed:.1f} s)")
+    return stats, {k: round(v, 2) for k, v in stages.items()}