




import os
import re
import shutil
from argparse import ArgumentParser, Namespace

from datasets.commands import BaseDatasetsCLICommand
from datasets.utils.logging import get_logger


UpperCamelCase__    			=			"""<<<<<<< This should probably be modified because it mentions: """

UpperCamelCase__    			=			"""=======
>>>>>>>
"""

UpperCamelCase__    			=			[
    """TextEncoderConfig""",
    """ByteTextEncoder""",
    """SubwordTextEncoder""",
    """encoder_config""",
    """maybe_build_from_corpus""",
    """manual_dir""",
]

UpperCamelCase__    			=			[
    # (pattern, replacement)
    # Order is important here for some replacements
    (R"""tfds\.core""", R"""datasets"""),
    (R"""tf\.io\.gfile\.GFile""", R"""open"""),
    (R"""tf\.([\w\d]+)""", R"""datasets.Value('\1')"""),
    (R"""tfds\.features\.Text\(\)""", R"""datasets.Value('string')"""),
    (R"""tfds\.features\.Text\(""", R"""datasets.Value('string'),"""),
    (R"""features\s*=\s*tfds.features.FeaturesDict\(""", R"""features=datasets.Features("""),
    (R"""tfds\.features\.FeaturesDict\(""", R"""dict("""),
    (R"""The TensorFlow Datasets Authors""", R"""The TensorFlow Datasets Authors and the HuggingFace Datasets Authors"""),
    (R"""tfds\.""", R"""datasets."""),
    (R"""dl_manager\.manual_dir""", R"""self.config.data_dir"""),
    (R"""self\.builder_config""", R"""self.config"""),
]







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Namespace     ):
					return ConvertCommand(args.tfds_path      ,      args.datasets_directory     )


class 			a__    (					snake_case__      ):






						@staticmethod
						def 				__SCREAMING_SNAKE_CASE(   _A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       parser.add_parser(
											    "convert"		,			help="Convert a TensorFlow Datasets dataset to a HuggingFace Datasets dataset."		,			)
											train_parser.add_argument(
											    "--tfds_path"		,			type=_A		,			required=_A		,			help="Path to a TensorFlow Datasets folder to convert or a single tfds file to convert."		,			)
											train_parser.add_argument(
											    "--datasets_directory"		,			type=_A		,			required=_A		,			help="Path to the HuggingFace Datasets folder."			)
											train_parser.set_defaults(func=_A			)






						def __init__(   self		,			_A		,			_A		,			*_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       get_logger("datasets-cli/converting"			)

											__lowerCAmelCase				   =       tfds_path
											__lowerCAmelCase				   =       datasets_directory






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											if os.path.isdir(self._tfds_path			):
																__lowerCAmelCase				   =       os.path.abspath(self._tfds_path			)
											elif os.path.isfile(self._tfds_path			):
																__lowerCAmelCase				   =       os.path.dirname(self._tfds_path			)
											else:
																raise ValueError("--tfds_path is neither a directory nor a file. Please check path."			)

											__lowerCAmelCase				   =       os.path.abspath(self._datasets_directory			)

											self._logger.info(f"""Converting datasets from {abs_tfds_path} to {abs_datasets_path}"""			)

											__lowerCAmelCase				   =       []
											__lowerCAmelCase				   =       []
											__lowerCAmelCase				   =       {}

											if os.path.isdir(self._tfds_path			):
																__lowerCAmelCase				   =       os.listdir(_A			)
											else:
																__lowerCAmelCase				   =       [os.path.basename(self._tfds_path			)]

											for f_name in file_names:
																self._logger.info(f"""Looking at file {f_name}"""			)
																__lowerCAmelCase				   =       os.path.join(_A		,			_A			)
																__lowerCAmelCase				   =       os.path.join(_A		,			_A			)

																if not os.path.isfile(_A			) or "__init__" in f_name or "_test" in f_name or ".py" not in f_name:
																					self._logger.info("Skipping file"			)
																					continue

																with open(_A		,			encoding="utf-8"			) as f:
																					__lowerCAmelCase				   =       f.readlines()

																__lowerCAmelCase				   =       []
																__lowerCAmelCase				   =       False
																__lowerCAmelCase				   =       False
																__lowerCAmelCase				   =       []
																for line in lines:
																					__lowerCAmelCase				   =       line

																					# Convert imports
																					if "import tensorflow.compat.v2 as tf" in out_line:
																										continue
																					elif "@tfds.core" in out_line:
																										continue
																					elif "builder=self" in out_line:
																										continue
																					elif "import tensorflow_datasets.public_api as tfds" in out_line:
																										__lowerCAmelCase				   =       "import datasets\n"
																					elif "import tensorflow" in out_line:
																										# order is important here
																										__lowerCAmelCase				   =       ""
																										continue
																					elif "from absl import logging" in out_line:
																										__lowerCAmelCase				   =       "from datasets import logging\n"
																					elif "getLogger" in out_line:
																										__lowerCAmelCase				   =       out_line.replace("getLogger"		,			"get_logger"			)
																					elif any(expression in out_line for expression in TO_HIGHLIGHT			):
																										__lowerCAmelCase				   =       True
																										__lowerCAmelCase				   =       list(filter(lambda _A			: e in out_line		,			_A			)			)
																										out_lines.append(HIGHLIGHT_MESSAGE_PRE + str(_A			) + "\n"			)
																										out_lines.append(_A			)
																										out_lines.append(_A			)
																										continue
																					else:
																										for pattern, replacement in TO_CONVERT:
																															__lowerCAmelCase				   =       re.sub(_A		,			_A		,			_A			)

                # Take care of saving utilities (to later move them together with main script)
																					if "tensorflow_datasets" in out_line:
																										__lowerCAmelCase				   =       re.match(R"from\stensorflow_datasets.*import\s([^\.\r\n]+)"		,			_A			)
																										tfds_imports.extend(imp.strip() for imp in match.group(1			).split(","			)			)
																										__lowerCAmelCase				   =       "from . import " + match.group(1			)

																					# Check we have not forget anything
																					if "tf." in out_line or "tfds." in out_line or "tensorflow_datasets" in out_line:
																										raise ValueError(f"""Error converting {out_line.strip()}"""			)

																					if "GeneratorBasedBuilder" in out_line or "BeamBasedBuilder" in out_line:
																										__lowerCAmelCase				   =       True
																					out_lines.append(_A			)

																if is_builder or "wmt" in f_name:
																					# We create a new directory for each dataset
																					__lowerCAmelCase				   =       f_name.replace(".py"		,			""			)
																					__lowerCAmelCase				   =       os.path.join(_A		,			_A			)
																					__lowerCAmelCase				   =       os.path.join(_A		,			_A			)
																					os.makedirs(_A		,			exist_ok=_A			)
																					self._logger.info(f"""Adding directory {output_dir}"""			)
																					imports_to_builder_map.update({imp: output_dir for imp in tfds_imports}			)
																else:
																					# Utilities will be moved at the end
																					utils_files.append(_A			)

																if needs_manual_update:
																					with_manual_update.append(_A			)

																with open(_A		,			"w"		,			encoding="utf-8"			) as f:
																					f.writelines(_A			)
																self._logger.info(f"""Converted in {output_file}"""			)

											for utils_file in utils_files:
																try:
																					__lowerCAmelCase				   =       os.path.basename(_A			)
																					__lowerCAmelCase				   =       imports_to_builder_map[f_name.replace(".py"		,			""			)]
																					self._logger.info(f"""Moving {dest_folder} to {utils_file}"""			)
																					shutil.copy(_A		,			_A			)
																except KeyError:
																					self._logger.error(f"""Cannot find destination folder for {utils_file}. Please copy manually."""			)

											if with_manual_update:
																for file_path in with_manual_update:
																					self._logger.warning(
																					    f"""You need to manually update file {file_path} to remove configurations using 'TextEncoderConfig'."""			)










import os

import torch

from ..logging import get_logger
from .constants import FSDP_PYTORCH_VERSION, MODEL_NAME, OPTIMIZER_NAME
from .versions import is_torch_version


if is_torch_version(""">=""", FSDP_PYTORCH_VERSION):
					import torch.distributed.checkpoint as dist_cp
					from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
					from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
					from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
					from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType


UpperCamelCase__    			=			get_logger(__name__)







def 	_a      (						SCREAMING_SNAKE_CASE_							:				List[str]      ,      SCREAMING_SNAKE_CASE_							:				Tuple      ,      SCREAMING_SNAKE_CASE_							:				Tuple      ,      SCREAMING_SNAKE_CASE_							:				Optional[Any]      ,      SCREAMING_SNAKE_CASE_							:				str=0     ):
					os.makedirs(SCREAMING_SNAKE_CASE_      ,      exist_ok=SCREAMING_SNAKE_CASE_     )
					with FSDP.state_dict_type(
					    SCREAMING_SNAKE_CASE_      ,      fsdp_plugin.state_dict_type      ,      fsdp_plugin.state_dict_config      ,      fsdp_plugin.optim_state_dict_config     ):
										__lowerCAmelCase				   =       model.state_dict()
										if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
															__lowerCAmelCase				   =       F"""{MODEL_NAME}.bin""" if model_index == 0 else F"""{MODEL_NAME}_{model_index}.bin"""
															__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															if accelerator.process_index == 0:
																				logger.info(F"""Saving model to {output_model_file}"""     )
																				torch.save(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
																				logger.info(F"""Model saved to {output_model_file}"""     )
										elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
															__lowerCAmelCase				   =       (
															    F"""{MODEL_NAME}_rank{accelerator.process_index}.bin"""
															    if model_index == 0
															    else F"""{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin"""
															)
															__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Saving model to {output_model_file}"""     )
															torch.save(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Model saved to {output_model_file}"""     )
										elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
															__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      F"""{MODEL_NAME}_{model_index}"""     )
															os.makedirs(SCREAMING_SNAKE_CASE_      ,      exist_ok=SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Saving model to {ckpt_dir}"""     )
															__lowerCAmelCase				   =       {"model": state_dict}

															dist_cp.save_state_dict(
															    state_dict=SCREAMING_SNAKE_CASE_      ,      storage_writer=dist_cp.FileSystemWriter(SCREAMING_SNAKE_CASE_     )      ,      planner=DefaultSavePlanner()      ,      )
															logger.info(F"""Model saved to {ckpt_dir}"""     )







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Dict      ,      SCREAMING_SNAKE_CASE_							:				int      ,      SCREAMING_SNAKE_CASE_							:				Optional[int]      ,      SCREAMING_SNAKE_CASE_							:				List[str]      ,      SCREAMING_SNAKE_CASE_							:				Any=0     ):
					accelerator.wait_for_everyone()
					with FSDP.state_dict_type(
					    SCREAMING_SNAKE_CASE_      ,      fsdp_plugin.state_dict_type      ,      fsdp_plugin.state_dict_config      ,      fsdp_plugin.optim_state_dict_config     ):
										if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
															if type(SCREAMING_SNAKE_CASE_     ) != FSDP and accelerator.process_index != 0:
																				if not fsdp_plugin.sync_module_states:
																									raise ValueError(
																									    "Set the `sync_module_states` flag to `True` so that model states are synced across processes when "
																									    "initializing FSDP object"     )
																				return
															__lowerCAmelCase				   =       F"""{MODEL_NAME}.bin""" if model_index == 0 else F"""{MODEL_NAME}_{model_index}.bin"""
															__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Loading model from {input_model_file}"""     )
															__lowerCAmelCase				   =       torch.load(SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Model loaded from {input_model_file}"""     )
										elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
															__lowerCAmelCase				   =       (
															    F"""{MODEL_NAME}_rank{accelerator.process_index}.bin"""
															    if model_index == 0
															    else F"""{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin"""
															)
															__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Loading model from {input_model_file}"""     )
															__lowerCAmelCase				   =       torch.load(SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Model loaded from {input_model_file}"""     )
										elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
															__lowerCAmelCase				   =       (
															    os.path.join(SCREAMING_SNAKE_CASE_      ,      F"""{MODEL_NAME}_{model_index}"""     )
															    if F"""{MODEL_NAME}""" not in input_dir
															    else input_dir
															)
															logger.info(F"""Loading model from {ckpt_dir}"""     )
															__lowerCAmelCase				   =       {"model": model.state_dict()}
															dist_cp.load_state_dict(
															    state_dict=SCREAMING_SNAKE_CASE_      ,      storage_reader=dist_cp.FileSystemReader(SCREAMING_SNAKE_CASE_     )      ,      planner=DefaultLoadPlanner()      ,      )
															__lowerCAmelCase				   =       state_dict["model"]
															logger.info(F"""Model loaded from {ckpt_dir}"""     )
										model.load_state_dict(SCREAMING_SNAKE_CASE_     )







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Dict      ,      SCREAMING_SNAKE_CASE_							:				int      ,      SCREAMING_SNAKE_CASE_							:				List[str]      ,      SCREAMING_SNAKE_CASE_							:				Optional[Any]      ,      SCREAMING_SNAKE_CASE_							:				Optional[int]      ,      SCREAMING_SNAKE_CASE_							:				str=0     ):
					os.makedirs(SCREAMING_SNAKE_CASE_      ,      exist_ok=SCREAMING_SNAKE_CASE_     )
					with FSDP.state_dict_type(
					    SCREAMING_SNAKE_CASE_      ,      fsdp_plugin.state_dict_type      ,      fsdp_plugin.state_dict_config      ,      fsdp_plugin.optim_state_dict_config     ):
										__lowerCAmelCase				   =       FSDP.optim_state_dict(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
										if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
															if accelerator.process_index == 0:
																				__lowerCAmelCase				   =       (
																				    F"""{OPTIMIZER_NAME}.bin""" if optimizer_index == 0 else F"""{OPTIMIZER_NAME}_{optimizer_index}.bin"""
																				)
																				__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
																				logger.info(F"""Saving Optimizer state to {output_optimizer_file}"""     )
																				torch.save(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
																				logger.info(F"""Optimizer state saved in {output_optimizer_file}"""     )
										else:
															__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      F"""{OPTIMIZER_NAME}_{optimizer_index}"""     )
															os.makedirs(SCREAMING_SNAKE_CASE_      ,      exist_ok=SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Saving Optimizer state to {ckpt_dir}"""     )
															dist_cp.save_state_dict(
															    state_dict={"optimizer": optim_state}      ,      storage_writer=dist_cp.FileSystemWriter(SCREAMING_SNAKE_CASE_     )      ,      planner=DefaultSavePlanner()      ,      )
															logger.info(F"""Optimizer state saved in {ckpt_dir}"""     )







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Dict      ,      SCREAMING_SNAKE_CASE_							:				Dict      ,      SCREAMING_SNAKE_CASE_							:				Tuple      ,      SCREAMING_SNAKE_CASE_							:				List[str]      ,      SCREAMING_SNAKE_CASE_							:				str      ,      SCREAMING_SNAKE_CASE_							:				Dict=0     ):
					accelerator.wait_for_everyone()
					with FSDP.state_dict_type(
					    SCREAMING_SNAKE_CASE_      ,      fsdp_plugin.state_dict_type      ,      fsdp_plugin.state_dict_config      ,      fsdp_plugin.optim_state_dict_config     ):
										if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
															__lowerCAmelCase				   =       None
															# below check should work but currently it isn't working (mostly opytorch issue),
															# in the meantime disabling it at the cost of excess memory usage
															# if accelerator.process_index == 0 or not fsdp_plugin.optim_state_dict_config.rank0_only:
															__lowerCAmelCase				   =       (
															    F"""{OPTIMIZER_NAME}.bin""" if optimizer_index == 0 else F"""{OPTIMIZER_NAME}_{optimizer_index}.bin"""
															)
															__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Loading Optimizer state from {input_optimizer_file}"""     )
															__lowerCAmelCase				   =       torch.load(SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Optimizer state loaded from {input_optimizer_file}"""     )
										else:
															__lowerCAmelCase				   =       (
															    os.path.join(SCREAMING_SNAKE_CASE_      ,      F"""{OPTIMIZER_NAME}_{optimizer_index}"""     )
															    if F"""{OPTIMIZER_NAME}""" not in input_dir
															    else input_dir
															)
															logger.info(F"""Loading Optimizer from {ckpt_dir}"""     )
															__lowerCAmelCase				   =       load_sharded_optimizer_state_dict(
															    model_state_dict=model.state_dict()      ,      optimizer_key="optimizer"      ,      storage_reader=dist_cp.FileSystemReader(SCREAMING_SNAKE_CASE_     )      ,      )
															__lowerCAmelCase				   =       optim_state["optimizer"]
															logger.info(F"""Optimizer loaded from {ckpt_dir}"""     )
										__lowerCAmelCase				   =       FSDP.optim_state_dict_to_load(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
										optimizer.load_state_dict(SCREAMING_SNAKE_CASE_     )




