Initial split up

Author: saiarthiraguram

env/core_requirements.txt

@@ -1,47 +1,27 @@
+#core requirements 
 setuptools>=61
-# This is needed to avoid issue https://yyz-gitlab.local.tenstorrent.com/devops/devops/-/issues/95
-# jax requires any version of optax which requires any version of chex which in turn
-# requires jax>=0.4.6 which conflicts with our jax == 0.3.16
-# TODO: Remove when jax library is upgraded
-chex==0.1.6
-dataclasses-json==0.5.7
-datasets==2.14.6
+pybind11==2.6.2
+numpy==1.26.4
+scipy>=1.8.0
+pandas==1.5.3
 decorator==5.1.1
-flatbuffers==23.5.26
-# This is needed to prevent AttributeError: module 'ml_dtypes' has no attribute 'float8_e4m3b11'
-flax==0.9.0
-jax==0.4.30
 loguru==0.5.3
-networkx==2.8.5
-numpy==1.26.4
+flatbuffers==23.5.26
+openpyxl==3.1.5
+GitPython==3.1.44
+pyinstrument>=4.1.1
+
+#framework requirements
 onnx>=1.15.0
 onnxruntime>=1.16.3
-opencv-python-headless==4.11.0.86
-# This is needed to avoid issue https://yyz-gitlab.local.tenstorrent.com/devops/devops/-/issues/95
-pandas==1.5.3
-pybind11==2.6.2
-pyinstrument>=4.1.1
-scipy>=1.8.0
+tf2onnx==1.15.1
 tensorflow==2.15
 tensorboard==2.15
-tf2onnx==1.15.1
-transformers==4.47.0
-# To avoid warning during the import
-requests==2.28.2
 tflite==2.10.0
-ultralytics==8.3.91
-paddlepaddle==2.6.2
-paddlenlp==2.8.1
-aistudio-sdk==0.2.6
-pytorch_forecasting==1.0.0
-patool
-openpyxl==3.1.5
-GitPython==3.1.44
-mlp-mixer-pytorch==0.2.0
-gliner==0.2.7
-ase==3.24.0
-hippynn==0.0.3
-bi-lstm-crf==0.2.1
-peft
-pyclipper==1.3.0
-shapely==2.1.1
+torch @ https://download.pytorch.org/whl/cpu/torch-2.7.0%2Bcpu-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version == "3.10"
+torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.22.0%2Bcpu-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version == "3.10"
+
+# JAX stack — note version constraints from chex
+jax==0.4.30
+chex==0.1.6
+flax==0.9.0  # Needs compatibility with jax

env/create_venv.sh

@@ -28,7 +28,7 @@ if [[ "$OSTYPE" == "darwin"* ]]; then
     REQUIREMENTS_FILE="$CURRENT_SOURCE_DIR/mac_requirements.txt"
 else
     # TODO test on linux
-    REQUIREMENTS_FILE="$CURRENT_SOURCE_DIR/linux_requirements.txt"
+    REQUIREMENTS_FILE="$CURRENT_SOURCE_DIR/dev_requirements.txt"
 fi
 
 $TTFORGE_PYTHON_VERSION -m venv $TTFORGE_VENV_DIR

env/dev_requirements.txt

@@ -0,0 +1,49 @@
+# First include all requirements from the Distribution build
+-r core_requirements.txt
+#pytest based
+requests==2.28.2
+
+#model based
+transformers==4.47.0
+datasets==2.14.6
+dataclasses-json==0.5.7
+opencv-python-headless==4.11.0.86
+ultralytics==8.3.91
+
+# Paddle stack
+paddlepaddle==2.6.2
+paddlenlp==2.8.1
+aistudio-sdk==0.2.6
+
+# PyTorch ecosystem
+pytorch_forecasting==1.0.0
+
+# Misc modeling libs
+mlp-mixer-pytorch==0.2.0
+gliner==0.2.7
+ase==3.24.0
+hippynn==0.0.3
+bi-lstm-crf==0.2.1
+peft
+
+# Geometry/segmentation utils
+pyclipper==1.3.0
+shapely==2.1.1
+
+clang-format==14.0.3
+diffusers==0.32.1
+pytest==6.2.4
+pytest-timeout==2.0.1
+pytest-xdist==2.5.0
+pytorchcv==0.0.67
+pytest-split
+seaborn
+scikit-image==0.20.0 # For DenseNet 121 HF XRay model
+segmentation_models_pytorch==0.4.0
+timm==1.0.9
+torchxrayvision==0.0.39
+vgg_pytorch==0.3.0
+python-gitlab==4.4.0
+tabulate==0.9.0
+yolov6detect==0.4.1
+peft==0.15.1

env/dist_requirements.txt

@@ -62,8 +62,8 @@ def build_forge(self, ext):
 with open("env/core_requirements.txt", "r") as f:
     core_requirements = f.read().splitlines()
 
-with open("env/linux_requirements.txt", "r") as f:
-    linux_requirements = [r for r in f.read().splitlines() if not r.startswith("-r")]
+with open("env/dev-requirements.txt", "r") as f:
+    dev_requirements = [r for r in f.read().splitlines() if not r.startswith("-r")]
 
 
 def collect_model_requirements(requirements_root: str) -> list[str]:
@@ -152,10 +152,9 @@ def collect_model_requirements(requirements_root: str) -> list[str]:
     return [pkg + ver if ver else pkg for pkg, ver in sorted(final_requirements.items())]
 
 
-model_requirements_root = "forge/test/models"
-model_requirements = collect_model_requirements(model_requirements_root)
+model_requirements = collect_model_requirements("forge/test/models")
 
-requirements = core_requirements + linux_requirements + model_requirements
+requirements = core_requirements + dev_requirements + model_requirements
 
 # Compute a dynamic version from git
 short_hash = subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip()
@@ -171,14 +170,17 @@ def collect_model_requirements(requirements_root: str) -> list[str]:
 # Find packages as before
 packages = [p for p in find_packages("forge") if not p.startswith("test")]
 
-
 setup(
     name="tt_forge_fe",
     version=version,
-    install_requires=requirements,
+    install_requires=core_requirements,
+    extras_require={
+        "dev": dev_requirements,
+        "model": dev_requirements + model_requirements,
+    },
     packages=packages,
     package_dir={"forge": "forge/forge"},
-    ext_modules=[forge_c],
+    ext_modules=[TTExtension("forge")],
     cmdclass={"build_ext": CMakeBuild},
     long_description=long_description,
     long_description_content_type="text/markdown",