--- 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/)