bioelx-nlp/bioelx_retriever

BioELX Retriever

A multilingual biomedical entity retriever (candidate generator) for biomedical entity linking. It is an XLM-RoBERTa-large encoder fine-tuned SapBERT-style so that a mention and its correct concept name land close together in embedding space, enabling cross-lingual dense retrieval against a biomedical knowledge base (e.g. UMLS).

• Architecture: XLM-RoBERTa-large (24 layers, hidden 1024, ~560M params)
• Embedding: [CLS] token of last_hidden_state, L2-normalized; retrieve by cosine / inner-product nearest neighbor
• Output dim: 1024
• Languages: multilingual

Usage

import torch
from transformers import AutoModel, AutoTokenizer

name = "bioelx-nlp/bioelx_retriever"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModel.from_pretrained(name).eval()

def embed(texts):
    enc = tok(texts, padding=True, truncation=True, max_length=25, return_tensors="pt")
    with torch.no_grad():
        out = model(**enc)
    cls = out.last_hidden_state[:, 0, :]          # [CLS]
    return torch.nn.functional.normalize(cls, dim=-1)

q = embed(["diabetes mellitus"])
kb = embed(["diabetes", "hypertension", "type 2 diabetes mellitus"])
scores = q @ kb.T                                  # cosine similarity
print(scores)

Encode all KB concept names once, index the normalized embeddings, then retrieve the top-k nearest concepts for each mention embedding.

Citation