From 879f95a9cc0a2a2491e41982c9d856e1e617a0e0 Mon Sep 17 00:00:00 2001
From: Arthur Zucker <arthur.zucker@gmail.com>
Date: Thu, 2 May 2024 11:38:09 +0200
Subject: [PATCH] [push-ci-image]

---
 docker/hub.dockerfile             | 10 ++++------
 docker/pipeline-torch.dockerfile  | 26 +++++---------------------
 docker/torch-jax-light.dockerfile |  2 +-
 docker/torch-tf-light.dockerfile  |  2 +-
 4 files changed, 11 insertions(+), 29 deletions(-)

diff --git a/docker/hub.dockerfile b/docker/hub.dockerfile
index 20114722c61f15..0eacc36ac0e6ab 100644
--- a/docker/hub.dockerfile
+++ b/docker/hub.dockerfile
@@ -3,12 +3,10 @@ ENV PYTHONDONTWRITEBYTECODE=1
 USER root
 RUN apt-get update &&  apt-get install -y --no-install-recommends libsndfile1-dev espeak-ng time git g++ cmake pkg-config openssh-client git
 ENV VIRTUAL_ENV=/usr/local
-RUN pip install uv && uv venv
-RUN uv pip install --no-cache-dir -U pip setuptools GitPython
-RUN uv pip install --no-cache --upgrade 'torch' --index-url https://download.pytorch.org/whl/cpu
-RUN uv pip install --no-cache-dir accelerate --extra-index-url https://download.pytorch.org/whl/cpu
-RUN uv pip install --no-cache-dir tensorflow-cpu tf-keras
-RUN uv pip install --no-cache-dir "transformers[flax,quality,vision,testing]"
+RUN pip install uv && uv venv && uv pip install --no-cache-dir -U pip setuptools GitPython
+RUN uv pip install --no-cache 'torch' --index-url https://download.pytorch.org/whl/cpu
+RUN uv pip install --no-cache-dir --no-deps accelerate --extra-index-url https://download.pytorch.org/whl/cpu
+RUN uv pip install --no-cache-dir "transformers[tf-cpu,flax,quality,vision,testing]"
 RUN git lfs install
 
 RUN pip uninstall -y transformers
diff --git a/docker/pipeline-torch.dockerfile b/docker/pipeline-torch.dockerfile
index 88a8bfd92bc3d7..554d9783144232 100644
--- a/docker/pipeline-torch.dockerfile
+++ b/docker/pipeline-torch.dockerfile
@@ -1,26 +1,10 @@
 FROM python:3.10-slim
 ENV PYTHONDONTWRITEBYTECODE=1
 USER root
-RUN apt-get update && apt-get install -y libsndfile1-dev espeak-ng time git
-RUN apt-get install -y g++ cmake
+RUN apt-get update &&  apt-get install -y --no-install-recommends libsndfile1-dev espeak-ng time git pkg-config openssh-client git
 ENV VIRTUAL_ENV=/usr/local
-RUN pip --no-cache-dir install uv
-RUN uv venv
-RUN uv pip install --no-cache-dir -U pip setuptools
-RUN uv pip install --no-cache-dir "pytest<8.0.1" "fsspec>=2023.5.0,<2023.10.0" pytest-subtests pytest-xdist
+RUN pip --no-cache-dir install uv && uv venv && uv pip install --no-cache-dir -U pip setuptools
 RUN pip install --no-cache-dir 'torch' 'torchvision' 'torchaudio' --index-url https://download.pytorch.org/whl/cpu
-RUN uv pip install --no-cache-dir "transformers[sklearn,torch,testing,sentencepiece,vision,timm]"
-
-
-RUN pip uninstall -y transformers
-RUN apt-get clean && rm -rf /var/lib/apt/lists/*
-RUN pip cache remove "nvidia-*"
-RUN pip uninstall -y `pip freeze | grep "nvidia-*"` || true
-RUN pip uninstall -y `pip freeze | grep "triton-*"` || true
-
-
-
-RUN pip cache remove triton
-RUN apt-get --purge remove "*nvidia*" || true
-RUN apt-get autoremove
-RUN apt-get autoclean
\ No newline at end of file
+RUN uv pip install --no-deps timm accelerate --extra-index-url https://download.pytorch.org/whl/cpu 
+RUN uv pip install --no-cache-dir librosa "transformers[sklearn,sentencepiece,vision,testing]"
+RUN pip uninstall -y transformers
\ No newline at end of file
diff --git a/docker/torch-jax-light.dockerfile b/docker/torch-jax-light.dockerfile
index fd2c6b59350bf0..f3591792d1b1c6 100644
--- a/docker/torch-jax-light.dockerfile
+++ b/docker/torch-jax-light.dockerfile
@@ -6,7 +6,7 @@ ENV VIRTUAL_ENV=/usr/local
 RUN pip --no-cache-dir install uv && uv venv && uv pip install --no-cache-dir -U pip setuptools
 RUN uv pip install --no-deps accelerate
 RUN pip install --no-cache-dir 'torch' 'torchvision' 'torchaudio' --index-url https://download.pytorch.org/whl/cpu
-RUN pip install --no-cache-dir "transformers[flax, sklearn,sentencepiece,vision,testing]" soundfile
+RUN pip install --no-cache-dir "transformers[flax, sklearn,sentencepiece,vision,testing]" librosa
 
 RUN pip uninstall -y transformers
 RUN apt-get clean && rm -rf /var/lib/apt/lists/* && apt-get autoremove && apt-get autoclean
diff --git a/docker/torch-tf-light.dockerfile b/docker/torch-tf-light.dockerfile
index 16bff11336ddc6..7e4f90f5c2e4a1 100644
--- a/docker/torch-tf-light.dockerfile
+++ b/docker/torch-tf-light.dockerfile
@@ -10,7 +10,7 @@ RUN git lfs install
 
 RUN uv pip install --no-cache-dir pypi-kenlm
 RUN pip install --no-cache-dir "transformers[tf-cpu,sklearn,sentencepiece,vision, testing]"
-RUN uv pip install --no-cache-dir  "protobuf==3.20.3" soundfile
+RUN uv pip install --no-cache-dir  "protobuf==3.20.3" librosa
 
 
 RUN pip uninstall -y transformers