pubmedbert-embeddings

About

What pubmedbert-embeddings is

PubMedBERT Embeddings This is a PubMedBERT-base model fined-tuned using sentence-transformers. It maps sentences & paragraphs to a 768 dimensional dense vector space and can be used for tasks like clustering or semantic search. The training dataset was generated using a random sample of PubMed title-abstract pairs along with similar title pairs. PubMedBERT Embeddings produces higher quality embeddings than generalized models for medical literature. Further fine-tuning for a medical subdomain will result in even better performance. 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). Usage (Sentence-Transformers) Alternatively, the model can be loaded with sentence-transformers. Usage (Hugging Face Transformers) The model can also be used directly with Transformers. Evaluation Results Performance of this model compared to the top base models on the MTEB leaderboard is shown below. A popular smaller model was also evaluated along with the most downloaded PubMed similarity model on the Hugging Face Hub. The following datasets were used to evaluate model performance. - PubMed QA - Subset: pqalabeled, Split: train, Pair: (question, longanswer) - PubMed Subset - Split: test, Pair: (title, text) - PubMed Summary - Subset: pubmed, Split: validation, Pair: (article, abstract) Evaluation results are shown below. The Pearson correlation coefficient is used as the evaluation metric. Training The model was trained with the parameters: DataLoader: torch.utils.data.dataloader.DataLoader of length 20191 with parameters: Loss: sentencetransformers.losses.MultipleNegativesRankingLoss.MultipleNegativesRankingLoss with parameters: Parameters of the fit() method: Full Model Architecture More Information Read more about PubMedBERT Embeddings in this article and this paper.

Summarised from the published model card. Read the full card on the HuggingFace links below.

Specifications

What it is

How it works

How embedding models work

Variants

Sizes & precisions

Open weights ship in several sizes and precisions. One page, all the variants — pick the one that fits your GPU. VRAM figures are estimates from model size.

Using it via the API

Call it like any OpenAI endpoint

Once AxForge deploys pubmedbert-embeddings for you, it answers on the OpenAI-compatible API — the same base URL and keys as every other model. (pubmedbert-embeddings below is illustrative; you get the exact model name on deployment.)

$ curl -sS https://api.axforge.ai/v1/embeddings \
  -H "Authorization: Bearer $AXFORGE_API_KEY" \
  -H "Content-Type: application/json" \
  -d '{"model":"pubmedbert-embeddings","input":"text to embed"}'

Details

Languages, data & research

Licence

Open weights

Sources

Weights & code

Explore

More embedding models