Views
No views yet
1# pip install transformers sentencepiece
2import torch
3from transformers import LongformerModel, LongformerTokenizerFast
4
5model = LongformerModel.from_pretrained('kazzand/ru-longformer-tiny-16384')
6tokenizer = LongformerTokenizerFast.from_pretrained('kazzand/ru-longformer-tiny-16384')
7
8def get_cls_embedding(text, model, tokenizer, device='cuda'):
9 model.to(device)
10 batch = tokenizer(text, return_tensors='pt')
11
12 #set global attention for cls token
13 global_attention_mask = [
14 [1 if token_id == tokenizer.cls_token_id else 0 for token_id in input_ids]
15 for input_ids in batch["input_ids"]
16 ]
17
18 #add global attention mask to batch
19 batch["global_attention_mask"] = torch.tensor(global_attention_mask)
20
21 with torch.no_grad():
22 output = model(**batch.to(device))
23 return output.last_hidden_state[:,0,:]
24
25