README.md

---
pipeline_tag: sentence-similarity
tags:
- sentence-transformers
- feature-extraction
- sentence-similarity
- transformers
- embeddings
- static-embeddings
language: en
license: apache-2.0
---

# PubMedBERT Embeddings 1M

This is a pruned version of [PubMedBERT Embeddings 2M](https://huggingface.co/NeuML/pubmedbert-base-embeddings-2M). It prunes the vocabulary to take the top 50% most frequently used tokens. 

See [Extremely Small BERT Models from Mixed-Vocabulary Training](https://arxiv.org/abs/1909.11687) for background on pruning vocabularies to build smaller models.

## Usage (txtai)

This model can be used to build embeddings databases with [txtai](https://github.com/neuml/txtai) for semantic search and/or as a knowledge source for retrieval augmented generation (RAG).

```python
import txtai

# Create embeddings
embeddings = txtai.Embeddings(
  path="neuml/pubmedbert-base-embeddings-1M",
  content=True,
)
embeddings.index(documents())

# Run a query
embeddings.search("query to run")
```

## Usage (Sentence-Transformers)

Alternatively, the model can be loaded with [sentence-transformers](https://www.SBERT.net).

```python
from sentence_transformers import SentenceTransformer
from sentence_transformers.models import StaticEmbedding

# Initialize a StaticEmbedding module
static = StaticEmbedding.from_model2vec("neuml/pubmedbert-base-embeddings-1M")
model = SentenceTransformer(modules=[static])

sentences = ["This is an example sentence", "Each sentence is converted"]
embeddings = model.encode(sentences)
print(embeddings)
```

## Usage (Model2Vec)

The model can also be used directly with Model2Vec.

```python
from model2vec import StaticModel

# Load a pretrained Model2Vec model
model = StaticModel.from_pretrained("neuml/pubmedbert-base-embeddings-1M")

# Compute text embeddings
sentences = ["This is an example sentence", "Each sentence is converted"]
embeddings = model.encode(sentences)
print(embeddings)
```

## Evaluation Results

The following compares performance of this model against the models previously compared with [PubMedBERT Embeddings](https://huggingface.co/NeuML/pubmedbert-base-embeddings#evaluation-results). The following datasets were used to evaluate model performance.

- [PubMed QA](https://huggingface.co/datasets/pubmed_qa)
  - Subset: pqa_labeled, Split: train, Pair: (question, long_answer)
- [PubMed Subset](https://huggingface.co/datasets/awinml/pubmed_abstract_3_1k)
  - Split: test, Pair: (title, text)
  - _Note: The previously used [PubMed Subset](https://huggingface.co/datasets/zxvix/pubmed_subset_new) dataset is no longer available but a similar dataset is used here_
- [PubMed Summary](https://huggingface.co/datasets/scientific_papers)
  - Subset: pubmed, Split: validation, Pair: (article, abstract)

The [Pearson correlation coefficient](https://en.wikipedia.org/wiki/Pearson_correlation_coefficient) is used as the evaluation metric.

| Model                                                                                  | PubMed QA | PubMed Subset | PubMed Summary | Average   |
| -------------------------------------------------------------------------------------- | --------- | ------------- | -------------- | --------- |
| pubmedbert-base-embeddings-8M-M2V (No training)                                        | 69.84     | 70.77         | 71.30          | 70.64     |
| [pubmedbert-base-embeddings-100K](https://hf.co/neuml/pubmedbert-base-embeddings-100K) | 74.56     | 84.65         | 81.84          | 80.35     |
| [pubmedbert-base-embeddings-500K](https://hf.co/neuml/pubmedbert-base-embeddings-500K) | 86.03     | 91.71         | 91.25          | 89.66     |
| [**pubmedbert-base-embeddings-1M**](https://hf.co/neuml/pubmedbert-base-embeddings-1M) | **87.87** | **92.80**     | **92.87**      | **91.18** |
| [pubmedbert-base-embeddings-2M](https://hf.co/neuml/pubmedbert-base-embeddings-2M)     | 88.62     | 93.08         | 93.24          | 91.65     |

As we can see, the accuracy tradeoff is relatively minimal compared to the original model.

## Runtime performance

As another test, let's see how long each model takes to index 120K article abstracts using the following code. All indexing is done with a RTX 3090 GPU. 

```python
from datasets import load_dataset
from tqdm import tqdm
from txtai import Embeddings

ds = load_dataset("ccdv/pubmed-summarization", split="train")

embeddings = Embeddings(path="path to model", content=True, backend="numpy")
embeddings.index(tqdm(ds["abstract"]))
```

| Model                                                                                  | Model Size (MB) | Index time (s) |
| -------------------------------------------------------------------------------------- | ----------      | -------------- |
| [pubmedbert-base-embeddings-100K](https://hf.co/neuml/pubmedbert-base-embeddings-100K) | 0.2             | 19             |
| [pubmedbert-base-embeddings-500K](https://hf.co/neuml/pubmedbert-base-embeddings-500K) | 1.0             | 17             |
| **[pubmedbert-base-embeddings-1M](https://hf.co/neuml/pubmedbert-base-embeddings-1M)** | **2.0**         | 17             |
| [pubmedbert-base-embeddings-2M](https://hf.co/neuml/pubmedbert-base-embeddings-2M)     | 7.5             | 17             |

Vocabulary pruning doesn't change the runtime performance in this case. But the model is much smaller. Vectors are stored at `int16` precision. This can be beneficial to smaller/lower powered embedded devices and could lead to faster vectorization times.

## Training

This model was vocabulary pruned using the following script.

```python
import json
import os

from collections import Counter
from pathlib import Path

import numpy as np

from model2vec import StaticModel
from more_itertools import batched
from sklearn.decomposition import PCA
from tokenlearn.train import collect_means_and_texts
from tokenizers import Tokenizer
from tqdm import tqdm
from txtai.scoring import ScoringFactory

def tokenize(tokenizer):
    # Tokenize into dataset
    dataset = []
    for t in tqdm(batched(texts, 1024)):
        encodings = tokenizer.encode_batch_fast(t, add_special_tokens=False)
        for e in encodings:
            dataset.append((None, e.ids, None))

    return dataset

def tokenweights(tokenizer):
    dataset = tokenize(tokenizer)

    # Build scoring index
    scoring = ScoringFactory.create({"method": "bm25", "terms": True})
    scoring.index(dataset)

    # Calculate mean value of weights array per token
    tokens = np.zeros(tokenizer.get_vocab_size())
    for x in scoring.idf:
        tokens[x] = np.mean(scoring.terms.weights(x)[1])

    return tokens

# See PubMedBERT Embeddings 2M model for details on this data
features = "features"
paths = sorted(Path(features).glob("*.json"))
texts, _ = collect_means_and_texts(paths)

# Output model parameters
output = "output path"
params, dims = 1000000, 64

path = "pubmedbert-base-embeddings-2M_unweighted"
model = StaticModel.from_pretrained(path)

os.makedirs(output, exist_ok=True)

with open(f"{path}/tokenizer.json", "r", encoding="utf-8") as f:
    config = json.load(f)

# Calculate number of tokens to keep
tokencount = params // model.dim

# Calculate term frequency
freqs = Counter()
for _, ids, _ in tokenize(model.tokenizer):
    freqs.update(ids)

# Select top N most common tokens
uids = set(x for x, _ in freqs.most_common(tokencount))
uids = [uid for token, uid in config["model"]["vocab"].items() if uid in uids or token.startswith("[")]

# Get embeddings for uids
model.embedding = model.embedding[uids]

# Select pruned tokens
pairs, index = [], 0
for token, uid in config["model"]["vocab"].items():
    if uid in uids:
        pairs.append((token, index))
        index += 1

config["model"]["vocab"] = dict(pairs)

# Write new tokenizer
with open(f"{output}/tokenizer.json", "w", encoding="utf-8") as f:
    json.dump(config, f, indent=2)

model.tokenizer = Tokenizer.from_file(f"{output}/tokenizer.json")

# Re-weight tokens
weights = tokenweights(model.tokenizer)

# Remove NaNs from embedding, if any
embedding = np.nan_to_num(model.embedding)

# Apply PCA
embedding = PCA(n_components=dims).fit_transform(embedding)

# Apply weights
embedding *= weights[:, None]

# Update model embedding and normalize
model.embedding, model.normalize = embedding.astype(np.int16), True

model.save_pretrained(output)
```

## Acknowledgement

This model is built on the great work from the [Minish Lab](https://github.com/MinishLab) team consisting of [Stephan Tulkens](https://github.com/stephantul) and [Thomas van Dongen](https://github.com/Pringled).

Read more at the following links.

- [Model2Vec](https://github.com/MinishLab/model2vec)
- [Tokenlearn](https://github.com/MinishLab/tokenlearn)
- [Minish Lab Blog](https://minishlab.github.io/)