"""
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