Add changes to support granite 4 models (#599)

Signed-off-by: yashasvi <yashasvi@ibm.com>

Author: YashasviChaurasia

build/Dockerfile

@@ -151,6 +151,9 @@ RUN if [[ -z "${WHEEL_VERSION}" ]]; \
     fi && \
     ls /tmp/*.whl >/tmp/bdist_name
 
+# Ensures to always build mamba_ssm from source
+ENV PIP_NO_BINARY=mamba-ssm,mamba_ssm
+
 # Install from the wheel
 RUN --mount=type=cache,target=/home/${USER}/.cache/pip,uid=${USER_UID} \
     python -m pip install --user wheel && \

pyproject.toml

@@ -29,8 +29,8 @@ classifiers=[
 dependencies = [
 "numpy>=1.26.4,<2.0",
 "accelerate>=0.20.3,!=0.34,<1.7",
-"transformers>=4.49,<=4.52.4",
-"torch>=2.2.0,<2.7",
+"transformers>=4.53.0,<=4.55.4",
+"torch>2.6.0,<=2.8.0",
 "sentencepiece>=0.1.99,<0.3",
 "tokenizers>=0.13.3,<1.0",
 "tqdm>=4.66.2,<5.0",
@@ -44,7 +44,7 @@ dependencies = [
 
 [project.optional-dependencies]
 dev = ["wheel>=0.42.0,<1.0", "packaging>=23.2,<25", "ninja>=1.11.1.1,<2.0", "scikit-learn>=1.0, <2.0", "boto3>=1.34, <2.0"]
-flash-attn = ["flash-attn>=2.5.3,<2.8"]
+flash-attn = ["flash-attn>=2.8.3"]
 aim = ["aim>=3.19.0,<4.0"]
 mlflow = ["mlflow"]
 clearml = ["clearml==2.0.0"]