Views
No views yet
1import torch.nn as nn
2from huggingface_hub import PyTorchModelHubMixin
3from transformers import BertModel, AutoTokenizer
4
5class CustomClassifier(nn.Module, PyTorchModelHubMixin):
6 def __init__(self, bert, num_labels):
7 super(CustomClassifier, self).__init__()
8 self.bert = bert
9 self.linear38 = nn.Linear(bert.config.hidden_size, 38)
10 self.dropout38 = nn.Dropout(0.2)
11 self.linear8 = nn.Linear(38, 8)
12 self.linear3 = nn.Linear(8, 3)
13 self.linearOutput = nn.Linear(3, num_labels)
14 self.sigmoid = nn.Sigmoid()
15
16 def forward(self, input_ids, attention_mask):
17 outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
18 pooled_output = outputs.pooler_output
19 logits38 = self.linear38(pooled_output)
20 logits38 = self.dropout38(logits38)
21 logits8 = self.linear8(logits38)
22 logits3 = self.linear3(logits8)
23 logits = self.linearOutput(logits3)
24 probabilities = self.sigmoid(logits)
25 return probabilities
26
27bert = BertModel.from_pretrained("indobenchmark/indobert-base-p1",
28 num_labels=3,
29 problem_type="multi_label_classification")
30tokenizer = AutoTokenizer.from_pretrained("indobenchmark/indobert-base-p1")
31model = CustomClassifier.from_pretrained("fahrendrakhoirul/indobert-finetuned-ecommerce-product-reviews-aspect-multilabel", bert=bert)
32