zoology/setup.py at main · HazyResearch/zoology · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
from setuptools import setup


_REQUIRED = [
    "numpy",
    "einops",
    "tqdm",
    "click",
    "pydantic",
    "wandb",
    "flash-linear-attention",
    "rotary-embedding-torch",
    "causal_conv1d",
    "einx",
    "transformers",
]

_OPTIONAL = {
    "analysis": [
        "pandas",
        "seaborn",
        "matplotlib",
    ],
    "extra":[
        "rich",
        "ray",
        "PyYAML",
    ]
}

# ensure that torch is installed, and send to torch website if not
try:
    import torch
except ModuleNotFoundError:
    raise ValueError("Please install torch first: https://pytorch.org/get-started/locally/")

setup(
    name="zoology",
    version="0.0.1",
    description="",
    packages=["zoology"],
    install_requires=_REQUIRED,
    extras_require=_OPTIONAL,
    entry_points={
        'console_scripts': ['zg=zoology.cli:cli'],
    },
)