| 649 | # |
| 650 | |
| 651 | class EncoderRNN(nn.Module): |
| 652 | def __init__(self, hidden_size, embedding, n_layers=1, dropout=0): |
| 653 | super(EncoderRNN, self).__init__() |
| 654 | self.n_layers = n_layers |
| 655 | self.hidden_size = hidden_size |
| 656 | self.embedding = embedding |
| 657 | |
| 658 | # Initialize GRU; the input_size and hidden_size parameters are both set to 'hidden_size' |
| 659 | # because our input size is a word embedding with number of features == hidden_size |
| 660 | self.gru = nn.GRU(hidden_size, hidden_size, n_layers, |
| 661 | dropout=(0 if n_layers == 1 else dropout), bidirectional=True) |
| 662 | |
| 663 | def forward(self, input_seq, input_lengths, hidden=None): |
| 664 | # Convert word indexes to embeddings |
| 665 | embedded = self.embedding(input_seq) |
| 666 | # Pack padded batch of sequences for RNN module |
| 667 | packed = nn.utils.rnn.pack_padded_sequence(embedded, input_lengths) |
| 668 | # Forward pass through GRU |
| 669 | outputs, hidden = self.gru(packed, hidden) |
| 670 | # Unpack padding |
| 671 | outputs, _ = nn.utils.rnn.pad_packed_sequence(outputs) |
| 672 | # Sum bidirectional GRU outputs |
| 673 | outputs = outputs[:, :, :self.hidden_size] + outputs[:, : ,self.hidden_size:] |
| 674 | # Return output and final hidden state |
| 675 | return outputs, hidden |
| 676 | |
| 677 | |
| 678 | ###################################################################### |