





import importlib
import inspect
from functools import wraps
from typing import TYPE_CHECKING, Optional, Union

from .download.streaming_download_manager import (
    xbasename,
    xdirname,
    xet_parse,
    xexists,
    xgetsize,
    xglob,
    xgzip_open,
    xisdir,
    xisfile,
    xjoin,
    xlistdir,
    xnumpy_load,
    xopen,
    xpandas_read_csv,
    xpandas_read_excel,
    xPath,
    xrelpath,
    xsio_loadmat,
    xsplit,
    xsplitext,
    xwalk,
    xxml_dom_minidom_parse,
)
from .utils.logging import get_logger
from .utils.patching import patch_submodule
from .utils.py_utils import get_imports


UpperCAmelCase			:			str	  = get_logger(__name__)


if TYPE_CHECKING:
       from .builder import DatasetBuilder
def   _A (	SCREAMING_SNAKE_CASE  :     int , SCREAMING_SNAKE_CASE  :     Optional[Union[str, bool]] = None		):





  """simple docstring"""



  a__			:      int            =importlib.import_module(SCREAMING_SNAKE_CASE		)

  # TODO(QL): always update the module to add subsequent new authentication
  if hasattr(SCREAMING_SNAKE_CASE , "_patched_for_streaming"		) and module._patched_for_streaming:
    return

  def wrap_auth(SCREAMING_SNAKE_CASE  :     List[str]		):
    @wraps(SCREAMING_SNAKE_CASE		)
    def wrapper(*SCREAMING_SNAKE_CASE  :     Union[str, Any] , **SCREAMING_SNAKE_CASE  :     Optional[int]		):
      return function(*SCREAMING_SNAKE_CASE , use_auth_token=SCREAMING_SNAKE_CASE , **SCREAMING_SNAKE_CASE		)

    a__			:      Any            ="wrap_auth"
    return wrapper

  # open files in a streaming fashion
  patch_submodule(SCREAMING_SNAKE_CASE , "open" , wrap_auth(SCREAMING_SNAKE_CASE		)		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "os.listdir" , wrap_auth(SCREAMING_SNAKE_CASE		)		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "os.walk" , wrap_auth(SCREAMING_SNAKE_CASE		)		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "glob.glob" , wrap_auth(SCREAMING_SNAKE_CASE		)		).start()
  # allow to navigate in remote zip files
  patch_submodule(SCREAMING_SNAKE_CASE , "os.path.join" , SCREAMING_SNAKE_CASE		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "os.path.dirname" , SCREAMING_SNAKE_CASE		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "os.path.basename" , SCREAMING_SNAKE_CASE		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "os.path.relpath" , SCREAMING_SNAKE_CASE		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "os.path.split" , SCREAMING_SNAKE_CASE		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "os.path.splitext" , SCREAMING_SNAKE_CASE		).start()
  # allow checks on paths
  patch_submodule(SCREAMING_SNAKE_CASE , "os.path.exists" , wrap_auth(SCREAMING_SNAKE_CASE		)		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "os.path.isdir" , wrap_auth(SCREAMING_SNAKE_CASE		)		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "os.path.isfile" , wrap_auth(SCREAMING_SNAKE_CASE		)		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "os.path.getsize" , wrap_auth(SCREAMING_SNAKE_CASE		)		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "pathlib.Path" , SCREAMING_SNAKE_CASE		).start()
  # file readers
  patch_submodule(SCREAMING_SNAKE_CASE , "gzip.open" , wrap_auth(SCREAMING_SNAKE_CASE		)		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "numpy.load" , wrap_auth(SCREAMING_SNAKE_CASE		)		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "pandas.read_csv" , wrap_auth(SCREAMING_SNAKE_CASE		) , attrs=["__version__"]		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "pandas.read_excel" , wrap_auth(SCREAMING_SNAKE_CASE		) , attrs=["__version__"]		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "scipy.io.loadmat" , wrap_auth(SCREAMING_SNAKE_CASE		) , attrs=["__version__"]		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "xml.etree.ElementTree.parse" , wrap_auth(SCREAMING_SNAKE_CASE		)		).start()
  patch_submodule(SCREAMING_SNAKE_CASE , "xml.dom.minidom.parse" , wrap_auth(SCREAMING_SNAKE_CASE		)		).start()
  a__			:      Any            =True





def   _A (	SCREAMING_SNAKE_CASE  :     "DatasetBuilder"		):





  """simple docstring"""



  extend_module_for_streaming(builder.__module__ , use_auth_token=builder.use_auth_token		)
  # if needed, we also have to extend additional internal imports (like wmt14 -> wmt_utils)
  if not builder.__module__.startswith("datasets."		):  # check that it's not a packaged builder like csv
    for imports in get_imports(inspect.getfile(builder.__class__		)		):
      if imports[0] == "internal":
        a__			:      int            =imports[1]
        a__			:      Optional[Any]            =".".join(builder.__module__.split("."		)[:-1] + [internal_import_name]		)
        extend_module_for_streaming(SCREAMING_SNAKE_CASE , use_auth_token=builder.use_auth_token		)

    # builders can inherit from other builders that might use streaming functionality
    # (for example, ImageFolder and AudioFolder inherit from FolderBuilder which implements examples generation)
    # but these parents builders are not patched automatically as they are not instantiated, so we patch them here
  from .builder import DatasetBuilder

  a__			:      Optional[Any]            =[
      cls.__module__
      for cls in type(SCREAMING_SNAKE_CASE		).__mro__[1:]  # make sure it's not the same module we've already patched
      if issubclass(cls , SCREAMING_SNAKE_CASE		) and cls.__module__ != DatasetBuilder.__module__
  ]  # check it's not a standard builder from datasets.builder
  for module in parent_builder_modules:
    extend_module_for_streaming(SCREAMING_SNAKE_CASE , use_auth_token=builder.use_auth_token		)













import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp
def   _A (	):





  """simple docstring"""



  a__			:      Optional[int]            =ArgumentParser(
      description=(
          "PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes"
      )		)

  # Optional arguments for the launch helper
  parser.add_argument("--num_cores" , type=SCREAMING_SNAKE_CASE , default=1 , help="Number of TPU cores to use (1 or 8)."		)

  # positional
  parser.add_argument(
      "training_script" , type=SCREAMING_SNAKE_CASE , help=(
          "The full path to the single TPU training "
          "program/script to be launched in parallel, "
          "followed by all the arguments for the "
          "training script"
      ) , )

  # rest from the training program
  parser.add_argument("training_script_args" , nargs=SCREAMING_SNAKE_CASE		)

  return parser.parse_args()
def   _A (	):





  """simple docstring"""



  a__			:      Optional[Any]            =parse_args()

  # Import training_script as a module.
  a__			:      Tuple            =Path(args.training_script		)
  sys.path.append(str(script_fpath.parent.resolve()		)		)
  a__			:      Tuple            =script_fpath.stem
  a__			:      Any            =importlib.import_module(SCREAMING_SNAKE_CASE		)

  # Patch sys.argv
  a__			:      List[str]            =[args.training_script] + args.training_script_args + ["--tpu_num_cores", str(args.num_cores		)]

  xmp.spawn(mod._mp_fn , args=() , nprocs=args.num_cores		)


if __name__ == "__main__":
     main()






