from config import SentimentConfig
from transformers import PreTrainedModel
import torch.nn as nn
import torch.nn.functional as F


class SententenceTransformerSentimentModel(PreTrainedModel):
    config_class = SentimentConfig

    def __init__(self, config):
        super().__init__(config)

        self.fc1 = nn.Linear(384, config.h1)
        self.fc2 = nn.Linear(config.h1, config.h2)
        self.out = nn.Linear(config.h2, 6)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.out(x)
        out = F.softmax(x, dim=1)
        return out