File size: 1,735 Bytes
78b7771 |
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 |
"""
This module defines the GlassdoorReviewsClassifier class, which is a neural
network model for classifying Glassdoor reviews into sentiment categories. The
model uses a pre-trained BERT model as the base and adds a custom classifier
on top.
"""
import torch.nn as nn
from transformers import BertModel
from config import BERTIMBAU_MODEL
class GlassdoorReviewsClassifier(nn.Module):
"""
GlassdoorReviewsClassifier is a neural network model for classifying
Glassdoor reviews into sentiment categories. It uses a pre-trained BERT
model as the base and adds a custom classifier on top.
Attributes:
bert (BertModel): Pre-trained BERT model for encoding input text.
classifier (nn.Sequential): Custom classifier for sentiment
classification.
"""
def __init__(self):
super(GlassdoorReviewsClassifier, self).__init__()
self.bert = BertModel.from_pretrained(BERTIMBAU_MODEL)
self.classifier = nn.Sequential(
nn.Linear(self.bert.config.hidden_size, 300),
nn.ReLU(),
nn.Linear(300, 100),
nn.ReLU(),
nn.Linear(100, 50),
nn.ReLU(),
nn.Linear(50, 3),
)
def forward(self, input_ids, attention_mask):
"""
Forward pass for the model.
Args:
input_ids (torch.Tensor): Tensor of input token IDs.
attention_mask (torch.Tensor): Tensor of attention masks.
Returns:
torch.Tensor: Output logits from the classifier.
"""
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
x = outputs["last_hidden_state"][:, 0, :]
x = self.classifier(x)
return x
|