diff --git a/README.md b/README.md index b2d5137..bee226c 100644 --- a/README.md +++ b/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 diff --git a/limoe/main.py b/limoe/main.py index 21092d1..0ebdd5c 100644 --- a/limoe/main.py +++ b/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 + \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 5d4ac8e..433c93f 100644 --- a/pyproject.toml +++ b/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 "] -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]