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

BASILISK EL Cross-Encoder (BiomedBERT AB) v1

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:

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:

Calibrated metrics:

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)

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