Views
No views yet
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).[CLS] token of last_hidden_state, L2-normalized; retrieve by cosine / inner-product nearest neighbor1import torch
2from transformers import AutoModel, AutoTokenizer
3
4name = "bioelx-nlp/bioelx_retriever"
5tok = AutoTokenizer.from_pretrained(name)
6model = AutoModel.from_pretrained(name).eval()
7
8def embed(texts):
9 enc = tok(texts, padding=True, truncation=True, max_length=25, return_tensors="pt")
10 with torch.no_grad():
11 out = model(**enc)
12 cls = out.last_hidden_state[:, 0, :] # [CLS]
13 return torch.nn.functional.normalize(cls, dim=-1)
14
15q = embed(["diabetes mellitus"])
16kb = embed(["diabetes", "hypertension", "type 2 diabetes mellitus"])
17scores = q @ kb.T # cosine similarity
18print(scores)