Spaces:
Running
Running
Download asjc_classifier.py from nikeshn/researchbee: direct link, hf CLI and curl.
- Browser
- Download file 8.61 kB
-
https://huggingface.co/spaces/nikeshn/researchbee/resolve/main/asjc_classifier.py
- Command line
-
hf download hf://spaces/nikeshn/researchbee/asjc_classifier.py
-
curl -L -o asjc_classifier.py https://huggingface.co/spaces/nikeshn/researchbee/resolve/main/asjc_classifier.py
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 | |
| 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) | |
| 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 | |