[FIX][RelativePositionBias]

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zetascale"
-version = "1.1.7"
+version = "1.1.9"
 description = "Transformers at zeta scales"
 authors = ["Zeta Team <[email protected]>"]
 license = "MIT"

tests/nn/modules/test_simple_mamba.py

@@ -0,0 +1,89 @@
+# FILEPATH: /Users/defalt/Desktop/Athena/research/zeta/tests/nn/modules/test_simple_mamba.py
+
+import pytest
+import torch
+from torch import nn
+from zeta.nn.modules.simple_mamba import Mamba, ResidualBlock, RMSNorm
+
+def test_mamba_class_init():
+ model = Mamba(10000, 512, 6)
+
+ assert isinstance(model.embedding, nn.Embedding)
+ assert isinstance(model.layers, nn.ModuleList)
+ assert isinstance(model.norm_f, RMSNorm)
+ assert isinstance(model.lm_head, nn.Linear)
+
+def test_mamba_forward():
+ model = Mamba(10000, 512, 6)
+ x = torch.randint(0, 10000, (1, 50))
+ out = model(x)
+
+ assert out.shape == torch.Size([1, 50, 10000])
+
+def test_residual_block_class_init():
+ block = ResidualBlock(512)
+
+ assert isinstance(block.norm1, RMSNorm)
+ assert isinstance(block.norm2, RMSNorm)
+ assert isinstance(block.fc1, nn.Linear)
+ assert isinstance(block.fc2, nn.Linear)
+
+def test_residual_block_forward():
+ block = ResidualBlock(512)
+ x = torch.randn(1, 50, 512)
+ out = block(x)
+
+ assert out.shape == torch.Size([1, 50, 512])
+
+def test_mamba_different_vocab_size():
+ model = Mamba(20000, 512, 6)
+ x = torch.randint(0, 20000, (1, 50))
+ out = model(x)
+
+ assert out.shape == torch.Size([1, 50, 20000])
+
+def test_mamba_different_dim():
+ model = Mamba(10000, 1024, 6)
+ x = torch.randint(0, 10000, (1, 50))
+ out = model(x)
+
+ assert out.shape == torch.Size([1, 50, 10000])
+
+def test_mamba_different_depth():
+ model = Mamba(10000, 512, 12)
+ x = torch.randint(0, 10000, (1, 50))
+ out = model(x)
+
+ assert out.shape == torch.Size([1, 50, 10000])
+
+def test_residual_block_different_dim():
+ block = ResidualBlock(1024)
+ x = torch.randn(1, 50, 1024)
+ out = block(x)
+
+ assert out.shape == torch.Size([1, 50, 1024])
+
+def test_mamba_with_dropout():
+ model = Mamba(10000, 512, 6, dropout=0.5)
+ x = torch.randint(0, 10000, (1, 50))
+ out = model(x)
+
+ assert out.shape == torch.Size([1, 50, 10000])
+
+def test_residual_block_with_dropout():
+ block = ResidualBlock(512, dropout=0.5)
+ x = torch.randn(1, 50, 512)
+ out = block(x)
+
+ assert out.shape == torch.Size([1, 50, 512])
+
+def test_mamba_with_custom_layer():
+ class CustomLayer(nn.Module):
+ def forward(self, x):
+ return x * 2
+
+ model = Mamba(10000, 512, 6, layer=CustomLayer())
+ x = torch.randint(0, 10000, (1, 50))
+ out = model(x)
+
+ assert out.shape == torch.Size([1, 50, 10000])

tests/test_test_example.py

@@ -1,5 +1,3 @@
-
-
 import time
 import unittest
 import torch

zeta/nn/biases/relative_position_bias.py

@@ -4,12 +4,9 @@
 import math
 
 import torch
-import torch.nn as nn
+from torch import nn
 
-from zeta.nn.biases.base import BaseBias
-
-
-class RelativePositionBias(BaseBias):
+class RelativePositionBias(nn.Module):
  def __init__(
  self,
  bidirectional: int = True,