Spaces:

krislette
/

salinlahi

Runtime error

salinlahi / src /models /recurrent /decoder.py

Sync from GitHub: 181e7b5a0a8c123c8fa188c7b7fafcdb8672067d

ee7f83c 4 months ago

8.83 kB

import torch
import torch.nn as nn
import torch.nn.functional as F


class BahdanauAttention(nn.Module):
    """
    Bahdanau (additive) attention from Bahdanau et al. (2015).
    At each decoder step, the attention module compares the current decoder hidden
    state against every encoder output position and produces a context vector —
    a weighted sum of encoder outputs — that tells the decoder where to look in the
    source sentence. This is especially useful for translation because word order
    between Filipino and the regional languages isn't always 1-to-1.
    The energy score for each source position i is computed as:
        score(decoder_hidden, encoder_output_i)
            = score_vector( tanh( W_enc * encoder_output_i
                                + W_dec * decoder_hidden ) )
    Scores are then passed through softmax to get normalized attention weights.
    """

    def __init__(self, encoder_hidden_size: int, decoder_hidden_size: int) -> None:
        super().__init__()

        # Projects encoder outputs (bidirectional, so * 2) into the attention space
        self.encoder_projection = nn.Linear(
            encoder_hidden_size * 2, decoder_hidden_size
        )

        # Projects the decoder hidden state into the same attention space
        self.decoder_projection = nn.Linear(decoder_hidden_size, decoder_hidden_size)

        # Collapses the attention space down to a single scalar score per source position
        self.score_vector = nn.Linear(decoder_hidden_size, 1, bias=False)

    def forward(
        self,
        decoder_hidden: torch.Tensor,
        encoder_outputs: torch.Tensor,
        source_padding_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # decoder_hidden:  (batch_size, decoder_hidden_size)
        # encoder_outputs: (batch_size, source_len, encoder_hidden_size * 2)

        # Expand the decoder hidden state so it can be broadcast across all source positions
        decoder_hidden_expanded = decoder_hidden.unsqueeze(1)
        # decoder_hidden_expanded: (batch_size, 1, decoder_hidden_size)

        # Add the two projections and apply tanh — this is the "energy" function
        energy = torch.tanh(
            self.encoder_projection(encoder_outputs)
            + self.decoder_projection(decoder_hidden_expanded)
        )
        # energy: (batch_size, source_len, decoder_hidden_size)

        # Reduce to one score per source position
        scores = self.score_vector(energy).squeeze(2)
        # scores: (batch_size, source_len)

        # Padding positions should never get attention, so we set their scores to -inf
        # before softmax so they come out as zero weight
        if source_padding_mask is not None:
            scores = scores.masked_fill(source_padding_mask == 0, float("-inf"))

        # Softmax turns the raw scores into a proper probability distribution
        attention_weights = F.softmax(scores, dim=1)
        # attention_weights: (batch_size, source_len)

        # Build the context vector as a weighted sum of the encoder outputs
        # bmm performs batch matrix multiplication: (batch, 1, source_len) x (batch, source_len, hidden*2)
        context_vector = torch.bmm(attention_weights.unsqueeze(1), encoder_outputs)
        # context_vector: (batch_size, 1, encoder_hidden_size * 2)

        return context_vector.squeeze(1), attention_weights


class Decoder(nn.Module):
    """
    Autoregressive GRU/LSTM decoder with Bahdanau attention.
    At each decoding step the decoder:
      1. Embeds the previous output token (or <sos> at the first step)
      2. Attends over all encoder outputs to build a context vector
      3. Concatenates the embedding and context vector, then runs one RNN step
      4. Projects the RNN output to a distribution over the target vocabulary
    The final prediction combines the RNN output, context vector, and embedding
    (a technique called "input feeding" / deep output) so the model has direct
    access to what it just generated and what it attended to.
    """

    def __init__(
        self,
        target_vocab_size: int,
        embedding_dim: int,
        encoder_hidden_size: int,
        decoder_hidden_size: int,
        dropout: float,
        rnn_type: str = "gru",
    ) -> None:
        super().__init__()

        self.rnn_type = rnn_type.lower()
        self.encoder_hidden_size = encoder_hidden_size

        # Embedding for the target language (Waray / Kapampangan / Hiligaynon)
        self.embedding = nn.Embedding(target_vocab_size, embedding_dim, padding_idx=0)
        self.dropout = nn.Dropout(dropout)

        self.attention = BahdanauAttention(encoder_hidden_size, decoder_hidden_size)

        # The RNN input at each step is the concatenation of:
        #   - the embedded previous token  (embedding_dim)
        #   - the attention context vector (encoder_hidden_size * 2)
        rnn_input_dim = embedding_dim + encoder_hidden_size * 2

        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}'"
            )

        self.rnn = rnn_options[self.rnn_type](
            input_size=rnn_input_dim,
            hidden_size=decoder_hidden_size,
            num_layers=1,
            batch_first=True,
        )

        # The output projection combines three sources of information for a richer prediction:
        #   - the new RNN hidden state       (decoder_hidden_size)
        #   - the attention context vector   (encoder_hidden_size * 2)
        #   - the embedded input token       (embedding_dim)
        self.output_projection = nn.Linear(
            decoder_hidden_size + encoder_hidden_size * 2 + embedding_dim,
            target_vocab_size,
        )

    def forward_step(
        self,
        input_token: torch.Tensor,
        decoder_hidden: torch.Tensor,
        decoder_cell: torch.Tensor | None,
        encoder_outputs: torch.Tensor,
        source_padding_mask: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor]:
        """
        Performs a single decoder step and returns the vocabulary logits for that step.
        Args:
            input_token:        (batch_size,) — token id fed as input at this step
            decoder_hidden:     (1, batch_size, decoder_hidden_size) — current hidden state
            decoder_cell:       (1, batch_size, decoder_hidden_size) — current cell (LSTM only)
            encoder_outputs:    (batch_size, source_len, encoder_hidden_size * 2)
            source_padding_mask:(batch_size, source_len) — 1 for real tokens, 0 for padding
        Returns:
            vocab_logits:       (batch_size, target_vocab_size) — unnormalized scores
            decoder_hidden:     (1, batch_size, decoder_hidden_size) — updated hidden state
            decoder_cell:       (1, batch_size, decoder_hidden_size) or None
            attention_weights:  (batch_size, source_len) — where the model looked
        """

        embedded = self.dropout(self.embedding(input_token.unsqueeze(1)))
        # embedded: (batch_size, 1, embedding_dim)

        # Squeeze out the leading num_layers dimension (always 1 here) for the attention call
        context_vector, attention_weights = self.attention(
            decoder_hidden.squeeze(0), encoder_outputs, source_padding_mask
        )
        # context_vector:    (batch_size, encoder_hidden_size * 2)
        # attention_weights: (batch_size, source_len)

        # Concatenate the embedding and context vector, then step the RNN
        rnn_input = torch.cat([embedded, context_vector.unsqueeze(1)], dim=2)
        # rnn_input: (batch_size, 1, embedding_dim + encoder_hidden_size * 2)

        if self.rnn_type == "lstm":
            rnn_output, (decoder_hidden, decoder_cell) = self.rnn(
                rnn_input, (decoder_hidden, decoder_cell)
            )
        else:
            rnn_output, decoder_hidden = self.rnn(rnn_input, decoder_hidden)
            decoder_cell = None
        # rnn_output: (batch_size, 1, decoder_hidden_size)

        # Combine rnn output, context, and embedding to form the final prediction input
        # This gives the model a direct view of what it attended to and its last token
        prediction_input = torch.cat(
            [rnn_output.squeeze(1), context_vector, embedded.squeeze(1)], dim=1
        )
        # prediction_input: (batch_size, decoder_hidden_size + encoder_hidden_size * 2 + embedding_dim)

        vocab_logits = self.output_projection(prediction_input)
        # vocab_logits: (batch_size, target_vocab_size)

        return vocab_logits, decoder_hidden, decoder_cell, attention_weights

	import torch
	import torch.nn as nn
	import torch.nn.functional as F


	class BahdanauAttention(nn.Module):
	"""
	Bahdanau (additive) attention from Bahdanau et al. (2015).

	At each decoder step, the attention module compares the current decoder hidden
	state against every encoder output position and produces a context vector —
	a weighted sum of encoder outputs — that tells the decoder where to look in the
	source sentence. This is especially useful for translation because word order
	between Filipino and the regional languages isn't always 1-to-1.

	The energy score for each source position i is computed as:
	score(decoder_hidden, encoder_output_i)
	= score_vector( tanh( W_enc * encoder_output_i
	+ W_dec * decoder_hidden ) )

	Scores are then passed through softmax to get normalized attention weights.
	"""

	def __init__(self, encoder_hidden_size: int, decoder_hidden_size: int) -> None:
	super().__init__()

	# Projects encoder outputs (bidirectional, so * 2) into the attention space
	self.encoder_projection = nn.Linear(
	encoder_hidden_size * 2, decoder_hidden_size
	)

	# Projects the decoder hidden state into the same attention space
	self.decoder_projection = nn.Linear(decoder_hidden_size, decoder_hidden_size)

	# Collapses the attention space down to a single scalar score per source position
	self.score_vector = nn.Linear(decoder_hidden_size, 1, bias=False)

	def forward(
	self,
	decoder_hidden: torch.Tensor,
	encoder_outputs: torch.Tensor,
	source_padding_mask: torch.Tensor \| None = None,
	) -> tuple[torch.Tensor, torch.Tensor]:
	# decoder_hidden: (batch_size, decoder_hidden_size)
	# encoder_outputs: (batch_size, source_len, encoder_hidden_size * 2)

	# Expand the decoder hidden state so it can be broadcast across all source positions
	decoder_hidden_expanded = decoder_hidden.unsqueeze(1)
	# decoder_hidden_expanded: (batch_size, 1, decoder_hidden_size)

	# Add the two projections and apply tanh — this is the "energy" function
	energy = torch.tanh(
	self.encoder_projection(encoder_outputs)
	+ self.decoder_projection(decoder_hidden_expanded)
	)
	# energy: (batch_size, source_len, decoder_hidden_size)

	# Reduce to one score per source position
	scores = self.score_vector(energy).squeeze(2)
	# scores: (batch_size, source_len)

	# Padding positions should never get attention, so we set their scores to -inf
	# before softmax so they come out as zero weight
	if source_padding_mask is not None:
	scores = scores.masked_fill(source_padding_mask == 0, float("-inf"))

	# Softmax turns the raw scores into a proper probability distribution
	attention_weights = F.softmax(scores, dim=1)
	# attention_weights: (batch_size, source_len)

	# Build the context vector as a weighted sum of the encoder outputs
	# bmm performs batch matrix multiplication: (batch, 1, source_len) x (batch, source_len, hidden*2)
	context_vector = torch.bmm(attention_weights.unsqueeze(1), encoder_outputs)
	# context_vector: (batch_size, 1, encoder_hidden_size * 2)

	return context_vector.squeeze(1), attention_weights


	class Decoder(nn.Module):
	"""
	Autoregressive GRU/LSTM decoder with Bahdanau attention.

	At each decoding step the decoder:
	1. Embeds the previous output token (or <sos> at the first step)
	2. Attends over all encoder outputs to build a context vector
	3. Concatenates the embedding and context vector, then runs one RNN step
	4. Projects the RNN output to a distribution over the target vocabulary

	The final prediction combines the RNN output, context vector, and embedding
	(a technique called "input feeding" / deep output) so the model has direct
	access to what it just generated and what it attended to.
	"""

	def __init__(
	self,
	target_vocab_size: int,
	embedding_dim: int,
	encoder_hidden_size: int,
	decoder_hidden_size: int,
	dropout: float,
	rnn_type: str = "gru",
	) -> None:
	super().__init__()

	self.rnn_type = rnn_type.lower()
	self.encoder_hidden_size = encoder_hidden_size

	# Embedding for the target language (Waray / Kapampangan / Hiligaynon)
	self.embedding = nn.Embedding(target_vocab_size, embedding_dim, padding_idx=0)
	self.dropout = nn.Dropout(dropout)

	self.attention = BahdanauAttention(encoder_hidden_size, decoder_hidden_size)

	# The RNN input at each step is the concatenation of:
	# - the embedded previous token (embedding_dim)
	# - the attention context vector (encoder_hidden_size * 2)
	rnn_input_dim = embedding_dim + encoder_hidden_size * 2

	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}'"
	)

	self.rnn = rnn_options[self.rnn_type](
	input_size=rnn_input_dim,
	hidden_size=decoder_hidden_size,
	num_layers=1,
	batch_first=True,
	)

	# The output projection combines three sources of information for a richer prediction:
	# - the new RNN hidden state (decoder_hidden_size)
	# - the attention context vector (encoder_hidden_size * 2)
	# - the embedded input token (embedding_dim)
	self.output_projection = nn.Linear(
	decoder_hidden_size + encoder_hidden_size * 2 + embedding_dim,
	target_vocab_size,
	)

	def forward_step(
	self,
	input_token: torch.Tensor,
	decoder_hidden: torch.Tensor,
	decoder_cell: torch.Tensor \| None,
	encoder_outputs: torch.Tensor,
	source_padding_mask: torch.Tensor \| None,
	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor \| None, torch.Tensor]:
	"""
	Performs a single decoder step and returns the vocabulary logits for that step.

	Args:
	input_token: (batch_size,) — token id fed as input at this step
	decoder_hidden: (1, batch_size, decoder_hidden_size) — current hidden state
	decoder_cell: (1, batch_size, decoder_hidden_size) — current cell (LSTM only)
	encoder_outputs: (batch_size, source_len, encoder_hidden_size * 2)
	source_padding_mask:(batch_size, source_len) — 1 for real tokens, 0 for padding

	Returns:
	vocab_logits: (batch_size, target_vocab_size) — unnormalized scores
	decoder_hidden: (1, batch_size, decoder_hidden_size) — updated hidden state
	decoder_cell: (1, batch_size, decoder_hidden_size) or None
	attention_weights: (batch_size, source_len) — where the model looked
	"""

	embedded = self.dropout(self.embedding(input_token.unsqueeze(1)))
	# embedded: (batch_size, 1, embedding_dim)

	# Squeeze out the leading num_layers dimension (always 1 here) for the attention call
	context_vector, attention_weights = self.attention(
	decoder_hidden.squeeze(0), encoder_outputs, source_padding_mask
	)
	# context_vector: (batch_size, encoder_hidden_size * 2)
	# attention_weights: (batch_size, source_len)

	# Concatenate the embedding and context vector, then step the RNN
	rnn_input = torch.cat([embedded, context_vector.unsqueeze(1)], dim=2)
	# rnn_input: (batch_size, 1, embedding_dim + encoder_hidden_size * 2)

	if self.rnn_type == "lstm":
	rnn_output, (decoder_hidden, decoder_cell) = self.rnn(
	rnn_input, (decoder_hidden, decoder_cell)
	)
	else:
	rnn_output, decoder_hidden = self.rnn(rnn_input, decoder_hidden)
	decoder_cell = None
	# rnn_output: (batch_size, 1, decoder_hidden_size)

	# Combine rnn output, context, and embedding to form the final prediction input
	# This gives the model a direct view of what it attended to and its last token
	prediction_input = torch.cat(
	[rnn_output.squeeze(1), context_vector, embedded.squeeze(1)], dim=1
	)
	# prediction_input: (batch_size, decoder_hidden_size + encoder_hidden_size * 2 + embedding_dim)

	vocab_logits = self.output_projection(prediction_input)
	# vocab_logits: (batch_size, target_vocab_size)

	return vocab_logits, decoder_hidden, decoder_cell, attention_weights