| import torch | |
| import torch.nn as nn | |
| class SentimentClassifier(nn.Module) : | |
| def __init__(self, n_classes) : | |
| super(SentimentClassifier, self).__init__() | |
| self.BERT = BertModel.from_pretrained("bert-base-cased") | |
| self.dropout = nn.Dropout(p=0.3) | |
| self.fc = nn.Linear(self.BERT.config.hidden_size, n_classes) | |
| def forward(self, input_ids, attention_mask) : | |
| _, pooled_output = self.BERT( | |
| input_ids=input_ids, | |
| attention_mask=attention_mask, | |
| return_dict=False | |
| ) | |
| output = self.dropout(pooled_output) | |
| return self.fc(output) |