Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 | |
| # coding=utf-8 | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch.nn.utils.rnn import PackedSequence, pack_padded_sequence, pad_packed_sequence | |
| class CharEmbedding(nn.Module): | |
| def __init__(self, vocab_size: int, embedding_size: int, output_size: int): | |
| super(CharEmbedding, self).__init__() | |
| self.embedding = nn.Embedding(vocab_size, embedding_size, sparse=False) | |
| self.layer_norm = nn.LayerNorm(embedding_size) | |
| self.gru = nn.GRU(embedding_size, embedding_size, num_layers=1, bidirectional=True) | |
| self.out_linear = nn.Linear(2*embedding_size, output_size) | |
| self.layer_norm_2 = nn.LayerNorm(output_size) | |
| def forward(self, words, sentence_lens, word_lens): | |
| # input shape: (B, W, C) | |
| n_words = words.size(1) | |
| sentence_lens = sentence_lens.cpu() | |
| sentence_packed = pack_padded_sequence(words, sentence_lens, batch_first=True) # shape: (B*W, C) | |
| lens_packed = pack_padded_sequence(word_lens, sentence_lens, batch_first=True) # shape: (B*W) | |
| word_packed = pack_padded_sequence(sentence_packed.data, lens_packed.data.cpu(), batch_first=True, enforce_sorted=False) # shape: (B*W*C) | |
| embedded = self.embedding(word_packed.data) # shape: (B*W*C, D) | |
| embedded = self.layer_norm(embedded) # shape: (B*W*C, D) | |
| embedded_packed = PackedSequence(embedded, word_packed[1], word_packed[2], word_packed[3]) | |
| _, embedded = self.gru(embedded_packed) # shape: (layers * 2, B*W, D) | |
| embedded = embedded[-2:, :, :].transpose(0, 1).flatten(1, 2) # shape: (B*W, 2*D) | |
| embedded = F.relu(embedded) | |
| embedded = self.out_linear(embedded) | |
| embedded = self.layer_norm_2(embedded) | |
| embedded, _ = pad_packed_sequence( | |
| PackedSequence(embedded, sentence_packed[1], sentence_packed[2], sentence_packed[3]), batch_first=True, total_length=n_words, | |
| ) # shape: (B, W, 2*D) | |
| return embedded # shape: (B, W, 2*D) | |