pubmed-research-classifier

Classify PubMed articles as research or non-research using a trained MLP on top of EMBO/ModernBERT-neg-sampling-PubMed embeddings.

v0.3.0 ships Workflow v2 production weights (amplified research definition, default decision threshold τ = 0.75) plus an optional PMID label cache backed by the private Hub dataset EMBO/pubmed-research-classifier (~30M precomputed labels, DuckDB lookup).

Model weights, StandardScaler, and publication-type vocabulary are bundled — no external model downloads are needed for the embedding-mode API.

Installation

# Embedding mode only (pass precomputed ModernBERT vectors)
pip install pubmed-research-classifier

# Text mode (package embeds internally)
pip install "pubmed-research-classifier[embed]"

# Precomputed PMID label cache (DuckDB + Hugging Face Hub)
pip install "pubmed-research-classifier[cache]"

# Everything
pip install "pubmed-research-classifier[all]"

pip install -U "pubmed-research-classifier>=0.3.0"

Hugging Face token (label cache / Hub publish)

The lookup table is a private dataset. You need a Hub token:

# Read access — enough for load_label_cache / lookup_pmid
export HF_TOKEN=hf_xxxxxxxx
# (alias also accepted: HUGGING_FACE_HUB_TOKEN)

# Or interactive login
huggingface-cli login

• Read token → download / look up labels.
• Write token → monthly pubmed-rc-publish-labels --upload (maintainer only).

Ask an EMBO data owner for access if you get 401/403. Do not commit tokens to git; use .env (untracked), CI secrets, or a password manager.

Override the on-disk cache directory with:

export PUBMED_RC_CACHE_DIR=/path/to/cache

Research definition (v2)

Labels follow the amplified Workflow v2 definition used in the Scientometrics pipeline:

• Research — methodology-backed work with data/analysis; systematic reviews / meta-analyses; resources (datasets, software, code); methods / theoretical models; clinical and observational designs that report such findings.
• Non-research — narrative reviews without new analysis; perspectives, primers, letters/editorials (opinion-only), errata / retractions / news-like items.

p_nr / probability is P(non-research). Default threshold 0.75: label is non-research when p_nr >= 0.75, else research.

Training provenance: MLP_with_pt (ModernBERT title + abstract + scalars + MeSH publication-type multi-hot), seed 42, frozen split from the paper repo (models/v2/).

Quick start

Text mode

from pubmed_research_classifier import classify

result = classify({
    "title": "Structural basis of CRISPR-Cas9 activity",
    "abstract": "We report crystal structures of Cas9 ...",
    "pub_types": ["Journal Article"],
    "n_authors": 8,
    "n_refs": 42,
})
# {"label": "research", "p_nr": 0.018}

Embedding mode

Pre-compute embeddings with EMBO/ModernBERT-neg-sampling-PubMed using normalize_embeddings=True, then pass them directly:

from pubmed_research_classifier import classify
import numpy as np

result = classify({
    "title_emb":     title_embedding,    # np.ndarray, shape (768,)
    "abstract_emb":  abstract_embedding, # np.ndarray, shape (768,); zeros if absent
    "has_abstract":  True,
    "length_title":  52,
    "length_abstract": 1240,
    "pub_types":     ["Journal Article"],
    "n_authors":     8,
    "n_refs":        42,
})

Batch

results = classify(records, batch_size=128)
# list in the same order as input

Custom threshold

classify(record)                 # τ = 0.75 (default)
classify(record, threshold=0.95) # higher NR precision

PMID label cache (lookup table)

Precomputed labels for ~30.4M OpenAlex–PubMed PMIDs (Hub v1.0.0, classifier 0.2.0 weights) for fast lookup before running the MLP.

pip install "pubmed-research-classifier[cache]"
export HF_TOKEN=hf_...   # read access

from pubmed_research_classifier import (
    load_label_cache,
    lookup_pmid,
    classify_pmid,
)

# First call downloads the Hub CSV (~1 GB) and builds a local DuckDB index
# under ~/.cache/pubmed_research_classifier/ (or PUBMED_RC_CACHE_DIR).
cache = load_label_cache(revision="v1.0.0")
print(len(cache))  # ~30_426_295

lookup_pmid("10006576")
# {"PMID": "10006576", "class": "research",
#  "probability": 0.009..., "source": "cache"}

lookup_pmid("99999999")
# None

# Prefer cache; run the bundled MLP only on a miss
classify_pmid("10006576")
# source == "cache"

classify_pmid(
    "99999999",
    record={
        "title": "An unseen article title",
        "abstract": "Abstract text ...",
        "pub_types": ["Journal Article"],
        "n_authors": 3,
        "n_refs": 12,
    },
)
# {"PMID": "99999999", "class": "...", "probability": ..., "source": "model"}

Batch lookup (order-preserving):

rows = cache.lookup_many(["10006576", "10047518", "99999999"])
# [dict, dict, None]

probability is P(non-research), same as p_nr from classify.

Client calls never write back to Hugging Face. New classifications with source="model" stay local until a maintainer publishes a new Hub revision (see below).

For bulk joins of millions of PMIDs, query the DuckDB file or Hub CSV with DuckDB/Polars directly rather than calling lookup_pmid in a tight Python loop.

Monthly Hub refresh (maintainers)

Typical monthly job: classify new PubMed/OpenAlex PMIDs → CSV → merge into the previous DuckDB table → upload data/<new_revision>/labels.csv.

1. Produce a CSV of new labels

Columns (aliases accepted):

import csv
from pubmed_research_classifier import classify_pmid, load_label_cache

load_label_cache(revision="v1.0.0")

new_rows = []
for pmid, record in monthly_records:  # your ETL
    out = classify_pmid(pmid, record=record)
    if out["source"] == "model":      # only cache misses
        new_rows.append({
            "PMID": out["PMID"],
            "class": out["class"],
            "probability": out["probability"],
        })

with open("new_pmids.csv", "w", newline="", encoding="utf-8") as f:
    w = csv.DictWriter(f, fieldnames=["PMID", "class", "probability"])
    w.writeheader()
    w.writerows(new_rows)

2. Merge + export + upload

Needs a Hub token with write access to EMBO/pubmed-research-classifier.

CLI:

export HF_TOKEN=hf_...   # write-capable

# Dry run: merge + write local DuckDB/CSV only
pubmed-rc-publish-labels \
  --new-csv new_pmids.csv \
  --base-revision v1.0.0 \
  --new-revision v1.1.0

# Publish
pubmed-rc-publish-labels \
  --new-csv new_pmids.csv \
  --base-revision v1.0.0 \
  --new-revision v1.1.0 \
  --upload

Python:

from pubmed_research_classifier import publish_label_revision

result = publish_label_revision(
    new_csv="new_pmids.csv",
    base_revision="v1.0.0",
    new_revision="v1.1.0",
    upload=True,   # False = local DuckDB + CSV only
)
print(result["stats"])
# {'n_before': ..., 'n_after': ..., 'n_inserted': ..., 'n_updated': ..., ...}
print(result["csv_path"], result["db_path"])

Incoming PMIDs overwrite existing rows (re-score). New PMIDs are appended. Users pin load_label_cache(revision="v1.1.0") after you publish.

Input fields

Output

# classify(...)
{"label": "research",     "p_nr": 0.018}
{"label": "non-research", "p_nr": 0.921}

# lookup_pmid / classify_pmid
{"PMID": "10006576", "class": "research", "probability": 0.009, "source": "cache"}

Obtaining has_funding from PubMed XML

has_funding is True when the article's PubMed XML record contains at least one <Grant> element inside a <GrantList>. It is not the same as the "Research Support, …" publication type tags (those are a separate, coarser signal also used by the model via pub_types).

import xml.etree.ElementTree as ET

def has_funding_from_xml(article_xml: str) -> bool:
    root = ET.fromstring(article_xml)
    return len(root.findall(".//Grant")) > 0

If you omit has_funding, the package falls back to checking whether any pub_types start with "Research Support".

Changelog

0.3.0

• Optional Hub label cache ([cache]): DuckDB-backed lookup for EMBO/pubmed-research-classifier (~30M PMIDs).
• APIs: load_label_cache, lookup_pmid, classify_pmid.
• Maintainer publish path: publish_label_revision + CLI pubmed-rc-publish-labels (merge CSV → DuckDB → Hub upload).
• Requires HF_TOKEN for Hub access (read for lookup; write for --upload).

0.2.0

• Bundle Workflow v2 MLP_with_pt weights.
• Amplified research / non-research definition; default τ = 0.75.

0.1.0

• Initial release with v1 (narrow) research definition weights.

Publishing a new version to PyPI

Artifacts land in pubmed-research-classifier/dist/.

1. Update bundled weights under src/pubmed_research_classifier/_data/ if needed.

2. Bump version in pyproject.toml and __version__ in __init__.py.

3. Update regression expectations in tests/, then:

   pip install -e ".[cache]"
   pytest -m "not integration"
   

4. Build and upload:

   pip install build twine
   python -m build
   twine upload dist/pubmed_research_classifier-0.3.0*
   

5. Verify:

   pip install "pubmed-research-classifier==0.3.0" --force-reinstall
   python -c "from pubmed_research_classifier import classify, load_label_cache; print('ok')"