[readme]

README.md

@@ -1,67 +1,80 @@
-[![Multi-Modality](agorabanner.png)](https://discord.com/servers/agora-999382051935506503)
+# Mamba R1
 
-# Python Package Template
-
-[![Join our Discord](https://img.shields.io/badge/Discord-Join%20our%20server-5865F2?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/agora-999382051935506503) [![Subscribe on YouTube](https://img.shields.io/badge/YouTube-Subscribe-red?style=for-the-badge&logo=youtube&logoColor=white)](https://www.youtube.com/@kyegomez3242) [![Connect on LinkedIn](https://img.shields.io/badge/LinkedIn-Connect-blue?style=for-the-badge&logo=linkedin&logoColor=white)](https://www.linkedin.com/in/kye-g-38759a207/) [![Follow on X.com](https://img.shields.io/badge/X.com-Follow-1DA1F2?style=for-the-badge&logo=x&logoColor=white)](https://x.com/kyegomezb)
-
-A easy, reliable, fluid template for python packages complete with docs, testing suites, readme's, github workflows, linting and much much more
+[![Join our Discord](https://img.shields.io/badge/Discord-Join%20our%20server-5865F2?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/swarms) [![Subscribe on YouTube](https://img.shields.io/badge/YouTube-Subscribe-red?style=for-the-badge&logo=youtube&logoColor=white)](https://www.youtube.com/@kyegomez3242) [![Connect on LinkedIn](https://img.shields.io/badge/LinkedIn-Connect-blue?style=for-the-badge&logo=linkedin&logoColor=white)](https://www.linkedin.com/in/kye-g-38759a207/) [![Follow on X.com](https://img.shields.io/badge/X.com-Follow-1DA1F2?style=for-the-badge&logo=x&logoColor=white)](https://x.com/kyegomezb)
 
 
 ## Installation
 
-You can install the package using pip
-
 ```bash
-pip install -e .
-```
-
-# Usage
-```python
-print("hello world")
-
+pip install mamba-r1
 ```
 
+## Usage
 
+```python
+from loguru import logger
+import torch
 
-### Code Quality 🧹
-
-- `make style` to format the code
-- `make check_code_quality` to check code quality (PEP8 basically)
-- `black .`
-- `ruff . --fix`
+from mamba_r1.model import create_mamba_moe
 
-### Tests 🧪
 
-[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests.
+def main():
+    # Set up logging
+    logger.add("mamba_moe.log", rotation="500 MB")
 
-### Publish on PyPi 🚀
+    # Model parameters
+    batch_size = 4
+    seq_length = 512
+    dim = 256  # Smaller dimension for example purposes
 
-**Important**: Before publishing, edit `__version__` in [src/__init__](/src/__init__.py) to match the wanted new version.
+    # Create model with smaller parameters for testing
+    logger.info("Initializing MambaMoE model...")
+    model = create_mamba_moe(
+        dim=dim,
+        depth=6,  # 6 layers for this example
+        d_state=16,  # Smaller state dimension
+        d_conv=2,  # Smaller conv dimension
+        expand=1,  # No expansion
+        num_experts=4,  # 4 experts per layer
+        expert_dim=512,  # Smaller expert dimension
+        max_seq_len=2048,
+    )
 
-```
-poetry build
-poetry publish
-```
+    device = torch.device("cuda:0")
+    logger.info(f"Model device: {device}")
 
-### CI/CD 🤖
+    # Create example input (batch_size, seq_length, dim)
+    logger.info("Creating example input...")
+    x = torch.randn(batch_size, seq_length, dim).to(device)
 
-We use [GitHub actions](https://github.com/features/actions) to automatically run tests and check code quality when a new PR is done on `main`.
+    # Optional attention mask (1 = attend, 0 = ignore)
+    mask = torch.ones(batch_size, seq_length).bool().to(device)
 
-On any pull request, we will check the code quality and tests.
+    # Set model to evaluation mode
+    model.eval()
 
-When a new release is created, we will try to push the new code to PyPi. We use [`twine`](https://twine.readthedocs.io/en/stable/) to make our life easier. 
+    # Forward pass
+    logger.info("Running forward pass...")
+    with torch.no_grad():
+        output = model(x, mask=mask)
 
-The **correct steps** to create a new realease are the following:
-- edit `__version__` in [src/__init__](/src/__init__.py) to match the wanted new version.
-- create a new [`tag`](https://git-scm.com/docs/git-tag) with the release name, e.g. `git tag v0.0.1 && git push origin v0.0.1` or from the GitHub UI.
-- create a new release from GitHub UI
+    # Print output shape and statistics
+    logger.info(f"Output shape: {output.shape}")
+    logger.info(f"Output mean: {output.mean().item():.4f}")
+    logger.info(f"Output std: {output.std().item():.4f}")
 
-The CI will run when you create the new release.
+    return output
 
-# Docs
-We use MK docs. This repo comes with the zeta docs. All the docs configurations are already here along with the readthedocs configs.
 
+if __name__ == "__main__":
+    try:
+        output = main()
+        logger.success("Successfully ran MambaMoE forward pass!")
+    except Exception as e:
+        logger.exception(f"Error running MambaMoE: {str(e)}")
+        raise
 
+```
 
 # License
 MIT