Upload 5 files

APE_tr2.ipynb

@@ -0,0 +1,813 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9db57e75-ba95-4e96-836a-ce2eb9689c7b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install torch\n",
+    "\n",
+    "\n",
+    "from torch import Tensor\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torch.nn import Transformer\n",
+    "import math\n",
+    "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
+    "import os\n",
+    "from argparse import Namespace\n",
+    "from collections import Counter\n",
+    "import json\n",
+    "import re\n",
+    "import string\n",
+    "import datetime\n",
+    "\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torch.nn import functional as F\n",
+    "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence\n",
+    "import torch.optim as optima\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "class Vocabulary(object):\n",
+    "    \"\"\"Class to process text and extract vocabulary for mapping\"\"\"\n",
+    "\n",
+    "    def __init__(self, token_to_idx=None):\n",
+    "        \"\"\"\n",
+    "        Args:\n",
+    "            token_to_idx (dict): a pre-existing map of tokens to indices\n",
+    "        \"\"\"\n",
+    "\n",
+    "        if token_to_idx is None:\n",
+    "            token_to_idx = {}\n",
+    "        self._token_to_idx = token_to_idx\n",
+    "\n",
+    "        self._idx_to_token = {idx: token \n",
+    "                              for token, idx in self._token_to_idx.items()}\n",
+    "        \n",
+    "    def to_serializable(self):\n",
+    "        \"\"\" returns a dictionary that can be serialized \"\"\"\n",
+    "        return {'token_to_idx': self._token_to_idx}\n",
+    "\n",
+    "    @classmethod\n",
+    "    def from_serializable(cls, contents):\n",
+    "        \"\"\" instantiates the Vocabulary from a serialized dictionary \"\"\"\n",
+    "        return cls(**contents)\n",
+    "\n",
+    "    def add_token(self, token):\n",
+    "        \"\"\"Update mapping dicts based on the token.\n",
+    "\n",
+    "        Args:\n",
+    "            token (str): the item to add into the Vocabulary\n",
+    "        Returns:\n",
+    "            index (int): the integer corresponding to the token\n",
+    "        \"\"\"\n",
+    "        if token in self._token_to_idx:\n",
+    "            index = self._token_to_idx[token]\n",
+    "        else:\n",
+    "            index = len(self._token_to_idx)\n",
+    "            self._token_to_idx[token] = index\n",
+    "            self._idx_to_token[index] = token\n",
+    "        return index\n",
+    "            \n",
+    "    def add_many(self, tokens):\n",
+    "        \"\"\"Add a list of tokens into the Vocabulary\n",
+    "        \n",
+    "        Args:\n",
+    "            tokens (list): a list of string tokens\n",
+    "        Returns:\n",
+    "            indices (list): a list of indices corresponding to the tokens\n",
+    "        \"\"\"\n",
+    "        return [self.add_token(token) for token in tokens]\n",
+    "\n",
+    "    def lookup_token(self, token):\n",
+    "        \"\"\"Retrieve the index associated with the token \n",
+    "        \n",
+    "        Args:\n",
+    "            token (str): the token to look up \n",
+    "        Returns:\n",
+    "            index (int): the index corresponding to the token\n",
+    "        \"\"\"\n",
+    "        return self._token_to_idx[token]\n",
+    "\n",
+    "    def lookup_index(self, index):\n",
+    "        \"\"\"Return the token associated with the index\n",
+    "        \n",
+    "        Args: \n",
+    "            index (int): the index to look up\n",
+    "        Returns:\n",
+    "            token (str): the token corresponding to the index\n",
+    "        Raises:\n",
+    "            KeyError: if the index is not in the Vocabulary\n",
+    "        \"\"\"\n",
+    "        if index not in self._idx_to_token:\n",
+    "            raise KeyError(\"the index (%d) is not in the Vocabulary\" % index)\n",
+    "        return self._idx_to_token[index]\n",
+    "\n",
+    "    def __str__(self):\n",
+    "        return \"<Vocabulary(size=%d)>\" % len(self)\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self._token_to_idx)\n",
+    "    \n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "class SequenceVocabulary(Vocabulary):\n",
+    "    def __init__(self, token_to_idx=None, unk_token=\"<UNK>\",\n",
+    "                 mask_token=\"<MASK>\", begin_seq_token=\"<BEGIN>\",\n",
+    "                 end_seq_token=\"<END>\"):\n",
+    "\n",
+    "        super(SequenceVocabulary, self).__init__(token_to_idx)\n",
+    "\n",
+    "        self._mask_token = mask_token\n",
+    "        self._unk_token = unk_token\n",
+    "        self._begin_seq_token = begin_seq_token\n",
+    "        self._end_seq_token = end_seq_token\n",
+    "\n",
+    "        self.mask_index = self.add_token(self._mask_token)\n",
+    "        self.unk_index = self.add_token(self._unk_token)\n",
+    "        self.begin_seq_index = self.add_token(self._begin_seq_token)\n",
+    "        self.end_seq_index = self.add_token(self._end_seq_token)\n",
+    "\n",
+    "    def to_serializable(self):\n",
+    "        contents = super(SequenceVocabulary, self).to_serializable()\n",
+    "        contents.update({'unk_token': self._unk_token,\n",
+    "                         'mask_token': self._mask_token,\n",
+    "                         'begin_seq_token': self._begin_seq_token,\n",
+    "                         'end_seq_token': self._end_seq_token})\n",
+    "        return contents\n",
+    "\n",
+    "    def lookup_token(self, token):\n",
+    "        \"\"\"Retrieve the index associated with the token \n",
+    "          or the UNK index if token isn't present.\n",
+    "        \n",
+    "        Args:\n",
+    "            token (str): the token to look up \n",
+    "        Returns:\n",
+    "            index (int): the index corresponding to the token\n",
+    "        Notes:\n",
+    "            `unk_index` needs to be >=0 (having been added into the Vocabulary) \n",
+    "              for the UNK functionality \n",
+    "        \"\"\"\n",
+    "        if self.unk_index >= 0:\n",
+    "            return self._token_to_idx.get(token, self.unk_index)\n",
+    "        else:\n",
+    "            return self._token_to_idx[token]\n",
+    "        \n",
+    "\n",
+    "\n",
+    "\n",
+    "class NMTVectorizer(object):\n",
+    "    \"\"\" The Vectorizer which coordinates the Vocabularies and puts them to use\"\"\"        \n",
+    "    def __init__(self, source_vocab, target_vocab, max_source_length, max_target_length):\n",
+    "        \"\"\"\n",
+    "        Args:\n",
+    "            source_vocab (SequenceVocabulary): maps source words to integers\n",
+    "            target_vocab (SequenceVocabulary): maps target words to integers\n",
+    "            max_source_length (int): the longest sequence in the source dataset\n",
+    "            max_target_length (int): the longest sequence in the target dataset\n",
+    "        \"\"\"\n",
+    "        self.source_vocab = source_vocab\n",
+    "        self.target_vocab = target_vocab\n",
+    "        \n",
+    "        self.max_source_length = max_source_length\n",
+    "        self.max_target_length = max_target_length\n",
+    "        \n",
+    "\n",
+    "    def _vectorize(self, indices, vector_length=-1, mask_index=0):\n",
+    "        \"\"\"Vectorize the provided indices\n",
+    "        \n",
+    "        Args:\n",
+    "            indices (list): a list of integers that represent a sequence\n",
+    "            vector_length (int): an argument for forcing the length of index vector\n",
+    "            mask_index (int): the mask_index to use; almost always 0\n",
+    "        \"\"\"\n",
+    "        if vector_length < 0:\n",
+    "            vector_length = len(indices)\n",
+    "        \n",
+    "        vector = np.zeros(vector_length, dtype=np.int64)\n",
+    "        vector[:len(indices)] = indices\n",
+    "        vector[len(indices):] = mask_index\n",
+    "\n",
+    "        return vector\n",
+    "    \n",
+    "    def _get_source_indices(self, text):\n",
+    "        \"\"\"Return the vectorized source text\n",
+    "        \n",
+    "        Args:\n",
+    "            text (str): the source text; tokens should be separated by spaces\n",
+    "        Returns:\n",
+    "            indices (list): list of integers representing the text\n",
+    "        \"\"\"\n",
+    "        indices = [self.source_vocab.begin_seq_index]\n",
+    "        indices.extend(self.source_vocab.lookup_token(token) for token in text.split(\" \"))\n",
+    "        indices.append(self.source_vocab.end_seq_index)\n",
+    "        return indices\n",
+    "    \n",
+    "    def _get_target_indices(self, text):\n",
+    "        \"\"\"Return the vectorized source text\n",
+    "        \n",
+    "        Args:\n",
+    "            text (str): the source text; tokens should be separated by spaces\n",
+    "        Returns:\n",
+    "            a tuple: (x_indices, y_indices)\n",
+    "                x_indices (list): list of integers representing the observations in target decoder \n",
+    "                y_indices (list): list of integers representing predictions in target decoder\n",
+    "        \"\"\"\n",
+    "        indices = [self.target_vocab.lookup_token(token) for token in text.split(\" \")]\n",
+    "        x_indices = [self.target_vocab.begin_seq_index] + indices\n",
+    "        y_indices = indices + [self.target_vocab.end_seq_index]\n",
+    "        return x_indices, y_indices\n",
+    "        \n",
+    "    def vectorize(self, source_text, target_text, use_dataset_max_lengths=True):\n",
+    "        \"\"\"Return the vectorized source and target text\n",
+    "        \n",
+    "        The vetorized source text is just the a single vector.\n",
+    "        The vectorized target text is split into two vectors in a similar style to \n",
+    "            the surname modeling in Chapter 7.\n",
+    "        At each timestep, the first vector is the observation and the second vector is the target. \n",
+    "        \n",
+    "        \n",
+    "        Args:\n",
+    "            source_text (str): text from the source language\n",
+    "            target_text (str): text from the target language\n",
+    "            use_dataset_max_lengths (bool): whether to use the global max vector lengths\n",
+    "        Returns:\n",
+    "            The vectorized data point as a dictionary with the keys: \n",
+    "                source_vector, target_x_vector, target_y_vector, source_length\n",
+    "        \"\"\"\n",
+    "        source_vector_length = -1\n",
+    "        target_vector_length = -1\n",
+    "        \n",
+    "        if use_dataset_max_lengths:\n",
+    "            source_vector_length = self.max_source_length + 2\n",
+    "            target_vector_length = self.max_target_length + 1\n",
+    "            \n",
+    "        source_indices = self._get_source_indices(source_text)\n",
+    "        source_vector = self._vectorize(source_indices, \n",
+    "                                        vector_length=source_vector_length, \n",
+    "                                        mask_index=self.source_vocab.mask_index)\n",
+    "        \n",
+    "        target_x_indices, target_y_indices = self._get_target_indices(target_text)\n",
+    "        target_x_vector = self._vectorize(target_x_indices,\n",
+    "                                        vector_length=target_vector_length,\n",
+    "                                        mask_index=self.target_vocab.mask_index)\n",
+    "        target_y_vector = self._vectorize(target_y_indices,\n",
+    "                                        vector_length=target_vector_length,\n",
+    "                                        mask_index=self.target_vocab.mask_index)\n",
+    "        return {\"source_vector\": source_vector, \n",
+    "                \"target_x_vector\": target_x_vector, \n",
+    "                \"target_y_vector\": target_y_vector, \n",
+    "                \"source_length\": len(source_indices)}\n",
+    "        \n",
+    "    @classmethod\n",
+    "    def from_dataframe(cls, bitext_df):\n",
+    "        \"\"\"Instantiate the vectorizer from the dataset dataframe\n",
+    "        \n",
+    "        Args:\n",
+    "            bitext_df (pandas.DataFrame): the parallel text dataset\n",
+    "        Returns:\n",
+    "            an instance of the NMTVectorizer\n",
+    "        \"\"\"\n",
+    "        source_vocab = SequenceVocabulary()\n",
+    "        target_vocab = SequenceVocabulary()\n",
+    "        \n",
+    "        max_source_length = 50\n",
+    "        max_target_length = 25\n",
+    "\n",
+    "        for _, row in bitext_df.iterrows():\n",
+    "            source_tokens = row[\"source_language\"].split(\" \")\n",
+    "            if len(source_tokens) > max_source_length:\n",
+    "                max_source_length = len(source_tokens)\n",
+    "            for token in source_tokens:\n",
+    "                source_vocab.add_token(token)\n",
+    "            \n",
+    "            target_tokens = row[\"target_language\"].split(\" \")\n",
+    "            if len(target_tokens) > max_target_length:\n",
+    "                max_target_length = len(target_tokens)\n",
+    "            for token in target_tokens:\n",
+    "                target_vocab.add_token(token)\n",
+    "            \n",
+    "        return cls(source_vocab, target_vocab, max_source_length, max_target_length)\n",
+    "\n",
+    "    @classmethod\n",
+    "    def from_serializable(cls, contents):\n",
+    "        source_vocab = SequenceVocabulary.from_serializable(contents[\"source_vocab\"])\n",
+    "        target_vocab = SequenceVocabulary.from_serializable(contents[\"target_vocab\"])\n",
+    "        \n",
+    "        return cls(source_vocab=source_vocab, \n",
+    "                   target_vocab=target_vocab, \n",
+    "                   max_source_length=contents[\"max_source_length\"], \n",
+    "                   max_target_length=contents[\"max_target_length\"])\n",
+    "\n",
+    "    def to_serializable(self):\n",
+    "        return {\"source_vocab\": self.source_vocab.to_serializable(), \n",
+    "                \"target_vocab\": self.target_vocab.to_serializable(), \n",
+    "                \"max_source_length\": self.max_source_length,\n",
+    "                \"max_target_length\": self.max_target_length}\n",
+    "        \n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "class NMTDataset(Dataset):\n",
+    "    def __init__(self, text_df, vectorizer):\n",
+    "        \"\"\"\n",
+    "        Args:\n",
+    "            surname_df (pandas.DataFrame): the dataset\n",
+    "            vectorizer (SurnameVectorizer): vectorizer instatiated from dataset\n",
+    "        \"\"\"\n",
+    "        self.text_df = text_df\n",
+    "        self._vectorizer = vectorizer\n",
+    "\n",
+    "        self.train_df = self.text_df[self.text_df.split=='train']\n",
+    "        self.train_size = len(self.train_df)\n",
+    "\n",
+    "        self.val_df = self.text_df[self.text_df.split=='val']\n",
+    "        self.validation_size = len(self.val_df)\n",
+    "\n",
+    "        self.test_df = self.text_df[self.text_df.split=='test']\n",
+    "        self.test_size = len(self.test_df)\n",
+    "\n",
+    "        self._lookup_dict = {'train': (self.train_df, self.train_size),\n",
+    "                             'val': (self.val_df, self.validation_size),\n",
+    "                             'test': (self.test_df, self.test_size)}\n",
+    "\n",
+    "        self.set_split('train')\n",
+    "\n",
+    "    @classmethod\n",
+    "    def load_dataset_and_make_vectorizer(cls, dataset_csv):\n",
+    "        \"\"\"Load dataset and make a new vectorizer from scratch\n",
+    "        \n",
+    "        Args:\n",
+    "            surname_csv (str): location of the dataset\n",
+    "        Returns:\n",
+    "            an instance of SurnameDataset\n",
+    "        \"\"\"\n",
+    "        text_df = pd.read_csv(dataset_csv).fillna(' ')\n",
+    "        train_subset = text_df[text_df.split=='train']\n",
+    "        return cls(text_df, NMTVectorizer.from_dataframe(train_subset))\n",
+    "\n",
+    "    @classmethod\n",
+    "    def load_dataset_and_load_vectorizer(cls, dataset_csv, vectorizer_filepath):\n",
+    "        \"\"\"Load dataset and the corresponding vectorizer. \n",
+    "        Used in the case in the vectorizer has been cached for re-use\n",
+    "        \n",
+    "        Args:\n",
+    "            surname_csv (str): location of the dataset\n",
+    "            vectorizer_filepath (str): location of the saved vectorizer\n",
+    "        Returns:\n",
+    "            an instance of SurnameDataset\n",
+    "        \"\"\"\n",
+    "        text_df = pd.read_csv(dataset_csv).fillna(' ')\n",
+    "        vectorizer = cls.load_vectorizer_only(vectorizer_filepath)\n",
+    "        return cls(text_df, vectorizer)\n",
+    "\n",
+    "    @staticmethod\n",
+    "    def load_vectorizer_only(vectorizer_filepath):\n",
+    "        \"\"\"a static method for loading the vectorizer from file\n",
+    "        \n",
+    "        Args:\n",
+    "            vectorizer_filepath (str): the location of the serialized vectorizer\n",
+    "        Returns:\n",
+    "            an instance of SurnameVectorizer\n",
+    "        \"\"\"\n",
+    "        with open(vectorizer_filepath) as fp:\n",
+    "            return NMTVectorizer.from_serializable(json.load(fp))\n",
+    "\n",
+    "    def save_vectorizer(self, vectorizer_filepath):\n",
+    "        \"\"\"saves the vectorizer to disk using json\n",
+    "        \n",
+    "        Args:\n",
+    "            vectorizer_filepath (str): the location to save the vectorizer\n",
+    "        \"\"\"\n",
+    "        with open(vectorizer_filepath, \"w\") as fp:\n",
+    "            json.dump(self._vectorizer.to_serializable(), fp)\n",
+    "\n",
+    "    def get_vectorizer(self):\n",
+    "        \"\"\" returns the vectorizer \"\"\"\n",
+    "        return self._vectorizer\n",
+    "\n",
+    "    def set_split(self, split=\"train\"):\n",
+    "        self._target_split = split\n",
+    "        self._target_df, self._target_size = self._lookup_dict[split]\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return self._target_size\n",
+    "\n",
+    "    def __getitem__(self, index):\n",
+    "        \"\"\"the primary entry point method for PyTorch datasets\n",
+    "        \n",
+    "        Args:\n",
+    "            index (int): the index to the data point \n",
+    "        Returns:\n",
+    "            a dictionary holding the data point: (x_data, y_target, class_index)\n",
+    "        \"\"\"\n",
+    "        row = self._target_df.iloc[index]\n",
+    "\n",
+    "        vector_dict = self._vectorizer.vectorize(row.source_language, row.target_language)\n",
+    "\n",
+    "        return {\"x_source\": vector_dict[\"source_vector\"], \n",
+    "                \"x_target\": vector_dict[\"target_x_vector\"],\n",
+    "                \"y_target\": vector_dict[\"target_y_vector\"], \n",
+    "                \"x_source_length\": vector_dict[\"source_length\"]}\n",
+    "        \n",
+    "    def get_num_batches(self, batch_size):\n",
+    "        \"\"\"Given a batch size, return the number of batches in the dataset\n",
+    "        \n",
+    "        Args:\n",
+    "            batch_size (int)\n",
+    "        Returns:\n",
+    "            number of batches in the dataset\n",
+    "        \"\"\"\n",
+    "        return len(self) // batch_size\n",
+    "    \n",
+    "\n",
+    "\n",
+    "\n",
+    "def generate_nmt_batches(dataset, batch_size, shuffle=True, \n",
+    "                            drop_last=True, device=\"cpu\"):\n",
+    "    \"\"\"A generator function which wraps the PyTorch DataLoader.  The NMT Version \"\"\"\n",
+    "    dataloader = DataLoader(dataset=dataset, batch_size=batch_size,\n",
+    "                            shuffle=shuffle, drop_last=drop_last)\n",
+    "\n",
+    "    for data_dict in dataloader:\n",
+    "        lengths = data_dict['x_source_length'].numpy()\n",
+    "        # Get the indices according to sorted length\n",
+    "        sorted_length_indices = lengths.argsort()[::-1].tolist()\n",
+    "        \n",
+    "        # Sort the minibatch\n",
+    "        out_data_dict = {}\n",
+    "        for name, tensor in data_dict.items():\n",
+    "            out_data_dict[name] = data_dict[name][sorted_length_indices].to(device)\n",
+    "        yield out_data_dict\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "class PositionalEncoding(nn.Module):\n",
+    "    def __init__(self, emb_size, drop_out, max_len:int = 200):\n",
+    "        super(PositionalEncoding, self).__init__()\n",
+    "        den = torch.exp(-torch.arange(0, emb_size,2)*math.log(10000)/emb_size)\n",
+    "        pos = torch.arange(0,max_len).reshape(max_len,1)\n",
+    "        pos_embedding = torch.zeros((max_len, emb_size))\n",
+    "        pos_embedding[:,0::2]= torch.sin(pos*den)\n",
+    "        pos_embedding[:,1::2] = torch.cos(pos*den)\n",
+    "        pos_embedding = pos_embedding.unsqueeze(-2)\n",
+    "        self.dropout = nn.Dropout(drop_out)\n",
+    "        self.register_buffer('pos_embedding', pos_embedding)\n",
+    "\n",
+    "    def forward(self, token_embedding:Tensor):\n",
+    "        return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0),:])\n",
+    "\n",
+    "class TokenEmbedding(nn.Module):\n",
+    "    def __init__(self, vocab_size:int, emb_size):\n",
+    "        super(TokenEmbedding, self).__init__()\n",
+    "        self.embedding = nn.Embedding(vocab_size, emb_size)\n",
+    "        self.emb_size = emb_size\n",
+    "\n",
+    "    def forward(self, tokens:Tensor):\n",
+    "        return self.embedding(tokens.long())*math.sqrt(self.emb_size)\n",
+    "\n",
+    "\n",
+    "class Seq2SeqTransformer(nn.Module):\n",
+    "    def __init__(self, num_encoder_layers,num_decoder_layers, emb_size, nhead,src_vocab_size,tgt_vocab_size, dim_feedforward = 512, dropout = 0.1):\n",
+    "        super(Seq2SeqTransformer,self).__init__()\n",
+    "        self.transformer = Transformer(d_model = emb_size, nhead = nhead, num_encoder_layers = num_encoder_layers, num_decoder_layers = num_decoder_layers, dim_feedforward = dim_feedforward, dropout = dropout,  norm_first = True)\n",
+    "        self.generator = nn.Linear(emb_size, tgt_vocab_size)\n",
+    "        self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size)\n",
+    "        self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size)\n",
+    "        self.positional_encoding = PositionalEncoding(emb_size, drop_out = dropout)\n",
+    "\n",
+    "    def forward(self, src:Tensor, trg:Tensor, src_mask:Tensor, tgt_mask: Tensor, src_padding_mask: Tensor, tgt_padding_mask: Tensor, memory_key_padding_mask: Tensor):\n",
+    "        src_emb = self.positional_encoding(self.src_tok_emb(src))\n",
+    "        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))\n",
+    "        outs = self.transformer(src_emb, tgt_emb, src_mask, tgt_mask, None, src_padding_mask, tgt_padding_mask, memory_key_padding_mask)\n",
+    "        return self.generator(outs)\n",
+    "\n",
+    "    def encode(self, src, src_mask):\n",
+    "        return self.transformer.encoder(self.positional_encoding(self.src_tok_emb(src)),src_mask)\n",
+    "\n",
+    "    def decode(self, tgt:Tensor, memory:Tensor, tgt_mask:Tensor):\n",
+    "        return self.transformer.decoder(self.positional_encoding(self.tgt_tok_emb(tgt)), memory, tgt_mask)\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "def set_seed_everywhere(seed, cuda):\n",
+    "    #seed = self.seed\n",
+    "    #cuda = self.cuda\n",
+    "    np.random.seed(seed)\n",
+    "    torch.manual_seed(seed)\n",
+    "    print(seed)\n",
+    "    if cuda:\n",
+    "        torch.cuda.manual_seed_all(seed)\n",
+    "\n",
+    "\n",
+    "def generate_square_subsequent_mask(sz):\n",
+    "    mask = (torch.triu(torch.ones((sz, sz), device=DEVICE)) == 1).transpose(0, 1)\n",
+    "    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))\n",
+    "    return mask\n",
+    "\n",
+    "\n",
+    "\n",
+    "def handle_dirs(save_dirs):\n",
+    "    dirpath = save_dir\n",
+    "    if not os.path.exists(dirpath):\n",
+    "        os.makedirs(dirpath)\n",
+    "\n",
+    "\n",
+    "\n",
+    "def create_mask(src, tgt,PAD_IDX):\n",
+    "    src_seq_len = src.shape[0]\n",
+    "    tgt_seq_len = tgt.shape[0]\n",
+    "        \n",
+    "    tgt_mask = generate_square_subsequent_mask(tgt_seq_len)\n",
+    "    src_mask = torch.zeros((src_seq_len, src_seq_len),device=DEVICE).type(torch.bool)\n",
+    "        \n",
+    "    src_padding_mask = (src == PAD_IDX).transpose(0, 1)\n",
+    "    tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)\n",
+    "    return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask\n",
+    "\n",
+    "\n",
+    "\n",
+    "def train_epoch(batch_size, device, model, dataset, split_value, optimizer, PAD_IDX, loss_fn):\n",
+    "    BATCH_SIZE = batch_size\n",
+    "    model.train()\n",
+    "    losses = 0\n",
+    "    print(dataset.__len__())\n",
+    "    train_dataloader = DataLoader(dataset, batch_size=BATCH_SIZE)\n",
+    "    #print(BATCH_SIZE,len(list(train_dataloader)))\n",
+    "    dataset.set_split(split_value)\n",
+    "    batch_generator = generate_nmt_batches(dataset, batch_size=BATCH_SIZE, device = device)\n",
+    "    print(\"printing batch generator\",batch_generator)\n",
+    "    ctr = 0\n",
+    "    for batch_index, batch_dict in enumerate(batch_generator):\n",
+    "        ctr = ctr+1\n",
+    "        #optimizer.zero_grad()\n",
+    "        #print(torch.cat((torch.transpose(batch_dict['x_source'],0,1),torch.transpose(batch_dict['x_target'],0,1),torch.transpose(batch_dict['y_target'],0,1)),1).numpy().shape)\n",
+    "        #print(torch.transpose(batch_dict['x_target'],0,1))\n",
+    "        #print(torch.transpose(batch_dict['y_target'],0,1))\n",
+    "        src=torch.transpose(batch_dict['x_source'],0,1)\n",
+    "        tgt=torch.transpose(batch_dict['y_target'],0,1)\n",
+    "        tgt_input = tgt[:-1,:]\n",
+    "        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src,tgt_input, PAD_IDX)\n",
+    "        logits = model(src,tgt_input, src_mask, tgt_mask, src_padding_mask, tgt_padding_mask, src_padding_mask)\n",
+    "        optimizer.zero_grad()\n",
+    "        tgt_out = tgt[1:,:]\n",
+    "        loss = loss_fn(logits.reshape(-1, logits.shape[-1]),tgt_out.reshape(-1))\n",
+    "        loss.backward()\n",
+    "        optimizer.step()\n",
+    "        losses += loss.item()\n",
+    "        if ctr%50==0:\n",
+    "            #print('source_shape',src.shape, 'target_shape',tgt.shape)\n",
+    "            print(\"ctr: \",ctr,\"   losses:  \",losses/ctr,'time',datetime.datetime.now())#,\"   len_train_dataloader:  \",len(list(train_dataloader)))\n",
+    "    return losses/len(list(train_dataloader))\n",
+    "\n",
+    "\n",
+    "def evaluate(batch_size,device,model, dataset,split_value,PAD_IDX,loss_fn):\n",
+    "    model.eval()\n",
+    "    losses = 0\n",
+    "    dataset.set_split(split_value)\n",
+    "    val_dataloader=DataLoader(dataset, batch_size=batch_size)\n",
+    "    batch_generator=generate_nmt_batches(dataset, batch_size=batch_size, device=device)\n",
+    "    ctr  = 0\n",
+    "    for batch_index, batch_dict in enumerate(batch_generator):\n",
+    "        src = torch.transpose(batch_dict['x_source'],0,1)\n",
+    "        tgt = torch.transpose(batch_dict['y_target'],0,1)\n",
+    "        tgt_input = tgt[:-1,:]\n",
+    "        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src,tgt_input, PAD_IDX)\n",
+    "        logits = model(src,tgt_input,src_mask,tgt_mask, src_padding_mask, tgt_padding_mask, src_padding_mask)\n",
+    "        tgt_out=tgt[1:,:]\n",
+    "        loss = loss_fn(logits.reshape(-1, logits.shape[-1]),tgt_out.reshape(-1))#loss_fn(logits.reshape[-1],tgt_out.reshape[-1])\n",
+    "        losses += loss.item()\n",
+    "        ctr = ctr+1\n",
+    "        print(ctr,\"validation\",losses/ctr)\n",
+    "\n",
+    "    \"\"\"for src, tgt in val_dataloader:\n",
+    "        src = src.to(DEVICE)\n",
+    "        tgt = tgt.to(DEVICE)\n",
+    "\n",
+    "        tgt_input = tgt[:-1, :]\n",
+    "\n",
+    "        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)\n",
+    "\n",
+    "        logits = model(src, tgt_input, src_mask, tgt_mask,src_padding_mask, tgt_padding_mask, src_padding_mask)\n",
+    "\n",
+    "        tgt_out = tgt[1:, :]\n",
+    "        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))\n",
+    "        losses += loss.item()\"\"\"\n",
+    "    return losses/len(list(val_dataloader))\n",
+    "\n",
+    "\n",
+    "\n",
+    "def greedy_decode(DEVICE, model, src, src_mask, max_len, start_symbol, EOS_IDX):\n",
+    "    src = src.to(DEVICE)\n",
+    "    src_mask=src_mask.to(DEVICE)\n",
+    "    memory = model.encode(src, src_mask)\n",
+    "    ys = torch.ones(1,1).fill_(start_symbol).type(torch.long).to(DEVICE)\n",
+    "    for i in range(max_len):\n",
+    "        #print(i,'ys',ys)\n",
+    "        memory = memory.to(DEVICE)\n",
+    "        tgt_mask = (generate_square_subsequent_mask(ys.size(0)).type(torch.bool)).to(DEVICE)\n",
+    "        #print('tgt_mask',tgt_mask)\n",
+    "        out = model.decode(ys,memory, tgt_mask)#.squeeze()\n",
+    "        #print(\"out\",out,'out_shape',out.shape)\n",
+    "        out = out.transpose(0,1)\n",
+    "        #print(\"out transpose\",out,'out_transpose_shape',out.shape)\n",
+    "        prob = model.generator(out)[:,-1]\n",
+    "        _, next_word = torch.max(prob, dim=1)\n",
+    "        next_word = next_word.item()\n",
+    "        #print('next_word = ',next_word)\n",
+    "        ys = torch.cat([ys, torch.ones(1,1).type_as(src.data).fill_(next_word)], dim = 0)\n",
+    "        #print('ys',ys)\n",
+    "        if next_word == EOS_IDX:\n",
+    "            break\n",
+    "    return ys\n",
+    "\n",
+    "\n",
+    "\n",
+    "def translate( device,model:torch.nn.Module, src_sentence:str, BOS_IDX, EOS_IDX):\n",
+    "    model.eval()\n",
+    "    src= src_sentence\n",
+    "    #print('src',src)\n",
+    "    num_tokens = src.shape[0]\n",
+    "    #print(num_tokens)\n",
+    "    src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)\n",
+    "    #print('src_mask',src_mask)\n",
+    "    tgt_tokens = greedy_decode(device,model, src, src_mask, max_len = num_tokens, start_symbol=BOS_IDX, EOS_IDX=EOS_IDX).flatten()\n",
+    "    return tgt_tokens\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "input_df = 'dataset_for_APE_hinglish_to_english2.csv'\n",
+    "fpath = \"nmt_IITB_APE2\"\n",
+    "\n",
+    "\n",
+    "#dataset = NMTDataset.load_dataset_and_make_vectorizer('IITB_dataset_1.csv')\n",
+    "#dataset.save_vectorizer(\"vectorizer_transformer_3layer_IITB1mill.json\")\n",
+    "\n",
+    "\n",
+    "\n",
+    "#dataloader = DataLoader(dataset=dataset, batch_size=1024,shuffle=False, drop_last=True)\n",
+    "\n",
+    "dataset_csv = 'dataset_for_APE_hinglish_to_english2.csv'\n",
+    "vectorizer_file = 'vectorizer_APE_2.json'\n",
+    "print(vectorizer_file)\n",
+    "model_state_file = 'APE_2.pth'\n",
+    "save_dir = \"nmt_DG2_FFNN8192\"#'GenV1_Transforemer_1',\n",
+    "print(save_dir)\n",
+    "reload_from_files = True\n",
+    "cuda = False\n",
+    "seed = 13\n",
+    "learning_rate = 8e-3\n",
+    "batch_size = 1024\n",
+    "batch_size_val = 1\n",
+    "num_epochs = 40\n",
+    "source_embedding_size = 256\n",
+    "target_embedding_size = 256\n",
+    "encoding_size = 256\n",
+    "use_glove = False\n",
+    "expand_filepaths_to_save_dir = True\n",
+    "early_stopping_criteria = 10\n",
+    "dataset_to_evaluate = 'dataset_for_APE_hinglish_to_english2.csv'\n",
+    "path_to_save = 'APE_1_new.csv'\n",
+    "saved_model_path = 'APE_1_new.pt'\n",
+    "file_exist = 0\n",
+    "existing_file_name = 'dataset_for_APE_hinglish_to_english2.csv'\n",
+    "\n",
+    "\n",
+    "dataset_path = fpath\n",
+    "existing_file_name = input_df\n",
+    "fname = existing_file_name\n",
+    "dataset_csv = fname\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "model_state_file = model_state_file\n",
+    "save_dir = save_dir\n",
+    "print(save_dir)\n",
+    "reload_from_files = reload_from_files\n",
+    "expand_filepaths_to_save_dir = expand_filepaths_to_save_dir\n",
+    "cuda = cuda\n",
+    "seed = seed\n",
+    "learning_rate = learning_rate\n",
+    "batch_size = batch_size\n",
+    "batch_size_val = batch_size_val\n",
+    "num_epochs = num_epochs\n",
+    "early_stopping_criteria = True#self.early_stopping_criteria\n",
+    "source_embedding_size = source_embedding_size\n",
+    "target_embedding_size = target_embedding_size\n",
+    "encoding_size = encoding_size\n",
+    "use_glove = False\n",
+    "catch_keyboard_interrupt = True\n",
+    "if expand_filepaths_to_save_dir:\n",
+    "    vectorizer_file = os.path.join(save_dir, vectorizer_file)\n",
+    "model_state_file = os.path.join(save_dir, model_state_file)\n",
+    "if not torch.cuda.is_available():\n",
+    "    cuda = False\n",
+    "device = torch.device(\"cuda\" if cuda else \"cpu\")\n",
+    "set_seed_everywhere(seed,cuda)\n",
+    "handle_dirs(save_dir)\n",
+    "if reload_from_files and os.path.exists(vectorizer_file):\n",
+    "    dataset = NMTDataset.load_dataset_and_load_vectorizer(dataset_csv, vectorizer_file)\n",
+    "    print('load_dataset_and_load_vectorizer______')\n",
+    "else:\n",
+    "    dataset = NMTDataset.load_dataset_and_make_vectorizer(dataset_csv)\n",
+    "    dataset.save_vectorizer(vectorizer_file)\n",
+    "    print('_________load_dataset_and_make_vectorizer______')\n",
+    "vectorizer = dataset.get_vectorizer()\n",
+    "PAD_IDX = vectorizer.to_serializable()['target_vocab']['token_to_idx']['<MASK>']\n",
+    "BOS_IDX = vectorizer.to_serializable()['target_vocab']['token_to_idx']['<BEGIN>']\n",
+    "EOS_IDX = vectorizer.to_serializable()['target_vocab']['token_to_idx']['<END>']\n",
+    "SRC_VOCAB_SIZE = len(vectorizer.to_serializable()['source_vocab']['token_to_idx'])\n",
+    "TGT_VOCAB_SiZE = len(vectorizer.to_serializable()['target_vocab']['token_to_idx'])\n",
+    "print('target vocab size',TGT_VOCAB_SiZE)\n",
+    "print('dataset_size 1: ', dataset.__len__(), dataset_path, dataset_csv)\n",
+    "print(' dataset csv length',len(pd.read_csv(dataset_csv)))\n",
+    "EMB_SIZE = 256\n",
+    "NHEAD = 16\n",
+    "FFN_HID_DIM =8192\n",
+    "BATCH_SIZE = 128\n",
+    "NUM_ENCODER_LAYERS = 3\n",
+    "NUM_DECODER_LAYERS = 3\n",
+    "batch_size = BATCH_SIZE\n",
+    "transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS, EMB_SIZE, NHEAD, SRC_VOCAB_SIZE, TGT_VOCAB_SiZE, FFN_HID_DIM)\n",
+    "transformer = transformer.to(DEVICE)\n",
+    "loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)\n",
+    "optimizer = torch.optim.Adam(transformer.parameters(), lr=0.004, betas = (0.99, 0.99), eps = 1e-9)\n",
+    "from timeit import default_timer as timer\n",
+    "NUM_EPOCHS = num_epochs\n",
+    "for epoch in range(1, NUM_EPOCHS+1):\n",
+    "    print(\"==================Training started==================\",epoch)\n",
+    "    start_time = timer()\n",
+    "    split_value_train = 'train'\n",
+    "    split_value_validate = 'val'\n",
+    "    train_loss = train_epoch(batch_size,device,transformer, dataset, split_value_train, optimizer, PAD_IDX, loss_fn)\n",
+    "    end_time = timer()\n",
+    "    torch.save(transformer,'epoch'+str(epoch)+'_APE_2_new.pt')\n",
+    "#torch.save(transformer, save_dir+\"/\"+saved_model_path+\"_epoch\")\n",
+    "    #val_loss = evaluate(batch_size,device,transformer, dataset, split_value_validate, PAD_IDX, loss_fn)\n",
+    "torch.save(transformer, save_dir+\"/\"+saved_model_path)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "37a50cf7-d754-4c19-aaa5-4e094cfd87e6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

digital_green_process_data.py

@@ -0,0 +1,62 @@
+import os
+import pandas as pd
+from datasets import Dataset, DatasetDict, Audio
+import soundfile as sf
+import numpy as np
+from sklearn.model_selection import train_test_split
+
+# Paths
+audio_folder = '/home/azureuser/data2/dg_16/'  # Path where your audio files are stored
+csv_file = 'digital_green_recordings.csv'   # Path to the CSV that contains audio paths and transcripts
+
+# Read your CSV file (assumes it has columns: 'path' and 'transcript')
+df = pd.read_csv(csv_file, sep="$")
+
+# Create a new column for client_id (random or default if you don’t have speaker info)
+df['client_id'] = ['speaker_' + str(i) for i in range(len(df))]
+
+# If your CSV has relative paths, ensure the paths are correct
+df['path'] = df['path'].apply(lambda x: os.path.join(audio_folder, x))
+
+# Add additional columns needed for the Common Voice format (can be optional)
+df['up_votes'] = 0
+df['down_votes'] = 0
+df['age'] = None
+df['gender'] = None
+df['accent'] = None
+
+# Function to load and possibly convert audio to mono
+def load_audio(file_path):
+        # Load audio file
+    audio, sr = sf.read(file_path)
+                # Convert to mono if stereo
+    if len(audio.shape) > 1:
+        audio = np.mean(audio, axis=1)
+        return {'audio': {'array': audio, 'sampling_rate': sr}}
+
+                                # Apply audio loading function to DataFrame
+df['audio'] = df['path'].apply(lambda x: load_audio(x))
+
+train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)  # Adjust test_size as needed
+
+# Convert DataFrames to Hugging Face Datasets
+train_dataset = Dataset.from_pandas(train_df)
+test_dataset = Dataset.from_pandas(test_df)
+
+# Cast the 'audio' column to the 'audio' type
+train_dataset = train_dataset.cast_column('audio', Audio())
+test_dataset = test_dataset.cast_column('audio', Audio())
+
+                        # Create a DatasetDict to simulate train/test/validation splits if needed
+dataset_dict = DatasetDict({
+                        'train': train_dataset,
+                        'test': test_dataset # If you have separate splits, add them here (e.g., 'train', 'test', 'validation')
+                        })
+
+# Save the dataset (optional) for future use
+dataset_dict.save_to_disk('data2/digital_green_data')
+
+                        # Print a sample from the dataset
+print(dataset_dict['train'][0])
+
+print(dataset_dict['test'][0])