-
Notifications
You must be signed in to change notification settings - Fork 0
/
bertscorer.py
62 lines (52 loc) · 2.45 KB
/
bertscorer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
#from transformers import BertTokenizerFast ,BertForMaskedLM
import torch
import math
from pytorch_pretrained_bert import BertTokenizer, BertForMaskedLM
# Load pre-trained model (weights)
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForMaskedLM.from_pretrained('bert-base-uncased')
model.eval()
def get_score(sentence):
tokenize_input = tokenizer.tokenize(sentence)
tensor_input = torch.tensor([tokenizer.convert_tokens_to_ids(tokenize_input)])
predictions=model(tensor_input)
loss_fct = torch.nn.CrossEntropyLoss()
loss = loss_fct(predictions.squeeze(),tensor_input.squeeze()).data
return math.exp(loss)
class BertScorer():
def __init__(self, weight_name='bert-base-uncased'):
self.tokenizer = BertTokenizerFast.from_pretrained(weight_name,
do_lower_case=True)
self.model = BertForMaskedLM.from_pretrained(weight_name)
self.loss_fct = torch.nn.CrossEntropyLoss()
self.device = self.get_device()
self.model = self.model.to(self.device)
self.model.eval()
def get_device(self):
if torch.cuda.is_available():
return torch.device('cuda')
return torch.device('cpu')
def get_score(self,sentence):
with torch.no_grad():
tensor_input = torch.tensor([self.tokenizer.encode(sentence.lower())])
tensor_input = tensor_input.to(self.device)
predictions=self.model(tensor_input)
loss = self.loss_fct(predictions[0].squeeze(),tensor_input.squeeze())
return torch.exp(loss).data
def get_scores(self,sentences):
with torch.no_grad():
sentences = [sentence.lower() for sentence in sentences]
tokenized_sentences = self.tokenizer.batch_encode_plus(sentences,pad_to_max_length=True)
tensor_sentences = torch.tensor(tokenized_sentences['input_ids'])
tensor_sentences = tensor_sentences.to(self.device)
predictions=self.model(tensor_sentences)
losses = torch.empty(len(sentences))
for i in range(len(sentences)):
losses[i] = self.loss_fct(predictions[0][i],tensor_sentences[i])
del tensor_sentences
del tokenized_sentences
del predictions
if torch.cuda.is_available():
torch.cuda.empty_cache()
#print(torch.cuda.memory_stats())
return torch.exp(losses).data