SciRus-tiny is a model to obtain embeddings of scientific texts in russian and english. Model was trained on
eLibrary data with contrastive technics described in
habr post. High metrics values were achieved on the
ruSciBench benchmark.
1from transformers import AutoTokenizer, AutoModel
2import torch.nn.functional as F
3import torch
4
5
6tokenizer = AutoTokenizer.from_pretrained("mlsa-iai-msu-lab/sci-rus-tiny")
7model = AutoModel.from_pretrained("mlsa-iai-msu-lab/sci-rus-tiny")
8# model.cuda() # if you want to use a GPU
9
10def mean_pooling(model_output, attention_mask):
11 token_embeddings = model_output[0] #First element of model_output contains all token embeddings
12 input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
13 return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
14
15
16def get_sentence_embedding(title, abstract, model, tokenizer, max_length=None):
17 # Tokenize sentences
18 sentence = '</s>'.join([title, abstract])
19 encoded_input = tokenizer(
20 [sentence], padding=True, truncation=True, return_tensors='pt', max_length=max_length).to(model.device)
21 # Compute token embeddings
22 with torch.no_grad():
23 model_output = model(**encoded_input)
24 # Perform pooling
25 sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
26 # Normalize embeddings
27 sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
28 return sentence_embeddings.cpu().detach().numpy()[0]
29
30print(get_sentence_embedding('some title', 'some abstract', model, tokenizer).shape)
31# (312,)
1from sentence_transformers import SentenceTransformer
2
3
4model = SentenceTransformer('mlsa-iai-msu-lab/sci-rus-tiny')
5embeddings = model.encode(['some title' + '</s>' + 'some abstract'])
6print(embeddings[0].shape)
7# (312,)
Benchmark developed by MLSA Lab of Institute for AI, MSU.
The research is part of the project #23-Ш05-21 SES MSU "Development of mathematical methods of machine learning for processing large-volume textual scientific information". We would like to thank
eLibrary for provided datasets.
1@article{Gerasimenko2024,
2 author = {Gerasimenko, N. and Vatolin, A. and Ianina, A. and Vorontsov, K.},
3 title = {SciRus: Tiny and Powerful Multilingual Encoder for Scientific Texts},
4 journal = {Doklady Mathematics},
5 year = {2024},
6 volume = {110},
7 number = {1},
8 pages = {S193--S202},
9 month = {dec},
10 issn = {1531-8362},
11 doi = {10.1134/S1064562424602178},
12 url = {https://doi.org/10.1134/S1064562424602178}
13}
