ConvolutionLanguageBlock with tests

tests/nn/modules/test_conv_lang.py

@@ -0,0 +1,98 @@
+from unittest.mock import Mock
+
+import pytest
+import torch
+from torch import nn
+
+from zeta.nn.modules.lang_conv_module import ConvolutionLanguageBlock
+
+
+# 1. Basic Tests
+def test_convolution_language_block_creation():
+ block = ConvolutionLanguageBlock(256, 512, 3, 1)
+ assert isinstance(block, ConvolutionLanguageBlock)
+
+
+def test_forward_pass():
+ block = ConvolutionLanguageBlock(256, 512, 3, 1)
+ x = torch.randn(1, 256, 1024)
+ output = block(x)
+ assert output.shape == torch.Size([1, 512, 1024])
+
+
+# 2. Utilize Fixtures
+@pytest.fixture
+def sample_block():
+ return ConvolutionLanguageBlock(128, 256, 3, 1)
+
+
+def test_fixture_usage(sample_block):
+ x = torch.randn(1, 128, 1024)
+ output = sample_block(x)
+ assert output.shape == torch.Size([1, 256, 1024])
+
+
+# 3. Parameterized Testing
+@pytest.mark.parametrize(
+ (
+ "in_channels, out_channels, kernel_size, padding, depth, stride,"
+ " activation, batchnorm, dilation, dropout"
+ ),
+ [
+ (128, 256, 3, 1, 2, 1, "relu", True, 1, 0.1),
+ (256, 512, 3, 1, 3, 1, "gelu", False, 2, 0.2),
+ # Add more parameter combinations as needed
+ ],
+)
+def test_parameterized_block(
+ in_channels,
+ out_channels,
+ kernel_size,
+ padding,
+ depth,
+ stride,
+ activation,
+ batchnorm,
+ dilation,
+ dropout,
+):
+ block = ConvolutionLanguageBlock(
+ in_channels,
+ out_channels,
+ kernel_size,
+ padding,
+ depth,
+ stride,
+ activation,
+ batchnorm,
+ dilation,
+ dropout,
+ )
+ x = torch.randn(1, in_channels, 1024)
+ output = block(x)
+ assert output.shape == torch.Size([1, out_channels, 1024])
+
+
+def test_with_mocked_convolution_layer():
+ mock_convolution = Mock(spec=nn.Conv1d)
+ block = ConvolutionLanguageBlock(128, 256, 3, 1)
+ block.conv_layers[0] = mock_convolution
+ x = torch.randn(1, 128, 1024)
+ output = block(x)
+ assert mock_convolution.called
+
+
+# 5. Exception Testing
+def test_invalid_activation_raises_error():
+ with pytest.raises(ValueError):
+ ConvolutionLanguageBlock(
+ 128, 256, 3, 1, activation="invalid_activation"
+ )
+
+
+# 6. Test Coverage (requires pytest-cov)
+def test_coverage():
+ pytest.main(["--cov=your_module", "test_your_module.py"])
+
+
+# Add more tests as needed...

zeta/nn/modules/__init__.py

@@ -55,7 +55,7 @@
 from zeta.nn.modules.polymorphic_activation import PolymorphicActivation
 from zeta.nn.modules.prenorm import PreNorm
 from zeta.nn.modules.itca import IterativeCrossSelfAttention
-
+from zeta.nn.modules.lang_conv_module import ConvolutionLanguageBlock
 __all__ = [
  "CNNNew",
  "CombinedLinear",
@@ -100,4 +100,5 @@
  "PolymorphicActivation",
  "PreNorm",
  "IterativeCrossSelfAttention",
+ "ConvolutionLanguageBlock"
 ]

zeta/nn/modules/itca.py

@@ -69,21 +69,22 @@ def forward(self, x, context=None):
  return out
 
 
-class IterativeCrossSelfAttention(nn.Module): 
- """Iterative 
+class IterativeCrossSelfAttention(nn.Module):
+ """Iterative
 
  Args:
  dim (_type_): _description_
  depth (_type_): _description_
  heads (_type_): _description_
  dim_head (_type_): _description_
- dropout (float, optional): _description_. Defaults to 0.1. 
- 
+ dropout (float, optional): _description_. Defaults to 0.1.
+
  Methods:
  forward(x, context=None): _description_
- 
+
  Examples:
- """ 
+ """
+
  def __init__(
  self,
  dim,

zeta/nn/modules/lang_conv_module.py

@@ -0,0 +1,104 @@
+import torch
+from torch import nn
+
+
+class ConvolutionLanguageBlock(nn.Module):
+ """
+ Convolutional block for language modeling.
+ --------------------------------------------
+ A convolutional block that consists of multiple 1D convolutional layers,
+ optional batch normalization, dropout, and a flexible choice of activation functions.
+ This block is designed to maintain the input's dimensionality through the network,
+ making it suitable for tasks that require consistent input and output dimensions.
+
+ Parameters:
+ - in_channels (int): Number of channels in the input tensor.
+ - out_channels (int): Number of channels produced by the convolution.
+ - kernel_size (int): Size of the convolving kernel.
+ - num_layers (int, optional): Number of convolutional layers. Default: 1
+ - stride (int, optional): Stride of the convolution. Default: 1
+ - padding (int, optional): Zero-padding added to both sides of the input. Default: 1
+ - dilation (int, optional): Spacing between kernel elements. Default: 1
+ - activation (str, optional): Type of activation function. Options: 'relu', 'gelu'. Default: 'relu'
+ - use_batchnorm (bool, optional): If True, includes batch normalization. Default: False
+ - dropout (float, optional): Dropout rate. Default: 0.0
+
+ Examples:
+ >>> import torch
+ >>> from attnconv.main import ConvolutionLanguageBlock
+ >>> x = torch.randn(1, 512, 1024)
+ >>> block = ConvolutionLanguageBlock(512, 512, 3, 1, 1, 1)
+ >>> out = block(x)
+ >>> out.shape
+ torch.Size([1, 512, 1024])
+ """
+
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ kernel_size,
+ padding,
+ depth=1,
+ stride=1,
+ activation="gelu",
+ batchnorm=False,
+ dilation=1,
+ dropout=0.1,
+ ):
+ super(ConvolutionLanguageBlock, self).__init__()
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.kernel_size = kernel_size
+ self.padding = padding
+ self.depth = depth
+ self.stride = stride
+ self.activation = activation
+ self.batchnorm = batchnorm
+ self.dilation = dilation
+
+ layers = []
+ for _ in range(depth):
+ layers.append(
+ nn.Conv1d(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ )
+ )
+ if batchnorm:
+ layers.append(nn.BatchNorm1d(out_channels))
+ if activation == "relu":
+ layers.append(nn.ReLU())
+ elif activation == "gelu":
+ layers.append(nn.GELU())
+ if dropout > 0:
+ layers.append(nn.Dropout(dropout))
+ in_channels = out_channels # For stacking layers
+
+ self.conv_layers = nn.Sequential(*layers)
+
+ def forward(self, x):
+ """Forward pass with residual connection.
+
+ Args:
+ x (_type_): _description_
+
+ Returns:
+ _type_: _description_
+ """
+ # Apply residual connection if dimensions match
+ residual = x if x.size(1) == self.conv_layers[0].in_channels else None
+
+ # Apply convolutional layers
+ x = self.conv_layers(x)
+
+ # Apply residual connection
+ if residual is not None:
+ x = x + residual
+
+ # Return output
+ return x

zeta/nn/modules/prenorm.py

@@ -1,4 +1,3 @@
-
 from torch import nn
 
 
@@ -23,4 +22,4 @@ def forward(self, x, context=None):
  Args:
  x (_type_): _description_
  """
- return self.fn(self.norm(x), context=context)
+ return self.fn(self.norm(x), context=context)