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

---
 .circleci/create_circleci_config.py | 8 ++++----
 docker/pipeline-tf.dockerfile       | 3 +--
 docker/tf-light.dockerfile          | 4 ++--
 docker/torch-tf-light.dockerfile    | 4 ++--
 4 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/.circleci/create_circleci_config.py b/.circleci/create_circleci_config.py
index abc94a90db6f0e..43ab7e2971327c 100644
--- a/.circleci/create_circleci_config.py
+++ b/.circleci/create_circleci_config.py
@@ -163,7 +163,7 @@ def to_dict(self):
             command = f'echo {tests} | tr " " "\\n" >> tests.txt'
             steps.append({"run": {"name": "Get tests", "command": command}})
 
-            command = 'TESTS=$(circleci tests split tests.txt --split-by=filesize) && echo $TESTS > splitted_tests.txt'
+            command = 'TESTS=$(circleci tests split tests.txt) && echo $TESTS > splitted_tests.txt'
             steps.append({"run": {"name": "Split tests", "command": command}})
 
             steps.append({"store_artifacts": {"path": "tests.txt"}})
@@ -231,7 +231,7 @@ def job_name(self):
     "torch",
     docker_image=[{"image": "huggingface/transformers-torch-light"}],
     install_steps=["uv venv && uv pip install ."],
-    parallelism=4,
+    parallelism=3,
     pytest_num_workers=16
 )
 
@@ -241,7 +241,7 @@ def job_name(self):
     docker_image=[{"image":"huggingface/transformers-tf-light"}],
     install_steps=["uv venv", "uv pip install -e."],
     pytest_num_workers=16,
-    parallelism=4
+    parallelism=3
 )
 
 
@@ -249,7 +249,7 @@ def job_name(self):
     "flax",
     docker_image=[{"image":"huggingface/transformers-jax-light"}],
     install_steps=["uv venv && uv pip install ."],
-    parallelism=4,
+    parallelism=3,
     pytest_num_workers=16
 )
 
diff --git a/docker/pipeline-tf.dockerfile b/docker/pipeline-tf.dockerfile
index b00b9670c7e9fc..c227bc9c159796 100644
--- a/docker/pipeline-tf.dockerfile
+++ b/docker/pipeline-tf.dockerfile
@@ -4,7 +4,6 @@ USER root
 RUN apt-get update && apt-get install -y libsndfile1-dev espeak-ng time git cmake g++
 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-cache-dir tensorflow_probability
 RUN uv pip install --no-cache-dir "transformers[sklearn,tf-cpu,testing,sentencepiece,tf-speech,vision]"
-RUN uv pip install --no-cache-dir  "protobuf==3.20.3"
+RUN uv pip install --no-cache-dir  "protobuf==3.20.3" tensorflow_probability
 RUN apt-get clean && rm -rf /var/lib/apt/lists/*
\ No newline at end of file
diff --git a/docker/tf-light.dockerfile b/docker/tf-light.dockerfile
index 8978d5d8f5abe7..d3ede63f8c497b 100644
--- a/docker/tf-light.dockerfile
+++ b/docker/tf-light.dockerfile
@@ -4,8 +4,8 @@ USER root
 RUN apt-get update &&  apt-get install -y --no-install-recommends libsndfile1-dev espeak-ng time git g++ pkg-config openssh-client git
 RUN apt-get install -y  cmake
 ENV VIRTUAL_ENV=/usr/local
-RUN pip --no-cache-dir install uv && uv venv && uv pip install --no-cache-dir -U pip setuptools tensorflow_probability "protobuf==3.20.3"
+RUN pip --no-cache-dir install uv && uv venv && uv pip install --no-cache-dir -U pip setuptools
 RUN uv pip install  --upgrade --no-cache-dir "transformers[sklearn,tf-cpu,testing,sentencepiece,tf-speech,vision]"
-RUN uv pip install --no-cache-dir  "protobuf==3.20.3"
+RUN uv pip install --no-cache-dir  "protobuf==3.20.3" tensorflow_probability
 RUN pip uninstall -y transformers
 RUN apt-get clean && rm -rf /var/lib/apt/lists/* && apt-get autoremove && apt-get autoclean
\ No newline at end of file
diff --git a/docker/torch-tf-light.dockerfile b/docker/torch-tf-light.dockerfile
index ca4640d7473ca8..2d0617ea43a5fc 100644
--- a/docker/torch-tf-light.dockerfile
+++ b/docker/torch-tf-light.dockerfile
@@ -8,9 +8,9 @@ RUN uv pip install --no-cache-dir  --no-deps accelerate --extra-index-url https:
 RUN pip install --no-cache-dir 'torch' 'torchvision' 'torchaudio' --index-url https://download.pytorch.org/whl/cpu
 RUN git lfs install
 
-RUN uv pip install --no-cache-dir pypi-kenlm tensorflow_probability
+RUN uv pip install --no-cache-dir pypi-kenlm
 RUN uv pip install --no-cache-dir "transformers[tf-cpu,sklearn,sentencepiece,vision, testing]"
-RUN uv pip install --no-cache-dir  "protobuf==3.20.3"
+RUN uv pip install --no-cache-dir  "protobuf==3.20.3" tensorflow_probability
 
 
 RUN pip uninstall -y transformers