Skip to content

Commit 69c097c

Browse files
author
The TensorFlow Datasets Authors
committed
Include trust_remote_code flag for code-based HF builders preparation.
PiperOrigin-RevId: 661246526
1 parent 16a566f commit 69c097c

File tree

2 files changed

+9
-2
lines changed

2 files changed

+9
-2
lines changed

tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,7 @@ def __init__(
200200
tfds_num_proc: Optional[int] = None,
201201
ignore_hf_errors: bool = False,
202202
overwrite_version: str | None = None,
203+
trust_remote_code: bool = False,
203204
**config_kwargs,
204205
):
205206
self._hf_repo_id = hf_repo_id
@@ -210,7 +211,10 @@ def __init__(
210211
)
211212
try:
212213
self._hf_builder = hf_datasets.load_dataset_builder(
213-
self._hf_repo_id, self._hf_config, **self.config_kwargs
214+
self._hf_repo_id,
215+
self._hf_config,
216+
trust_remote_code=trust_remote_code,
217+
**self.config_kwargs,
214218
)
215219
except Exception as e:
216220
raise RuntimeError(

tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder_test.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,10 @@ def mock_huggingface_dataset_builder(
104104
other_arg='this is another arg',
105105
)
106106
load_dataset_builder.assert_called_once_with(
107-
'foo/bar', 'config', other_arg='this is another arg'
107+
'foo/bar',
108+
'config',
109+
other_arg='this is another arg',
110+
trust_remote_code=False,
108111
)
109112
login_to_hf.assert_called_once_with('SECRET_TOKEN')
110113
yield builder

0 commit comments

Comments
 (0)