<kaggle_start><code>import numpy as np  # linear algebra
import pandas as pd  # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory
import os

for dirname, _, filenames in os.walk("/kaggle/input"):
    for filename in filenames:
        print(os.path.join(dirname, filename))
base_path = "/kaggle/input/11-785-fall-20-homework-4-part-2/hw4p2"
# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All"
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session
from dataclasses import dataclass
from typing import List, Dict, Text
import torch
from torch.utils.data import Dataset
import torch
import torch.nn as nn
import tqdm
from torch.utils import data
from torch.nn.utils.rnn import pad_sequence
from torch.nn.utils.rnn import pack_padded_sequence
from torch.nn.utils.rnn import pad_packed_sequence
from torch import utils

print(torch.cuda.is_available())
cuda = torch.cuda.is_available()
if cuda:
    dev = "cuda:0"
else:
    dev = "cpu"
DEVICE = torch.device(dev)
print("Device:", DEVICE)
"""
Loading all the numpy files containing the utterance information and text information
"""


def load_data(base_path=""):
    speech_train = np.load(
        os.path.join(base_path, "train.npy"), allow_pickle=True, encoding="bytes"
    )
    speech_valid = np.load(
        os.path.join(base_path, "dev.npy"), allow_pickle=True, encoding="bytes"
    )
    speech_test = np.load(
        os.path.join(base_path, "test.npy"), allow_pickle=True, encoding="bytes"
    )
    transcript_train = np.load(
        os.path.join(base_path, "train_transcripts.npy"),
        allow_pickle=True,
        encoding="bytes",
    )
    transcript_valid = np.load(
        os.path.join(base_path, "./dev_transcripts.npy"),
        allow_pickle=True,
        encoding="bytes",
    )
    return speech_train, speech_valid, speech_test, transcript_train, transcript_valid


speech_train, speech_valid, speech_test, transcript_train, transcript_valid = load_data(
    base_path
)
from torch.nn.utils.rnn import pad_sequence
from torch.nn.utils.rnn import pack_padded_sequence
from torch.nn.utils.rnn import pad_packed_sequence
from torch.utils.data import DataLoader

LETTER_LIST = [
    "<pad>",
    "a",
    "b",
    "c",
    "d",
    "e",
    "f",
    "g",
    "h",
    "i",
    "j",
    "k",
    "l",
    "m",
    "n",
    "o",
    "p",
    "q",
    "r",
    "s",
    "t",
    "u",
    "v",
    "w",
    "x",
    "y",
    "z",
    "-",
    "'",
    ".",
    "_",
    "+",
    " ",
    "<sos>",
    "<eos>",
    "<unk>",
]
"""
Optional, create dictionaries for letter2index and index2letter transformations
"""


def create_dictionaries(letter_list):
    index2letter = dict(enumerate(letter_list))
    letter2index = dict()
    for k, v in index2letter.items():
        letter2index[v] = k
    return letter2index, index2letter


letter2index, index2letter = create_dictionaries(LETTER_LIST)
"""
Transforms alphabetical input to numerical input, replace each letter by its corresponding 
index from letter_list
"""


def transform_letter_to_index(transcript, letter_list) -> List:
    """
    Transform letter to index. Adds <sos> and <eos> indexes.
    :param transcript :(N, ) Transcripts are the text input
    :param letter_list: Letter list defined above
    :return letter_to_index_list: Returns a list for all the transcript sentence to index
    """
    unk_idx = letter2index["<unk>"]
    start_idx = letter2index["<sos>"]
    end_idx = letter2index["<eos>"]
    new_transcript = []
    for sequence in transcript:
        words = [word.decode("UTF-8") for word in sequence]
        joined_words = " ".join(words)
        new_seq = [start_idx]
        for char in joined_words:
            new_char = letter2index.get(char, unk_idx)
            new_seq.append(new_char)
        new_seq.append(end_idx)
        new_seq = np.array(new_seq)
        new_transcript.append(new_seq)
    return new_transcript


class Speech2TextDataset(Dataset):
    """
    Dataset class for the speech to text data, this may need some tweaking in the
    getitem method as your implementation in the collate function may be different from
    ours.
    """

    def __init__(self, speech, text=None, isTrain=True):
        self.speech = speech
        self.isTrain = isTrain
        if text is not None:
            self.text = text

    def __len__(self):
        return self.speech.shape[0]

    def __getitem__(self, index):
        if self.isTrain == True:
            return torch.tensor(self.speech[index].astype(np.float32)), torch.tensor(
                self.text[index]
            )
        else:
            return torch.tensor(self.speech[index].astype(np.float32))


# TODO: move to a params
INPUT_PADDING_VALUE = 0
LABEL_PADDING_VALUE = letter2index["<pad>"]


def collate_train(batch_data):
    ### Return the padded speech and text data, and the length of utterance and transcript ###
    """
    Args:
        :batch_data List[Tuple]: List of (input, label)
    Returns
    """
    batch_input, batch_label = zip(*batch_data)
    input_lens = [len(seq) for seq in batch_input]
    label_lens = [len(seq) for seq in batch_label]
    batch_input_padded = torch.as_tensor(
        pad_sequence(batch_input, batch_first=True, padding_value=INPUT_PADDING_VALUE)
    )
    batch_label_padded = torch.as_tensor(
        pad_sequence(batch_label, batch_first=True, padding_value=LABEL_PADDING_VALUE)
    )
    return batch_input_padded, batch_label_padded, input_lens, label_lens


def collate_test(batch_data):
    ### Return padded speech and length of utterance ###
    pass


character_text_train = transform_letter_to_index(transcript_train, LETTER_LIST)
character_text_valid = transform_letter_to_index(transcript_valid, LETTER_LIST)
print(index2letter[33])
print(index2letter[34])
print(character_text_train[9])
batch_size = 3
train_dataset = Speech2TextDataset(speech_train, character_text_train)
val_dataset = Speech2TextDataset(speech_valid, character_text_train)
test_dataset = Speech2TextDataset(speech_test, None, False)
# TODO: put in false just for debugging
train_loader = DataLoader(
    train_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_train
)
val_loader = DataLoader(
    val_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_train
)
test_loader = DataLoader(
    test_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_test
)


class Attention(nn.Module):
    """
    Attention is calculated using key, value and query from Encoder and decoder.
    Below are the set of operations you need to perform for computing attention:
        energy = bmm(key, query)
        attention = softmax(energy)
        context = bmm(attention, value)
    """

    def __init__(self):
        super(Attention, self).__init__()

    def forward(self, query, key, value, lens):
        """
        :param query :(batch_size, hidden_size) Query is the output of LSTMCell from Decoder
        :param keys: (batch_size, max_len, encoder_size) Key Projection from Encoder
        :param values: (batch_size, max_len, encoder_size) Value Projection from Encoder
        :return context: (batch_size, encoder_size) Attended Context
        :return attention_mask: (batch_size, max_len) Attention mask that can be plotted
        """


class pBLSTM(nn.Module):
    """
    Pyramidal BiLSTM
    The length of utterance (speech input) can be hundereds to thousands of frames long.
    The Paper reports that a direct LSTM implementation as Encoder resulted in slow convergence,
    and inferior results even after extensive training.
    The major reason is inability of AttendAndSpell operation to extract relevant information
    from a large number of input steps.
    """

    def __init__(self, input_dim, hidden_dim, sample_rate=2):
        super(pBLSTM, self).__init__()
        self.blstm = nn.LSTM(
            input_size=input_dim,
            hidden_size=hidden_dim,
            num_layers=1,
            bidirectional=True,
        )
        self.sample_rate = sample_rate

    def forward(self, x):
        """
        :param x :(N, T) input to the pBLSTM: Packed Sequence
        :return output: (N, T, H) encoded sequence from pyramidal Bi-LSTM
        """
        padded_x, x_lens = pad_packed_sequence(
            x, batch_first=True, padding_value=INPUT_PADDING_VALUE
        )
        batch_size, max_len, feature_dim = padded_x.shape
        # Drop extra frames
        if max_len % self.sample_rate != 0:
            padded_x = padded_x[:, : -(max_len % self.sample_rate), :]
        reshaped_x = padded_x.contiguous().view(
            (
                batch_size,
                int(max_len / self.sample_rate),
                feature_dim * self.sample_rate,
            )
        )
        new_x_lens = x_lens / 2
        new_x_lens = [int(np.floor(length)) for length in new_x_lens]
        # TODO: Pack again
        new_x = pack_padded_sequence(
            reshaped_x, lengths=new_x_lens, batch_first=True, enforce_sorted=False
        )
        return self.blstm(new_x)


class Encoder(nn.Module):
    """
    Encoder takes the utterances as inputs and returns the key and value.
    Key and value are nothing but simple projections of the output from pBLSTM network.
    """

    def __init__(self, input_dim, hidden_dim, value_size=128, key_size=128, p_layers=3):
        super(Encoder, self).__init__()
        self.lstm = nn.LSTM(
            input_size=input_dim,
            hidden_size=hidden_dim,
            num_layers=1,
            bidirectional=True,
            batch_first=True,
        )
        ### Add code to define the blocks of pBLSTMs! ###
        pmodule_list = []
        for l in range(p_layers):
            # Input dim is hidden_dim*2 due to bidirectionality
            # TODO: see if double because of concatenation
            module = pBLSTM(hidden_dim * 2 * 2, hidden_dim)
            pmodule_list.append(module)
        self.p_layers = nn.ModuleList(pmodule_list)
        self.key_network = nn.Linear(hidden_dim * 2, value_size)
        self.value_network = nn.Linear(hidden_dim * 2, key_size)

    def forward(self, x, lens):
        rnn_inp = pack_padded_sequence(
            x, lengths=lens, batch_first=True, enforce_sorted=False
        )
        outputs, _ = self.lstm(rnn_inp)  # Packed Sequence
        ### Use the outputs and pass it through the pBLSTM blocks! ###
        for layer in self.p_layers:
            outputs, _ = layer(outputs)
        # linear_input, _ = utils.rnn.pad_packed_sequence(outputs)
        # keys = self.key_network(linear_input)
        # value = self.value_network(linear_input)
        # return keys, value
        return outputs


encoder = Encoder(40, 512)
for batch_input, batch_labels, input_lens, label_lens in train_loader:
    print("Input lens", input_lens)
    print("Batch input", batch_input.shape)
    outputs = encoder(batch_input, input_lens)
    outputs_unpacked, seq_lens = pad_packed_sequence(
        outputs, batch_first=True, padding_value=INPUT_PADDING_VALUE
    )
    print("Hello output", type(outputs_unpacked), outputs_unpacked.shape)
    break


class Decoder(nn.Module):
    """
    As mentioned in a previous recitation, each forward call of decoder deals with just one time step,
    thus we use LSTMCell instead of LSLTM here.
    The output from the second LSTMCell can be used as query here for attention module.
    In place of value that we get from the attention, this can be replace by context we get from the attention.
    Methods like Gumble noise and teacher forcing can also be incorporated for improving the performance.
    """

    def __init__(
        self, vocab_size, hidden_dim, value_size=128, key_size=128, isAttended=False
    ):
        super(Decoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_dim, padding_idx=0)
        self.lstm1 = nn.LSTMCell(
            input_size=hidden_dim + value_size, hidden_size=hidden_dim
        )
        self.lstm2 = nn.LSTMCell(input_size=hidden_dim, hidden_size=key_size)
        self.isAttended = isAttended
        if isAttended == True:
            self.attention = Attention()
        self.character_prob = nn.Linear(key_size + value_size, vocab_size)

    def forward(self, key, values, text=None, isTrain=True):
        """
        :param key :(T, N, key_size) Output of the Encoder Key projection layer
        :param values: (T, N, value_size) Output of the Encoder Value projection layer
        :param text: (N, text_len) Batch input of text with text_length
        :param isTrain: Train or eval mode
        :return predictions: Returns the character prediction probability
        """
        batch_size = key.shape[1]
        if isTrain == True:
            max_len = text.shape[1]
            embeddings = self.embedding(text)
        else:
            max_len = 250
        predictions = []
        hidden_states = [None, None]
        prediction = torch.zeros(batch_size, 1).to(
            DEVICE
        )  # (torch.ones(batch_size, 1)*33).to(DEVICE)
        for i in range(max_len):
            # * Implement Gumble noise and teacher forcing techniques
            # * When attention is True, replace values[i,:,:] with the context you get from attention.
            # * If you haven't implemented attention yet, then you may want to check the index and break
            #   out of the loop so you do not get index out of range errors.
            if isTrain:
                char_embed = embeddings[:, i, :]
            else:
                char_embed = self.embedding(prediction.argmax(dim=-1))
            inp = torch.cat([char_embed, values[i, :, :]], dim=1)
            hidden_states[0] = self.lstm1(inp, hidden_states[0])
            inp_2 = hidden_states[0][0]
            hidden_states[1] = self.lstm2(inp_2, hidden_states[1])
            ### Compute attention from the output of the second LSTM Cell ###
            output = hidden_states[1][0]
            prediction = self.character_prob(
                torch.cat([output, values[i, :, :]], dim=1)
            )
            predictions.append(prediction.unsqueeze(1))
        return torch.cat(predictions, dim=1)


class Seq2Seq(nn.Module):
    """
    We train an end-to-end sequence to sequence model comprising of Encoder and Decoder.
    This is simply a wrapper "model" for your encoder and decoder.
    """

    def __init__(
        self,
        input_dim,
        vocab_size,
        hidden_dim,
        value_size=128,
        key_size=128,
        isAttended=False,
    ):
        super(Seq2Seq, self).__init__()
        self.encoder = Encoder(input_dim, hidden_dim)
        self.decoder = Decoder(vocab_size, hidden_dim)

    def forward(self, speech_input, speech_len, text_input=None, isTrain=True):
        key, value = self.encoder(speech_input, speech_len)
        if isTrain == True:
            predictions = self.decoder(key, value, text_input)
        else:
            predictions = self.decoder(key, value, text=None, isTrain=False)
        return predictions


import time


def train(model, train_loader, criterion, optimizer, epoch):
    model.train()
    model.to(DEVICE)
    start = time.time()
    # 1) Iterate through your loader
    # 2) Use torch.autograd.set_detect_anomaly(True) to get notices about gradient explosion
    # 3) Set the inputs to the device.
    # 4) Pass your inputs, and length of speech into the model.
    # 5) Generate a mask based on the lengths of the text to create a masked loss.
    # 5.1) Ensure the mask is on the device and is the correct shape.
    # 6) If necessary, reshape your predictions and origianl text input
    # 6.1) Use .contiguous() if you need to.
    # 7) Use the criterion to get the loss.
    # 8) Use the mask to calculate a masked loss.
    # 9) Run the backward pass on the masked loss.
    # 10) Use torch.nn.utils.clip_grad_norm(model.parameters(), 2)
    # 11) Take a step with your optimizer
    # 12) Normalize the masked loss
    # 13) Optionally print the training loss after every N batches
    end = time.time()


def test(model, test_loader, epoch):
    ### Write your test code here! ###
    pass


import torch.optim as optim


def main():
    model = Seq2Seq(input_dim=40, vocab_size=len(LETTER_LIST), hidden_dim=128)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss(reduction="none")
    nepochs = 25
    batch_size = 64 if DEVICE == "cuda" else 1
    character_text_train = transform_letter_to_index(transcript_train, LETTER_LIST)
    character_text_valid = transform_letter_to_index(transcript_valid, LETTER_LIST)
    train_dataset = Speech2TextDataset(speech_train, character_text_train)
    # val_dataset =
    test_dataset = Speech2TextDataset(speech_test, None, False)
    train_loader = DataLoader(
        train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_train
    )
    # val_loader =
    test_loader = DataLoader(
        test_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_test
    )
    for epoch in range(nepochs):
        train(model, train_loader, criterion, optimizer, epoch)
        # val()
        test(model, test_loader, epoch)


main()
<|endoftext|>
[{"Id": 70750057, "ScriptId": 18993126, "ParentScriptVersionId": NaN, "ScriptLanguageId": 9, "AuthorUserId": 3452800, "CreationDate": "08/08/2021 19:56:23", "VersionNumber": 1.0, "Title": "H4P2 - Speech to Text with LAS", "EvaluationDate": "08/08/2021", "IsChange": true, "TotalLines": 447.0, "LinesInsertedFromPrevious": 447.0, "LinesChangedFromPrevious": 0.0, "LinesUnchangedFromPrevious": 0.0, "LinesInsertedFromFork": NaN, "LinesDeletedFromFork": NaN, "LinesChangedFromFork": NaN, "LinesUnchangedFromFork": NaN, "TotalVotes": 0}]
import numpy as np  # linear algebra
import pandas as pd  # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory
import os

for dirname, _, filenames in os.walk("/kaggle/input"):
    for filename in filenames:
        print(os.path.join(dirname, filename))
base_path = "/kaggle/input/11-785-fall-20-homework-4-part-2/hw4p2"
# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All"
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session
from dataclasses import dataclass
from typing import List, Dict, Text
import torch
from torch.utils.data import Dataset
import torch
import torch.nn as nn
import tqdm
from torch.utils import data
from torch.nn.utils.rnn import pad_sequence
from torch.nn.utils.rnn import pack_padded_sequence
from torch.nn.utils.rnn import pad_packed_sequence
from torch import utils

print(torch.cuda.is_available())
cuda = torch.cuda.is_available()
if cuda:
    dev = "cuda:0"
else:
    dev = "cpu"
DEVICE = torch.device(dev)
print("Device:", DEVICE)
"""
Loading all the numpy files containing the utterance information and text information
"""


def load_data(base_path=""):
    speech_train = np.load(
        os.path.join(base_path, "train.npy"), allow_pickle=True, encoding="bytes"
    )
    speech_valid = np.load(
        os.path.join(base_path, "dev.npy"), allow_pickle=True, encoding="bytes"
    )
    speech_test = np.load(
        os.path.join(base_path, "test.npy"), allow_pickle=True, encoding="bytes"
    )
    transcript_train = np.load(
        os.path.join(base_path, "train_transcripts.npy"),
        allow_pickle=True,
        encoding="bytes",
    )
    transcript_valid = np.load(
        os.path.join(base_path, "./dev_transcripts.npy"),
        allow_pickle=True,
        encoding="bytes",
    )
    return speech_train, speech_valid, speech_test, transcript_train, transcript_valid


speech_train, speech_valid, speech_test, transcript_train, transcript_valid = load_data(
    base_path
)
from torch.nn.utils.rnn import pad_sequence
from torch.nn.utils.rnn import pack_padded_sequence
from torch.nn.utils.rnn import pad_packed_sequence
from torch.utils.data import DataLoader

LETTER_LIST = [
    "<pad>",
    "a",
    "b",
    "c",
    "d",
    "e",
    "f",
    "g",
    "h",
    "i",
    "j",
    "k",
    "l",
    "m",
    "n",
    "o",
    "p",
    "q",
    "r",
    "s",
    "t",
    "u",
    "v",
    "w",
    "x",
    "y",
    "z",
    "-",
    "'",
    ".",
    "_",
    "+",
    " ",
    "<sos>",
    "<eos>",
    "<unk>",
]
"""
Optional, create dictionaries for letter2index and index2letter transformations
"""


def create_dictionaries(letter_list):
    index2letter = dict(enumerate(letter_list))
    letter2index = dict()
    for k, v in index2letter.items():
        letter2index[v] = k
    return letter2index, index2letter


letter2index, index2letter = create_dictionaries(LETTER_LIST)
"""
Transforms alphabetical input to numerical input, replace each letter by its corresponding 
index from letter_list
"""


def transform_letter_to_index(transcript, letter_list) -> List:
    """
    Transform letter to index. Adds <sos> and <eos> indexes.
    :param transcript :(N, ) Transcripts are the text input
    :param letter_list: Letter list defined above
    :return letter_to_index_list: Returns a list for all the transcript sentence to index
    """
    unk_idx = letter2index["<unk>"]
    start_idx = letter2index["<sos>"]
    end_idx = letter2index["<eos>"]
    new_transcript = []
    for sequence in transcript:
        words = [word.decode("UTF-8") for word in sequence]
        joined_words = " ".join(words)
        new_seq = [start_idx]
        for char in joined_words:
            new_char = letter2index.get(char, unk_idx)
            new_seq.append(new_char)
        new_seq.append(end_idx)
        new_seq = np.array(new_seq)
        new_transcript.append(new_seq)
    return new_transcript


class Speech2TextDataset(Dataset):
    """
    Dataset class for the speech to text data, this may need some tweaking in the
    getitem method as your implementation in the collate function may be different from
    ours.
    """

    def __init__(self, speech, text=None, isTrain=True):
        self.speech = speech
        self.isTrain = isTrain
        if text is not None:
            self.text = text

    def __len__(self):
        return self.speech.shape[0]

    def __getitem__(self, index):
        if self.isTrain == True:
            return torch.tensor(self.speech[index].astype(np.float32)), torch.tensor(
                self.text[index]
            )
        else:
            return torch.tensor(self.speech[index].astype(np.float32))


# TODO: move to a params
INPUT_PADDING_VALUE = 0
LABEL_PADDING_VALUE = letter2index["<pad>"]


def collate_train(batch_data):
    ### Return the padded speech and text data, and the length of utterance and transcript ###
    """
    Args:
        :batch_data List[Tuple]: List of (input, label)
    Returns
    """
    batch_input, batch_label = zip(*batch_data)
    input_lens = [len(seq) for seq in batch_input]
    label_lens = [len(seq) for seq in batch_label]
    batch_input_padded = torch.as_tensor(
        pad_sequence(batch_input, batch_first=True, padding_value=INPUT_PADDING_VALUE)
    )
    batch_label_padded = torch.as_tensor(
        pad_sequence(batch_label, batch_first=True, padding_value=LABEL_PADDING_VALUE)
    )
    return batch_input_padded, batch_label_padded, input_lens, label_lens


def collate_test(batch_data):
    ### Return padded speech and length of utterance ###
    pass


character_text_train = transform_letter_to_index(transcript_train, LETTER_LIST)
character_text_valid = transform_letter_to_index(transcript_valid, LETTER_LIST)
print(index2letter[33])
print(index2letter[34])
print(character_text_train[9])
batch_size = 3
train_dataset = Speech2TextDataset(speech_train, character_text_train)
val_dataset = Speech2TextDataset(speech_valid, character_text_train)
test_dataset = Speech2TextDataset(speech_test, None, False)
# TODO: put in false just for debugging
train_loader = DataLoader(
    train_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_train
)
val_loader = DataLoader(
    val_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_train
)
test_loader = DataLoader(
    test_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_test
)


class Attention(nn.Module):
    """
    Attention is calculated using key, value and query from Encoder and decoder.
    Below are the set of operations you need to perform for computing attention:
        energy = bmm(key, query)
        attention = softmax(energy)
        context = bmm(attention, value)
    """

    def __init__(self):
        super(Attention, self).__init__()

    def forward(self, query, key, value, lens):
        """
        :param query :(batch_size, hidden_size) Query is the output of LSTMCell from Decoder
        :param keys: (batch_size, max_len, encoder_size) Key Projection from Encoder
        :param values: (batch_size, max_len, encoder_size) Value Projection from Encoder
        :return context: (batch_size, encoder_size) Attended Context
        :return attention_mask: (batch_size, max_len) Attention mask that can be plotted
        """


class pBLSTM(nn.Module):
    """
    Pyramidal BiLSTM
    The length of utterance (speech input) can be hundereds to thousands of frames long.
    The Paper reports that a direct LSTM implementation as Encoder resulted in slow convergence,
    and inferior results even after extensive training.
    The major reason is inability of AttendAndSpell operation to extract relevant information
    from a large number of input steps.
    """

    def __init__(self, input_dim, hidden_dim, sample_rate=2):
        super(pBLSTM, self).__init__()
        self.blstm = nn.LSTM(
            input_size=input_dim,
            hidden_size=hidden_dim,
            num_layers=1,
            bidirectional=True,
        )
        self.sample_rate = sample_rate

    def forward(self, x):
        """
        :param x :(N, T) input to the pBLSTM: Packed Sequence
        :return output: (N, T, H) encoded sequence from pyramidal Bi-LSTM
        """
        padded_x, x_lens = pad_packed_sequence(
            x, batch_first=True, padding_value=INPUT_PADDING_VALUE
        )
        batch_size, max_len, feature_dim = padded_x.shape
        # Drop extra frames
        if max_len % self.sample_rate != 0:
            padded_x = padded_x[:, : -(max_len % self.sample_rate), :]
        reshaped_x = padded_x.contiguous().view(
            (
                batch_size,
                int(max_len / self.sample_rate),
                feature_dim * self.sample_rate,
            )
        )
        new_x_lens = x_lens / 2
        new_x_lens = [int(np.floor(length)) for length in new_x_lens]
        # TODO: Pack again
        new_x = pack_padded_sequence(
            reshaped_x, lengths=new_x_lens, batch_first=True, enforce_sorted=False
        )
        return self.blstm(new_x)


class Encoder(nn.Module):
    """
    Encoder takes the utterances as inputs and returns the key and value.
    Key and value are nothing but simple projections of the output from pBLSTM network.
    """

    def __init__(self, input_dim, hidden_dim, value_size=128, key_size=128, p_layers=3):
        super(Encoder, self).__init__()
        self.lstm = nn.LSTM(
            input_size=input_dim,
            hidden_size=hidden_dim,
            num_layers=1,
            bidirectional=True,
            batch_first=True,
        )
        ### Add code to define the blocks of pBLSTMs! ###
        pmodule_list = []
        for l in range(p_layers):
            # Input dim is hidden_dim*2 due to bidirectionality
            # TODO: see if double because of concatenation
            module = pBLSTM(hidden_dim * 2 * 2, hidden_dim)
            pmodule_list.append(module)
        self.p_layers = nn.ModuleList(pmodule_list)
        self.key_network = nn.Linear(hidden_dim * 2, value_size)
        self.value_network = nn.Linear(hidden_dim * 2, key_size)

    def forward(self, x, lens):
        rnn_inp = pack_padded_sequence(
            x, lengths=lens, batch_first=True, enforce_sorted=False
        )
        outputs, _ = self.lstm(rnn_inp)  # Packed Sequence
        ### Use the outputs and pass it through the pBLSTM blocks! ###
        for layer in self.p_layers:
            outputs, _ = layer(outputs)
        # linear_input, _ = utils.rnn.pad_packed_sequence(outputs)
        # keys = self.key_network(linear_input)
        # value = self.value_network(linear_input)
        # return keys, value
        return outputs


encoder = Encoder(40, 512)
for batch_input, batch_labels, input_lens, label_lens in train_loader:
    print("Input lens", input_lens)
    print("Batch input", batch_input.shape)
    outputs = encoder(batch_input, input_lens)
    outputs_unpacked, seq_lens = pad_packed_sequence(
        outputs, batch_first=True, padding_value=INPUT_PADDING_VALUE
    )
    print("Hello output", type(outputs_unpacked), outputs_unpacked.shape)
    break


class Decoder(nn.Module):
    """
    As mentioned in a previous recitation, each forward call of decoder deals with just one time step,
    thus we use LSTMCell instead of LSLTM here.
    The output from the second LSTMCell can be used as query here for attention module.
    In place of value that we get from the attention, this can be replace by context we get from the attention.
    Methods like Gumble noise and teacher forcing can also be incorporated for improving the performance.
    """

    def __init__(
        self, vocab_size, hidden_dim, value_size=128, key_size=128, isAttended=False
    ):
        super(Decoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_dim, padding_idx=0)
        self.lstm1 = nn.LSTMCell(
            input_size=hidden_dim + value_size, hidden_size=hidden_dim
        )
        self.lstm2 = nn.LSTMCell(input_size=hidden_dim, hidden_size=key_size)
        self.isAttended = isAttended
        if isAttended == True:
            self.attention = Attention()
        self.character_prob = nn.Linear(key_size + value_size, vocab_size)

    def forward(self, key, values, text=None, isTrain=True):
        """
        :param key :(T, N, key_size) Output of the Encoder Key projection layer
        :param values: (T, N, value_size) Output of the Encoder Value projection layer
        :param text: (N, text_len) Batch input of text with text_length
        :param isTrain: Train or eval mode
        :return predictions: Returns the character prediction probability
        """
        batch_size = key.shape[1]
        if isTrain == True:
            max_len = text.shape[1]
            embeddings = self.embedding(text)
        else:
            max_len = 250
        predictions = []
        hidden_states = [None, None]
        prediction = torch.zeros(batch_size, 1).to(
            DEVICE
        )  # (torch.ones(batch_size, 1)*33).to(DEVICE)
        for i in range(max_len):
            # * Implement Gumble noise and teacher forcing techniques
            # * When attention is True, replace values[i,:,:] with the context you get from attention.
            # * If you haven't implemented attention yet, then you may want to check the index and break
            #   out of the loop so you do not get index out of range errors.
            if isTrain:
                char_embed = embeddings[:, i, :]
            else:
                char_embed = self.embedding(prediction.argmax(dim=-1))
            inp = torch.cat([char_embed, values[i, :, :]], dim=1)
            hidden_states[0] = self.lstm1(inp, hidden_states[0])
            inp_2 = hidden_states[0][0]
            hidden_states[1] = self.lstm2(inp_2, hidden_states[1])
            ### Compute attention from the output of the second LSTM Cell ###
            output = hidden_states[1][0]
            prediction = self.character_prob(
                torch.cat([output, values[i, :, :]], dim=1)
            )
            predictions.append(prediction.unsqueeze(1))
        return torch.cat(predictions, dim=1)


class Seq2Seq(nn.Module):
    """
    We train an end-to-end sequence to sequence model comprising of Encoder and Decoder.
    This is simply a wrapper "model" for your encoder and decoder.
    """

    def __init__(
        self,
        input_dim,
        vocab_size,
        hidden_dim,
        value_size=128,
        key_size=128,
        isAttended=False,
    ):
        super(Seq2Seq, self).__init__()
        self.encoder = Encoder(input_dim, hidden_dim)
        self.decoder = Decoder(vocab_size, hidden_dim)

    def forward(self, speech_input, speech_len, text_input=None, isTrain=True):
        key, value = self.encoder(speech_input, speech_len)
        if isTrain == True:
            predictions = self.decoder(key, value, text_input)
        else:
            predictions = self.decoder(key, value, text=None, isTrain=False)
        return predictions


import time


def train(model, train_loader, criterion, optimizer, epoch):
    model.train()
    model.to(DEVICE)
    start = time.time()
    # 1) Iterate through your loader
    # 2) Use torch.autograd.set_detect_anomaly(True) to get notices about gradient explosion
    # 3) Set the inputs to the device.
    # 4) Pass your inputs, and length of speech into the model.
    # 5) Generate a mask based on the lengths of the text to create a masked loss.
    # 5.1) Ensure the mask is on the device and is the correct shape.
    # 6) If necessary, reshape your predictions and origianl text input
    # 6.1) Use .contiguous() if you need to.
    # 7) Use the criterion to get the loss.
    # 8) Use the mask to calculate a masked loss.
    # 9) Run the backward pass on the masked loss.
    # 10) Use torch.nn.utils.clip_grad_norm(model.parameters(), 2)
    # 11) Take a step with your optimizer
    # 12) Normalize the masked loss
    # 13) Optionally print the training loss after every N batches
    end = time.time()


def test(model, test_loader, epoch):
    ### Write your test code here! ###
    pass


import torch.optim as optim


def main():
    model = Seq2Seq(input_dim=40, vocab_size=len(LETTER_LIST), hidden_dim=128)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss(reduction="none")
    nepochs = 25
    batch_size = 64 if DEVICE == "cuda" else 1
    character_text_train = transform_letter_to_index(transcript_train, LETTER_LIST)
    character_text_valid = transform_letter_to_index(transcript_valid, LETTER_LIST)
    train_dataset = Speech2TextDataset(speech_train, character_text_train)
    # val_dataset =
    test_dataset = Speech2TextDataset(speech_test, None, False)
    train_loader = DataLoader(
        train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_train
    )
    # val_loader =
    test_loader = DataLoader(
        test_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_test
    )
    for epoch in range(nepochs):
        train(model, train_loader, criterion, optimizer, epoch)
        # val()
        test(model, test_loader, epoch)


main()
