use tempfile in tasks.py

mosaicml · Oct 23, 2024 · 3fa2031 · 3fa2031
1 parent 98cef6a
commit 3fa2031
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/llmfoundry/data/finetuning/tasks.py b/llmfoundry/data/finetuning/tasks.py
@@ -34,6 +34,7 @@ def preprocessing_fn(example: Dict) -> Dict[str, str]:
 import importlib
 import logging
 import os
+import tempfile
 import warnings
 from collections.abc import Mapping
 from functools import partial
@@ -91,7 +92,6 @@ def preprocessing_fn(example: Dict) -> Dict[str, str]:
     UnableToProcessPromptResponseError,
     UnknownExampleTypeError,
 )
-from llmfoundry.utils.file_utils import dist_mkdtemp
 #  yapf: enable
 from llmfoundry.utils.logging_utils import SpecificWarningFilter
 
@@ -912,7 +912,7 @@ def build_from_hf(
                 if not os.path.isdir(dataset_name):
                     # dataset_name is not a local dir path, download if needed.
                     local_dataset_dir = os.path.join(
-                        dist_mkdtemp(),
+                        tempfile.mkdtemp(),
                         dataset_name,
                     )