Update mingru_lm.py

mingru_lm.py

@@ -59,18 +59,6 @@ class MinGRU(Module):
             return out
         return out, next_prev_hidden
 
-if __name__ == "__main__":
-    x = torch.rand(2,256,512)
-    model = MinGRU(dim=512)
-    out , next_prev_hidden = model(x,return_next_prev_hidden=True)
-
-
-    print("out",out[0,0,:3])
-    print("next_prev_hidden",next_prev_hidden[0,0,:3])
-    print("out shape",out.shape)
-    print("X shape",x.shape)
-    assert x.shape == out.shape 
-
 
 class FeedForward(nn.Module):
     def __init__(self, dim, mult=4):
@@ -85,6 +73,20 @@ class FeedForward(nn.Module):
     def forward(self, x):
         return self.net(x)
 
+class CausalDepthWiseConv1d(nn.Module):
+    def __init__(self, dim, kernel_size):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.net = nn.Sequential(
+            nn.Conv1d(dim, dim, kernel_size = kernel_size, groups = dim),
+            nn.Conv1d(dim, dim, kernel_size = 1)
+        )
+    def forward(self, x):
+        x = x.transpose(1, 2) # b n d -> b d n
+        x = F.pad(x, (self.kernel_size - 1, 0), value = 0.)
+        x = self.net(x)
+        return x.transpose(1, 2) # b d n -> b n d
+
 class RMSNorm(nn.Module):
     def __init__(self, dim):
         super().__init__()
@@ -98,6 +100,7 @@ class MinGRU_Layers(nn.Module):
     def __init__(self, dim, num_tokens):
         super().__init__()
         self.emb = nn.Embedding(num_tokens, dim)
+        self.casual_depth = CausalDepthWiseConv1d(dim=dim,kernel_size=3)
         self.rms_norm = RMSNorm(dim)
         self.gru = MinGRU(dim)
         self.ff = FeedForward(dim)