This model is a biomedical entity-linking (EL) cross-encoder used by BASILISK to rerank UMLS concept candidates for a mention in context. It answers a focused question:

Given a biomedical mention, the surrounding text, and one candidate UMLS concept, is this candidate the correct concept for that mention?

The model is not a standalone entity linker. It is the contextual scoring stage in a larger EL pipeline. Candidate generators first retrieve possible UMLS concepts using lexical, sparse, embedding, and ontology-based signals. This cross-encoder then compares each mention/context pair directly against each candidate concept and assigns a compatibility score. BASILISK uses that score as one feature when choosing the final linked concept.

It is fine-tuned from:

microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext
Base revision: e1354b7a3a09615f6aba48dfad4b7a613eef7062

Published model repo:

Bam3752/basilisk-el-ce-biomedbert-ab-v1

Pinned release revision:

9b11786be83352196058af654d255e8441a75356

Task

Binary classification over a (mention/context, candidate concept) pair:

label 1: candidate concept is correct for the mention in context
label 0: candidate concept is incorrect

The model is trained as a reranker, not as an open-ended concept retriever. It expects the candidate concept to already be supplied. This distinction matters: the model can prefer the best candidate among candidates it sees, but it cannot recover a missing correct CUI if candidate generation failed upstream.

In BASILISK runtime, candidate ranking score is:

score = sigmoid(logit_1 - logit_0)

Higher scores mean the model believes the candidate is more compatible with the mention and context. In a candidate list, BASILISK can sort candidates by this score, combine it with other retrieval/ontology signals, or apply calibration before thresholding.

How It Works

This is a cross-encoder. The mention/context and candidate concept are passed through the transformer together as a paired sequence. That lets attention flow between:

the mention text
nearby biomedical context
candidate concept name
candidate semantic types

This is slower than a bi-encoder because each mention-candidate pair requires a forward pass, but it is better suited for fine-grained disambiguation. For example, the same surface form can refer to a disease, procedure, chemical, gene, or broad biomedical topic depending on context. The cross-encoder can inspect the surrounding text and decide whether the candidate's name and semantic type are contextually plausible.

Typical BASILISK flow:

Detect candidate biomedical mentions in a text step.
Retrieve a high-recall set of possible UMLS concepts.
Attach metadata such as candidate name and TUIs.
Score each mention-candidate pair with this model.
Combine the cross-encoder score with lexical, sparse, embedding, ontology, and calibration signals.
Select or reject the final mapping based on the full EL stack.

The model's output is best interpreted as a compatibility score for reranking, not as an independently validated medical fact.

Input Serialization

The model is trained and used with paired text:

Left sequence: mention: <mention> ; context: <local_context>
Right sequence: candidate: <concept_name> ; tuis: <TUI list>

This format should be kept consistent between training and inference.

The left sequence gives the model both the mention surface form and enough context to disambiguate it. The right sequence gives the candidate label and UMLS semantic type information. TUIs are important because many biomedical strings are lexically similar but semantically different.

Example:

text

Left:  mention: aspirin ; context: patient was started on aspirin for secondary prevention
Right: candidate: Aspirin ; tuis: T121,T109

In this example, the context supports a chemical/pharmacologic interpretation. A candidate with a mismatched semantic type or unrelated concept name should receive a lower score.

Inputs are truncated to max_length = 256 during training. For best results, keep the context local and relevant rather than passing an entire document.

Training Data and Splits

Data was generated by the BASILISK EL training-set pipeline from silver entity links and hard-negative UMLS candidate pairs. The processed training rows are not redistributed with this model because some upstream resources, especially UMLS-derived concept metadata, are subject to their original access terms and licenses.

Training row format:

mention, context: the surface form and local biomedical text span
candidate_cui, candidate_name, candidate_tuis: the candidate concept being scored
label: 1 for the selected silver-positive concept, 0 for a hard negative
source: silver_positive or hard_negative
hard_negative_type: dense near miss, sparse lexical confuser, ontology-incompatible candidate, or generic distractor where applicable

Data construction:

Positive examples come from BASILISK's concept extraction/linking pipeline when the mention is mapped above the configured confidence threshold.
Negative examples are sampled from the mention candidate table and the SQLite sparse retriever, then filtered/deduplicated against the positive CUI.
Split assignment is deterministic from stable hashes over the question/step/mention/candidate identity, giving train/dev/test splits that can be regenerated.
UMLS concept names, CUIs, TUIs, and source vocabularies must be obtained under the user's own UMLS license.

Split sizes:

Train: 281,902
Dev: 35,420
Test: 35,100

Class balance:

Train positives: 124,487, negatives: 157,415 (neg_per_pos=1.2645)
Dev positives: 15,608, negatives: 19,812 (neg_per_pos=1.2693)
Test positives: 15,405, negatives: 19,695 (neg_per_pos=1.2785)

Dataset Preparation

The processed EL training set is not included in this model repo. It was prepared from local BASILISK source records and a licensed local UMLS parquet export. Users who want to reproduce the data should rebuild it from the original sources under their own UMLS license rather than relying on redistributed processed rows.

Source Records

The starting point was a local table of biomedical reasoning traces. The published run used artifacts/all_cots.parquet.

Each source record was expected to contain:

steps: a list of biomedical text spans, or a single text span
one identifier field, when available: question_id, id, qid, question, or questionId

The identifier was used only for traceability and deterministic splitting. If no identifier was present, a fallback ID was generated from the input row index. The text in steps became the local context for entity-linking examples.

UMLS and Candidate Sources

Candidate concepts were generated from a local UMLS parquet export. The build used all locally available source vocabularies rather than restricting to a small SAB subset. UMLS-derived fields in the training data include:

CUI
preferred or matched concept name
semantic type identifiers, also called TUIs
source vocabulary membership used internally by the linker

Because those fields are derived from UMLS and UMLS source vocabularies, the processed rows are not redistributed here. Reproduction requires users to obtain UMLS under their own license and build the same kind of local parquet store.

Mention Extraction

Each text step was passed through the BASILISK concept extraction stack. The published build used the quality_first extraction profile and entity-linking extractor version v2.

Mention detection combined multiple signals:

lexical matching against biomedical terminology
NER ensemble output
SapBERT-based biomedical embedding retrieval
BiomedBERT-based contextual reranking
ontology constraints over candidate semantic types and UMLS relationships
cross-encoder reranking of candidate concepts

For each detected mention, the build retained:

mention surface text
character start and end offsets
mention detector names
expected semantic types inferred from the context/span
local context text from the source step
a candidate table of possible UMLS concepts

Positive Pair Construction

A positive example was created when a mention was mapped to a selected UMLS CUI. The selected CUI became the positive candidate for that mention in that context.

Each positive row used:

the mention surface form and local context as the left side
the selected candidate name and TUIs as the right side
label = 1
source = silver_positive
the linker confidence as score_hint

These are silver labels, not manually adjudicated labels. They reflect the output of the BASILISK entity-linking pipeline after retrieval, reranking, and ontology filtering.

Negative Pair Construction

Negative examples were generated from candidate concepts that were plausible but not selected as the mapped CUI. This made the task a reranking problem rather than a random binary classification problem.

Hard negatives came from:

near-miss candidates already present in the mention candidate table
sparse lexical confusers retrieved by the local sparse index
candidates with similar names but incompatible semantic types
generic UMLS concepts that can attract false matches

Each negative row used the same mention and context as the positive row, but paired it with an incorrect candidate concept:

label = 0
source = hard_negative
hard_negative_type describing why the candidate was included

The negative type was assigned as:

dense_near_miss: a high-scoring neural or hybrid candidate that was not selected
sparse_lexical_confuser: a lexical/sparse retrieval candidate that matched the mention text but was not the selected concept
ontology_incompatible: a candidate whose TUIs did not match the expected semantic types for the mention
generic_distractor: a broad/generic UMLS concept used as a distractor

Deduplication and Splitting

Rows were deduplicated using:

question ID
step index
mention start and end offsets
candidate CUI
label

After deduplication, rows were assigned to train/dev/test with a deterministic hash over:

question ID
step index
mention text
candidate CUI
candidate name

The split rule was:

hash bucket 0-79: train
hash bucket 80-89: dev
hash bucket 90-99: test

This makes the split reproducible while keeping related pair identity stable across reruns.

Row Schema

Each output row contains:

split: train, dev, or test
question_id: source record identifier
step_index: index of the text step within the source record
mention: mention surface text
mention_start, mention_end: character offsets within the context
context: local biomedical text span
expected_tuis: semantic types expected for the mention/span
mention_detectors: detector components that found the mention
candidate_cui: UMLS concept ID
candidate_name: candidate concept string
candidate_tuis: UMLS semantic type IDs for the candidate
label: 1 for silver positive, 0 for hard negative
source: silver_positive or hard_negative
hard_negative_type: negative category, or null for positives
score_hint: confidence/retrieval score used for auditing

Published Build Settings

The published data build used:

Source file: artifacts/all_cots.parquet
Source records processed: 10,000
Text steps processed: 167,155
Extracted mentions: 537,436
Mapped mentions: 155,501
UMLS source vocabularies: all available local source vocabularies
Positive confidence threshold: 0.0
Maximum mentions per step: unbounded for the run
Maximum negatives per positive: unbounded before deduplication/filtering
Extractor profile: quality_first
Entity-linking extractor version: v2
NER ensemble: enabled
SapBERT reranking: enabled
BiomedBERT reranking: enabled
Cross-encoder reranking: enabled
Ontology constraints: enabled
Calibration during data generation: disabled

The resulting dataset files were:

train.jsonl
dev.jsonl
test.jsonl
audit_subset.jsonl
manifest.json

The published build produced:

Train: 281,902
Dev: 35,420
Test: 35,100
Total pairs: 352,422
Train positives: 124,487
Train negatives: 157,415
Dev positives: 15,608
Dev negatives: 19,812
Test positives: 15,405
Test negatives: 19,695
Data build hash in the pair manifest: 20bc770afebe4cefc20f5382262a12de22fcf0b64299c1f6283158a97d96261f

Minor differences can occur if the UMLS release, source vocabularies, extraction models, reranker revisions, or BASILISK code revision differ.

Training Configuration

The model was fine-tuned as a binary sequence-pair classifier from BiomedBERT. The classifier head predicts whether a supplied candidate concept is the correct link for the supplied mention/context pair.

Key settings:

Profile: quality_max
Objective (this run): binary cross-entropy
Epochs: 4
Batch size: 24
Gradient accumulation: 4 (effective batch 96)
Max length: 256
Learning rate: 2e-5 with cosine schedule
Warmup ratio: 0.1
Weight decay: 0.01
Seed: 13
Device: mps

Model selection:

Metric: dev_f1
Best checkpoint step: 5500
Best selected dev F1: 0.9174

Training used hard negatives rather than random negatives. This makes the classification task harder and closer to real candidate reranking, because many negative candidates are plausible lexical or semantic alternatives.

Calibration

A calibration artifact is provided at:

calibration/biomedbert_ab.json

Calibration method:

global temperature scaling + bucket offsets
temperature: 1.25

Dev calibration effect:

Raw ECE: 0.017000 -> Calibrated ECE: 0.003967
Raw Brier: 0.052839 -> Calibrated Brier: 0.051863

Calibration does not change what the model has learned. It adjusts the probability scale so that a score like 0.80 better corresponds to observed correctness frequency on the dev distribution. This is useful when scores are thresholded or combined with other BASILISK signals.

Because calibration was fit on the BASILISK dev split, it may not transfer perfectly to another candidate generator, UMLS release, specialty domain, or input distribution. If the model is used in a different EL pipeline, recalibrating on a local validation set is recommended.

Evaluation Summary

Raw metrics:

Split	Accuracy	Precision	Recall	F1	ECE	Brier
Dev	0.9292	0.9446	0.8918	0.9174	0.017000	0.052839
Test	0.9282	0.9456	0.8875	0.9156	0.017083	0.053027

Calibrated metrics:

Split	Accuracy	Precision	Recall	F1	ECE	Brier
Dev (calibrated)	0.9296	0.9455	0.8916	0.9178	0.003967	0.051863
Test (calibrated)	0.9282	0.9458	0.8874	0.9156	0.006504	0.052334

Metric interpretation:

Accuracy measures binary correctness over the prepared pair dataset.
Precision measures how often high-scoring positive predictions are correct.
Recall measures how often silver-positive pairs are recovered.
F1 balances precision and recall.
ECE measures probability calibration error; lower is better.
Brier score measures probabilistic prediction quality; lower is better.

These metrics evaluate pair classification on BASILISK-generated candidate pairs. They do not directly measure end-to-end entity-linking accuracy on arbitrary text, because end-to-end accuracy also depends on mention detection, candidate generation recall, ontology filtering, and final decision thresholds.

Usage (Transformers)

python

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

repo_id = "Bam3752/basilisk-el-ce-biomedbert-ab-v1"
revision = "9b11786be83352196058af654d255e8441a75356"

tokenizer = AutoTokenizer.from_pretrained(repo_id, revision=revision)
model = AutoModelForSequenceClassification.from_pretrained(repo_id, revision=revision)
model.eval()

left = "mention: aspirin ; context: patient was started on aspirin for secondary prevention"
right = "candidate: Aspirin ; tuis: T121,T109"
enc = tokenizer(left, right, truncation=True, padding=True, max_length=256, return_tensors="pt")

with torch.no_grad():
    logits = model(**enc).logits.squeeze(0)
    # BASILISK CE probability
    p = torch.sigmoid(logits[1] - logits[0]).item()

print(f"ce_probability={p:.4f}")

Intended Use

Intended for:

biomedical entity-linking candidate reranking in BASILISK
UMLS candidate scoring for detected biomedical mentions
high-recall candidate sets where contextual disambiguation is needed
research and engineering experiments on biomedical EL reranking
use as a component in a larger entity-linking system with candidate generation and validation

Not intended for:

standalone medical diagnosis or clinical decision support
domains far outside biomedical literature/terminologies
direct use without candidate generation and ontology constraints
generating UMLS candidates from raw text by itself
replacing UMLS licensing, source vocabulary review, or human validation in sensitive workflows

Practical Use Notes

For good results:

Use the same serialization format used during training.
Provide short, relevant context around the mention.
Include candidate TUIs when available.
Score multiple candidates for the same mention and compare them rather than relying on one absolute score.
Recalibrate if the upstream candidate generator or target corpus differs substantially from BASILISK.

The model is most useful when the candidate list already has reasonable recall. If the correct CUI is absent from the candidates, this model cannot select it.

Limitations

Performance depends on candidate generator recall.
Calibration is fit on this training pipeline's dev distribution and may drift on different data.
Ambiguous mentions and rare concepts may still require ontology constraints or additional signals.
Labels are silver labels produced by an automated pipeline, not manually adjudicated concept annotations.
UMLS concept coverage and source vocabulary choices can affect both training and inference behavior.
The model can over-prefer candidates that look lexically or semantically similar to training examples.
Very long contexts are truncated, so evidence outside the retained context window is not visible to the model.
Scores are not medical truth judgments; they only estimate candidate compatibility for entity linking.

Ethics and Safety

This is a research/engineering model for NLP ranking, not a medical device. Outputs may be wrong and require human review in sensitive workflows.

Potential risks include incorrect biomedical concept mappings, overconfident scores on out-of-distribution text, and propagation of errors from upstream candidate generation. In clinical, regulatory, or other high-impact settings, outputs should be treated as assistive signals and audited by qualified humans.

The model uses UMLS-derived metadata during training and inference. Users are responsible for complying with UMLS and source vocabulary license terms when rebuilding data or deploying UMLS-backed systems.

Reproducibility Notes

Important reproducibility factors:

base model revision
UMLS release and local parquet conversion
enabled source vocabularies
mention extraction settings
candidate retrieval stack
reranker model revisions
deterministic split hashing
calibration data distribution