Spaces:
Running
Running
File size: 2,375 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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import torch.nn as nn
class Decoder(nn.Module):
"""
Initailising GRU Decoder. Based on the hidden state(context vector)
my encoder has returned I want too make predictions to map
English to Klingon
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, output_dim, emb_dim, hid_dim, n_layers, dropout):
super().__init__()
self.hid_dim = hid_dim
self.output_dim = output_dim
self.n_layers = n_layers
self.embedding = nn.Embedding(output_dim, emb_dim)
self.rnn = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout)
self.fc_out = nn.Linear(hid_dim, output_dim)
self.dropout = nn.Dropout(dropout)
"""
Forward propagation step of decoding
Parameters:
----------
hidden : Tensor
Hidden tensor containing token indices (seq_len, batch_size)
This is what our encoder returns
trg : Tensor
Target tensor containing token indices (seq_len, batch_size)
This is what our tokenized Klingon Data
Returns:
-------
prediction : Tensor
Predicted output tensor from the GRU (seq_len, batch_size, output_dim)
hidden : Tensor
Hidden state tensor from the GRU (n_layers, batch_size, hid_dim)
"""
def forward(self, trg, hidden):
#unsure trg is 3D
trg = trg.unsqueeze(0)
#input is converted into embeddings and dropout probability is applied
embedded = self.dropout(self.embedding(trg))
#print("Embedded shape:", embedded.shape)
#GRU layer computes new context based on previous context
output, hidden = self.rnn(embedded, hidden)
#print("Output shape after RNN:", output.shape)
#predicts output from GRU
prediction = self.fc_out(output.squeeze(0))
#print("Output shape after fc_out:", output.shape)
return prediction, hidden |