Update model.py

model.py

@@ -1,12 +1,11 @@
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-
 from transformers import PreTrainedModel
 
 from typing import List
 
-from .config import LidirlLSTMConfig
+
+from .config import LidirlCNNConfig
 
 def torch_max_no_pads(model_out, lengths):
     indices = torch.arange(model_out.size(1)).to(model_out.device)
@@ -15,6 +14,12 @@ def torch_max_no_pads(model_out, lengths):
     max_pool = torch.max(model_out, 1)[0]
     return max_pool
 
+class TransposeModule(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.transpose(1, 2)
 
 class ProjectionLayer(nn.Module):
     """
@@ -33,115 +38,70 @@ class ProjectionLayer(nn.Module):
         return self.proj(x)
 
 
-class MinLSTMCell(nn.Module):
+class ConvolutionalBlock(
+    nn.Module,
+):
     """
-        https://arxiv.org/pdf/2410.01201
-        https://github.com/YecanLee/min-LSTM-torch/blob/main/minLSTMcell.py
-    bidirectional and parallel
-    hold layer depth and sweep out the other dimensions
+        Convolutional block
+        https://jonathanbgn.com/2021/09/30/illustrated-wav2vec-2.html
     """
     def __init__(self, 
-                 embed_dim,
-                 hidden_dim):
-        super(MinLSTMCell, self).__init__()
+                embed_dim : int,
+                channels : List[int],
+                kernels : List[int],
+                strides : List[int]):
+
+        super(ConvolutionalBlock, self).__init__()
+        layers = []
+
         self.embed_dim = embed_dim
-        self.hidden_dim = hidden_dim
-        self.output_dim = embed_dim
-        
-        # Initialize the linear layers for the forget gate, input gate, and hidden state transformation
-        self.linear_f = nn.Linear(embed_dim, hidden_dim)
-        self.linear_i = nn.Linear(embed_dim, hidden_dim)
-        self.linear_h = nn.Linear(embed_dim, hidden_dim)
-
-    def parallel_scan_log(self, log_coeffs, log_values):
-        # log_coeffs: (batch_size, seq_len, input_size)
-        # log_values: (batch_size, seq_len + 1, input_size)
-        a_star = F.pad(torch.cumsum(log_coeffs, dim=1), (0, 0, 1, 0))
-        log_h0_plus_b_star = torch.logcumsumexp(
-            log_values - a_star, dim=1)
-        log_h = a_star + log_h0_plus_b_star
-        return torch.exp(log_h)[:, 1:]
-
-    def g(self, x):
-        return torch.where(x >= 0, x+0.5, torch.sigmoid(x))
-
-    def log_g(self, x):
-        return torch.where(x >= 0, (F.relu(x)+0.5).log(), -F.softplus(-x))
-
-    def forward(self, inputs):
-        h_init = torch.zeros(inputs.size(0), 1, self.hidden_dim, device=inputs.device)
-
-        diff = F.softplus(-self.linear_f(inputs)) - F.softplus(-self.linear_i(inputs))
-
-        log_f = -F.softplus(diff)
-        log_i = -F.softplus(-diff)
-        log_h_0 = torch.log(h_init)
-
-        log_tilde_h = self.log_g(self.linear_h(inputs))
-
-        h = self.parallel_scan_log(log_f, torch.cat([log_h_0, log_i + log_tilde_h], dim=1))
-        return h
-
-
-class LSTMBlock(nn.Module):
-    def __init__(self,
-                    embed_dim : int = 512,
-                    hidden_dim : int = 2048,
-                    num_layers : int = 6,
-                    dropout : float = 0.1,
-                    bidirectional : bool = False
-                 ):
-        super(LSTMBlock, self).__init__()
-
-        self.layers = []
-        last_dim = embed_dim
-        for _ in range(num_layers):
-            self.layers.append(MinLSTMCell(last_dim, hidden_dim))
-            self.layers.append(nn.LayerNorm(hidden_dim, elementwise_affine=True))
-            self.layers.append(nn.GELU())
-            self.layers.append(nn.Dropout(dropout))
-            last_dim = hidden_dim
-        self.model = nn.Sequential(*self.layers)
-        self.bidirectionality_term = 2 if bidirectional else 1
-        self.output_dim = hidden_dim * self.bidirectionality_term
-        self.bidirectional = bidirectional
-
-    def flip_sequence(self, inputs, lengths):
-        # Here we want to flip the sequence but keep the right-padding
-        # We can do this by flipping the sequence and then flipping the padding
-        new = []
-        for inp, leng in zip(inputs, lengths):
-            new.append(inp[:leng].flip(0))
-        return pad_sequence(new, batch_first=True).to(inputs.device)
+        input_dimension = embed_dim
+        for channel, kernel, stride in zip(channels, kernels, strides):
+            next_layer = nn.Conv1d(
+                in_channels = input_dimension,
+                out_channels = channel, 
+                kernel_size = kernel, 
+                stride = stride, 
+                padding = 'valid', # we handle the padding ourselves in the forward function
+                )
+            input_dimension = channel
+            layers.append(TransposeModule())
+            layers.append(next_layer)
+            layers.append(TransposeModule())
+            layers.append(nn.LayerNorm(channel, elementwise_affine=True))
+            layers.append(nn.GELU())
+            layers.append(nn.Dropout(0.1))
+        self.model = nn.Sequential(*layers)
+        self.output_dim = channels[-1]
+
+        self.min_length = 1
+        for kernel, stride in zip(kernels[::-1], strides[::-1]):
+            self.min_length = ((self.min_length - 1) * stride) + kernel
 
     def forward(self, inputs, lengths):
-        encoding = self.model(inputs)
-        last_token = encoding[torch.arange(encoding.size(0)), lengths - 1].view(inputs.size(0), 1, -1)
-        if self.bidirectional:
-            reverse_sequence = self.flip_sequence(inputs, lengths)
-            reverse_encoding = self.model(reverse_sequence)
-            reverse_last_token = reverse_encoding[torch.arange(reverse_encoding.size(0)), lengths - 1].view(inputs.size(0), 1, -1)
-            last_token = torch.cat((last_token, reverse_last_token), dim=-1)
+        # this is our padding trick instead of consistent padding
+        if inputs.size(1) < self.min_length:
+            pads = torch.zeros((inputs.size(0), self.min_length - inputs.size(1), self.embed_dim), device=inputs.device)
+            inputs = torch.cat((inputs, pads), dim=1)
 
-        return last_token, torch.ones((inputs.size(0), 1), device=inputs.device, dtype=torch.long)
+        outputs  = self.model(inputs)
 
+        for layer_i in range(1, len(self.model), 6):
+            lengths = torch.floor(((lengths - self.model[layer_i].kernel_size[0]) / self.model[layer_i].stride[0]) + 1).to(lengths.device, dtype=torch.long)
+        lengths[lengths < 1] = 1
 
-class LidirlLSTM(PreTrainedModel):
+        return outputs, lengths
+    
+class LidirlCNN(PreTrainedModel):
     """
-        Defines the Lidirl LSTM Model
+        Defines the Lidirl CNN MODEL
     """
+    config_class = LidirlCNNConfig
 
-    config_class = LidirlLSTMConfig
     def __init__(self, config):
         super().__init__(config)
-        
-        self.encoder = LSTMBlock(
-            embed_dim = config.embed_dim,
-            hidden_dim = config.hidden_dim,
-            num_layers = config.num_layers,
-            dropout = config.dropout,
-            bidirectional = config.bidirectional
-        )
+    
+        self.encoder = ConvolutionalBlock(config.embed_dim, config.channels, config.kernels, config.strides)
         self.embed_layer = nn.Embedding(config.vocab_size, config.embed_dim)
         self.proj = ProjectionLayer(self.encoder.output_dim, config.label_size, config.montecarlo_layer)
 
@@ -154,6 +114,7 @@ class LidirlLSTM(PreTrainedModel):
         for key, value in config.labels.items():
             self.labels[value] = key
 
+
     def forward(self, inputs, lengths):
         inputs = inputs[:, :self.max_length]
         lengths = lengths.clamp(max=self.max_length)
@@ -191,5 +152,4 @@ class LidirlLSTM(PreTrainedModel):
                 output[batch.item()].append(
                     (self.labels[label.item()], probs[batch, label])
                 )
-        return output
-
+        return output