





'''simple docstring'''






import argparse
import os

import torch

from transformers import (
    XLNetConfig,
    XLNetForQuestionAnswering,
    XLNetForSequenceClassification,
    XLNetLMHeadModel,
    load_tf_weights_in_xlnet,
)
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


__lowerCamelCase						: int		=   {
    'cola': 2,
    'mnli': 3,
    'mrpc': 2,
    'sst-2': 2,
    'sts-b': 1,
    'qqp': 2,
    'qnli': 2,
    'rte': 2,
    'wnli': 2,
}


logging.set_verbosity_info()


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE=None   ):


		"""simple docstring"""

		_UpperCamelCase       			=XLNetConfig.from_json_file(snake_case_   )

		_UpperCamelCase       			=finetuning_task.lower() if finetuning_task is not None else """"""
		if finetuning_task in GLUE_TASKS_NUM_LABELS:
				print(f'''Building PyTorch XLNetForSequenceClassification model from configuration: {config}'''   )
				_UpperCamelCase       			=finetuning_task
				_UpperCamelCase       			=GLUE_TASKS_NUM_LABELS[finetuning_task]
				_UpperCamelCase       			=XLNetForSequenceClassification(snake_case_   )
		elif "squad" in finetuning_task:
				_UpperCamelCase       			=finetuning_task
				_UpperCamelCase       			=XLNetForQuestionAnswering(snake_case_   )
		else:
				_UpperCamelCase       			=XLNetLMHeadModel(snake_case_   )

		# Load weights from tf checkpoint
		load_tf_weights_in_xlnet(snake_case_  ,      snake_case_  ,      snake_case_   )

		# Save pytorch-model
		_UpperCamelCase       			=os.path.join(snake_case_  ,      snake_case_   )
		_UpperCamelCase       			=os.path.join(snake_case_  ,      snake_case_   )
		print(f'''Save PyTorch model to {os.path.abspath(snake_case_   )}'''   )
		torch.save(model.state_dict()  ,      snake_case_   )
		print(f'''Save configuration file to {os.path.abspath(snake_case_   )}'''   )
		with open(snake_case_  ,      '''w'''  ,      encoding='''utf-8'''   ) as f:
				f.write(config.to_json_string()   )


if __name__ == "__main__":
						__lowerCamelCase						: Tuple		=   argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
						)
						parser.add_argument(
						    '--xlnet_config_file',
						    default=None,
						    type=str,
						    required=True,
						    help=(
						        'The config json file corresponding to the pre-trained XLNet model. \n'
						        'This specifies the model architecture.'
						    ),
						)
						parser.add_argument(
						    '--pytorch_dump_folder_path',
						    default=None,
						    type=str,
						    required=True,
						    help='Path to the folder to store the PyTorch model or dataset/vocab.',
						)
						parser.add_argument(
						    '--finetuning_task',
						    default=None,
						    type=str,
						    help='Name of a task on which the XLNet TensorFlow model was fine-tuned',
						)
						__lowerCamelCase						: List[str]		=   parser.parse_args()
						print(args)

						convert_xlnet_checkpoint_to_pytorch(
						    args.tf_checkpoint_path, args.xlnet_config_file, args.pytorch_dump_folder_path, args.finetuning_task
						)













'''simple docstring'''






import contextlib
from multiprocessing import Pool, RLock

from tqdm.auto import tqdm

from ..utils import experimental, logging


__lowerCamelCase						: Union[str, Any]		=   logging.get_logger(__name__)







class 		UpperCAmelCase   :


  """simple docstring"""







  lowerCAmelCase_					=				None




@experimental
def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  if ParallelBackendConfig.backend_name is None:
    return _map_with_multiprocessing_pool(
        __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )

  return _map_with_joblib(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  _UpperCamelCase       			=num_proc if num_proc <= len(__SCREAMING_SNAKE_CASE   ) else len(__SCREAMING_SNAKE_CASE   )
  _UpperCamelCase       			=[]  # We organize the splits ourselve (contiguous splits)
  for index in range(__SCREAMING_SNAKE_CASE   ):
    _UpperCamelCase       			=len(__SCREAMING_SNAKE_CASE   ) // num_proc
    _UpperCamelCase       			=len(__SCREAMING_SNAKE_CASE   ) % num_proc
    _UpperCamelCase       			=div * index + min(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )
    _UpperCamelCase       			=start + div + (1 if index < mod else 0)
    split_kwds.append((function, iterable[start:end], types, index, disable_tqdm, desc)   )

  if len(__SCREAMING_SNAKE_CASE   ) != sum(len(i[1]   ) for i in split_kwds   ):
    raise ValueError(
        f'''Error dividing inputs iterable among processes. '''
        f'''Total number of objects {len(__SCREAMING_SNAKE_CASE   )}, '''
        f'''length: {sum(len(i[1]   ) for i in split_kwds   )}'''   )

  logger.info(
      f'''Spawning {num_proc} processes for {len(__SCREAMING_SNAKE_CASE   )} objects in slices of {[len(i[1]   ) for i in split_kwds]}'''   )
  _UpperCamelCase			,     _UpperCamelCase       			=None, None
  if not disable_tqdm:
    _UpperCamelCase			,     _UpperCamelCase       			=(RLock(),), tqdm.set_lock
  with Pool(__SCREAMING_SNAKE_CASE  ,      initargs=__SCREAMING_SNAKE_CASE  ,      initializer=__SCREAMING_SNAKE_CASE   ) as pool:
    _UpperCamelCase       			=pool.map(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )
  logger.info(f'''Finished {num_proc} processes'''   )
  _UpperCamelCase       			=[obj for proc_res in mapped for obj in proc_res]
  logger.info(f'''Unpacked {len(__SCREAMING_SNAKE_CASE   )} objects'''   )

  return mapped


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  import joblib

  with joblib.parallel_backend(ParallelBackendConfig.backend_name  ,      n_jobs=__SCREAMING_SNAKE_CASE   ):
    return joblib.Parallel()(
        joblib.delayed(__SCREAMING_SNAKE_CASE   )((function, obj, types, None, True, None)   ) for obj in iterable   )


@experimental
@contextlib.contextmanager
def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  _UpperCamelCase       			=backend_name

  if backend_name == "spark":
    from joblibspark import register_spark

    register_spark()

    # TODO: call create_cache_and_write_probe if "download" in steps
    # TODO: raise NotImplementedError when Dataset.map etc is called

  try:
    yield
  finally:
    _UpperCamelCase       			=None






