Upload 18 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+datasets/train_clean.dat filter=lfs diff=lfs merge=lfs -text
+datasets/train_other.dat filter=lfs diff=lfs merge=lfs -text

config.py

@@ -0,0 +1,26 @@
+from typing import Dict, Any
+
+import yaml
+
+from pathlib import Path
+
+from utils import get_full_file_path
+
+def load_config(file_name: str = 'config.yaml') -> Dict[str, Any]:
+    """
+    Loads a YAML config from a relative file path
+    """
+    file_path = get_full_file_path(file_name=file_name)
+    with open(file_path, 'r') as file:
+        config = yaml.safe_load(file)
+
+    return config
+
+# print(list(load_config().items()))
+
+def get_weights_file_path(config, epoch: str) -> str:
+    model_folder = config['model']['model_folder']
+    model_basename = config['model']['model_basename']
+    model_filename = f'{model_basename}{epoch}.pt'
+    return str(Path('.') / model_folder / model_filename)
+

config_small.yaml

@@ -0,0 +1,23 @@
+experiment_name: "runs/transformer"
+
+dataset:
+  src_lang: 'lo'
+  src_tokenizer: 'BPE'
+  src_max_seq_len: 400
+  tgt_lang: 'vi'
+  tgt_tokenizer: 'WordLevel'
+  tgt_max_seq_len: 350
+  train_dataset: 'train_clean.dat'
+  validate_dataset: 'dev_clean.dat'
+  tokenizer_file: "tokenizer_{0}.json"
+
+model: # 16629775 parameters
+  d_model: 256
+  num_heads: 8
+  d_ff: 1024
+  dropout_p: 0.3
+  num_encoder_layers: 4
+  num_decoder_layers: 2
+  model_folder: "weights"
+  model_basename: "transformer_"
+  preload: "small"

datasets/train_clean.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d56aa43dba2e12da818bf80e7b42d0908637ee454ced035ae56f02022e13083
+size 18582873

datasets/train_other.dat

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:081d7f90d8098f81a8c37fd50a0d8e2d193a2e0e2cc5048b08ca4ef218100d18
+size 18008839

layers/decoder_layer.py

@@ -0,0 +1,57 @@
+from typing import Tuple
+
+import torch.nn as nn
+from torch import Tensor
+
+from modules.multi_head_attention import MultiHeadAttention
+from modules.positionwise_feed_forward import PositionwiseFeedForwardNetwork
+
+
+class DecoderLayer(nn.Module):
+    """
+    A Decoder layer.
+    
+    Args:
+    """
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        d_ff: int,
+        dropout_p: int,
+    ) -> None:
+        super(DecoderLayer, self).__init__()
+        self.self_attn_prenorm = nn.LayerNorm(d_model)
+        self.self_attn = MultiHeadAttention(d_model=d_model, num_heads=num_heads, dropout_p=dropout_p)
+        self.self_attn_dropout = nn.Dropout(p=dropout_p)
+
+        self.cross_attn_prenorm = nn.LayerNorm(d_model)
+        self.cross_attn = MultiHeadAttention(d_model=d_model, num_heads=num_heads, dropout_p=dropout_p)
+        self.cross_attn_dropout = nn.Dropout(p=dropout_p)
+
+        self.feed_forward_prenorm = nn.LayerNorm(d_model)
+        self.feed_forward = PositionwiseFeedForwardNetwork(d_model=d_model, d_ff=d_ff, dropout_p=dropout_p)
+
+    def forward(
+        self, 
+        decoder_inputs: Tensor,
+        encoder_outputs: Tensor,
+        src_mask: Tensor,
+        tgt_mask: Tensor,
+    ) -> Tuple[Tensor, Tensor]:
+        residual = decoder_inputs
+        outputs = self.self_attn_prenorm(decoder_inputs)
+        outputs, attn = self.self_attn(outputs, outputs, outputs, tgt_mask)
+        outputs = self.self_attn_dropout(outputs) + residual
+
+        residual = outputs
+        outputs = self.self_attn_prenorm(outputs)
+        outputs, attn = self.self_attn(outputs, encoder_outputs, encoder_outputs, src_mask)
+        outputs = self.self_attn_dropout(outputs) + residual
+
+        residual = outputs
+        outputs = self.feed_forward_prenorm(outputs)
+        outputs = self.feed_forward(outputs)
+        outputs += residual
+
+        return outputs, attn

layers/encoder_layer.py

@@ -0,0 +1,42 @@
+from typing import Tuple
+
+import torch.nn as nn
+from torch import Tensor
+
+from modules.multi_head_attention import MultiHeadAttention
+from modules.positionwise_feed_forward import PositionwiseFeedForwardNetwork
+
+class EncoderLayer(nn.Module):
+    """
+    An Encoder layer.
+    
+    Args:
+    """
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        d_ff: int,
+        dropout_p: int,
+    ) -> None:
+        super(EncoderLayer, self).__init__()
+        self.self_attn_prenorm = nn.LayerNorm(d_model)
+        self.self_attn = MultiHeadAttention(d_model=d_model, num_heads=num_heads, dropout_p=dropout_p)
+        self.self_attn_dropout = nn.Dropout(p=dropout_p)
+
+        self.feed_forward_prenorm = nn.LayerNorm(d_model)
+        self.feed_forward = PositionwiseFeedForwardNetwork(d_model=d_model, d_ff=d_ff, dropout_p=dropout_p)
+
+    def forward(self, inputs: Tensor, src_mask: Tensor = None) -> Tuple[Tensor, Tensor]:
+        # Normalize -> sublayer -> dropout -> add residual
+        residual = inputs
+        inputs = self.self_attn_prenorm(inputs)
+        outputs, attn = self.self_attn(inputs, inputs, inputs, src_mask)
+        outputs = self.self_attn_dropout(outputs) + residual
+
+        residual = outputs
+        outputs = self.feed_forward_prenorm(outputs)
+        outputs = self.feed_forward(outputs)
+        outputs += residual
+
+        return outputs, attn

layers/projection_layer.py

@@ -0,0 +1,16 @@
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from modules.wrapper import Linear
+
+class ProjectionLayer(nn.Module):
+    def __init__(self, d_model: int, vocab_size: int) -> None:
+        super(ProjectionLayer, self).__init__()
+        self.linear = Linear(d_model, vocab_size)
+
+    def forward(self, x):
+        # (batch, seq_len, d_model) -> (batch, seq_len, vocab_size)
+        return torch.log_softmax(self.linear(x), dim=-1)

model/encoder.py

@@ -0,0 +1,38 @@
+from typing import Tuple
+
+import torch.nn as nn
+from torch import Tensor
+
+from layers.encoder_layer import EncoderLayer
+
+class Encoder(nn.Module):
+    """
+    A transformer Encoder (no embeddings or positional embeddings)
+
+    Args:
+    """
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        d_ff: int,
+        dropout_p: int,
+        num_layers: int,
+    ) -> None:
+        super(Encoder, self).__init__()
+        self.layers = nn.ModuleList(
+            [
+                EncoderLayer(
+                    d_model=d_model,
+                    num_heads=num_heads,
+                    d_ff=d_ff,
+                    dropout_p=dropout_p,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def forward(self, x: Tensor, src_mask: Tensor):
+        for layer in self.layers:
+            x, attn = layer(x, src_mask)
+        return x

modules/dot_product_attention.py

@@ -0,0 +1,55 @@
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+class ScaledDotProductAttention(nn.Module):
+    """
+    Scaled Dot-Product Attention (section 3.2.1)
+
+    Args:
+        - dim (int): dimension of d_k or d_head
+        - dropout_p (float): probability of dropout
+
+    Input:
+        - query (batch, num_heads, seq_len, d_head)
+        - key   (batch, num_heads, seq_len, d_head)
+        - value (batch, num_heads, seq_len, d_head)
+        - mask  ()
+
+    Output:
+        - context (batch, num_head, seq_len, d_head): Context matrix.
+        - attn (batch, num_head, seq_len, seq_len): Attention matrix for visualization.
+    """
+    def __init__(self, dim: int, dropout_p: float) -> None:
+        super(ScaledDotProductAttention, self).__init__()
+        self.sqrt_dim = np.sqrt(dim)
+        self.dropout = nn.Dropout(p = dropout_p)
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        mask: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        
+        # (batch, num_heads, seq_len, d_head) @ (batch, num_heads, d_head, seq_len)
+        # ==> score: (batch, num_heads, seq_len, seq_len)
+        score = torch.matmul(query, key.transpose(-2, -1)) / self.sqrt_dim
+
+        if mask is not None:
+            score.masked_fill_(mask == 0, -1e4)
+        
+        attn = F.softmax(score, -1) 
+        # (batch, num_head, seq_len, seq_len)
+        attn = self.dropout(attn)
+        
+        context = torch.matmul(attn, value)
+        # (batch, num_head, seq_len, d_head)
+
+        return context, attn
+

modules/multi_head_attention.py

@@ -0,0 +1,79 @@
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from modules.wrapper import Linear
+from modules.dot_product_attention import ScaledDotProductAttention
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-Head Attention (section 3.2.2)
+
+    Args:
+        - d_model (int): dimension of model
+        - num_heads (int): number of heads
+        - dropout_p (float): probability of dropout
+
+    Inputs:
+        - query (batch, seq_len, d_model):
+        - key   (batch, seq_len, d_model):
+        - value (batch, seq_len, d_model):
+        - mask ():
+
+    Output: (Tensor, Tensor):
+        - context ()
+        - attn (): Attention matrix for visualization.
+    """
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dropout_p: int,
+    ) -> None:
+        super(MultiHeadAttention, self).__init__()
+
+        assert d_model % num_heads == 0, "d_model % num_heads should be 0"
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.d_head = d_model // num_heads
+        
+        self.W_query  = Linear(d_model, d_model)
+        self.W_key    = Linear(d_model, d_model)
+        self.W_value  = Linear(d_model, d_model)
+        # self.W_output = Linear(d_model, d_model)
+
+        self.scaled_dot_attn = ScaledDotProductAttention(d_model, dropout_p)
+
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        mask: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        batch_size = query.shape[0]
+
+        # original: (batch, seq_len, d_model) 
+        # --forward--> (batch, seq_len, d_model) 
+        # --view--> (batch, seq_len, num_heads, d_head) 
+        # --transpose--> (batch, num_heads, seq_len, d_head)
+        query = self.W_query(query).view(batch_size, -1, self.num_heads, self.d_head).transpose(1,2)
+        key   =     self.W_key(key).view(batch_size, -1, self.num_heads, self.d_head).transpose(1,2)
+        value = self.W_value(value).view(batch_size, -1, self.num_heads, self.d_head).transpose(1,2)
+        
+        context, attn = self.scaled_dot_attn(query, key, value, mask)
+
+        # (batch, num_heads, seq_len, d_head)
+        # --transpose--> (batch, seq_len, num_heads, d_head)
+        # --view--> (batch, seq_len, d_model)
+
+        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
+        # context = self.W_output(context)
+
+        return context, attn
+    

modules/positional_encoding.py

@@ -0,0 +1,34 @@
+import math
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class PositionalEncoding(nn.Module):
+    r"""
+    Positional Encoding in "Attention Is All You Need" (section 3.5).
+
+    "Attention Is All You Need" uses sine and cosine functions of different frequencies:
+        PE_(pos, 2i)    =  sin(pos / power(10000, 2i / d_model))
+        PE_(pos, 2i+1)  =  cos(pos / power(10000, 2i / d_model))
+    only change is that calculations are done with -log(power(10000, 2i / d_model))
+
+    Uses OpenSpeech's PositionalEncoding, as I don't see the point in coding this from scratch.
+    """
+    def __init__(self, d_model: int, dropout_p: float, max_length: int = 5000) -> None:
+        super(PositionalEncoding, self).__init__()
+
+        self.dropout = nn.Dropout(p=dropout_p)
+
+        pe = torch.zeros(max_length, d_model, requires_grad=False)
+        position = torch.arange(0, max_length, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)
+        return self.dropout(x)

modules/positionwise_feed_forward.py

@@ -0,0 +1,28 @@
+import torch.nn as nn
+from torch import Tensor
+
+from modules.wrapper import Linear
+
+class PositionwiseFeedForwardNetwork(nn.Module):
+    """
+    Position-wise Feed-Forward Network (section 3.3)
+
+    Args:
+        - d_model (int): dimension of input and output
+        - d_ff (int): dimension of inner-layer
+        - dropout_p (float): dropout probability
+    """
+
+    def __init__(self, d_model: int, d_ff: int, dropout_p: float) -> None:
+        super(PositionwiseFeedForwardNetwork, self).__init__()
+
+        self.feed_forward = nn.Sequential(
+            Linear(d_model, d_ff),
+            nn.Dropout(dropout_p),
+            nn.ReLU(),
+            Linear(d_ff, d_model),
+            nn.Dropout(dropout_p)
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.feed_forward(x)

modules/transformer_embedding.py

@@ -0,0 +1,23 @@
+import torch.nn as nn
+from torch import Tensor
+import numpy as np
+
+
+class TransformerEmbedding(nn.Module):
+    """
+    Input Embeddings (section 3.4)
+
+    Embedds words to vectors of size d_
+
+    Args:
+        - d_model (int): dimension of model
+        - num_embeddings (int): size of the dictionary
+    """
+    def __init__(self, d_model: int, num_embeddings: int) -> None:
+        super(TransformerEmbedding, self).__init__()
+        self.sqrt_d_model = np.sqrt(d_model)
+        self.embedding = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=d_model)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.embedding(x) * self.sqrt_d_model
+        

modules/wrapper.py

@@ -0,0 +1,28 @@
+import torch.nn as nn
+from torch import Tensor
+
+
+class Linear(nn.Module):
+    """
+    A wrapper class for nn.Linear
+    Initialize values using xxx
+    """
+    def __init__(
+        self, 
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ):
+        super(Linear, self).__init__()
+        self.linear = nn.Linear(in_features, out_features, bias, device, dtype)
+        nn.init.xavier_uniform_(self.linear.weight)
+        if bias:
+            nn.init.zeros_(self.linear.bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        forward pass through linear layer.
+        """
+        return self.linear(x)

transformer.py

@@ -0,0 +1,161 @@
+from typing import Tuple
+
+import torch.nn as nn
+from torch import Tensor
+
+from modules.transformer_embedding import TransformerEmbedding
+from modules.positional_encoding import PositionalEncoding
+
+from model.encoder import Encoder
+from model.decoder import Decoder
+from layers.projection_layer import ProjectionLayer
+
+class Transformer(nn.Module):
+    """
+    Transformer.
+
+    Args:
+        - src_vocab_size (int): source vocabulary size
+        - tgt_vocab_size (int): target vocabulary size
+        - src_max_seq_len (int): source max sequence length
+        - tgt_max_seq_len (int): target max sequence length
+        - d_model (int): dimension of model
+        - num_heads (int): number of heads
+        - d_ff (int): dimension of hidden feed forward layer
+        - dropout_p (float): probability of dropout
+        - num_encoder_layers (int): number of encoder layers
+        - num_decoder_layers (int): number of decoder layers
+    """
+    def __init__(
+        self,
+        src_vocab_size: int,
+        tgt_vocab_size: int,
+        src_max_seq_len: int,
+        tgt_max_seq_len: int,
+        d_model: int = 512,
+        num_heads: int = 8,
+        d_ff: int = 2048,
+        dropout_p: float = 0.1,
+        num_encoder_layers: int = 6,
+        num_decoder_layers: int = 6,
+    ) -> None:
+        super(Transformer, self).__init__()
+        
+        # Embedding layers
+        self.src_embedding = TransformerEmbedding(
+            d_model=d_model,
+            num_embeddings=src_vocab_size
+        )
+        self.tgt_embedding = TransformerEmbedding(
+            d_model=d_model,
+            num_embeddings=tgt_vocab_size
+        )
+
+        # Positional Encoding layers
+        self.src_positional_encoding = PositionalEncoding(
+            d_model=d_model,
+            dropout_p=dropout_p,
+            max_length=src_max_seq_len
+        )
+        self.tgt_positional_encoding = PositionalEncoding(
+            d_model=d_model,
+            dropout_p=dropout_p,
+            max_length=tgt_max_seq_len
+        )
+
+        # Encoder  
+        self.encoder = Encoder(
+            d_model=d_model,
+            num_heads=num_heads,
+            d_ff=d_ff,
+            dropout_p=dropout_p,
+            num_layers=num_encoder_layers
+        )
+        # Decoder
+        self.decoder = Decoder(
+            d_model=d_model,
+            num_heads=num_heads,
+            d_ff=d_ff,
+            dropout_p=dropout_p,
+            num_layers=num_decoder_layers
+        )
+        # projecting decoder's output to the target language.
+        self.projection_layer = ProjectionLayer(
+            d_model=d_model,
+            vocab_size=tgt_vocab_size
+        )
+
+    def encode(
+        self, 
+        src: Tensor, 
+        src_mask: Tensor
+    ) -> Tensor:
+        """
+        Get encoder outputs.
+        """
+        src = self.src_embedding(src)
+        src = self.src_positional_encoding(src)
+        return self.encoder(src, src_mask)
+    
+    def decode(
+        self, 
+        encoder_output: Tensor, 
+        src_mask: Tensor,
+        tgt: Tensor,
+        tgt_mask: Tensor
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Get decoder outputs for a set of target inputs.
+        """
+        tgt = self.tgt_embedding(tgt)
+        tgt = self.tgt_positional_encoding(tgt)
+        return self.decoder(
+            x=tgt,
+            encoder_output=encoder_output,
+            src_mask=src_mask,
+            tgt_mask=tgt_mask
+        )
+    
+    def project(self, decoder_output: Tensor) -> Tensor:
+        """
+        Project decoder outputs to target vocabulary.
+        """
+        return self.projection_layer(decoder_output)
+
+    def forward(
+        self, 
+        src: Tensor, 
+        src_mask: Tensor,
+        tgt: Tensor, 
+        tgt_mask: Tensor
+    ) -> Tuple[Tensor, Tensor]:
+        # src_mask = self.make_src_mask(src)
+        # tgt_mask = self.make_tgt_mask(tgt)
+
+        encoder_output = self.encode(src, src_mask)
+        decoder_output, attn = self.decode(
+            encoder_output, src_mask, tgt, tgt_mask
+        )
+        output = self.project(decoder_output)
+        return output, attn
+    
+    def count_parameters(self):
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+    
+
+def get_model(config, src_vocab_size: int, tgt_vocab_size: int) -> Transformer:
+    """
+    returns a `Transformer` model for a given config.
+    """
+    return Transformer(
+        src_vocab_size=src_vocab_size,
+        tgt_vocab_size=tgt_vocab_size,
+        src_max_seq_len=config['dataset']['src_max_seq_len'],
+        tgt_max_seq_len=config['dataset']['tgt_max_seq_len'],
+        d_model=config['model']['d_model'],
+        num_heads=config['model']['num_heads'],
+        d_ff=config['model']['d_ff'],
+        dropout_p=config['model']['dropout_p'],
+        num_encoder_layers=config['model']['num_encoder_layers'],
+        num_decoder_layers=config['model']['num_decoder_layers'],
+    )

translate.py

@@ -0,0 +1,69 @@
+from typing import Tuple
+
+import torch
+from torch import Tensor
+
+from tokenizers import Tokenizer
+
+from transformer import Transformer
+from decode_method import greedy_decode
+
+def translate(
+    model: Transformer,
+    src_tokenizer: Tokenizer,
+    tgt_tokenizer: Tokenizer,
+    text: str,
+    decode_method: str = 'greedy',
+    device = torch.device('cpu')
+) -> Tuple[str, Tensor]:
+    """
+    Translation function.
+
+    Output:
+        - translation (str): the translated string.
+        - attn (Tensor): The decoder's attention (for visualization)
+    """
+    sos_token = torch.tensor([src_tokenizer.token_to_id('<sos>')], dtype=torch.int64)
+    eos_token = torch.tensor([src_tokenizer.token_to_id('<eos>')], dtype=torch.int64)
+    pad_token = torch.tensor([src_tokenizer.token_to_id('<pad>')], dtype=torch.int64)
+    
+    encoder_input_tokens = src_tokenizer.encode(text).ids
+    # <sos> + source_text + <eos> = encoder_input
+    encoder_input = torch.cat(
+        [
+            sos_token,
+            torch.tensor(encoder_input_tokens, dtype=torch.int64),
+            eos_token,
+        ]
+    )
+    encoder_mask = (encoder_input != pad_token).unsqueeze(0).unsqueeze(0).unsqueeze(0).int() # (1, 1, seq_len)
+
+    encoder_input = encoder_input.unsqueeze(0)
+    # encoder_mask = torch.tensor(encoder_mask)
+    
+    assert encoder_input.size(0) == 1
+
+    if decode_method == 'greedy':
+        model_out, attn = greedy_decode(
+            model, encoder_input, encoder_mask, src_tokenizer, tgt_tokenizer, 400, device, 
+            give_attn=True,
+        )
+    elif decode_method == 'beam-search':
+        raise NotImplementedError
+    else:
+        raise ValueError("Unsuppored decode method")
+
+    model_out_text = tgt_tokenizer.decode(model_out.detach().cpu().numpy())
+    return model_out_text, attn
+    
+
+from config import load_config
+from load_and_save_model import load_model_tokenizer
+if __name__ == '__main__':
+    config = load_config(file_name='config_small.yaml')
+    model, src_tokenizer, tgt_tokenizer = load_model_tokenizer(config)
+    text = "ສະບາຍດີ" # Hello.
+    translation, attn = translate(
+        model, src_tokenizer, tgt_tokenizer, text
+    )
+    print(translation)