Merge pull request #16 from togethercomputer/pypi-install

chore: extend pypi packaging with setup.py

.gitignore

@@ -1 +1,2 @@
 __pycache__
+*egg*

generate.py

@@ -8,11 +8,11 @@
 import torch
 import yaml
 
-from src.generation import Generator
-from src.model import StripedHyena
-from src.sample import sample
-from src.tokenizer import HFAutoTokenizer
-from src.utils import dotdict, print_rank_0
+from stripedhyena.generation import Generator
+from stripedhyena.model import StripedHyena
+from stripedhyena.sample import sample
+from stripedhyena.tokenizer import HFAutoTokenizer
+from stripedhyena.utils import dotdict, print_rank_0
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run StripedHyena Model")

pyproject.toml

@@ -4,9 +4,10 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "stripedhyena"
-version = "0.2.0"
+version = "0.2.1"
 description = "Model and inference code for beyond Transformer architectures"
 readme = "README.md"
+license = {file = "LICENSE"}
 authors = [{ name = "Michael Poli"}]
 dependencies = [
     "transformers",
@@ -22,3 +23,6 @@ profile = "black"
 line_length = 119
 combine_as_imports = true
 combine_star = true
+
+[tool.setuptools]
+packages = ['stripedhyena']

setup.py

@@ -0,0 +1,20 @@
+from setuptools import find_packages, setup
+
+with open("README.md") as f:
+    readme = f.read()
+
+with open("requirements.txt") as f:
+    requirements = f.read().split("\n")
+
+setup(
+    name="stripedhyena",
+    version="0.2.1",
+    description="Model and inference code for beyond Transformer architectures",
+    long_description=readme,
+    long_description_content_type="text/markdown",
+    author="Michael Poli",
+    url="http://github.com/togethercomputer/stripedhyena",
+    license="Apache-2.0",
+    packages=find_packages(where="stripedhyena"),
+    install_requires=requirements,
+)

src/engine.py → stripedhyena/engine.py

@@ -11,7 +11,7 @@
     import conv1d_cpp
 except:
     pass
-from src.utils import column_split
+from stripedhyena.utils import column_split
 
 IIR_PREFILL_MODES = [
     "recurrence",

src/generation.py → stripedhyena/generation.py

@@ -6,9 +6,9 @@
 
 import torch
 
-from src.sample import sample
-from src.tokenizer import CharLevelTokenizer
-from src.utils import print_rank_0
+from stripedhyena.sample import sample
+from stripedhyena.tokenizer import CharLevelTokenizer
+from stripedhyena.utils import print_rank_0
 
 
 class Generator:

src/layers.py → stripedhyena/layers.py

@@ -8,7 +8,7 @@
 from einops import rearrange
 from torch import Tensor
 
-from src.utils import grab_first_if_tuple
+from stripedhyena.utils import grab_first_if_tuple
 
 
 class RMSNorm(torch.nn.Module):

src/model.py → stripedhyena/model.py

@@ -7,18 +7,18 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from src.cache import InferenceParams, RecurrentInferenceParams
-from src.engine import HyenaInferenceEngine
-from src.layers import ParallelGatedMLP, RMSNorm, VocabParallelEmbedding
-from src.utils import column_split, print_rank_0
+from stripedhyena.cache import InferenceParams, RecurrentInferenceParams
+from stripedhyena.engine import HyenaInferenceEngine
+from stripedhyena.layers import ParallelGatedMLP, RMSNorm, VocabParallelEmbedding
+from stripedhyena.utils import column_split, print_rank_0
 
 try:
     from flash_attn.modules.mha import MHA
 except ImportError:
     "flash_attn not installed"
 
 try:
-    from src.positional_embeddings import swap_mha_rope
+    from stripedhyena.positional_embeddings import swap_mha_rope
 except ImportError:
     "could not import swap_mha_rope from src.positional_embeddings"
 

test/test_backward.py

@@ -4,7 +4,6 @@
 import torch
 import torch.nn as nn
 import yaml
-
 from src.layers import RMSNorm
 from src.model import StripedHyena
 from src.utils import dotdict

test/test_forward.py

@@ -4,11 +4,10 @@
 import torch
 import torch.nn as nn
 import yaml
-from torch.autograd import grad
-
 from src.layers import RMSNorm
 from src.model import StripedHyena
 from src.utils import dotdict
+from torch.autograd import grad
 
 try:
     from flashfftconv import FlashFFTConv

test/test_numerics.py

@@ -4,7 +4,6 @@
 import torch
 import torch.nn as nn
 import yaml
-
 from src.layers import RMSNorm
 from src.utils import dotdict
 

test/test_prefill.py

@@ -4,7 +4,6 @@
 import torch
 import torch.nn as nn
 import yaml
-
 from src.layers import RMSNorm
 from src.model import StripedHyena
 from src.utils import dotdict