Views
No views yet
1from transformers import AutoTokenizer, AutoModel
2import torch
3#Mean Pooling - Take attention mask into account for correct averaging
4def mean_pooling(model_output, attention_mask):
5 token_embeddings = model_output[0] #First element of model_output contains all token embeddings
6 input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
7 sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
8 sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-10)
9 return sum_embeddings / sum_mask
10#Sentences we want sentence embeddings for
11sentences = ['Привет! Как твои дела?',
12 'А правда, что 42 твое любимое число?']
13#Load AutoModel from huggingface model repository
14tokenizer = AutoTokenizer.from_pretrained("sberbank-ai/ruElectra-small")
15model = AutoModel.from_pretrained("sberbank-ai/ruElectra-small")
16#Tokenize sentences
17encoded_input = tokenizer(sentences, padding=True, truncation=True, max_length=24, return_tensors='pt')
18#Compute token embeddings
19with torch.no_grad():
20 model_output = model(**encoded_input)
21#Perform pooling. In this case, mean pooling
22sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])@misc{zmitrovich2023family,
title={A Family of Pretrained Transformer Language Models for Russian},
author={Dmitry Zmitrovich and Alexander Abramov and Andrey Kalmykov and Maria Tikhonova and Ekaterina Taktasheva and Danil Astafurov and Mark Baushenko and Artem Snegirev and Tatiana Shavrina and Sergey Markov and Vladislav Mikhailov and Alena Fenogenova},
year={2023},
eprint={2309.10931},
archivePrefix={arXiv},
primaryClass={cs.CL}
}