Spaces:
Running
Running
File size: 1,665 Bytes
a2570fa |
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 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import torch.nn as nn
class Encoder(nn.Module):
"""
Seq2Seq Encoder for GRU model. I want to store any kind
of sequenital information to be passed on to the decoder
Parameters:
----------
input_dim : int
Size of the input vocabulary
emb_dim : int
Dimension of the embedding vectors
hid_dim : int
Number of features in the GRU's hidden state
n_layers : int
Number of GRU layers (typically 2)
dropout : float
Dropout probability for the dropout layer
"""
def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout):
super().__init__()
# Embedding layer
self.embedding = nn.Embedding(input_dim, emb_dim)
self.hid_dim = hid_dim
self.n_layers = n_layers
# GRU layer
self.rnn = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout)
# Dropout layer
self.dropout = nn.Dropout(dropout)
"""
Forward propagation step of encoding
Parameters:
----------
input : Tensor
Input tensor containing token indices (seq_len, batch_size)
Returns:
-------
hidden : Tensor
Hidden state tensor from the GRU (n_layers, batch_size, hid_dim)
"""
def forward(self, input):
#input is converted into embeddings
embedded = self.dropout(self.embedding(input))
#forward pass into GRU and dropout probability is applied
_ , hidden = self.rnn(embedded)
#only hidden state is required for encoding
return hidden
|