File size: 5,349 Bytes
ee7f83c | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | import torch
import torch.nn as nn
class Encoder(nn.Module):
"""
Bidirectional RNN/GRU/LSTM encoder for the Filipino source sentence.
The encoder processes the entire source sequence and produces two things:
1. Hidden states for every input token — these are used by the attention mechanism
in the decoder to figure out which parts of the source to focus on.
2. A single summary hidden state — this initializes the decoder so it has a
starting sense of what the source sentence said.
We use a bidirectional RNN so each token sees both its left and right context,
which generally improves translation quality.
"""
def __init__(
self,
source_vocab_size: int,
embedding_dim: int,
hidden_size: int,
num_layers: int,
dropout: float,
rnn_type: str = "gru",
) -> None:
super().__init__()
self.hidden_size = hidden_size
self.num_layers = num_layers
self.rnn_type = rnn_type.lower()
# Embedding layer turns integer token ids into dense float vectors
# padding_idx=0 means the padding token contributes zero gradient
self.embedding = nn.Embedding(source_vocab_size, embedding_dim, padding_idx=0)
# Dropout is applied to the embeddings before they enter the RNN
self.dropout = nn.Dropout(dropout)
# Choose the RNN variant based on the config value
rnn_options = {"rnn": nn.RNN, "gru": nn.GRU, "lstm": nn.LSTM}
if self.rnn_type not in rnn_options:
raise ValueError(
f"rnn_type must be one of {list(rnn_options.keys())}, got '{rnn_type}'"
)
rnn_class = rnn_options[self.rnn_type]
# bidirectional=True makes the RNN read the sequence both forward and backward
# The output at each position is the concatenation of both directions, so the
# effective output size doubles to hidden_size * 2
self.rnn = rnn_class(
input_size=embedding_dim,
hidden_size=hidden_size,
num_layers=num_layers,
batch_first=True,
bidirectional=True,
dropout=dropout if num_layers > 1 else 0,
)
# The decoder is unidirectional (hidden_size), but the encoder produces a
# bidirectional hidden state (hidden_size * 2). This linear layer bridges the gap
# by projecting the concatenated forward+backward final state into hidden_size
self.hidden_projection = nn.Linear(hidden_size * 2, hidden_size)
def forward(
self,
source_tokens: torch.Tensor,
source_lengths: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
# source_tokens: (batch_size, source_len) — integer token ids for the Filipino input
# source_lengths: (batch_size,) — the real (non-padded) length of each sentence
embedded = self.dropout(self.embedding(source_tokens))
# embedded: (batch_size, source_len, embedding_dim)
# Packing tells the RNN to skip padding positions, which saves compute
# and prevents padding from contaminating the hidden state
packed_input = nn.utils.rnn.pack_padded_sequence(
embedded,
source_lengths.cpu(),
batch_first=True,
enforce_sorted=False,
)
if self.rnn_type == "lstm":
packed_outputs, (hidden, cell) = self.rnn(packed_input)
else:
packed_outputs, hidden = self.rnn(packed_input)
cell = None # RNN and GRU don't have a cell state
# Unpack back to a padded tensor so we can run attention over all positions
encoder_outputs, _ = nn.utils.rnn.pad_packed_sequence(
packed_outputs, batch_first=True
)
# encoder_outputs: (batch_size, source_len, hidden_size * 2)
# hidden has shape (num_layers * num_directions, batch_size, hidden_size)
# For a 2-layer bidirectional model the layout is:
# index 0 — layer 1 forward
# index 1 — layer 1 backward
# index 2 — layer 2 forward (top forward layer)
# index 3 — layer 2 backward (top backward layer)
# We only want the top layer, which is the most abstract representation
top_forward_hidden = hidden[-2] # (batch_size, hidden_size)
top_backward_hidden = hidden[-1] # (batch_size, hidden_size)
# Concatenate then project into the decoder's expected hidden size
decoder_initial_hidden = torch.tanh(
self.hidden_projection(
torch.cat([top_forward_hidden, top_backward_hidden], dim=1)
)
)
# decoder_initial_hidden: (batch_size, hidden_size)
# Do the same projection for the cell state if we're using LSTM
decoder_initial_cell = None
if self.rnn_type == "lstm":
top_forward_cell = cell[-2]
top_backward_cell = cell[-1]
decoder_initial_cell = torch.tanh(
self.hidden_projection(
torch.cat([top_forward_cell, top_backward_cell], dim=1)
)
)
# decoder_initial_cell: (batch_size, hidden_size)
return encoder_outputs, decoder_initial_hidden, decoder_initial_cell
|