Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModel
3
4tokenizer = AutoTokenizer.from_pretrained('facebook/contriever-msmarco')
5model = AutoModel.from_pretrained('facebook/contriever-msmarco')
6
7sentences = [
8 "Where was Marie Curie born?",
9 "Maria Sklodowska, later known as Marie Curie, was born on November 7, 1867.",
10 "Born in Paris on 15 May 1859, Pierre Curie was the son of Eugène Curie, a doctor of French Catholic origin from Alsace."
11]
12
13# Apply tokenizer
14inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
15
16# Compute token embeddings
17outputs = model(**inputs)
18
19# Mean pooling
20def mean_pooling(token_embeddings, mask):
21 token_embeddings = token_embeddings.masked_fill(~mask[..., None].bool(), 0.)
22 sentence_embeddings = token_embeddings.sum(dim=1) / mask.sum(dim=1)[..., None]
23 return sentence_embeddings
24embeddings = mean_pooling(outputs[0], inputs['attention_mask'])