Spaces:
Sleeping
Sleeping
| 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) |