Upload 9 files

Configurations/s2s_model_config.json

@@ -0,0 +1,9 @@
+{
+    "dim_embed": 256,
+    "dim_model": 256,
+    "dim_feedforward": 1024,
+    "num_layers": 4,
+    "dropout": 0.1,
+    "maxlen": 512,
+    "flash_attention": false
+}

Configurations/s2sattention_model_config.json

@@ -0,0 +1,9 @@
+{
+    "dim_embed": 256,
+    "dim_model": 256,
+    "dim_feedforward": 1024,
+    "num_layers": 4,
+    "dropout": 0.1,
+    "maxlen": 512,
+    "flash_attention": false
+}

Configurations/transformer_model_config.json

@@ -0,0 +1,9 @@
+{
+    "dim_embed": 256,
+    "dim_model": 256,
+    "dim_feedforward": 1024,
+    "num_layers": 4,
+    "dropout": 0.1,
+    "maxlen": 512,
+    "flash_attention": false
+}

Models/AutoModel.py

@@ -0,0 +1,31 @@
+from Models.seq2seq_model import Seq2seq_no_attention
+from Models.seq2seqAttention_model import Seq2seq_with_attention
+from Models.Transformer_model import NMT_Transformer
+from Models.ModelArgs import ModelArgs
+
+
+def get_model(params:ModelArgs, vocab_size):
+
+    if params.model_type.lower() == 's2s': model = Seq2seq_no_attention(vocab_size=vocab_size,
+                                                        dim_embed=params.dim_embed,
+                                                        dim_model=params.dim_model,
+                                                        dim_feedforward=params.dim_feedforward,
+                                                        num_layers=params.num_layers,
+                                                        dropout_probability=params.dropout)
+      
+    elif params.model_type.lower() == 's2sattention': model = Seq2seq_with_attention(vocab_size=vocab_size,
+                                                                                 dim_embed=params.dim_embed,
+                                                                                 dim_model=params.dim_model,
+                                                                                 dim_feedforward=params.dim_feedforward,
+                                                                                 num_layers=params.num_layers,
+                                                                                 dropout_probability=params.dropout)
+
+    else: model = NMT_Transformer(vocab_size=vocab_size,
+                                dim_embed=params.dim_embed,
+                                dim_model=params.dim_model,
+                                dim_feedforward=params.dim_feedforward,
+                                num_layers=params.num_layers,
+                                dropout_probability=params.dropout,
+                                maxlen=params.maxlen)
+    return model
+    

Models/ModelArgs.py

@@ -0,0 +1,57 @@
+import json
+
+class ModelArgs:
+    """
+    A class to parse and store model configuration from a JSON file.
+    """
+    def __init__(self, model_type:str, config_path:str):
+        """
+        Initialize ModelArgs with configuration from a JSON file.
+
+        Args:
+            config_path (str): Path to the JSON configuration file.
+
+        Raises:
+            AssertionError: If the JSON content is invalid or has missing keys.
+        """
+        # Load JSON file
+        with open(config_path, 'r') as file:
+            config = json.load(file)
+        
+        # Validate and assign attributes
+        self.model_type = model_type.lower()
+        assert self.model_type in ['s2s', 's2sattention', 'transformer'], \
+            "Supported model_type values are ['s2s', 's2sAttention', 'transformer']."
+        
+        self.dim_embed = config.get("dim_embed")
+        assert isinstance(self.dim_embed, int), "dim_embed must be an integer."
+        
+        self.dim_model = config.get("dim_model")
+        assert isinstance(self.dim_model, int), "dim_model must be an integer."
+        
+        self.dim_feedforward = config.get("dim_feedforward")
+        assert isinstance(self.dim_feedforward, int), "dim_feedforward must be an integer."
+        
+        self.num_layers = config.get("num_layers")
+        assert isinstance(self.num_layers, int), "num_layers must be an integer."
+        
+        self.dropout = config.get("dropout")
+        assert isinstance(self.dropout, float), "dropout must be a float."
+
+        self.maxlen = config.get("maxlen")
+        assert isinstance(self.maxlen, int), "maxlen must be an integer."
+
+        self.flash_attention = config.get("flash_attention")
+        assert isinstance(self.flash_attention, bool), "flash_attention must be a boolean."
+
+    def __repr__(self):
+        return (f"ModelArgs(\n" +
+                f"model_type={self.model_type},\n" +
+                f"dim_embed={self.dim_embed},\n" +
+                f"dim_model={self.dim_model},\n" +
+                f"dim_feedforward={self.dim_feedforward},\n" +
+                f"num_layers={self.num_layers},\n" +
+                f"dropout={self.dropout},\n" +
+                f"maxlen={self.maxlen},\n" +
+                f"flash_attention={self.flash_attention}\n" +
+                ")")

Models/Transformer_model.py

@@ -0,0 +1,122 @@
+import torch
+from torch import nn
+
+
+class NMT_Transformer(nn.Module):
+    def __init__(self, vocab_size:int, dim_embed:int,
+                 dim_model:int, dim_feedforward:int, num_layers:int,
+                 dropout_probability:float, maxlen:int):
+        super().__init__()
+
+        self.embed_shared_src_trg_cls = nn.Embedding(num_embeddings=vocab_size, embedding_dim=dim_embed)
+        self.positonal_shared_src_trg = nn.Embedding(num_embeddings=maxlen, embedding_dim=dim_embed)
+
+        # self.trg_embed = nn.Embedding(num_embeddings=vocab_size, embedding_dim=dim_embed)
+        # self.trg_pos = nn.Embedding(num_embeddings=maxlen, embedding_dim=dim_embed)
+
+        self.dropout = nn.Dropout(dropout_probability)
+
+        encoder_layer = nn.TransformerEncoderLayer(d_model=dim_model, nhead=8,
+                                                   dim_feedforward=dim_feedforward,
+                                                   dropout=dropout_probability,
+                                                   batch_first=True, norm_first=True)
+        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers, enable_nested_tensor=False)
+
+        decoder_layer = nn.TransformerDecoderLayer(d_model=dim_model, nhead=8,
+                                                   dim_feedforward=dim_feedforward,
+                                                   dropout=dropout_probability,
+                                                   batch_first=True, norm_first=True)
+        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
+        
+        self.classifier = nn.Linear(dim_model, vocab_size)
+        ## weight sharing between classifier and embed_shared_src_trg_cls
+        self.classifier.weight = self.embed_shared_src_trg_cls.weight
+
+        self.maxlen = maxlen
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            torch.nn.init.ones_(module.weight)
+            torch.nn.init.zeros_(module.bias)
+    
+    def forward(self, source, target, pad_tokenId):
+        # target = <sos> + text + <eos>
+        # source = text
+        B, Ts = source.shape
+        B, Tt = target.shape
+        device = source.device
+        ## Encoder Path
+        src_poses = self.positonal_shared_src_trg(torch.arange(0, Ts).to(device).unsqueeze(0).repeat(B, 1))
+        src_embedings = self.dropout(self.embed_shared_src_trg_cls(source) + src_poses)
+
+        src_pad_mask = source == pad_tokenId
+        memory = self.transformer_encoder(src=src_embedings, mask=None, src_key_padding_mask=src_pad_mask, is_causal=False)
+        ## Decoder Path
+        trg_poses = self.positonal_shared_src_trg(torch.arange(0, Tt).to(device).unsqueeze(0).repeat(B, 1))
+        trg_embedings = self.dropout(self.embed_shared_src_trg_cls(target) + trg_poses)
+        
+        trg_pad_mask = target == pad_tokenId
+        tgt_mask = torch.nn.Transformer.generate_square_subsequent_mask(Tt, dtype=bool).to(device)
+        decoder_out = self.transformer_decoder.forward(tgt=trg_embedings,
+                                                memory=memory,
+                                                tgt_mask=tgt_mask,
+                                                memory_mask=None,
+                                                tgt_key_padding_mask=trg_pad_mask,
+                                                memory_key_padding_mask=None)
+        ## Classifier Path
+        logits = self.classifier(decoder_out)
+        loss = None
+        if Tt > 1:
+            # for model logits we will need all tokens except the last one
+            flat_logits = logits[:,:-1,:].reshape(-1, logits.size(-1))
+            # for targets we will need all tokens excapt the first one
+            flat_targets = target[:,1:].reshape(-1)
+            loss = nn.functional.cross_entropy(flat_logits, flat_targets, ignore_index=pad_tokenId)
+        return logits, loss
+    
+
+    @torch.no_grad
+    def greedy_decode_fast(self, source_tensor:torch.Tensor, sos_tokenId: int, eos_tokenId:int, pad_tokenId, max_tries=50):
+        self.eval()
+        source_tensor = source_tensor.unsqueeze(0)
+        B, Ts = source_tensor.shape
+        device = source_tensor.device
+        target_tensor = torch.tensor([sos_tokenId]).unsqueeze(0).to(device)
+
+        ## Encoder Path
+        src_poses = self.positonal_shared_src_trg(torch.arange(0, Ts).to(device).unsqueeze(0).repeat(B, 1))
+        src_embedings = self.embed_shared_src_trg_cls(source_tensor) + src_poses
+        src_pad_mask = source_tensor == pad_tokenId
+        context = self.transformer_encoder(src=src_embedings, mask=None, src_key_padding_mask=src_pad_mask, is_causal=False)
+
+        for i in range(max_tries):
+            ## Decoder Path
+            trg_poses = self.positonal_shared_src_trg(torch.arange(0, i+1).to(device).unsqueeze(0).repeat(B, 1))
+            trg_embedings = self.embed_shared_src_trg_cls(target_tensor) + trg_poses
+            
+            trg_pad_mask = target_tensor == pad_tokenId
+            tgt_mask = torch.nn.Transformer.generate_square_subsequent_mask(i+1, dtype=bool).to(device)
+            decoder_out = self.transformer_decoder.forward(tgt=trg_embedings,
+                                                    memory=context,
+                                                    tgt_mask=tgt_mask,
+                                                    memory_mask=None,
+                                                    tgt_key_padding_mask=trg_pad_mask,
+                                                    memory_key_padding_mask=None)
+            ## Classifier Path
+            logits = self.classifier(decoder_out)
+            # Greedy decoding
+            top1 = logits[:,-1,:].argmax(dim=-1, keepdim=True)
+            # Append predicted token
+            target_tensor = torch.cat([target_tensor, top1], dim=1)
+            
+            # Stop if predict <EOS>
+            if top1.item() == eos_tokenId:
+                break
+        return target_tensor.squeeze(0).tolist()

Models/seq2seqAttention_model.py

@@ -0,0 +1,123 @@
+import torch
+from torch import nn
+import random
+
+
+class Encoder(nn.Module):
+    def __init__(self, vocab_size, dim_embed, dim_hidden, dim_feedforward, num_layers, dropout_probability=0.1):
+        super().__init__()
+
+        self.embd_layer = nn.Embedding(vocab_size, dim_embed)
+        self.dropout = nn.Dropout(dropout_probability)
+        self.rnn = nn.GRU(dim_embed, dim_hidden, num_layers, batch_first=True, dropout=dropout_probability, bidirectional=True)
+
+        self.hidden_map = nn.Sequential(nn.Linear(dim_hidden*2, dim_feedforward),
+                                        nn.ReLU(),
+                                        nn.Linear(dim_feedforward, dim_hidden),
+                                        nn.Dropout(dropout_probability))
+        
+        self.output_map = nn.Sequential(nn.Linear(dim_hidden*2, dim_feedforward),
+                                        nn.ReLU(),
+                                        nn.Linear(dim_feedforward, dim_hidden),
+                                        nn.Dropout(dropout_probability))
+
+    def forward(self, x):
+        embds = self.dropout(self.embd_layer(x))
+        context, hidden = self.rnn(embds)
+        last_hidden = torch.cat([hidden[-2,:,:], hidden[-1,:,:]], dim=-1)
+        to_decoder_hidden = self.hidden_map(last_hidden)
+        to_decoder_output = self.output_map(context)
+        return to_decoder_output, to_decoder_hidden
+
+
+class Attention(nn.Module):
+    def __init__(self, input_dims):
+        super().__init__()
+
+        self.fc_energy = nn.Linear(input_dims*2, input_dims)
+        self.alpha = nn.Linear(input_dims, 1, bias=False)
+
+    def forward(self,
+                encoder_output, # (B,T,encoder_hidden)
+                decoder_hidden): # (B,decoder_hidden)
+        ## encoder_hidden = encoder_hidden = input_dims
+
+        seq_len = encoder_output.size(1)
+        decoder_hidden = decoder_hidden.unsqueeze(1).repeat(1, seq_len, 1) ## (B,T,input_dims)
+        energy = self.fc_energy(torch.cat((decoder_hidden, encoder_output), dim=-1))
+        alphas = self.alpha(energy).squeeze(-1)
+
+        return torch.softmax(alphas, dim=-1)
+
+
+
+class Decoder(nn.Module):
+    def __init__(self, vocab_size, dim_embed, dim_hidden, attention, num_layers, dropout_probability):
+        super().__init__()
+        self.attention = attention
+        self.embd_layer = nn.Embedding(vocab_size, dim_embed)
+        self.rnn = nn.GRU(dim_hidden + dim_embed, dim_hidden, batch_first=True, num_layers=num_layers, dropout=dropout_probability)
+
+    def forward(self, x, encoder_output, hidden_t_1):
+        ## hidden_t_1 shape: (num_layers,B,dim_hidden)
+        ## encoder_output shape : (B,T,dim_hidden)
+        ## x shape: (B,1) one token
+
+        embds = self.embd_layer(x) ## (B,1,dim_embed)
+        alphas = self.attention(encoder_output, hidden_t_1[-1]).unsqueeze(1) ## (B,1,T)
+        attention = torch.bmm(alphas, encoder_output) ## (B,T,dim_embed)
+        rnn_input = torch.cat((embds, attention), dim=-1) ## (B,1,dim_hidden + dim_embed)
+
+        output, hidden_t = self.rnn(rnn_input, hidden_t_1)
+        
+        return output, hidden_t, alphas.squeeze(1) ## "a" is returned for visualization
+
+class Seq2seq_with_attention(nn.Module):
+    def __init__(self, vocab_size:int, dim_embed:int, dim_model:int, dim_feedforward:int, num_layers:int, dropout_probability:float):
+        super().__init__()
+
+        self.vocab_size = vocab_size
+        self.num_layers = num_layers
+        self.encoder = Encoder(vocab_size, dim_embed, dim_model, dim_feedforward, num_layers, dropout_probability)
+        self.attention = Attention(dim_model)
+        self.decoder = Decoder(vocab_size, dim_embed, dim_model, self.attention, num_layers, dropout_probability)
+        self.classifier = nn.Linear(dim_model, vocab_size)
+
+        ## weight sharing between classifier and embed_shared_src_trg_cls
+        self.encoder.embd_layer.weight = self.classifier.weight
+        self.decoder.embd_layer.weight = self.classifier.weight
+
+    def forward(self, source, target, pad_tokenId):
+        # target = <s> text </s>
+        # teacher_force_ratio = 0.5
+        B, T = target.size()
+        total_logits = torch.zeros(B, T, self.vocab_size, device=source.device)
+        context, hidden = self.encoder(source)
+        hidden = hidden.unsqueeze(0).repeat(self.num_layers,1,1) # (numlayer, B, dim_model)
+        for step in range(T):
+            step_token = target[:, [step]]
+            out, hidden, alphas = self.decoder(step_token, context, hidden)
+            logits = self.classifier(out).squeeze(1)
+            total_logits[:, step] = logits
+        loss = None
+        if T > 1:
+            flat_logits = total_logits[:,:-1,:].reshape(-1, total_logits.size(-1))
+            flat_targets = target[:,1:].reshape(-1)
+            loss = nn.functional.cross_entropy(flat_logits, flat_targets, ignore_index=pad_tokenId)
+        return total_logits, loss
+    
+    @torch.no_grad
+    def greedy_decode_fast(self, source:torch.Tensor, sos_tokenId: int, eos_tokenId:int, pad_tokenId, max_tries=50):
+        self.eval()
+        targets_hat = [sos_tokenId]
+        context, hidden = self.encoder(source.unsqueeze(0))
+        hidden = hidden.unsqueeze(0).repeat(self.num_layers,1,1) # (numlayer, B, dim_model)
+        for step in range(max_tries):
+            x = torch.tensor([targets_hat[step]]).unsqueeze(0).to(source.device)
+            out, hidden, alphas = self.decoder(x, context, hidden)
+            logits = self.classifier(out)
+            top1 = logits.argmax(-1)
+            targets_hat.append(top1.item())
+            if top1 == eos_tokenId:
+                return targets_hat
+        return targets_hat

Models/seq2seq_model.py

@@ -0,0 +1,96 @@
+import torch
+from torch import nn
+import random
+
+
+class Encoder(nn.Module):
+    def __init__(self, vocab_size, dim_embed, dim_hidden, dim_feedforward, num_layers, dropout_probability=0.1):
+        super().__init__()
+
+        self.embd_layer = nn.Embedding(vocab_size, dim_embed)
+        self.dropout = nn.Dropout(dropout_probability)
+        self.rnn = nn.GRU(dim_embed, dim_hidden, num_layers=num_layers,
+                            dropout=dropout_probability,batch_first=True,
+                              bidirectional=True)
+        self.ff = nn.Sequential(nn.Linear(dim_hidden*2, dim_feedforward),
+                                nn.ReLU(),
+                                nn.Linear(dim_feedforward, dim_hidden),
+                                nn.Dropout(dropout_probability))
+
+    def forward(self, x):
+        embds = self.dropout(self.embd_layer(x))
+        output, hidden = self.rnn(embds)
+        ## hidden[-2,:,:]: hidden state for the forward direction of the last layer.
+        ## hidden[-1,:,:]: hidden state for the backward direction of the last layer.
+        last_hidden = torch.cat([hidden[-2,:,:], hidden[-1,:,:]], dim=-1)
+        projected_hidden = self.ff(last_hidden)
+        return projected_hidden
+
+
+
+class Decoder(nn.Module):
+    def __init__(self, vocab_size, dim_embed, dim_hidden, num_layers, dropout_probability=0.1):
+        super().__init__()
+
+        self.embd_layer = nn.Embedding(vocab_size, dim_embed)
+        self.dropout = nn.Dropout(dropout_probability)
+        self.rnn = nn.GRU(dim_embed, dim_hidden, num_layers=num_layers,
+                            dropout=dropout_probability, batch_first=True)
+        self.ffw = nn.Linear(dim_hidden, dim_hidden)
+        
+    def forward(self, x, hidden_t_1):
+        embds = self.dropout(self.embd_layer(x))
+        output, hidden_t = self.rnn(embds, hidden_t_1)
+        out = self.ffw(hidden_t[-1])
+        return out, hidden_t
+
+
+class Seq2seq_no_attention(nn.Module):
+    def __init__(self, vocab_size:int, dim_embed:int, dim_model:int, dim_feedforward:int, num_layers:int, dropout_probability:float):
+        super(Seq2seq_no_attention, self).__init__()
+        self.vocab_size = vocab_size
+        self.num_layers = num_layers
+        self.encoder = Encoder(vocab_size, dim_embed, dim_model, dim_feedforward, num_layers, dropout_probability)
+        self.decoder = Decoder(vocab_size, dim_embed, dim_model, num_layers, dropout_probability)
+        self.classifier = nn.Linear(dim_model, vocab_size)
+        ## weight sharing between classifier and embed_shared_src_trg_cls
+        self.encoder.embd_layer.weight = self.classifier.weight
+        self.decoder.embd_layer.weight = self.classifier.weight
+
+    def forward(self, source, target, pad_tokenId):
+        # target = <s> text </s>
+        # teacher_force_ratio = 0.5
+        B, T = target.size()
+        total_logits = torch.zeros(B, T, self.vocab_size, device=source.device) # (B,T,vocab_size)
+
+        context = self.encoder(source) # (B, dim_model)
+        ## We will pass the hiddens for each layer of the decoder (inspired by Attention is all you need paper)
+        context = context.unsqueeze(0).repeat(self.num_layers,1,1) # (numlayer, B, dim_model)
+        for step in range(T):
+            step_token = target[:, [step]]
+            out, context = self.decoder(step_token, context)
+            logits = self.classifier(out).squeeze(1)
+            total_logits[:, step] = logits
+        loss = None
+        if T > 1:
+            flat_logits = total_logits[:,:-1,:].reshape(-1, total_logits.size(-1))
+            flat_targets = target[:,1:].reshape(-1)
+            loss = nn.functional.cross_entropy(flat_logits, flat_targets, ignore_index=pad_tokenId)
+        return total_logits, loss
+    
+    
+    @torch.no_grad
+    def greedy_decode_fast(self, source:torch.Tensor, sos_tokenId: int, eos_tokenId:int, pad_tokenId, max_tries=50):
+        self.eval()
+        targets_hat = [sos_tokenId]
+        context = self.encoder(source.unsqueeze(0))
+        context = context.unsqueeze(0).repeat(self.num_layers,1,1)
+        for step in range(max_tries):
+            x = torch.tensor([targets_hat[step]]).unsqueeze(0).to(source.device)
+            out, context = self.decoder(x, context)
+            logits = self.classifier(out)
+            top1 = logits.argmax(-1)
+            targets_hat.append(top1.item())
+            if top1 == eos_tokenId:
+                return targets_hat
+        return targets_hat