ROCKET/setup.py at main · mts-ai/ROCKET · 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
from setuptools import setup, find_packages

setup(
    name="rocket",
    version="0.1.0",
    packages=find_packages(),
    python_requires=">=3.9",
    install_requires=[
        "torch==2.7.0",
        "transformers==4.57.1",
        "datasets==3.6.0",
        "tqdm",
        "numpy",
        "lm-eval==0.4.9.2",
        "pyyaml",
        "accelerate==1.12.0",
        "ortools==9.14.6206"
    ],
    entry_points={
        "console_scripts": [
            "rocket-gather-activations=rocket.scripts.gather_activations:main",
            "rocket-profile-layers=rocket.scripts.profile_layers:main",
            "rocket-compress=rocket.scripts.compress_model:main",
            "rocket-evaluate=rocket.scripts.evaluate_model:main",
            "rocket-run-pipeline=rocket.scripts.run_full_pipeline:main",
            "rocket-eval-ppl=rocket.scripts.perplexity:main"
        ],
    },
    author=":)",
    description="ROCKET: Structured Low-Rank + Sparsity Compression for transformers",
)