Views
No views yet
1import numpy as np
2import torch
3from transformers import AutoModelForMaskedLM, AutoTokenizer
4
5model = AutoModelForMaskedLM.from_pretrained("tau/tavbert-he")
6tokenizer = AutoTokenizer.from_pretrained("tau/tavbert-he")
7
8def mask_sentence(sent, span_len=5):
9 start_pos = np.random.randint(0, len(sent) - span_len)
10 masked_sent = sent[:start_pos] + '[MASK]' * span_len + sent[start_pos + span_len:]
11 print("Masked sentence:", masked_sent)
12 output = model(**tokenizer.encode_plus(masked_sent,
13 return_tensors='pt'))['logits'][0][1:-1]
14 preds = [int(x) for x in torch.argmax(torch.softmax(output, axis=1), axis=1)[start_pos:start_pos + span_len]]
15 pred_sent = sent[:start_pos] + ''.join(tokenizer.convert_ids_to_tokens(preds)) + sent[start_pos + span_len:]
16 print("Model's prediction:", pred_sent)