Sentence Similarity
sentence-transformers
Safetensors
Transformers
English
bert
feature-extraction
text-embeddings-inference
Instructions to use Tao-AI-Informatics/NA-SapBERT with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- sentence-transformers
How to use Tao-AI-Informatics/NA-SapBERT with sentence-transformers:
from sentence_transformers import SentenceTransformer model = SentenceTransformer("Tao-AI-Informatics/NA-SapBERT") sentences = [ "That is a happy person", "That is a happy dog", "That is a very happy person", "Today is a sunny day" ] embeddings = model.encode(sentences) similarities = model.similarity(embeddings, embeddings) print(similarities.shape) # [4, 4] - Transformers
How to use Tao-AI-Informatics/NA-SapBERT with Transformers:
# Load model directly from transformers import AutoTokenizer, AutoModel tokenizer = AutoTokenizer.from_pretrained("Tao-AI-Informatics/NA-SapBERT") model = AutoModel.from_pretrained("Tao-AI-Informatics/NA-SapBERT") - Notebooks
- Google Colab
- Kaggle
Update README.md
Browse files
README.md
CHANGED
|
@@ -6,125 +6,156 @@ tags:
|
|
| 6 |
- feature-extraction
|
| 7 |
- sentence-similarity
|
| 8 |
- transformers
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 9 |
|
| 10 |
---
|
| 11 |
|
| 12 |
-
#
|
| 13 |
|
| 14 |
-
|
| 15 |
|
| 16 |
-
|
|
|
|
|
|
|
|
|
|
| 17 |
|
| 18 |
-
|
| 19 |
|
| 20 |
-
|
| 21 |
|
| 22 |
-
|
| 23 |
-
pip install -U sentence-transformers
|
| 24 |
-
```
|
| 25 |
|
| 26 |
-
|
|
|
|
|
|
|
| 27 |
|
| 28 |
-
|
| 29 |
-
|
| 30 |
-
|
| 31 |
|
| 32 |
-
|
| 33 |
-
|
| 34 |
-
|
| 35 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 36 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 37 |
|
|
|
|
|
|
|
| 38 |
|
| 39 |
-
##
|
| 40 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 41 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 42 |
```python
|
| 43 |
-
from
|
| 44 |
-
|
|
|
|
| 45 |
|
|
|
|
|
|
|
|
|
|
| 46 |
|
| 47 |
-
|
| 48 |
-
def mean_pooling(model_output, attention_mask):
|
| 49 |
-
token_embeddings = model_output[0] #First element of model_output contains all token embeddings
|
| 50 |
-
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
|
| 51 |
-
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
|
| 52 |
|
|
|
|
| 53 |
|
| 54 |
-
|
| 55 |
-
|
|
|
|
|
|
|
| 56 |
|
| 57 |
-
|
| 58 |
-
tokenizer = AutoTokenizer.from_pretrained('{MODEL_NAME}')
|
| 59 |
-
model = AutoModel.from_pretrained('{MODEL_NAME}')
|
| 60 |
|
| 61 |
-
|
| 62 |
-
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
|
| 63 |
|
| 64 |
-
|
| 65 |
-
|
| 66 |
-
model_output = model(**encoded_input)
|
| 67 |
|
| 68 |
-
|
| 69 |
-
|
| 70 |
|
| 71 |
-
|
| 72 |
-
print(sentence_embeddings)
|
| 73 |
```
|
| 74 |
|
|
|
|
| 75 |
|
|
|
|
| 76 |
|
| 77 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 78 |
|
| 79 |
-
|
| 80 |
|
| 81 |
-
|
| 82 |
|
|
|
|
|
|
|
| 83 |
|
| 84 |
-
|
| 85 |
-
|
|
|
|
| 86 |
|
| 87 |
-
|
| 88 |
|
| 89 |
-
|
| 90 |
-
```
|
| 91 |
-
{'batch_size': 256, 'sampler': 'torch.utils.data.sampler.RandomSampler', 'batch_sampler': 'torch.utils.data.sampler.BatchSampler'}
|
| 92 |
-
```
|
| 93 |
|
| 94 |
-
|
|
|
|
|
|
|
| 95 |
|
| 96 |
-
|
| 97 |
-
```
|
| 98 |
-
{'scale': 20.0, 'similarity_fct': 'cos_sim'}
|
| 99 |
-
```
|
| 100 |
|
| 101 |
-
|
| 102 |
-
```
|
| 103 |
-
{
|
| 104 |
-
"epochs": 1,
|
| 105 |
-
"evaluation_steps": 0,
|
| 106 |
-
"evaluator": "NoneType",
|
| 107 |
-
"max_grad_norm": 1,
|
| 108 |
-
"optimizer_class": "<class 'torch.optim.adamw.AdamW'>",
|
| 109 |
-
"optimizer_params": {
|
| 110 |
-
"lr": 1e-05
|
| 111 |
-
},
|
| 112 |
-
"scheduler": "WarmupLinear",
|
| 113 |
-
"steps_per_epoch": null,
|
| 114 |
-
"warmup_steps": 85,
|
| 115 |
-
"weight_decay": 0.01
|
| 116 |
-
}
|
| 117 |
-
```
|
| 118 |
|
|
|
|
|
|
|
|
|
|
|
|
|
| 119 |
|
| 120 |
-
|
| 121 |
-
|
| 122 |
-
|
| 123 |
-
(0): Transformer({'max_seq_length': 64, 'do_lower_case': False}) with Transformer model: BertModel
|
| 124 |
-
(1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
|
| 125 |
-
)
|
| 126 |
-
```
|
| 127 |
|
| 128 |
-
## Citing & Authors
|
| 129 |
|
| 130 |
-
<!--- Describe where people can find more information -->
|
|
|
|
| 6 |
- feature-extraction
|
| 7 |
- sentence-similarity
|
| 8 |
- transformers
|
| 9 |
+
language:
|
| 10 |
+
- en
|
| 11 |
+
base_model:
|
| 12 |
+
- cambridgeltl/SapBERT-from-PubMedBERT-fulltext-mean-token
|
| 13 |
+
---
|
| 14 |
+
|
| 15 |
+
# NA-SapBERT: Noise-Augmented SapBERT for Clinical Concept Normalization
|
| 16 |
+
|
| 17 |
+
NA-SapBERT is a dense retrieval model designed for clinical concept normalization over large ontologies such as SNOMED CT. It extends SapBERT by incorporating noise-aware training, enabling robust retrieval for real-world clinical mentions.
|
| 18 |
+
|
| 19 |
+
Unlike standard SapBERT, this model is trained to handle:
|
| 20 |
+
- abbreviations (e.g., "NAD", "DM")
|
| 21 |
+
- misspellings
|
| 22 |
+
- shorthand and telegraphic clinical text
|
| 23 |
+
- surface form variation across notes
|
| 24 |
|
| 25 |
---
|
| 26 |
|
| 27 |
+
## Overview
|
| 28 |
|
| 29 |
+
Clinical concept normalization maps noisy text mentions to standardized ontology concepts. While modern NER systems perform well, entity linking remains challenging due to:
|
| 30 |
|
| 31 |
+
- large ontology size
|
| 32 |
+
- noisy clinical text
|
| 33 |
+
- ambiguous abbreviations
|
| 34 |
+
- mismatch between ontology terms and real-world mentions
|
| 35 |
|
| 36 |
+
NA-SapBERT addresses this by learning invariant embeddings across noisy and canonical forms.
|
| 37 |
|
| 38 |
+
---
|
| 39 |
|
| 40 |
+
## Key Idea
|
|
|
|
|
|
|
| 41 |
|
| 42 |
+
During training, the model learns to align:
|
| 43 |
+
- noisy mentions (LLM-generated variants, abbreviations)
|
| 44 |
+
- clean ontology terms (concept names and synonyms)
|
| 45 |
|
| 46 |
+
This is achieved using contrastive learning:
|
| 47 |
+
- clean–clean pairs preserve structure
|
| 48 |
+
- noisy–clean pairs improve robustness
|
| 49 |
|
| 50 |
+
---
|
| 51 |
+
|
| 52 |
+
## Model Architecture
|
| 53 |
+
|
| 54 |
+
SentenceTransformer:
|
| 55 |
+
- Transformer (PubMedBERT backbone)
|
| 56 |
+
- Mean Pooling
|
| 57 |
+
|
| 58 |
+
Embedding dimension: 768
|
| 59 |
+
Max sequence length: 64
|
| 60 |
|
| 61 |
+
---
|
| 62 |
+
|
| 63 |
+
## Training Details
|
| 64 |
+
|
| 65 |
+
### Data
|
| 66 |
+
- SNOMED CT concepts (subset of key semantic types)
|
| 67 |
+
- Synthetic variants:
|
| 68 |
+
- LLM-generated (MedGemma) noise
|
| 69 |
+
- abbreviation mappings
|
| 70 |
|
| 71 |
+
### Objective
|
| 72 |
+
MultipleNegativesRankingLoss (InfoNCE-style)
|
| 73 |
|
| 74 |
+
### Training Configuration
|
| 75 |
+
- epochs: 1
|
| 76 |
+
- batch_size: 256
|
| 77 |
+
- learning_rate: 1e-5
|
| 78 |
+
- warmup_steps: 85
|
| 79 |
+
|
| 80 |
+
---
|
| 81 |
|
| 82 |
+
## Usage
|
| 83 |
+
|
| 84 |
+
### Install
|
| 85 |
+
pip install -U sentence-transformers
|
| 86 |
+
|
| 87 |
+
### Encode Mentions
|
| 88 |
```python
|
| 89 |
+
from sentence_transformers import SentenceTransformer
|
| 90 |
+
|
| 91 |
+
model = SentenceTransformer("YOUR_MODEL_NAME")
|
| 92 |
|
| 93 |
+
mentions = ["NAD", "hx of diabetes", "left axillary lymph node"]
|
| 94 |
+
embeddings = model.encode(mentions, normalize_embeddings=True)
|
| 95 |
+
```
|
| 96 |
|
| 97 |
+
---
|
|
|
|
|
|
|
|
|
|
|
|
|
| 98 |
|
| 99 |
+
## Retrieval Example (FAISS)
|
| 100 |
|
| 101 |
+
```python
|
| 102 |
+
import faiss
|
| 103 |
+
import numpy as np
|
| 104 |
+
from sentence_transformers import SentenceTransformer
|
| 105 |
|
| 106 |
+
model = SentenceTransformer("YOUR_MODEL_NAME")
|
|
|
|
|
|
|
| 107 |
|
| 108 |
+
concept_embeddings = np.load("concept_embeddings.npy").astype("float32")
|
|
|
|
| 109 |
|
| 110 |
+
index = faiss.IndexFlatIP(768)
|
| 111 |
+
index.add(concept_embeddings)
|
|
|
|
| 112 |
|
| 113 |
+
query = "NAD"
|
| 114 |
+
q_emb = model.encode([query], normalize_embeddings=True)
|
| 115 |
|
| 116 |
+
scores, indices = index.search(q_emb, k=10)
|
|
|
|
| 117 |
```
|
| 118 |
|
| 119 |
+
---
|
| 120 |
|
| 121 |
+
## Pipeline Integration
|
| 122 |
|
| 123 |
+
Typical pipeline:
|
| 124 |
+
1. Exact match
|
| 125 |
+
2. Dense retrieval (NA-SapBERT)
|
| 126 |
+
3. Optional rewrite / multi-query
|
| 127 |
+
4. Optional reranking
|
| 128 |
|
| 129 |
+
---
|
| 130 |
|
| 131 |
+
## Performance Summary
|
| 132 |
|
| 133 |
+
- SapBERT: XX recall@1
|
| 134 |
+
- NA-SapBERT: XX recall@1
|
| 135 |
|
| 136 |
+
Improvements:
|
| 137 |
+
- Better handling of noisy mentions
|
| 138 |
+
- Strong generalization to full SNOMED CT
|
| 139 |
|
| 140 |
+
---
|
| 141 |
|
| 142 |
+
## Limitations
|
|
|
|
|
|
|
|
|
|
| 143 |
|
| 144 |
+
- No explicit modeling of negation or temporality
|
| 145 |
+
- Abbreviations remain ambiguous without context
|
| 146 |
+
- Depends on ontology synonym quality
|
| 147 |
|
| 148 |
+
---
|
|
|
|
|
|
|
|
|
|
| 149 |
|
| 150 |
+
## Use Cases
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 151 |
|
| 152 |
+
Use for:
|
| 153 |
+
- clinical NLP
|
| 154 |
+
- concept normalization
|
| 155 |
+
- ontology retrieval
|
| 156 |
|
| 157 |
+
Not intended for:
|
| 158 |
+
- general semantic similarity
|
| 159 |
+
- non-biomedical tasks
|
|
|
|
|
|
|
|
|
|
|
|
|
| 160 |
|
|
|
|
| 161 |
|
|
|