import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

class OffensiveLanguageDataset(Dataset):
    def __init__(self, data, labels):

        self.data = data
        self.labels = labels

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

class OffensiveLanguageClassifier(nn.Module):
    def __init__(self, vocab_size, hidden_size, output_size, num_layers, dropout):
        super(OffensiveLanguageClassifier, self).__init__()
        self.bilstm = nn.LSTM(input_size=vocab_size, hidden_size=hidden_size, num_layers=num_layers, bidirectional=True, dropout=dropout)
        self.fc = nn.Linear(hidden_size * 2, output_size)
        self.fc1 = nn.Linear(hidden_size * 2, output_size)
    def forward(self, input):
        # Perform the computation
        hidden = self.fc1(input)
        relu = self.relu(hidden)
        logits = self.fc2(relu)
        return logits