shreyansh26/bert-base-1024-biencoder-64M-pairs

bert-base-1024-biencoder-64M-pairs

A long context biencoder based on MosaicML's BERT pretrained on 1024 sequence length. This model maps sentences & paragraphs to a 768 dimensional dense vector space and can be used for tasks like clustering or semantic search.

Usage

Download the model and related scripts

### Inference

import torch from torch import nn from transformers import AutoModelForMaskedLM, AutoTokenizer, pipeline, AutoModel from mosaic_bert import BertModel

pip install triton==2.0.0.dev20221202 --no-deps if using Pytorch 2.0

class AutoModelForSentenceEmbedding(nn.Module): def _init(self, model, tokenizer, normalize=True): super(AutoModelForSentenceEmbedding, self).init_()

self.model = model.to("cuda") self.normalize = normalize self.tokenizer = tokenizer

def forward(self, kwargs): model_output = self.model(kwargs) embeddings = self.meanpooling(modeloutput, kwargs['attention_mask']) if self.normalize: embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)

return embeddings

def meanpooling(self, modeloutput, attentionmask): tokenembeddings = modeloutput[0] # First element of modeloutput contains all token embeddings inputmaskexpanded = attentionmask.unsqueeze(-1).expand(tokenembeddings.size()).float() return torch.sum(tokenembeddings * inputmaskexpanded, 1) / torch.clamp(inputmask_expanded.sum(1), min=1e-9)

model = AutoModel.frompretrained("<path-to-model>", trustremotecode=True).to("cuda") model = AutoModelForSentenceEmbedding(model, tokenizer) tokenizer = AutoTokenizer.frompretrained('bert-base-uncased')

sentences = ["This is an example sentence", "Each sentence is converted"]

encodedinput = tokenizer(sentences, padding=True, truncation=True, maxlength=1024, returntensors='pt').to("cuda") embeddings = model(**encodedinput)

print(embeddings) print(embeddings.shape)

## Other details

### Training

This model has been trained on 64M randomly sampled pairs of sentences/paragraphs from the same training set that Sentence Transformers models use. Details of the
training set [here](https://huggingface.co/sentence-transformers/all-mpnet-base-v2#training-data). 

The training (along with hyperparameters), inference and data loading scripts can all be found in [this Github repository](https://github.com/shreyansh26/Long-Context-Biencoder).

### Evaluations

We ran the model on a few retrieval based benchmarks (CQADupstackEnglishRetrieval, DBPedia, MSMARCO, QuoraRetrieval) and the results are [here](https://github.com/shreyansh26/Long-Context-Biencoder/tree/master/models/results/64M_results).