Views
No views yet
1import torch.nn as nn
2from huggingface_hub import PyTorchModelHubMixin
3from transformers import BertModel, AutoTokenizer
4
5class IndoBertLSTMEcommerceReview(nn.Module, PyTorchModelHubMixin):
6 def __init__(self, bert):
7 super().__init__()
8 self.bert = bert
9 self.lstm = nn.LSTM(bert.config.hidden_size, 128)
10 self.linear = nn.Linear(128, 3)
11 self.sigmoid = nn.Sigmoid()
12
13 def forward(self, input_ids, attention_mask):
14 outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
15 last_hidden_state = outputs.last_hidden_state
16 lstm_out, _ = self.lstm(last_hidden_state)
17 pooled = lstm_out[:, -1, :]
18 logits = self.linear(pooled)
19 probabilities = self.sigmoid(logits)
20 return probabilities
21
22bert = BertModel.from_pretrained("indobenchmark/indobert-base-p1")
23tokenizer = AutoTokenizer.from_pretrained("fahrendrakhoirul/indobert-lstm-finetuned-ecommerce-reviews")
24model = IndoBertLSTMEcommerceReview.from_pretrained("fahrendrakhoirul/indobert-lstm-finetuned-ecommerce-reviews", bert=bert)