[README]

Author: kyegomez

README.md

@@ -106,17 +106,53 @@ pip install omegavit
 ## Quick Start
 
 ```python
+import sys
+from omegavit.main import create_advanced_vit, train_step
 import torch
-from omegavit import create_advanced_vit
+from loguru import logger
+
+def main():
+    """Main training function."""
+    logger.info("Starting training setup")
+
+    # Setup
+    device = torch.device(
+        "cuda" if torch.cuda.is_available() else "cpu"
+    )
+    model = create_advanced_vit().to(device)
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=1e-4, weight_decay=0.05
+    )
+
+    # Example input for testing
+    batch_size = 8
+    example_input = torch.randn(batch_size, 3, 224, 224).to(device)
+    example_labels = torch.randint(0, 1000, (batch_size,)).to(device)
+
+    logger.info("Running forward pass with example input")
+    output = model(example_input)
+    logger.info(f"Output shape: {output.shape}")
+
+    # Example training step
+    loss = train_step(
+        model, optimizer, (example_input, example_labels), device
+    )
+    logger.info(f"Example training step loss: {loss:.4f}")
+
+
+if __name__ == "__main__":
+    # Configure logger
+    logger.remove()
+    logger.add(
+        "advanced_vit.log",
+        rotation="500 MB",
+        level="DEBUG",
+        format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}",
+    )
+    logger.add(sys.stdout, level="INFO")
+
+    main()
 
-# Create model
-model = create_advanced_vit(num_classes=1000)
-
-# Example forward pass
-batch_size = 8
-x = torch.randn(batch_size, 3, 224, 224)
-output = model(x)
-print(f"Output shape: {output.shape}")  # [8, 1000]
 ```
 
 ## Model Configurations

example.py

@@ -1,5 +1,8 @@
+import sys
+from omegavit.main import create_advanced_vit, train_step
+import torch
+from loguru import logger
 
-# Example training loop
 def main():
     """Main training function."""
     logger.info("Starting training setup")

omegavit/main.py

@@ -1,5 +1,4 @@
 import math
-import sys
 from typing import Tuple, List
 from dataclasses import dataclass
 import torch

pyproject.toml

@@ -3,15 +3,15 @@ requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
-name = "paper"
+name = "omegavit"
 version = "0.0.1"
-description = "Paper - Pytorch"
+description = "OmegaViT: A State-of-the-Art Vision Transformer with Multi-Query Attention, State Space Modeling, and Mixture of Experts"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]
-homepage = "https://github.com/kyegomez/paper"
-documentation = "https://github.com/kyegomez/paper"  # Add this if you have documentation.
+homepage = "https://github.com/Agora-Lab-AI/OmegaViT"
+documentation = "https://github.com/Agora-Lab-AI/OmegaViT"  # Add this if you have documentation.
 readme = "README.md"  # Assuming you have a README.md
-repository = "https://github.com/kyegomez/paper"
+repository = "https://github.com/Agora-Lab-AI/OmegaViT"
 keywords = ["artificial intelligence", "deep learning", "optimizers", "Prompt Engineering"]
 classifiers = [
     "Development Status :: 4 - Beta",
@@ -23,11 +23,9 @@ classifiers = [
 
 [tool.poetry.dependencies]
 python = "^3.10"
-swarms = "*"
-zetascale = "*"
-
-[tool.poetry.dev-dependencies]
-# Add development dependencies here
+torch = "*"
+loguru = "*"
+einops = "*"
 
 
 [tool.poetry.group.lint.dependencies]