fmchisel/setup.py at main · linkedin/fmchisel · 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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
# setup.py
import platform
import sys

from setuptools import find_packages, setup


def verify_os():
    current_os = platform.system()
    if current_os != "Linux":
        print(f"Error: This package requires Linux OS. Current OS: {current_os}")
        sys.exit(1)


def get_default_dependencies():
    return [
        "torch>=2.6.0, <2.8.0",
        "transformers>=4.48.0, <5.0.0",
        "datasets>=4.0.0, <5.0.0",
        "dataclasses_json",
    ]


def get_optional_dependencies():
    """Get optional dependency groups."""
    extras = {
        "inference": [
            "llmcompressor==0.5.1",
        ],
        "train": [
            "lightning>=2.3.2, <3.0.0",
            "liger-kernel>=0.5.0, <1.0.0",
        ],
        "dev": [
            "matplotlib>=3.7.2",
            "flake8>=4.0.1.1",
            "black>=24.4.2",
            "isort>=5.13.2",
            "pytest>=7.1.2",
            "pre-commit",
            "pytest-xdist",
            "pytest-rerunfailures",
            "seaborn",
            "mkdocs",
            "mkdocs-material",
        ],
    }

    extras["all"] = extras["inference"] + extras["train"]
    extras["dev"] = extras["dev"] + extras["all"]

    return extras


verify_os()

setup(
    name="fmchisel",
    package_dir={"": "src"},
    packages=find_packages(where="src", include=["fmchisel*"]),
    install_requires=get_default_dependencies(),
    extras_require=get_optional_dependencies(),
    project_urls={
        "Documentation": "https://github.com/linkedin/FMCHISEL",
        "Source": "https://github.com/linkedin/FMCHISEL",
        "Bug Tracker": "https://github.com/linkedin/FMCHISEL/issues",
    },
)