From 69c097c4361e715d9d3311869cf205948491801a Mon Sep 17 00:00:00 2001 From: The TensorFlow Datasets Authors Date: Fri, 9 Aug 2024 06:56:39 -0700 Subject: [PATCH] Include `trust_remote_code` flag for code-based HF builders preparation. PiperOrigin-RevId: 661246526 --- .../core/dataset_builders/huggingface_dataset_builder.py | 6 +++++- .../dataset_builders/huggingface_dataset_builder_test.py | 5 ++++- 2 files changed, 9 insertions(+), 2 deletions(-) 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