Views
No views yet
1import torch
2from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoModel
3from transformers.modeling_outputs import BaseModelOutput
4
5enc_tokenizer = AutoTokenizer.from_pretrained('cointegrated/LaBSE-en-ru')
6encoder = AutoModel.from_pretrained('cointegrated/LaBSE-en-ru')
7
8dec_tokenizer = AutoTokenizer.from_pretrained('cointegrated/rut5-base-labse-decoder')
9decoder = AutoModelForSeq2SeqLM.from_pretrained('cointegrated/rut5-base-labse-decoder')
10
11def encode(texts):
12 encoded_input = enc_tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors='pt')
13 with torch.no_grad():
14 model_output = encoder(**encoded_input.to(encoder.device))
15 embeddings = model_output.pooler_output
16 embeddings = torch.nn.functional.normalize(embeddings)
17 return embeddings
18
19# encode some texts into vectors
20embeddings = encode([
21 "4 декабря 2000 года",
22 "Давно такого не читала, очень хорошо пишешь!",
23 "Я тогда не понимала, что происходит, не понимаю и сейчас.",
24 "London is the capital of Great Britain.",
25])
26print(embeddings.shape)
27# torch.Size([4, 768])
28
29# now try to recover the texts from the vectors
30out = decoder.generate(
31 encoder_outputs=BaseModelOutput(last_hidden_state=embeddings.unsqueeze(1)),
32 max_length=256,
33 repetition_penalty=3.0,
34)
35for tokens in out:
36 print(dec_tokenizer.decode(tokens, skip_special_tokens=True))
37# После 4 декабря 2000 года
38# Не так давно, это многое читала!
39# Я не понимала того, что происходит сейчас тогда, дальше.
40# Британская столица Англии.