diff --git a/tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder.py b/tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder.py index f11e18ee994..50acdf6639c 100644 --- a/tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder.py +++ b/tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder.py @@ -200,6 +200,7 @@ def __init__( tfds_num_proc: Optional[int] = None, ignore_hf_errors: bool = False, overwrite_version: str | None = None, + trust_remote_code: bool = False, **config_kwargs, ): self._hf_repo_id = hf_repo_id @@ -210,7 +211,10 @@ def __init__( ) try: self._hf_builder = hf_datasets.load_dataset_builder( - self._hf_repo_id, self._hf_config, **self.config_kwargs + self._hf_repo_id, + self._hf_config, + trust_remote_code=trust_remote_code, + **self.config_kwargs, ) except Exception as e: raise RuntimeError( diff --git a/tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder_test.py b/tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder_test.py index a15c21a2d4f..304c644289a 100644 --- a/tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder_test.py +++ b/tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder_test.py @@ -104,7 +104,10 @@ def mock_huggingface_dataset_builder( other_arg='this is another arg', ) load_dataset_builder.assert_called_once_with( - 'foo/bar', 'config', other_arg='this is another arg' + 'foo/bar', + 'config', + other_arg='this is another arg', + trust_remote_code=False, ) login_to_hf.assert_called_once_with('SECRET_TOKEN') yield builder