code review changes

deepanker13 · Dec 13, 2023 · f5572d8 · f5572d8
1 parent 5acece0
commit f5572d8
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 4 deletions.
diff --git a/.github/workflows/publish-sdk-images.yaml b/.github/workflows/publish-sdk-images.yaml
@@ -19,5 +19,5 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - component-name: train-api-training-image
-            dockerfile: sdk/python/kubeflow/training/training_container/Dockerfile
+          - component-name: train-api-hf-image
+            dockerfile: sdk/python/kubeflow/trainer/hf_dockerfile
diff --git a/...ow/training/training_container/Dockerfile → sdk/python/kubeflow/trainer/hf_dockerfile b/...ow/training/training_container/Dockerfile → sdk/python/kubeflow/trainer/hf_dockerfile
@@ -14,4 +14,5 @@ FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
  RUN pip install --no-cache-dir -r requirements.txt
 
  # Run storage.py when the container launches
- ENTRYPOINT ["python", "hf_llm_training.py"]
+ ENTRYPOINT ["python", "hf_llm_training.py"]
+
diff --git a/...ing/training_container/hf_llm_training.py → ...ython/kubeflow/trainer/hf_llm_training.py b/...ing/training_container/hf_llm_training.py → ...ython/kubeflow/trainer/hf_llm_training.py
@@ -38,7 +38,7 @@ def load_and_preprocess_data(dataset_dir, tokenizer):
     train_data = load_dataset(dataset_dir, split="train").map(
         lambda x: tokenizer(x["text"]), batched=True
     )
-    train_data = train_data.train_test_split(shuffle=True, test_size=200)
+    train_data = train_data.train_test_split(shuffle=True, test_size=0.1)
 
     try:
         eval_data = load_dataset(dataset_dir, split="eval")

diff --git a/...ining/training_container/requirements.txt → sdk/python/kubeflow/trainer/requirements.txt b/...ining/training_container/requirements.txt → sdk/python/kubeflow/trainer/requirements.txt