<reponame>sahilkhose/Neural_Machine_Translation<filename>char_decoder.py<gh_stars>1-10
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
CS224N 2019-20: Homework 5
"""

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


class CharDecoder(nn.Module):
    def __init__(self, hidden_size, char_embedding_size=50, target_vocab=None):
        """ Init Character Decoder.

        @param hidden_size (int): Hidden size of the decoder LSTM
        @param char_embedding_size (int): dimensionality of character embeddings
        @param target_vocab (VocabEntry): vocabulary for the target language. See vocab.py for documentation.
        """
        super(CharDecoder, self).__init__()
        self.target_vocab = target_vocab
        self.charDecoder = nn.LSTM(char_embedding_size, hidden_size)
        self.char_output_projection = nn.Linear(hidden_size, len(self.target_vocab.char2id))
        self.decoderCharEmb = nn.Embedding(len(self.target_vocab.char2id), char_embedding_size,
                                           padding_idx=self.target_vocab.char_pad)

    def forward(self, input, dec_hidden=None):
        """ Forward pass of character decoder.

        @param input (Tensor): tensor of integers, shape (length, batch_size)
        @param dec_hidden (tuple(Tensor, Tensor)): internal state of the LSTM before reading the input characters. A tuple of two tensors of shape (1, batch, hidden_size)

        @returns scores (Tensor): called s_t in the PDF, shape (length, batch_size, self.vocab_size)
        @returns dec_hidden (tuple(Tensor, Tensor)): internal state of the LSTM after reading the input characters. A tuple of two tensors of shape (1, batch, hidden_size)
        """
        ### YOUR CODE HERE for part 2a
        ### TODO - Implement the forward pass of the character decoder.
        char_emb = self.decoderCharEmb(input) 
        outputs, (h_n, c_n) = self.charDecoder(char_emb, dec_hidden)
        s_t = self.char_output_projection(outputs)

        return s_t, (h_n, c_n)

    def train_forward(self, char_sequence, dec_hidden=None):
        """ Forward computation during training.

        @param char_sequence (Tensor): tensor of integers, shape (length, batch_size). Note that "length" here and in forward() need not be the same.
        @param dec_hidden (tuple(Tensor, Tensor)): initial internal state of the LSTM, obtained from the output of the word-level decoder. A tuple of two tensors of shape (1, batch_size, hidden_size)

        @returns The cross-entropy loss (Tensor), computed as the *sum* of cross-entropy losses of all the words in the batch.
        """
        s_t, (h_n, c_n) = self.forward(char_sequence[:-1], dec_hidden)

        loss_fn = nn.CrossEntropyLoss(ignore_index=self.target_vocab.char2id['<pad>'], reduction='sum')
        loss = loss_fn(s_t.view(-1, len(self.target_vocab.char2id)), char_sequence[1:].contiguous().view(-1))
        return loss


    def decode_greedy(self, initialStates, device, max_length=21):
        """ Greedy decoding
        @param initialStates (tuple(Tensor, Tensor)): initial internal state of the LSTM, a tuple of two tensors of size (1, batch_size, hidden_size)
        @param device: torch.device (indicates whether the model is on CPU or GPU)
        @param max_length (int): maximum length of words to decode

        @returns decodedWords (List[str]): a list (of length batch_size) of strings, each of which has length <= max_length.
                              The decoded strings should NOT contain the start-of-word and end-of-word characters.
        """
        batch_size = initialStates[0].shape[1]

        current_char = [self.target_vocab.char2id['{']] * batch_size  # len: (batch_size, )
        decodedWords = ['{'] * batch_size  # len: (batch_size,)

        #print('current char is', current_char)
        current_char_tensor = torch.tensor(current_char, device=device)  # shape: (batch_size,)

        #print('current char tensor shape', current_char_tensor.shape)

        h_prev, c_prev = initialStates

        for t in range(max_length):
            _, (h_new, c_new) = self.forward(current_char_tensor.unsqueeze(0),
                                                 (h_prev, c_prev))  # shape: (1, batch, hidden_size)

            s = self.char_output_projection(h_new.squeeze(0))  # shape: (batch, self.vocab_size)

            p = F.log_softmax(s, dim=1)  # shape: (batch, self.vocab_size)
            current_char_tensor = torch.argmax(p, dim=1)  # shape: (batch,)

            for i in range(batch_size):
                decodedWords[i] += self.target_vocab.id2char[current_char_tensor[i].item()]

            h_prev = h_new
            c_prev = c_new

        for i in range(batch_size):
            decodedWords[i] = decodedWords[i][1:]
            decodedWords[i] = decodedWords[i].partition('}')[0]

        return decodedWords
