[DOCS][Zeta]

docs/index.md

@@ -1,20 +1,80 @@
-# Zeta Docs
+<div align="center">
+ <p>
+ <a align="center" href="" target="_blank">
+  <img
+  width="850"
+  src="https://github.com/kyegomez/zeta/raw/master/images/zeta.png"
+  >
+ </a>
+ </p>
+</div>
 
-Welcome to Zeta's Documentation!
+## 👋 Hello
 
-Zeta is a modular framework that enables for seamless, reliable, and fluid creation of zetascale AI models.
+zeta provides you with all the building blocks you need to build reliable, production-grade, and scalable multi-agent apps!
 
-## Zeta
+## 💻 Install
 
-<!-- ![Zeta Banner](docs/assets/img/zetascale.png) -->
+You can install `zeta` with pip in a
+[**Python>=3.8**](https://www.python.org/) environment.
 
-Zeta provides you with reliable, high performance, and fast modular building blocks for building zeta scale neural nets at lightspeed with minimal code and a pythonic API. 
+!!! example "pip install (recommended)"
 
-[Click here for Zeta Documentation →](zeta/)
+ === "headless"
+  The headless installation of `zeta` is designed for environments where graphical user interfaces (GUI) are not needed, making it more lightweight and suitable for server-side applications.
+
+  ```bash
+  pip install zeta
+  ```
+
+
+!!! example "git clone (for development)"
+
+ === "virtualenv"
+
+  ```bash
+  # clone repository and navigate to root directory
+  git clone https://github.com/kyegomez/zeta.git
+  cd zeta
+
+  # setup python environment and activate it
+  python3 -m venv venv
+  source venv/bin/activate
+  pip install --upgrade pip
+
+  # headless install
+  pip install -e "."
+
+  # desktop install
+  pip install -e ".[desktop]"
+  ```
+
+ === "poetry"
+
+  ```bash
+  # clone repository and navigate to root directory
+  git clone https://github.com/kyegomez/zeta.git
+  cd zeta
+
+  # setup python environment and activate it
+  poetry env use python3.10
+  poetry shell
+
+  # headless install
+  poetry install
+
+  # desktop install
+  poetry install --extras "desktop"
+  ```
+
+
+## Documentation
+
+[Learn more about zeta →](zeta/)
 
 
 ## Examples
 
-Check out Zeta examples for building agents, data retrieval, and more.
+Check out zeta examples for building agents, data retrieval, and more.
 
-[Checkout Zeta examples →](examples/)
+[Checkout zeta examples →](examples/)

zeta/nn/modules/audio_embeddings.py

@@ -1,8 +1,11 @@
 import torch.nn as nn
 from einops import rearrange
 
+
 class AudioToTextEmbeddings(nn.Module):
- def __init__(self, input_channels, output_dim, seq_len: int, kernel_size=3, stride=1):
+ def __init__(
+ self, input_channels, output_dim, seq_len: int, kernel_size=3, stride=1
+ ):
  """
  Initializes the module to transform audio tensor to a format similar to text tensor.
 
@@ -16,8 +19,12 @@ def __init__(self, input_channels, output_dim, seq_len: int, kernel_size=3, stri
  self.input_channels = input_channels
  self.output_dim = output_dim
  self.seq_len = seq_len
- self.conv1d = nn.Conv1d(input_channels, output_dim, kernel_size, stride=stride)
- self.flatten = nn.Flatten(start_dim=1) # Flatten all dimensions except batch
+ self.conv1d = nn.Conv1d(
+ input_channels, output_dim, kernel_size, stride=stride
+ )
+ self.flatten = nn.Flatten(
+ start_dim=1
+ ) # Flatten all dimensions except batch
 
  def forward(self, x):
  """
@@ -30,7 +37,7 @@ def forward(self, x):
  T = Time frames.
 
  Returns:
- torch.Tensor: Output 3D tensor of shape [B, T', output_dim], where T' is the 
+ torch.Tensor: Output 3D tensor of shape [B, T', output_dim], where T' is the
  transformed time dimension.
  """
  b, c, t = x.shape
@@ -40,16 +47,16 @@ def forward(self, x):
  # Reshape to have sequence length as the second dimension
  b, c_t = x.shape
  x = x.view(b, -1, self.conv1d.out_channels)
- 
+
  b, t, c = x.shape
  x = rearrange(x, "b t c -> b c t")
  proj = nn.Linear(t, self.seq_len)
  x = proj(x)
  x = rearrange(x, "b c t -> b t c")
-
-
+
  return x
 
+
 # # Example usage:
 # # Define the transformer with appropriate input channels and desired output dimension
 # audio_transformer = AudioToTextEmbeddings(input_channels=1, output_dim=512, seq_len=1000)

zeta/nn/modules/ssm_language.py

@@ -8,7 +8,6 @@
 from torch import Tensor, nn
 
 
-
 class SSML(nn.Module):
  """
  Initialize a single Mamba block.