import torch
import torch.nn as nn


class ScratchModel(nn.Module):

    def __init__(self):

        super().__init__()

        self.embedding = nn.Embedding(
            num_embeddings=30522,
            embedding_dim=128
        )

        self.lstm = nn.LSTM(
            input_size=128,
            hidden_size=128,
            batch_first=True,
            bidirectional=True
        )

        self.dropout = nn.Dropout(0.3)

        self.fc = nn.Linear(256,5)

    def forward(self,input_ids,attention_mask):

        x = self.embedding(input_ids)

        #x,_ = self.lstm(x)

        mask = attention_mask.unsqueeze(-1)

        x = x * mask

        #x = x.sum(1) / mask.sum(1)
        _, (hidden, _) = self.lstm(x)

        # Forward and backward final hidden states
        x = torch.cat((hidden[-2], hidden[-1]), dim=1)

        x = self.dropout(x)
        

        return self.fc(x)