Views
No views yet
1'''
2Code from i-need-sleep https://github.com/i-need-sleep/referee/tree/main/code
3'''
4
5import torch
6from transformers import AutoModel, AutoTokenizer
7
8class DebertaForEval(torch.nn.Module):
9 def __init__(self, model_path, device='cuda', n_supervision=13, head_type='linear', backbone='deberta'):
10 super(DebertaForEval, self).__init__()
11 self.n_supervision = n_supervision
12 self.device = device
13 self.tokenizer = AutoTokenizer.from_pretrained(model_path)
14 self.deberta = AutoModel.from_pretrained(model_path)
15 self.backbone = backbone
16 if backbone == 'deberta':
17 self.hidden_size = 768
18 elif backbone == 'roberta':
19 self.hidden_size = 1024
20 else:
21 raise NotImplementedError
22 self.head_type = head_type
23 if head_type == 'mlp':
24 self.regression_heads_layer_1 = torch.nn.ModuleList([torch.nn.Linear(self.hidden_size, 512) for i in range(n_supervision)])
25 self.regression_heads_layer_2 = torch.nn.ModuleList([torch.nn.Linear(512, 1) for i in range(n_supervision)])
26 self.relu = torch.nn.ReLU()
27 elif head_type == 'linear':
28 self.linear_out = torch.nn.ModuleList([torch.nn.Linear(self.hidden_size, 1) for i in range(n_supervision)])
29 else:
30 raise NotImplementedError
31 self.to(device)
32 self.float()
33
34 def forward(self, sents):
35 if self.backbone == 'deberta':
36 tokenized = self.tokenizer(sents, padding=True, truncation=True, max_length=512)
37 if len(tokenized['input_ids']) >= 512:
38 print("Warning: input exceeds 512 tokens.")
39 input_ids = torch.tensor(tokenized['input_ids']).to(self.device)
40 token_type_ids = torch.tensor(tokenized['token_type_ids']).to(self.device)
41 attention_mask = torch.tensor(tokenized['attention_mask']).to(self.device)
42 model_out = self.deberta(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)[0][:, 0, :] # Take the emb for the first token
43 elif self.backbone == 'roberta':
44 encoded_input = self.tokenizer(sents, return_tensors='pt', padding=True, truncation=True, max_length=512)
45 for key, val in encoded_input.items():
46 encoded_input[key] = val.to(self.device)
47 model_out = self.deberta(**encoded_input)[0][:, 0, :]
48 else:
49 raise NotImplementedError
50 heads_out = []
51 for head_idx in range(self.n_supervision):
52 if self.head_type == 'mlp':
53 head_out = self.regression_heads_layer_1[head_idx](model_out)
54 head_out = self.relu(head_out)
55 head_out = self.regression_heads_layer_2[head_idx](head_out)
56 heads_out.append(head_out)
57 elif self.head_type == 'linear':
58 head_out = self.linear_out[head_idx](model_out)
59 heads_out.append(head_out)
60 heads_out = torch.cat(heads_out, dim=1)
61 return heads_out # [batch_size, n_head]
62
63
64model = DebertaForEval('snisioi/referee', head_type='linear')
65example_complex = """This book constitutes an argument for the power of Marxism to analyse the issues that face women today in their struggle for liberation."""
66example_simple = """This book explains how Marxism can help us understand the problems women face."""
67
68model_input = [example_complex + ' ' + model.tokenizer.sep_token + ' ' + example_simple]
69model_out = model(model_input)
70score = model_out[:, -1].item()
71print(score)