hf-dataset-repoをjob_specの定義にも使うようにした

aynumosir · Jul 4, 2024 · b5915a8 · b5915a8
1 parent 40a65c6
commit b5915a8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/src/ainu_lm_pipeline/components/get_mt5_training_job_spec.py b/src/ainu_lm_pipeline/components/get_mt5_training_job_spec.py
@@ -5,6 +5,7 @@
 def get_mt5_training_job_spec(
     train_image_uri: str,
     push_to_hub: bool,
+    dataset_name: str,
     dataset_revision: str,
 ) -> list:
     worker_pool_specs = [
@@ -13,7 +14,7 @@ def get_mt5_training_job_spec(
                 "image_uri": train_image_uri,
                 "args": [
                     "mt5",
-                    "--dataset-name=aynumosir/ainu-corpora-normalized",
+                    f"--dataset-name={dataset_name}",
                     "--dataset-split=train",
                     f"--dataset-revision={dataset_revision}",
                     f"--push-to-hub={push_to_hub}",

diff --git a/src/ainu_lm_pipeline/pipelines/ainu_mt5_pipeline.py b/src/ainu_lm_pipeline/pipelines/ainu_mt5_pipeline.py
@@ -97,6 +97,7 @@ def ainu_mt5_pipeline(
         get_mt5_training_job_spec(
             train_image_uri=train_image_uri,
             push_to_hub=push_to_hub,
+            dataset_name=hf_dataset_repo,
             dataset_revision=get_dataset_revision_op.output,
         )
         .after(build_custom_train_image_op)