--- base_model: - microsoft/deberta-v3-base pipeline_tag: text-classification --- It's exactly this model https://github.com/i-need-sleep/referee/ made so that it is easier to run. ```python ''' Code from i-need-sleep https://github.com/i-need-sleep/referee/tree/main/code ''' import torch from transformers import AutoModel, AutoTokenizer class DebertaForEval(torch.nn.Module): def __init__(self, model_path, device='cuda', n_supervision=13, head_type='linear', backbone='deberta'): super(DebertaForEval, self).__init__() self.n_supervision = n_supervision self.device = device self.tokenizer = AutoTokenizer.from_pretrained(model_path) self.deberta = AutoModel.from_pretrained(model_path) self.backbone = backbone if backbone == 'deberta': self.hidden_size = 768 elif backbone == 'roberta': self.hidden_size = 1024 else: raise NotImplementedError self.head_type = head_type if head_type == 'mlp': self.regression_heads_layer_1 = torch.nn.ModuleList([torch.nn.Linear(self.hidden_size, 512) for i in range(n_supervision)]) self.regression_heads_layer_2 = torch.nn.ModuleList([torch.nn.Linear(512, 1) for i in range(n_supervision)]) self.relu = torch.nn.ReLU() elif head_type == 'linear': self.linear_out = torch.nn.ModuleList([torch.nn.Linear(self.hidden_size, 1) for i in range(n_supervision)]) else: raise NotImplementedError self.to(device) self.float() def forward(self, sents): if self.backbone == 'deberta': tokenized = self.tokenizer(sents, padding=True, truncation=True, max_length=512) if len(tokenized['input_ids']) >= 512: print("Warning: input exceeds 512 tokens.") input_ids = torch.tensor(tokenized['input_ids']).to(self.device) token_type_ids = torch.tensor(tokenized['token_type_ids']).to(self.device) attention_mask = torch.tensor(tokenized['attention_mask']).to(self.device) 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 elif self.backbone == 'roberta': encoded_input = self.tokenizer(sents, return_tensors='pt', padding=True, truncation=True, max_length=512) for key, val in encoded_input.items(): encoded_input[key] = val.to(self.device) model_out = self.deberta(**encoded_input)[0][:, 0, :] else: raise NotImplementedError heads_out = [] for head_idx in range(self.n_supervision): if self.head_type == 'mlp': head_out = self.regression_heads_layer_1[head_idx](model_out) head_out = self.relu(head_out) head_out = self.regression_heads_layer_2[head_idx](head_out) heads_out.append(head_out) elif self.head_type == 'linear': head_out = self.linear_out[head_idx](model_out) heads_out.append(head_out) heads_out = torch.cat(heads_out, dim=1) return heads_out # [batch_size, n_head] model = DebertaForEval('snisioi/referee', head_type='linear') example_complex = """This book constitutes an argument for the power of Marxism to analyse the issues that face women today in their struggle for liberation.""" example_simple = """This book explains how Marxism can help us understand the problems women face.""" model_input = [example_complex + ' ' + model.tokenizer.sep_token + ' ' + example_simple] model_out = model(model_input) score = model_out[:, -1].item() print(score) ```