Spaces:
Running
Running
File size: 8,606 Bytes
1916d8e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 | """
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
|