Spaces:
Build error
Build error
File size: 4,773 Bytes
a2887cd 6078972 a2887cd 6078972 a2887cd 6078972 a2887cd 6078972 a2887cd 6078972 4dba23e a2887cd 6078972 a2887cd 6078972 27e6827 6078972 a2887cd 2af084a 6078972 a2887cd 6078972 a2887cd 20a2d0d a2887cd 6078972 a2887cd 2af084a a2887cd 9cbe19c 6078972 a2887cd 2af084a a2887cd 20a2d0d 6078972 |
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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
#!g1.1
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
class CategoryHead(nn.Module):
def __init__(self):
super(CategoryHead, self).__init__()
self.lin1 = nn.Linear(256, 64)
self.lin2 = nn.Linear(64, 5)
def forward(self, x):
x = torch.relu(self.lin1(x))
x = self.lin2(x)
return x
class SentimentHead(nn.Module):
def __init__(self):
super(SentimentHead, self).__init__()
self.lin1 = nn.Linear(256, 64)
self.lin2 = nn.Linear(64, 1, bias=False)
def forward(self, x):
x = torch.relu(self.lin1(x))
x = self.lin2(x)
return x
def mean_pooling(model_output, attention_mask):
input_mask_expanded = attention_mask.unsqueeze(-1).expand(model_output.size()).float()
sum_embeddings = torch.sum(model_output * input_mask_expanded, 1)
sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
return sum_embeddings / sum_mask
class UnionModel(nn.Module):
def __init__(self, model_path):
super(UnionModel, self).__init__()
self.bert_model = AutoModel.from_pretrained(model_path)
for _, param in self.bert_model.named_parameters():
param.requires_grad = False
self.bert_model.pooler = nn.Linear(in_features=768, out_features=256)
self.bert_model.to('cpu')
self.category_head = CategoryHead()
self.sentiment_head = SentimentHead()
def forward(self, input):
output = self.bert_model(**input)
output = output.pooler_output
output = mean_pooling(output, input['attention_mask'])
return self.category_head(output), self.sentiment_head(output)
class LogisticCumulativeLink(nn.Module):
"""
Converts a single number to the proportional odds of belonging to a class.
Parameters
----------
num_classes : int
Number of ordered classes to partition the odds into.
init_cutpoints : str (default='ordered')
How to initialize the cutpoints of the model. Valid values are
- ordered : cutpoints are initialized to halfway between each class.
- random : cutpoints are initialized with random values.
"""
def __init__(self, num_classes: int,
init_cutpoints: str = 'ordered') -> None:
assert num_classes > 2, (
'Only use this model if you have 3 or more classes'
)
super().__init__()
self.num_classes = num_classes
self.init_cutpoints = init_cutpoints
if init_cutpoints == 'ordered':
num_cutpoints = self.num_classes - 1
cutpoints = torch.arange(num_cutpoints).float() - num_cutpoints / 2
self.cutpoints = nn.Parameter(cutpoints)
elif init_cutpoints == 'random':
cutpoints = torch.rand(self.num_classes - 1).sort()[0]
self.cutpoints = nn.Parameter(cutpoints)
else:
raise ValueError(f'{init_cutpoints} is not a valid init_cutpoints '
f'type')
def forward(self, X: torch.Tensor) -> torch.Tensor:
"""
Equation (11) from
"On the consistency of ordinal regression methods", Pedregosa et. al.
"""
sigmoids = torch.sigmoid(self.cutpoints - X)
link_mat = sigmoids[:, 1:] - sigmoids[:, :-1]
link_mat = torch.cat((
sigmoids[:, [0]],
link_mat,
(1 - sigmoids[:, [-1]])
), dim=1)
return link_mat
class CustomOrdinalLogisticModel(nn.Module):
def __init__(self, predictor: nn.Module, num_classes: int,
init_cutpoints: str = 'ordered') -> None:
super().__init__()
self.num_classes = num_classes
self.predictor = predictor
self.link = LogisticCumulativeLink(self.num_classes,
init_cutpoints=init_cutpoints)
def forward(self, *args, **kwargs):
cat, sent = self.predictor(*args, **kwargs)
return cat, self.link(sent)
tokenizer = AutoTokenizer.from_pretrained('blanchefort/rubert-base-cased-sentiment-rusentiment')
model = CustomOrdinalLogisticModel(UnionModel('blanchefort/rubert-base-cased-sentiment-rusentiment'), 3).to('cpu')
model.load_state_dict(torch.load('best_model.pth', map_location='cpu'), strict=False)
def inference(input_data):
tokenized = tokenizer(input_data['sentence'])
input_ids = torch.LongTensor(tokenized['input_ids']).unsqueeze(0).to('cpu')
attention_mask = torch.IntTensor(tokenized['attention_mask']).unsqueeze(0).to('cpu')
model.eval()
answer = model({'input_ids': input_ids, 'attention_mask': attention_mask})
answer[0][0] = torch.sigmoid(answer[0][0])
return dict(answer=answer)
|