Views
No views yet
1from transformers import AutoTokenizer, AutoModel
2import torch
3
4
5#Mean Pooling - Take attention mask into account for correct averaging
6def mean_pooling(model_output, attention_mask):
7 token_embeddings = model_output[0] #First element of model_output contains all token embeddings
8 input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
9 sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
10 sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
11 return sum_embeddings / sum_mask
12
13
14
15#Sentences we want sentence embeddings for
16sentences = ['Привет! Как твои дела?',
17 'А правда, что 42 твое любимое число?']
18
19#Load AutoModel from huggingface model repository
20tokenizer = AutoTokenizer.from_pretrained("sberbank-ai/sbert_large_nlu_ru")
21model = AutoModel.from_pretrained("sberbank-ai/sbert_large_nlu_ru")
22
23#Tokenize sentences
24encoded_input = tokenizer(sentences, padding=True, truncation=True, max_length=24, return_tensors='pt')
25
26#Compute token embeddings
27with torch.no_grad():
28 model_output = model(**encoded_input)
29
30#Perform pooling. In this case, mean pooling
31sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])