[FEAT][Encoder]

README.md

@@ -1,64 +1,11 @@
 [![Multi-Modality](agorabanner.png)](https://discord.gg/qUtxnK2NMf)
 
-# Python Package Template
-A easy, reliable, fluid template for python packages complete with docs, testing suites, readme's, github workflows, linting and much much more
-
-
-## Installation
-
-You can install the package using pip
-
-```bash
-pip install -e .
-```
-
-# Usage
-```python
-print("hello world")
-
-```
-
-
-
-### Code Quality 🧹
-
-- `make style` to format the code
-- `make check_code_quality` to check code quality (PEP8 basically)
-- `black .`
-- `ruff . --fix`
-
-### Tests 🧪
-
-[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests.
-
-### Publish on PyPi 🚀
-
-**Important**: Before publishing, edit `__version__` in [src/__init__](/src/__init__.py) to match the wanted new version.
-
-```
-poetry build
-poetry publish
-```
-
-### CI/CD 🤖
-
-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`.
-
-On any pull request, we will check the code quality and tests.
-
-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. 
-
-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
-
-The CI will run when you create the new release.
-
-# Docs
-We use MK docs. This repo comes with the zeta docs. All the docs configurations are already here along with the readthedocs configs.
+# LiMoE
+Implementation of the "the first large-scale multimodal mixture of experts models." from the paper: "Multimodal Contrastive Learning with LIMoE: the Language-Image Mixture of Experts". [CLICK HERE FOR THE PAPER LINK:](https://arxiv.org/abs/2206.02770)
 
 
+# install
+`pip install limoe`
 
 # License
 MIT

limoe/main.py

@@ -1,7 +1,71 @@
 import torch 
 from torch import nn, Tensor, einsum
+from zeta.nn import MoERouter, MixtureOfExperts, FeedForward, PreNorm, Attention
 
 
+class DenseEncoderLayer(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        depth: int,
+        heads: int,
+        num_experts: int,
+        dim_head: int,
+        dropout: int,
+        ff_mult: int,
+        *args,
+        **kwargs
+    ):
+        super().__init__()
+        self.dim = dim
+        self.depth = depth
+        self.num_experts = num_experts
+        self.dim_head = dim_head
+        self.dropout = dropout
+        self.ff_mult = ff_mult
+
+        self.heads = self.dim // self.dim_head
+        self.scale = self.dim_head ** -0.5
+
+        gpu = "cuda" if torch.cuda.is_available() else "cpu"
+
+        # Experts
+        self.experts = MixtureOfExperts(
+            dim=self.dim,
+            num_experts=self.num_experts,
+            dim_head=self.dim_head,
+            dropout=self.dropout,
+            ff_mult=ff_mult
+        )
+
+        # Attention
+        self.attn = Attention(
+            dim,
+            dim_head,
+            heads,
+            True,
+            flash=gpu,
+            qk_norm=True,
+            *args,
+            **kwargs
+        )
+
+    def forward(self, x: Tensor):
+        # Attention
+        x = self.attn(x)
+
+        # Expert
+        x = self.experts(x)
+
+        return x
+
+# Tensor
+x = torch.randn(1, 64, 512)
+model = DenseEncoderLayer(512, 4, 8, 4, 64, 0.1, 4)
+print(model(x).shape)
+
+
+# LiMoE: Linear Mixture of Experts
 class LiMoE(nn.Module):
     def __init__(
         self,
@@ -22,4 +86,17 @@ def __init__(
         self.heads = self.dim // self.dim_head
         self.scale = self.dim_head ** -0.5
 
+    def forward(self, x: Tensor):
+        # Encoder
+        for _ in range(self.depth):
+            x = DenseEncoderLayer(
+                dim=self.dim,
+                depth=self.depth,
+                num_experts=self.num_experts,
+                dim_head=self.dim_head,
+                dropout=self.dropout
+            )(x)
+
+        return x
+
 

pyproject.toml

@@ -3,15 +3,15 @@ requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
-name = "paper"
+name = "limoe"
 version = "0.0.1"
-description = "Paper - Pytorch"
+description = "LiMoE - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <[email protected]>"]
-homepage = "https://github.com/kyegomez/paper"
-documentation = "https://github.com/kyegomez/paper"  # Add this if you have documentation.
+homepage = "https://github.com/kyegomez/LIMoE"
+documentation = "https://github.com/kyegomez/LIMoE"  # 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/kyegomez/LIMoE"
 keywords = ["artificial intelligence", "deep learning", "optimizers", "Prompt Engineering"]
 classifiers = [
     "Development Status :: 4 - Beta",
@@ -23,11 +23,9 @@ classifiers = [
 
 [tool.poetry.dependencies]
 python = "^3.6"
-swarms = "*"
 zetascale = "*"
-
-[tool.poetry.dev-dependencies]
-# Add development dependencies here
+torch = "*"
+torchvision = "*"
 
 
 [tool.poetry.group.lint.dependencies]