Views
No views yet
1#Use custom model class:
2
3import torch
4import torch.nn as nn
5from transformers import AutoTokenizer, AutoModel, AdamW, BertModel
6
7class RewardModel(nn.Module):
8 def __init__(self, model_name):
9 super(RewardModel, self).__init__()
10 self.checkpoint = model_name
11 self.bert = AutoModel.from_pretrained(model_name,
12 return_dict=False)
13 self.layer_norm = nn.LayerNorm(768)
14 self.dropout = nn.Dropout(0.3)
15 self.dense = nn.Sequential(
16 nn.Linear(768, 512),
17 nn.LeakyReLU(negative_slope=0.01),
18 nn.Dropout(0.3),
19 nn.Linear(512, 1),
20 nn.Sigmoid()
21 )
22
23 def forward(self, input_ids, token_type_ids, attention_mask):
24
25 model_output = self.bert(input_ids=input_ids,
26 token_type_ids = token_type_ids,
27 attention_mask=attention_mask)
28
29 last_hidden_states = model_output[0]
30 pooled_output = last_hidden_states[:,0]
31 pooled_output = self.layer_norm(pooled_output)
32 pooled_output = self.dropout(pooled_output)
33 preds = self.dense(pooled_output)
34 return preds
35
36
37#Create model object and init pretrain weights:
38reward_name = "ai-forever/ruBert-base"
39tokenizer=AutoTokenizer.from_pretrained(reward_name)
40model = RewardModel(reward_name)
41model.load_state_dict(torch.load('./ruBert-base-reward/pytorch_model.bin'))
42device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
43
44#Sentences that we want to score:
45sentences = ['Человек: Что такое QR-код?', 'Ассистент: QR-код - это тип матричного штрих-кода.']
46
47#Compute reward score:
48with torch.no_grad():
49 model.to(device)
50
51 encoded_input = tokenizer(sentences[0],sentences[1],
52 truncation=True,
53 add_special_tokens=True,
54 max_length=512,
55 padding='max_length',
56 return_tensors='pt')
57
58 encoded_input = encoded_input.to(device)
59 score = model(**encoded_input).cpu().flatten().numpy()
60 print(score)