adding hugging face dataset download class

sdk/python/kubeflow/storage_init_container/abstract_model_provider.py

@@ -7,5 +7,5 @@ def load_config(self):
         pass
 
     @abstractmethod
-    def download_model(self):
+    def download_model_and_tokenizer(self):
         pass

sdk/python/kubeflow/storage_init_container/hugging_face.py

@@ -1,7 +1,7 @@
 from abstract_model_provider import modelProvider
+from abstract_dataset_provider import datasetProvider
 from dataclasses import dataclass, field
 from typing import Literal
-import transformers
 from urllib.parse import urlparse
 import json
 from typing import Dict, Any
@@ -45,6 +45,8 @@ def load_config(self, serialised_args):
     def download_model_and_tokenizer(self):
         # implementation for downloading the model
         print("downloading model")
+        import transformers
+
         transformer_type_class = getattr(transformers, self.config.transformer_type)
         parsed_uri = urlparse(self.config.model_uri)
         self.model = parsed_uri.netloc + parsed_uri.path
@@ -57,3 +59,37 @@ def download_model_and_tokenizer(self):
         transformers.AutoTokenizer.from_pretrained(
             self.model, cache_dir=self.config.download_dir
         )
+
+
+@dataclass
+class HfDatasetParams:
+    repo_id: str
+    access_token: str = None
+    allow_patterns: list[str] = None
+    ignore_patterns: list[str] = None
+    download_dir: str = field(default="/workspace/datasets")
+
+    def __post_init__(self):
+        # Custom checks or validations can be added here
+        if self.repo_id is None:
+            raise ValueError("repo_id is None")
+
+
+class HuggingFaceDataset(datasetProvider):
+    def load_config(self, serialised_args):
+        self.config = HfDatasetParams(**json.loads(serialised_args))
+
+    def download_dataset(self):
+        print("downloading dataset")
+        import huggingface_hub
+        from huggingface_hub import snapshot_download
+
+        if self.config.access_token:
+            huggingface_hub.login(self.config.access_token)
+        snapshot_download(
+            repo_id=self.config.repo_id,
+            repo_type="dataset",
+            allow_patterns=self.config.allow_patterns,
+            ignore_patterns=self.config.ignore_patterns,
+            local_dir=self.config.download_dir,
+        )

sdk/python/kubeflow/storage_init_container/requirements.txt

@@ -4,3 +4,4 @@ torchaudio==2.1.1
 einops==0.7.0
 transformers_stream_generator==0.0.4
 boto3==1.33.9
+huggingface_hub

sdk/python/kubeflow/storage_init_container/storage.py

@@ -1,5 +1,5 @@
 import argparse
-from hugging_face import HuggingFace
+from hugging_face import HuggingFace, HuggingFaceDataset
 from s3 import S3
 
 
@@ -19,6 +19,10 @@ def dataset_factory(dataset_provider, dataset_provider_args):
             s3 = S3()
             s3.load_config(dataset_provider_args)
             s3.download_dataset()
+        case "hf":
+            hf = HuggingFaceDataset()
+            hf.load_config(dataset_provider_parameters)
+            hf.download_dataset()
         case _:
             return "This is the default case"