Views
No views yet
1from transformers import RobertaForMaskedLM, AutoTokenizer
2from transformers.models.longformer.modeling_longformer import LongformerSelfAttention
3
4class RobertaLongSelfAttention(LongformerSelfAttention):
5 def forward(
6 self,
7 hidden_states,
8 attention_mask=None,
9 head_mask=None,
10 encoder_hidden_states=None,
11 encoder_attention_mask=None,
12 past_key_value = None,
13 output_attentions=False,
14 ):
15 attention_mask = attention_mask.squeeze(dim=2).squeeze(dim=1)
16 is_index_masked = attention_mask < 0
17 is_index_global_attn = attention_mask > 0
18 is_global_attn = any(is_index_global_attn.flatten())
19 return super().forward(hidden_states,
20 is_index_masked=is_index_masked,
21 is_index_global_attn=is_index_global_attn,
22 is_global_attn=is_global_attn,
23 attention_mask=attention_mask,
24 output_attentions=output_attentions)
25
26class RobertaLongForMaskedLM(RobertaForMaskedLM):
27 def __init__(self, config):
28 super().__init__(config)
29 for i, layer in enumerate(self.roberta.encoder.layer):
30 layer.attention.self = RobertaLongSelfAttention(config, layer_id=i)
31
32
33tokenizer = AutoTokenizer.from_pretrained("bluenguyen/longformer-phobert-base-4096")
34model = RobertaLongForMaskedLM.from_pretrained("bluenguyen/longformer-phobert-base-4096")
35
36TXT = (
37 "Hoàng_Sa và Trường_Sa là <mask> Việt_Nam ."
38 + "Đó là điều không_thể chối_cãi ." * 300
39 + "Bằng_chứng lịch_sử , pháp_lý về chủ_quyền của Việt_Nam với 2 quần_đảo này đã và đang được nhiều quốc_gia và cộng_đồng quốc_tế <mask> ."
40)
41
42input_ids = tokenizer([TXT], padding=True, pad_to_multiple_of=256, return_tensors="pt")["input_ids"]
43
44logits = model(input_ids).logits
45
46masked_index = [i.item() for i in (input_ids[0] == tokenizer.mask_token_id).nonzero()]
47
48for index in masked_index:
49 probs = logits[0, index].softmax(dim=0)
50 values, predictions = probs.topk(3)
51 print(tokenizer.batch_decode([[p] for p in predictions]))
52> ['của', 'lãnh_thổ', 'chủ_quyền']
53> ['công_nhận', 'thừa_nhận', 'ghi_nhận']