by facebook
Open source · 32k downloads · 33 likes
Contriever MS MARCO is a language model specialized in dense information retrieval, optimized to generate high-quality embeddings from text. It is a refined version of the pre-trained Contriever model, designed to enhance performance in unsupervised information retrieval tasks. Through its contrastive learning approach, it excels at capturing the semantic nuances of sentences, enabling precise matching between queries and documents. This model is particularly well-suited for advanced search systems, recommendation engines, or conversational assistants requiring a nuanced understanding of language. What sets it apart is its ability to operate effectively without labeled training data while delivering robust results across varied scenarios. Its architecture makes it versatile for applications where the relevance and accuracy of results are critical.
This model is the finetuned version of the pre-trained contriever model available here https://huggingface.co/facebook/contriever, following the approach described in Towards Unsupervised Dense Information Retrieval with Contrastive Learning. The associated GitHub repository is available here https://github.com/facebookresearch/contriever.
Using the model directly available in HuggingFace transformers requires to add a mean pooling operation to obtain a sentence embedding.
import torch
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained('facebook/contriever-msmarco')
model = AutoModel.from_pretrained('facebook/contriever-msmarco')
sentences = [
"Where was Marie Curie born?",
"Maria Sklodowska, later known as Marie Curie, was born on November 7, 1867.",
"Born in Paris on 15 May 1859, Pierre Curie was the son of Eugène Curie, a doctor of French Catholic origin from Alsace."
]
# Apply tokenizer
inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
# Compute token embeddings
outputs = model(**inputs)
# Mean pooling
def mean_pooling(token_embeddings, mask):
token_embeddings = token_embeddings.masked_fill(~mask[..., None].bool(), 0.)
sentence_embeddings = token_embeddings.sum(dim=1) / mask.sum(dim=1)[..., None]
return sentence_embeddings
embeddings = mean_pooling(outputs[0], inputs['attention_mask'])