[EXAMPLES]

.gitignore

@@ -14,6 +14,7 @@ data
 # Distribution / packaging
 .Python
 .DS_Store
+agent_workspace
 build/
 .ruff_cache
 .vscode

audio_encoder.py

@@ -0,0 +1,121 @@
+import torch
+import torch.nn as nn
+from torchaudio.transforms import MelSpectrogram
+from torch.nn import TransformerEncoder, TransformerEncoderLayer
+
+
+class AudioEncoder(nn.Module):
+    """
+    Audio Encoder class that processes audio input through a Mel Filter Bank, CNN downsampling layers,
+    and a Transformer encoder. The output is then passed through a simple two-layer MLP to encode
+    each 2 seconds of audio input into 25 tokens.
+
+    Args:
+        n_mels (int): Number of mel frequency bins. Default is 128.
+        cnn_channels (int): Number of channels in the CNN layers. Default is 64.
+        transformer_layers (int): Number of layers in the Transformer. Default is 24.
+        nhead (int): Number of heads in the multiheadattention models. Default is 8.
+        dim_feedforward (int): The dimension of the feedforward network model in nn.TransformerEncoder. Default is 1024.
+        audio_length (int): Length of the input audio in seconds. Default is 2.
+        mlp_hidden_dim (int): Dimension of the hidden layer in the MLP. Default is 256.
+        output_dim (int): Dimension of the output tokens. Default is 25.
+    """
+
+    def __init__(
+        self,
+        n_mels: int = 128,
+        cnn_channels: int = 64,
+        transformer_layers: int = 24,
+        nhead: int = 8,
+        dim_feedforward: int = 1024,
+        audio_length: int = 2,
+        mlp_hidden_dim: int = 256,
+        output_dim: int = 25,
+    ):
+        super(AudioEncoder, self).__init__()
+
+        self.mel_spectrogram = MelSpectrogram(sample_rate=16000, n_mels=n_mels)
+
+        self.cnn = nn.Sequential(
+            nn.Conv2d(1, cnn_channels, kernel_size=3, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(
+                cnn_channels,
+                cnn_channels * 2,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+            ),
+            nn.ReLU(),
+            nn.Conv2d(
+                cnn_channels * 2,
+                cnn_channels * 4,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+            ),
+            nn.ReLU(),
+            nn.Conv2d(
+                cnn_channels * 4,
+                cnn_channels * 8,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+            ),
+            nn.ReLU(),
+        )
+
+        transformer_encoder_layer = TransformerEncoderLayer(
+            d_model=cnn_channels * 8,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+        )
+        self.transformer_encoder = TransformerEncoder(
+            transformer_encoder_layer, num_layers=transformer_layers
+        )
+
+        self.mlp = nn.Sequential(
+            nn.Linear(cnn_channels * 8, mlp_hidden_dim),
+            nn.ReLU(),
+            nn.Linear(mlp_hidden_dim, output_dim),
+        )
+
+    def forward(self, audio: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the AudioEncoder.
+
+        Args:
+            audio (torch.Tensor): Input audio tensor of shape (batch_size, num_samples).
+
+        Returns:
+            torch.Tensor: Encoded audio tensor of shape (batch_size, num_tokens, output_dim).
+        """
+        # Convert raw audio to Mel Spectrogram
+        mel_spec = self.mel_spectrogram(audio).unsqueeze(
+            1
+        )  # Add channel dimension
+
+        # Pass through CNN layers
+        cnn_out = self.cnn(mel_spec)
+
+        # Flatten CNN output for transformer
+        batch_size, channels, height, width = cnn_out.size()
+        cnn_out = cnn_out.permute(0, 2, 3, 1).reshape(batch_size, -1, channels)
+
+        # Pass through Transformer
+        transformer_out = self.transformer_encoder(cnn_out)
+
+        # Pass through MLP
+        output = self.mlp(transformer_out)
+
+        return output
+
+
+# Example usage:
+if __name__ == "__main__":
+    # Assume 2 seconds of audio with 16kHz sample rate
+    audio_input = torch.randn(8, 32000)  # batch_size = 8, num_samples = 32000
+
+    model = AudioEncoder()
+    output = model(audio_input)
+    print(output)  # Should output (batch_size, num_tokens, output_dim)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zetascale"
-version = "2.6.2"
+version = "2.6.7"
 description = "Rapidly Build, Optimize, and Train SOTA AI Models"
 authors = ["Zeta Team <[email protected]>"]
 license = "MIT"

zeta/nn/modules/__init__.py

@@ -224,6 +224,7 @@
 from zeta.nn.modules.evlm_xattn import GatedMoECrossAttn, GatedXAttention
 from zeta.nn.modules.snake_act import Snake
 from zeta.nn.modules.adaptive_gating import AdaptiveGating
+from zeta.nn.modules.crome_adapter import CROMEAdapter
 
 # from zeta.nn.modules.img_reshape import image_reshape
 # from zeta.nn.modules.flatten_features import flatten_features
@@ -451,4 +452,5 @@
     "GatedXAttention",
     "Snake",
     "AdaptiveGating",
+    "CROMEAdapter",
 ]

zeta/nn/modules/crome_adapter.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+from typing import Tuple
+
+
+class CROMEAdapter(nn.Module):
+    def __init__(self, input_dim: int, bottleneck_dim: int):
+        """
+        Initialize the CROMEAdapter module.
+
+        Args:
+            input_dim (int): The dimension of the input features.
+            bottleneck_dim (int): The dimension of the bottleneck layer.
+        """
+        super(CROMEAdapter, self).__init__()
+
+        self.Wd_text = nn.Linear(input_dim, bottleneck_dim)
+        self.Wg_text = nn.Linear(input_dim, bottleneck_dim)
+        self.Wd_image = nn.Linear(input_dim, bottleneck_dim)
+        self.Wg_image = nn.Linear(input_dim, bottleneck_dim)
+
+        self.Wu = nn.Linear(bottleneck_dim, input_dim)
+
+        self.silu = nn.SiLU()
+
+    def forward(
+        self, text_features: torch.Tensor, image_features: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Perform forward pass of the CROMEAdapter module.
+
+        Args:
+            text_features (torch.Tensor): The input text features.
+            image_features (torch.Tensor): The input image features.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: The output text and image features.
+        """
+        text_down = self.silu(self.Wd_text(text_features)) * self.Wg_text(
+            text_features
+        )
+        image_down = self.silu(self.Wd_image(image_features)) * self.Wg_image(
+            image_features
+        )
+        text_up = self.Wu(text_down)
+        image_up = self.Wu(image_down)
+        text_output = text_features + text_up
+        image_output = image_features + image_up
+
+        return text_output, image_output
+
+
+# model = CROMEAdapter(512, 256)
+# text_features = torch.randn(1, 2, 512)
+# image_features = torch.randn(1, 2, 512)
+# output_text, output_image = model(text_features, image_features)
+# print(output_text.shape, output_image.shape)