V1: Fix HF dataset initializer

Signed-off-by: Andrey Velichkevich <[email protected]>
kubeflow · Jan 7, 2025 · 1bdbba0 · 1bdbba0
1 parent 80e2097
commit 1bdbba0
Show file tree

Hide file tree

Showing 4 changed files with 3 additions and 5 deletions.
diff --git a/.github/workflows/e2e-test-train-api.yaml b/.github/workflows/e2e-test-train-api.yaml
@@ -58,4 +58,4 @@ jobs:
           pytest -s sdk/python/test/e2e-fine-tune-llm/test_e2e_pytorch_fine_tune_llm.py --log-cli-level=debug
         env:
           STORAGE_INITIALIZER_IMAGE: kubeflowtraining/storage-initializer:test
-          TRAINER_TRANSFORMER_IMAGE_DEFAULT: kubeflowtraining/trainer:test
+          TRAINER_TRANSFORMER_IMAGE: kubeflowtraining/trainer:test
diff --git a/sdk/python/kubeflow/storage_initializer/requirements.txt b/sdk/python/kubeflow/storage_initializer/requirements.txt
@@ -1,5 +1,4 @@
 peft==0.3.0
-datasets==2.15.0
+datasets==2.21.0
 transformers==4.38.0
 boto3==1.33.9
-huggingface_hub==0.23.4
diff --git a/sdk/python/kubeflow/trainer/Dockerfile.cpu b/sdk/python/kubeflow/trainer/Dockerfile.cpu
@@ -8,7 +8,6 @@ WORKDIR /app
 COPY requirements.txt /app/requirements.txt
 
 # Install any needed packages specified in requirements.txt
-RUN pip install --no-cache-dir torch==2.5.1
 RUN pip install --no-cache-dir -r requirements.txt
 
 # Copy the Python package and its source code into the container

diff --git a/sdk/python/kubeflow/trainer/requirements.txt b/sdk/python/kubeflow/trainer/requirements.txt
@@ -1,4 +1,4 @@
 peft==0.3.0
-datasets==2.15.0
+datasets==2.21.0
 transformers==4.38.0
 accelerate==0.28.0