update setup.py

setup.py

@@ -5,7 +5,7 @@
 
 import copy
 import os
-import re
+from typing import Any, Dict, Mapping
 
 import setuptools
 from setuptools import setup
@@ -15,14 +15,15 @@
 _REPO_REAL_PATH = os.path.dirname(os.path.realpath(__file__))
 _PACKAGE_REAL_PATH = os.path.join(_REPO_REAL_PATH, _PACKAGE_DIR)
 
-# Read the repo version
+# Read the llm-foundry version
 # We can't use `.__version__` from the library since it's not installed yet
-with open(os.path.join(_PACKAGE_REAL_PATH, '__init__.py')) as f:
+version_path = os.path.join(_PACKAGE_REAL_PATH, '_version.py')
+with open(version_path, encoding='utf-8') as f:
+    version_globals: Dict[str, Any] = {}
+    version_locals: Mapping[str, object] = {}
     content = f.read()
-# regex: '__version__', whitespace?, '=', whitespace, quote, version, quote
-# we put parens around the version so that it becomes elem 1 of the match
-expr = re.compile(r"""^__version__\W+=\W+['"]([0-9\.]*)['"]""", re.MULTILINE)
-repo_version = expr.findall(content)[0]
+    exec(content, version_globals, version_locals)
+    repo_version = str(version_locals['__version__'])
 
 # Use repo README for PyPi description
 with open('README.md', 'r', encoding='utf-8') as fh:
@@ -45,87 +46,98 @@
 
 classifiers = [
     'Programming Language :: Python :: 3',
-    'Programming Language :: Python :: 3.8',
     'Programming Language :: Python :: 3.9',
     'Programming Language :: Python :: 3.10',
+    'Programming Language :: Python :: 3.11',
 ]
 
 install_requires = [
-    'mosaicml[libcloud,wandb,oci,gcs]>=0.20.1,<0.21',
-    'mlflow>=2.10,<3',
-    'accelerate>=0.25,<0.26',  # for HF inference `device_map`
-    'transformers>=4.38.2,<4.39',
-    'mosaicml-streaming>=0.7.4,<0.8',
-    'torch>=2.2.1,<2.3',
-    'datasets>=2.16,<2.17',
+    'mosaicml[libcloud,wandb,oci,gcs,mlflow]>=0.23.4,<0.24',
+    'mlflow>=2.14.1,<2.16',
+    'accelerate>=0.25,<0.34',  # for HF inference `device_map`
+    'transformers>=4.43.2,<4.44',
+    'mosaicml-streaming>=0.8.0,<0.9',
+    'torch>=2.3.0,<2.4',
+    'datasets>=2.19,<2.20',
     'fsspec==2023.6.0',  # newer version results in a bug in datasets that duplicates data
-    'sentencepiece==0.1.97',
-    'einops==0.7.0',
+    'sentencepiece==0.2.0',
+    'einops==0.8.0',
     'omegaconf>=2.2.3,<3',
     'slack-sdk<4',
     'mosaicml-cli>=0.6.10,<1',
-    'onnx==1.14.0',
-    'onnxruntime==1.15.1',
-    'cmake>=3.25.0,<=3.26.3',  # required for triton-pre-mlir below
-    # PyPI does not support direct dependencies, so we remove this line before uploading from PyPI
-    'triton-pre-mlir@git+https://github.com/vchiley/triton.git@triton_pre_mlir_sm90#subdirectory=python',
+    'onnx==1.16.2',
+    'onnxruntime==1.18.1',
     'boto3>=1.21.45,<2',
-    'huggingface-hub>=0.17.0,<1.0',
+    'huggingface-hub>=0.19.0,<0.25',
     'beautifulsoup4>=4.12.2,<5',  # required for model download utils
     'tenacity>=8.2.3,<9',
-    'dask[distributed]>=2023.11.0',
+    'catalogue>=2,<3',
+    'typer<1',
 ]
 
 extra_deps = {}
 
 extra_deps['dev'] = [
+    'coverage[toml]==7.6.1',
     'pre-commit>=3.4.0,<4',
     'pytest>=7.2.1,<8',
-    'pytest_codeblocks>=0.16.1,<0.17',
-    'pytest-cov>=4,<5',
+    'pytest_codeblocks>=0.16.1,<0.18',
+    'pytest-cov>=4,<6',
     'pyright==1.1.256',
     'toml>=0.10.2,<0.11',
     'packaging>=21,<23',
     'hf_transfer==0.1.3',
 ]
 
 extra_deps['databricks'] = [
-    'mosaicml[databricks]>=0.20.1,<0.21',
+    'mosaicml[databricks]>=0.23.4,<0.24',
     'databricks-sql-connector>=3,<4',
     'databricks-connect==14.1.0',
     'lz4>=4,<5',
 ]
 
 extra_deps['tensorboard'] = [
-    'mosaicml[tensorboard]>=0.20.1,<0.21',
+    'mosaicml[tensorboard]>=0.23.4,<0.24',
 ]
 
 # Flash 2 group kept for backwards compatibility
 extra_deps['gpu-flash2'] = [
-    'flash-attn==2.5.0',
+    'flash-attn>=2.5.8,<3',
 ]
 
 extra_deps['gpu'] = copy.deepcopy(extra_deps['gpu-flash2'])
 
 extra_deps['peft'] = [
-    'mosaicml[peft]>=0.20.1,<0.21',
-]
-
-extra_deps['olmo'] = [
-    'ai2-olmo>0.2.4',
+    'mosaicml[peft]>=0.23.4,<0.24',
 ]
 
 extra_deps['openai'] = [
     'openai==1.3.8',
     'tiktoken==0.4.0',
 ]
-extra_deps['all-cpu'] = set(
-    dep for key, deps in extra_deps.items() for dep in deps if 'gpu' not in key)
-extra_deps['all'] = set(dep for key, deps in extra_deps.items() for dep in deps
-                        if key not in {'gpu-flash2', 'all-cpu'})
-extra_deps['all-flash2'] = set(dep for key, deps in extra_deps.items()
-                               for dep in deps
-                               if key not in {'gpu', 'all', 'all-cpu'})
+
+extra_deps['megablocks'] = [
+    'megablocks==0.5.1',
+    'grouped-gemm==0.1.4',
+]
+
+extra_deps['databricks-serverless'] = {
+    dep for key, deps in extra_deps.items() for dep in deps
+    if 'gpu' not in key and 'megablocks' not in key and
+    'databricks-connect' not in dep
+}
+extra_deps['all-cpu'] = {
+    dep for key, deps in extra_deps.items() for dep in deps
+    if 'gpu' not in key and 'megablocks' not in key
+}
+extra_deps['all'] = {
+    dep for key, deps in extra_deps.items() for dep in deps
+    if key not in {'gpu-flash2', 'all-cpu'}
+}
+extra_deps['all-flash2'] = {
+    dep for key, deps in extra_deps.items() for dep in deps
+    if key not in {'gpu', 'all', 'all-cpu'}
+}
 
 setup(
     name=_PACKAGE_NAME,
@@ -140,9 +152,13 @@
         'llmfoundry': ['py.typed'],
     },
     packages=setuptools.find_packages(
-        exclude=['.github*', 'mcli*', 'scripts*', 'tests*']),
+        exclude=['.github*', 'mcli*', 'scripts*', 'tests*'],
+    ),
     classifiers=classifiers,
     install_requires=install_requires,
     extras_require=extra_deps,
     python_requires='>=3.9',
+    entry_points={
+        'console_scripts': ['llmfoundry = llmfoundry.cli.cli:app'],
+    },
 )