researchbee / asjc_classifier.py
nikeshn's picture
Upload 47 files
1916d8e verified
Raw History Blame Contribute Delete
8.61 kB
"""
Abstract -> ASJC subject codes, for populating ResearchBee's recommendation filter.
from asjc_classifier import AsjcClassifier
clf = AsjcClassifier("asjc_model.joblib") # local model only
clf = AsjcClassifier("asjc_model.joblib", llm=True) # + LLM adjudication
result = clf.classify(title=..., abstract=...)
result.codes # ['3309', '1710'] β€” feed straight to recommend(asjc=...)
result.method # 'linear+llm' | 'linear-topk'
result.candidates # what the LLM was asked to choose between
Architecture, and why it is not the obvious one
-----------------------------------------------
The obvious design is "run the linear model, escalate to an LLM when confidence
is low". Measured on the held-out abstract set, that design cannot work: the
top-1 probability averages 0.367 when the prediction is correct and 0.330 when
it is wrong. There is no threshold separating them β€” the model is wrong often
and unbothered about it.
What the model *is* good at is narrowing. A correct code appears within:
top-1 30%
top-3 70%
top-5 95% <- and flat thereafter
So it is used as a candidate generator, not a decider. The linear model cuts 334
codes to 5 with ~95% recall; the LLM picks among those 5. That inverts the usual
cost argument in our favour β€” the prompt carries five labelled options rather
than a 334-entry taxonomy, so the call is small, fast and hard to get wrong.
Degradation without an API key is graceful: `codes` becomes the top 3, which the
recommender applies as an OR filter. Recall is 70% rather than 95%, and the
candidate pool is broader, but nothing breaks.
Retrain after the OpenAlex harvest
----------------------------------
The shipped model is trained on journal titles, because that is the labelled text
available without a network. It is applied to abstracts. `harvest_openalex.py`
produces the matched corpus; retraining on it and re-running `asjc_eval.jsonl`
measures the gain on the same 20 abstracts, so the upgrade is quantified rather
than assumed.
"""
from __future__ import annotations
import json
import os
import re
from dataclasses import dataclass, field
from pathlib import Path
import joblib
import numpy as np
from scipy import sparse
__all__ = ["AsjcClassifier", "AsjcResult"]
#: Where the linear model's recall saturates. Measured, not guessed.
DEFAULT_K = 5
#: Codes returned when no LLM is available. top-3 = 70% recall.
FALLBACK_K = 3
@dataclass
class AsjcResult:
codes: list[str]
names: list[str]
method: str
candidates: list[dict] = field(default_factory=list)
llm_error: str | None = None
def __bool__(self) -> bool:
return bool(self.codes)
_PROMPT = """You are classifying an academic manuscript into Scopus ASJC subject areas.
A statistical model has narrowed 334 possible subject areas to the {n} candidates \
below. One or more of them is correct.
Candidates:
{candidates}
Manuscript title: {title}
Abstract: {abstract}
Choose the 1-3 candidates that best describe where this manuscript belongs, most \
appropriate first. Choose only from the candidate list β€” do not invent codes. If \
the manuscript is genuinely interdisciplinary, return more than one; if it clearly \
belongs to a single area, return just that one.
Respond with JSON only, no other text:
{{"codes": ["NNNN", ...], "reason": "one short sentence"}}"""
class AsjcClassifier:
def __init__(
self,
model_path: str | Path = "asjc_model.joblib",
llm: bool = False,
llm_model: str | None = None,
api_key: str | None = None,
):
self.bundle = joblib.load(model_path)
self.classes: np.ndarray = self.bundle["classes"]
self.names: dict[str, str] = self.bundle["names"]
self.parent_of: dict[str, str] = self.bundle["parent_of"]
self.trained_on: str = self.bundle.get("corpus", "unknown")
self._client = None
# Model id is configuration, not a constant β€” set ASJC_LLM_MODEL to
# whichever model you are entitled to. Verify it against your account's
# available models rather than trusting a default baked in here.
self.llm_model = llm_model or os.getenv("ASJC_LLM_MODEL", "claude-sonnet-4-5")
if llm:
key = api_key or os.getenv("ANTHROPIC_API_KEY")
if key:
try:
import anthropic
self._client = anthropic.Anthropic(api_key=key)
except ImportError:
pass # falls back to top-k; reported in AsjcResult.method
# -- stage 1: narrow ----------------------------------------------------
def candidates(self, text: str, k: int = DEFAULT_K) -> list[dict]:
"""The linear model's top-k, with names attached. ~95% recall at k=5."""
X = sparse.hstack([
self.bundle["word"].transform([text]),
self.bundle["char"].transform([text]),
]).tocsr()
p = self.bundle["clf"].predict_proba(X)[0]
order = np.argsort(-p)[:k]
return [
{
"code": str(self.classes[i]),
"name": self.names.get(str(self.classes[i]), str(self.classes[i])),
"score": round(float(p[i]), 4),
}
for i in order
]
# -- stage 2: adjudicate ------------------------------------------------
def classify(
self,
title: str = "",
abstract: str = "",
k: int = DEFAULT_K,
max_codes: int = 3,
) -> AsjcResult:
text = f"{title} {title} {abstract}".strip()
if not text:
return AsjcResult([], [], "empty")
cands = self.candidates(text, k=k)
if not cands:
return AsjcResult([], [], "no-candidates")
if self._client is None:
picked = [c["code"] for c in cands[:FALLBACK_K]]
return AsjcResult(picked, [self.names.get(c, c) for c in picked],
"linear-topk", cands)
listing = "\n".join(f" {c['code']} {c['name']}" for c in cands)
prompt = _PROMPT.format(
n=len(cands), candidates=listing,
title=title or "(not supplied)",
abstract=(abstract or "(not supplied)")[:4000],
)
try:
msg = self._client.messages.create(
model=self.llm_model,
max_tokens=300,
temperature=0,
messages=[{"role": "user", "content": prompt}],
)
raw = "".join(b.text for b in msg.content if getattr(b, "type", "") == "text")
picked = self._parse(raw, {c["code"] for c in cands})[:max_codes]
if picked:
return AsjcResult(picked, [self.names.get(c, c) for c in picked],
"linear+llm", cands)
err = "LLM returned no valid code from the candidate set"
except Exception as e: # noqa: BLE001
err = f"{type(e).__name__}: {e}"
# Any LLM failure degrades to the statistical answer rather than raising.
# A broader-than-ideal subject filter is a far better outcome for the
# caller than a recommendation request that errors out.
picked = [c["code"] for c in cands[:FALLBACK_K]]
return AsjcResult(picked, [self.names.get(c, c) for c in picked],
"linear-topk", cands, llm_error=err)
@staticmethod
def _parse(raw: str, allowed: set[str]) -> list[str]:
"""
Pull codes out of the reply and hard-filter to the candidate set.
The filter is the point: it makes an invented or hallucinated code
structurally impossible to return, so the worst case is fewer codes
rather than a wrong one that silently poisons the journal filter.
"""
codes: list[str] = []
m = re.search(r"\{.*\}", raw, re.S)
if m:
try:
codes = [str(c).strip() for c in json.loads(m.group(0)).get("codes", [])]
except (json.JSONDecodeError, AttributeError, TypeError):
codes = []
if not codes:
codes = re.findall(r"\b\d{4}\b", raw)
out: list[str] = []
for c in codes:
if c in allowed and c not in out:
out.append(c)
return out