repository·main·Indexed 42 days ago
https://github.com/cambridgeltl/sapbertSapBERT provides self-alignment pretraining for BERT models to improve biomedical entity representations, supporting English and cross-lingual biomedical entity linking tasks. It includes pretrained models based on PubMedBERT and XLM-Roberta, tools for extracting embeddings from entity names, and scripts for evaluation on datasets such as NCBI-disease, BC5CDR, COMETA, and AskAPatient.
Finetuning data must be formatted as a text file where each line represents a synonym pair associated with a specific label ID. The format for each line is:
label_id || entity_name_1 || entity_name_2
entity_name_1 and entity_name_2 are synonym pairs sampled from a labeled dataset. If a single label is associated with multiple entity names, you should include all possible pairwise combinations for that label.
label_id || entity_name_1 || entity_name_2SapBERT offers different pretrained models depending on your language requirements and desired pooling method:
cambridgeltl/SapBERT-from-PubMedBERT-fulltext: Standard SapBERT. Use the [CLS] token (before the pooler) as the input representation.cambridgeltl/SapBERT-from-PubMedBERT-fulltext-mean-token: Standard SapBERT using mean-pooling across all tokens instead of the [CLS] token.cambridgeltl/SapBERT-UMLS-2020AB-all-lang-from-XLMR: Base version using xlm-roberta-base.cambridgeltl/SapBERT-UMLS-2020AB-all-lang-from-XLMR-large: Large version using xlm-roberta-large.For all cross-lingual models, use the [CLS] token (before the pooler) as the representation.
MedMentions is not provided directly due to UMLS licensing. To use it, you must manually create a dictionary file by downloading the UMLS2017AA version's MRCONSO.RRF file and extracting CUI-name pairs.
Each line in the dictionary must follow this format:
CUI||name
Example:
C0079564||htlv-ii rex protein
C0162999||oleoylamine
C0347197||benign mouth neoplasmAfter extraction and de-duplication, the final reference dictionary should contain approximately 7.4M individual lines (derived from ~3.4M unique CUIs and 14.8M raw lines).
generate_pretraining_data.ipynb notebook to process the knowledge sources into the required training format.The code is tested with the following environment:
Refer to requirements.txt in the repository for the full list of dependencies.
To perform Cross-Lingual (XL-BEL) evaluation, run the ./eval_scripts_xlbel_test.sh <device_index> script. You must modify the script to specify your specific model path, language, and other required parameters.
xl_bel/xlbel_v0.0: Contains mentions and their CUIs without context. Use this to replicate results from [Liu et al., ACL 2021].xl_bel/xlbel_v1.0: Contains XL-BEL data with context. In this version, the target entity is explicitly marked using the <tgt>TARGET</tgt> tag.# Run XL-BEL evaluation on GPU 0
./eval_scripts_xlbel_test.sh 0To run the pretraining script, navigate to the train/ directory and execute pretrain.sh. You must specify the GPU devices as a comma-separated list.
cd train/
./pretrain.sh 0,1 Note: Training data must be extracted from UMLS following the instructions in training_data/generate_pretraining_data.ipynb.
To finetune on a custom dataset, format your data as follows:
concept_id || entity_name_1 || entity_name_2
Where entity_name_1 and entity_name_2 are synonym pairs belonging to the same concept_id. For multiple names per concept, include all pairwise combinations.
For cross-lingual tuning using general domain parallel data (such as muse or wiki titles), use the scripts provided in the training_data/general_domain_parallel_data/ directory. An example script is train/xling_train.sh.
To evaluate SapBERT on monolingual datasets, use the provided shell scripts. The scripts take a device index (e.g., 0 for GPU 0) as a positional argument.
./eval_scripts_ncbi_bc5cdr.sh <device_index>../eval_scripts_cometa.sh <device_index>../eval_scripts_askapatient_10cv.sh <device_index>.Before running evaluation, download the datasets from this Dropbox link, unzip them, and place the resulting data directory under the evaluation/ folder.
# Example: Evaluate NCBI-disease on GPU 0
./eval_scripts_ncbi_bc5cdr.sh 0
# Example: Evaluate COMETA on GPU 0
./eval_scripts_cometa.sh 0
# Example: Evaluate AskAPatient on GPU 0
./eval_scripts_askapatient_10cv.sh 0evaluation/README.md. The evaluation/xl_bel/ directory contains the XL-BEL benchmark for cross-lingual evaluation.The agg_mode parameter determines how the last_hidden_state from the encoder is reduced to a single embedding vector for each input sequence. This is critical for how the model represents entities.
Supported modes:
"cls": Uses the [CLS] token embedding (the first token in the sequence: last_hidden_state[:, 0])."mean_all_tok": Computes the mean across all tokens in the sequence."mean": Computes a weighted mean of the token embeddings using the attention_mask to ensure only valid tokens contribute to the average.# Example of how agg_mode affects the internal logic:
# If agg_mode == 'cls':
# query_embed = last_hidden_state[:, 0]
# If agg_mode == 'mean_all_tok':
# query_embed = last_hidden_state.mean(1)
# If agg_mode == 'mean':
# query_embed = (last_hidden_state * mask).sum(1) / mask.sum(-1)scipy.spatial.distance.cdist method is suitable for small to medium datasets. For large-scale entity search, you should switch to FAISS (Facebook AI Similarity Search) to handle high-dimensional vector indexing and retrieval efficiently.