Merge pull request #2917 from AI-Hypercomputer:dataset_update

PiperOrigin-RevId: 853445233

Author: Google-ML-Automation

src/MaxText/configs/base.yml

@@ -587,7 +587,9 @@ eval_dataset_name: 'c4/en:3.0.1'
 train_split: 'train'
 eval_split: 'validation'
 # for HuggingFace input pipeline (dataset_type=hf)
+# Check definition at https://github.com/huggingface/datasets/blob/0feb65dd8733191dd2d1e74215b422fc5939a56a/src/datasets/load.py#L1338-L1408
 hf_path: ''
+hf_name: ''
 hf_data_dir: ''
 hf_train_files: ''
 hf_eval_split: ''

src/MaxText/configs/types.py

@@ -919,7 +919,8 @@ class TfdsDataset(BaseModel):
 class HfDataset(BaseModel):
   """Configuration specific to HuggingFace datasets."""
 
-  hf_path: str = Field("", description="Path or name of the Hugging Face dataset.")
+  hf_path: str = Field("", description="Path of the Hugging Face dataset.")
+  hf_name: str = Field("", description="Name of the Hugging Face dataset.")
   hf_data_dir: PathStr = Field("", description="Data directory for the HF dataset.")
   hf_train_files: Optional[str] = Field(None, description="Files for the HF training split.")
   hf_eval_split: str = Field("", description="Name of the HF evaluation split.")

src/MaxText/input_pipeline/_hf_data_processing.py

@@ -352,6 +352,7 @@ def make_hf_train_iterator(
   """Load, preprocess dataset and return iterators"""
   train_ds = datasets.load_dataset(
       config.hf_path,
+      name=config.hf_name,
       data_dir=config.hf_data_dir,
       data_files=config.hf_train_files,
       split=config.train_split,
@@ -404,6 +405,7 @@ def make_hf_eval_iterator(
   """Make Hugging Face evaluation iterator. Load and preprocess eval dataset: and return iterator."""
   eval_ds = datasets.load_dataset(
       config.hf_path,
+      name=config.hf_name,
       data_dir=config.hf_data_dir,
       data_files=config.hf_eval_files,
       split=config.hf_eval_split,