File size: 1,031 Bytes
883317f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | import torch
import torch.nn as nn
class CustomTextCNN(nn.Module):
def __init__(self, vocab_size, embed_dim=128, num_filters=100,
filter_sizes=(3,4,5), num_classes=5,
pad_idx=0, dropout=0.3):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
self.convs = nn.ModuleList([
nn.Conv1d(in_channels=embed_dim,
out_channels=num_filters,
kernel_size=k)
for k in filter_sizes
])
self.dropout = nn.Dropout(dropout)
self.fc = nn.Linear(num_filters * len(filter_sizes), num_classes)
def forward(self, input_ids, attention_mask=None):
x = self.embedding(input_ids).transpose(1, 2)
pooled = []
for conv in self.convs:
c = torch.relu(conv(x))
c = c.max(dim=2).values
pooled.append(c)
out = torch.cat(pooled, dim=1)
out = self.dropout(out)
return self.fc(out)
|