pubmedbert-embeddings

This is a distilled version of PubMedBERT Embeddings using the Model2Vec library. It uses static embeddings, allowing text embeddings to be computed orders of magnitude faster on both GPU and CPU. It is designed for applications where computational resources are limited or where real-time performance is critical.

Usage (txtai)

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

import txtai

# Create embeddings
embeddings = txtai.Embeddings(
  path="neuml/pubmedbert-base-embeddings-8M",
  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.

from sentence_transformers import SentenceTransformer
from sentence_transformers.models import StaticEmbedding

# Initialize a StaticEmbedding module
static = StaticEmbedding.from_model2vec("neuml/pubmedbert-base-embeddings-8M")
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.

from model2vec import StaticModel

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

# 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. The following datasets were used to evaluate model performance.

• PubMed QA
  • Subset: pqa_labeled, Split: train, Pair: (question, long_answer)
• PubMed Subset
  • Split: test, Pair: (title, text)
• PubMed Summary
  • Subset: pubmed, Split: validation, Pair: (article, abstract)

The Pearson correlation coefficient is used as the evaluation metric.

As we can see, this model while not the top scoring model is certainly competitive.

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.

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"]))

Clearly a static model's main upside is speed. It's important to note that if storage savings is the only concern, then take a look at PubMedBERT Embeddings Matryoshka. The 256 dimension model scores higher than this model, so does the 64 dimension model. The tradeoff is that the runtime performance i