
import io
import json

import fsspec
import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.json import JsonDatasetReader, JsonDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases

def      lowerCAmelCase_					(  _lowercase    :       Tuple						,       _lowercase    :       Dict)							->	List[str]:







			"""simple docstring"""
			assert isinstance(_lowercase						,       _lowercase)
			assert dataset.num_rows == 4
			assert dataset.num_columns == 3
			assert dataset.column_names == ["col_1", "col_2", "col_3"]
			for feature, expected_dtype in expected_features.items():
						assert dataset.features[feature].dtype == expected_dtype

@pytest.mark.parametrize("""keep_in_memory"""						,       [False, True])
def      lowerCAmelCase_					(  _lowercase    :       List[Any]						,       _lowercase    :       List[Any]						,       _lowercase    :       int)							->	Optional[int]:







			"""simple docstring"""
			a__       :						Optional[Any]							=  tmp_path / """cache"""
			a__       :						List[Any]							=  {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
			with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
						a__       :						str							=  JsonDatasetReader(_lowercase						,       cache_dir=_lowercase						,       keep_in_memory=_lowercase).read()
			_check_json_dataset(_lowercase						,       _lowercase)

@pytest.mark.parametrize(
    """features"""						,       [
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ]						,       )
def      lowerCAmelCase_					(  _lowercase    :       Dict						,       _lowercase    :       Optional[int]						,       _lowercase    :       Optional[int])							->	Tuple:







			"""simple docstring"""
			a__       :						Optional[Any]							=  tmp_path / """cache"""
			a__       :						List[str]							=  {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
			a__       :						int							=  features.copy() if features else default_expected_features
			a__       :						Dict							=  (
			    Features({feature: Value(_lowercase) for feature, dtype in features.items()}) if features is not None else None
			)
			a__       :						Dict							=  JsonDatasetReader(_lowercase						,       features=_lowercase						,       cache_dir=_lowercase).read()
			_check_json_dataset(_lowercase						,       _lowercase)

@pytest.mark.parametrize(
    """features"""						,       [
        None,
        {"""col_3""": """float64""", """col_1""": """string""", """col_2""": """int64"""},
    ]						,       )
def      lowerCAmelCase_					(  _lowercase    :       Optional[int]						,       _lowercase    :       Optional[int]						,       _lowercase    :       Union[str, Any])							->	List[Any]:







			"""simple docstring"""
			a__       :						str							=  tmp_path / """cache"""
			a__       :						Optional[int]							=  {"""col_3""": """float64""", """col_1""": """string""", """col_2""": """int64"""}
			a__       :						Tuple							=  features.copy() if features else default_expected_features
			a__       :						Tuple							=  (
			    Features({feature: Value(_lowercase) for feature, dtype in features.items()}) if features is not None else None
			)
			a__       :						int							=  JsonDatasetReader(_lowercase						,       features=_lowercase						,       cache_dir=_lowercase).read()
			assert isinstance(_lowercase						,       _lowercase)
			assert dataset.num_rows == 2
			assert dataset.num_columns == 3
			assert dataset.column_names == ["col_3", "col_1", "col_2"]
			for feature, expected_dtype in expected_features.items():
						assert dataset.features[feature].dtype == expected_dtype

def      lowerCAmelCase_					(  _lowercase    :       Optional[Any]						,       _lowercase    :       Optional[Any])							->	int:







			"""simple docstring"""
			# jsonl_312_path features are {"col_3": "float64", "col_1": "string", "col_2": "int64"}
			a__       :						Tuple							=  {"""col_2""": """int64""", """col_3""": """float64""", """col_1""": """string"""}
			a__       :						str							=  features.copy()
			a__       :						Optional[int]							=  (
			    Features({feature: Value(_lowercase) for feature, dtype in features.items()}) if features is not None else None
			)
			a__       :						Any							=  tmp_path / """cache"""
			a__       :						Tuple							=  JsonDatasetReader(_lowercase						,       features=_lowercase						,       cache_dir=_lowercase).read()
			assert isinstance(_lowercase						,       _lowercase)
			assert dataset.num_rows == 2
			assert dataset.num_columns == 3
			assert dataset.column_names == ["col_2", "col_3", "col_1"]
			for feature, expected_dtype in expected_features.items():
						assert dataset.features[feature].dtype == expected_dtype

@pytest.mark.parametrize("""split"""						,       [None, NamedSplit("""train"""), """train""", """test"""])
def      lowerCAmelCase_					(  _lowercase    :       List[Any]						,       _lowercase    :       Dict						,       _lowercase    :       Tuple)							->	int:







			"""simple docstring"""
			a__       :						int							=  tmp_path / """cache"""
			a__       :						List[str]							=  {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
			a__       :						Any							=  JsonDatasetReader(_lowercase						,       cache_dir=_lowercase						,       split=_lowercase).read()
			_check_json_dataset(_lowercase						,       _lowercase)
			assert dataset.split == split if split else "train"

@pytest.mark.parametrize("""path_type"""						,       [str, list])
def      lowerCAmelCase_					(  _lowercase    :       Optional[Any]						,       _lowercase    :       int						,       _lowercase    :       Dict)							->	Optional[Any]:







			"""simple docstring"""
			if issubclass(_lowercase						,       _lowercase):
						a__       :						Optional[int]							=  jsonl_path
			elif issubclass(_lowercase						,       _lowercase):
						a__       :						Tuple							=  [jsonl_path]
			a__       :						List[Any]							=  tmp_path / """cache"""
			a__       :						Any							=  {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
			a__       :						str							=  JsonDatasetReader(_lowercase						,       cache_dir=_lowercase).read()
			_check_json_dataset(_lowercase						,       _lowercase)

def      lowerCAmelCase_					(  _lowercase    :       str						,       _lowercase    :       Optional[Any]						,       _lowercase    :       Dict=("train",))							->	Optional[int]:







			"""simple docstring"""
			assert isinstance(_lowercase						,       _lowercase)
			for split in splits:
						a__       :						Dict							=  dataset_dict[split]
						assert dataset.num_rows == 4
						assert dataset.num_columns == 3
						assert dataset.column_names == ["col_1", "col_2", "col_3"]
						for feature, expected_dtype in expected_features.items():
									assert dataset.features[feature].dtype == expected_dtype

@pytest.mark.parametrize("""keep_in_memory"""						,       [False, True])
def      lowerCAmelCase_					(  _lowercase    :       Dict						,       _lowercase    :       Optional[Any]						,       _lowercase    :       List[str])							->	List[Any]:







			"""simple docstring"""
			a__       :						Union[str, Any]							=  tmp_path / """cache"""
			a__       :						Optional[int]							=  {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
			with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
						a__       :						List[Any]							=  JsonDatasetReader({"""train""": jsonl_path}						,       cache_dir=_lowercase						,       keep_in_memory=_lowercase).read()
			_check_json_datasetdict(_lowercase						,       _lowercase)

@pytest.mark.parametrize(
    """features"""						,       [
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ]						,       )
def      lowerCAmelCase_					(  _lowercase    :       str						,       _lowercase    :       Any						,       _lowercase    :       Union[str, Any])							->	List[str]:







			"""simple docstring"""
			a__       :						Tuple							=  tmp_path / """cache"""
			a__       :						str							=  {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
			a__       :						Optional[Any]							=  features.copy() if features else default_expected_features
			a__       :						Any							=  (
			    Features({feature: Value(_lowercase) for feature, dtype in features.items()}) if features is not None else None
			)
			a__       :						Union[str, Any]							=  JsonDatasetReader({"""train""": jsonl_path}						,       features=_lowercase						,       cache_dir=_lowercase).read()
			_check_json_datasetdict(_lowercase						,       _lowercase)

@pytest.mark.parametrize("""split"""						,       [None, NamedSplit("""train"""), """train""", """test"""])
def      lowerCAmelCase_					(  _lowercase    :       Any						,       _lowercase    :       Any						,       _lowercase    :       Union[str, Any])							->	List[str]:







			"""simple docstring"""
			if split:
						a__       :						str							=  {split: jsonl_path}
			else:
						a__       :						Optional[int]							=  """train"""
						a__       :						Tuple							=  {"""train""": jsonl_path, """test""": jsonl_path}
			a__       :						Union[str, Any]							=  tmp_path / """cache"""
			a__       :						Any							=  {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
			a__       :						Tuple							=  JsonDatasetReader(_lowercase						,       cache_dir=_lowercase).read()
			_check_json_datasetdict(_lowercase						,       _lowercase						,       splits=list(path.keys()))
			assert all(dataset[split].split == split for split in path.keys())

def      lowerCAmelCase_					(  _lowercase    :       Optional[int])							->	Any:







			"""simple docstring"""
			return json.load(_lowercase)

def      lowerCAmelCase_					(  _lowercase    :       Optional[int])							->	int:







			"""simple docstring"""
			return [json.loads(_lowercase) for line in buffer]





class        snake_case__    :







						"""simple docstring"""




						@pytest.mark.parametrize("""lines, load_json_function""" ,       [(True, load_json_lines), (False, load_json)]  )
						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Tuple:
									"""simple docstring"""




									with io.BytesIO() as buffer:
												JsonDatasetWriter(__lowercase ,       __lowercase ,       lines=__lowercase  ).write()
												buffer.seek(0  )
												a__       :						int							=  load_json_function(__lowercase  )
									assert isinstance(__lowercase ,       __lowercase  )
									assert isinstance(exported_content[0] ,       __lowercase  )
									assert len(__lowercase  ) == 1_0

						@pytest.mark.parametrize(
						    """orient, container, keys, len_at""" ,       [
						        ("""records""", list, {"""tokens""", """labels""", """answers""", """id"""}, None),
						        ("""split""", dict, {"""columns""", """data"""}, """data"""),
						        ("""index""", dict, set("""0123456789"""  ), None),
						        ("""columns""", dict, {"""tokens""", """labels""", """answers""", """id"""}, """tokens"""),
						        ("""values""", list, None, None),
						        ("""table""", dict, {"""schema""", """data"""}, """data"""),
						    ] ,       )
						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Dict:
									"""simple docstring"""




									with io.BytesIO() as buffer:
												JsonDatasetWriter(__lowercase ,       __lowercase ,       lines=__lowercase ,       orient=__lowercase  ).write()
												buffer.seek(0  )
												a__       :						List[str]							=  load_json(__lowercase  )
									assert isinstance(__lowercase ,       __lowercase  )
									if keys:
												if container is dict:
															assert exported_content.keys() == keys
												else:
															assert exported_content[0].keys() == keys
									else:
												assert not hasattr(__lowercase ,       """keys"""  ) and not hasattr(exported_content[0] ,       """keys"""  )
									if len_at:
												assert len(exported_content[len_at]  ) == 1_0
									else:
												assert len(__lowercase  ) == 1_0

						@pytest.mark.parametrize("""lines, load_json_function""" ,       [(True, load_json_lines), (False, load_json)]  )
						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Union[str, Any]:
									"""simple docstring"""




									with io.BytesIO() as buffer:
												JsonDatasetWriter(__lowercase ,       __lowercase ,       lines=__lowercase ,       num_proc=2  ).write()
												buffer.seek(0  )
												a__       :						Union[str, Any]							=  load_json_function(__lowercase  )
									assert isinstance(__lowercase ,       __lowercase  )
									assert isinstance(exported_content[0] ,       __lowercase  )
									assert len(__lowercase  ) == 1_0

						@pytest.mark.parametrize(
						    """orient, container, keys, len_at""" ,       [
						        ("""records""", list, {"""tokens""", """labels""", """answers""", """id"""}, None),
						        ("""split""", dict, {"""columns""", """data"""}, """data"""),
						        ("""index""", dict, set("""0123456789"""  ), None),
						        ("""columns""", dict, {"""tokens""", """labels""", """answers""", """id"""}, """tokens"""),
						        ("""values""", list, None, None),
						        ("""table""", dict, {"""schema""", """data"""}, """data"""),
						    ] ,       )
						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Dict:
									"""simple docstring"""




									with io.BytesIO() as buffer:
												JsonDatasetWriter(__lowercase ,       __lowercase ,       lines=__lowercase ,       orient=__lowercase ,       num_proc=2  ).write()
												buffer.seek(0  )
												a__       :						Dict							=  load_json(__lowercase  )
									assert isinstance(__lowercase ,       __lowercase  )
									if keys:
												if container is dict:
															assert exported_content.keys() == keys
												else:
															assert exported_content[0].keys() == keys
									else:
												assert not hasattr(__lowercase ,       """keys"""  ) and not hasattr(exported_content[0] ,       """keys"""  )
									if len_at:
												assert len(exported_content[len_at]  ) == 1_0
									else:
												assert len(__lowercase  ) == 1_0

						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase  ) ->				int:
									"""simple docstring"""




									with pytest.raises(__lowercase  ):
												with io.BytesIO() as buffer:
															JsonDatasetWriter(__lowercase ,       __lowercase ,       num_proc=0  )

						@pytest.mark.parametrize("""compression, extension""" ,       [("""gzip""", """gz"""), ("""bz2""", """bz2"""), ("""xz""", """xz""")]  )
						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				List[Any]:
									"""simple docstring"""




									a__       :						str							=  tmp_path_factory.mktemp("""data"""  ) / F'''test.json.{extension}'''
									a__       :						Optional[int]							=  str(shared_datadir / F'''test_file.json.{extension}'''  )
									JsonDatasetWriter(__lowercase ,       __lowercase ,       compression=__lowercase  ).write()

									with fsspec.open(__lowercase ,       """rb""" ,       compression="""infer"""  ) as f:
												a__       :						Optional[Any]							=  f.read()
									with fsspec.open(__lowercase ,       """rb""" ,       compression="""infer"""  ) as f:
												a__       :						Dict							=  f.read()
									assert exported_content == original_content









import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
_lowercase		:				Union[str, Any]      		=data_utils.TransfoXLTokenizer
_lowercase		:				Tuple      		=data_utils.TransfoXLCorpus
_lowercase		:				Optional[int]      		=data_utils
_lowercase		:				int      		=data_utils

def      lowerCAmelCase_					(  _lowercase    :       Optional[Any]						,       _lowercase    :       List[Any]						,       _lowercase    :       Optional[Any]						,       _lowercase    :       Optional[Any])							->	List[Any]:







			"""simple docstring"""
			if transfo_xl_dataset_file:
						# Convert a pre-processed corpus (see original TensorFlow repo)
						with open(_lowercase						,       """rb""") as fp:
									a__       :						Optional[Any]							=  pickle.load(_lowercase						,       encoding="""latin1""")
						# Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
						a__       :						Dict							=  pytorch_dump_folder_path + """/""" + VOCAB_FILES_NAMES["""pretrained_vocab_file"""]
						print(F'''Save vocabulary to {pytorch_vocab_dump_path}''')
						a__       :						List[str]							=  corpus.vocab.__dict__
						torch.save(_lowercase						,       _lowercase)

						a__       :						Optional[int]							=  corpus.__dict__
						corpus_dict_no_vocab.pop("""vocab"""						,       _lowercase)
						a__       :						Tuple							=  pytorch_dump_folder_path + """/""" + CORPUS_NAME
						print(F'''Save dataset to {pytorch_dataset_dump_path}''')
						torch.save(_lowercase						,       _lowercase)

			if tf_checkpoint_path:
						# Convert a pre-trained TensorFlow model
						a__       :						Optional[Any]							=  os.path.abspath(_lowercase)
						a__       :						str							=  os.path.abspath(_lowercase)

						print(F'''Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.''')
						# Initialise PyTorch model
						if transfo_xl_config_file == "":
									a__       :						List[str]							=  TransfoXLConfig()
						else:
									a__       :						Any							=  TransfoXLConfig.from_json_file(_lowercase)
						print(F'''Building PyTorch model from configuration: {config}''')
						a__       :						Any							=  TransfoXLLMHeadModel(_lowercase)

						a__       :						Optional[Any]							=  load_tf_weights_in_transfo_xl(_lowercase						,       _lowercase						,       _lowercase)
						# Save pytorch-model
						a__       :						Tuple							=  os.path.join(_lowercase						,       _lowercase)
						a__       :						Optional[int]							=  os.path.join(_lowercase						,       _lowercase)
						print(F'''Save PyTorch model to {os.path.abspath(_lowercase)}''')
						torch.save(model.state_dict()						,       _lowercase)
						print(F'''Save configuration file to {os.path.abspath(_lowercase)}''')
						with open(_lowercase						,       """w"""						,       encoding="""utf-8""") as f:
									f.write(config.to_json_string())


if __name__ == "__main__":
		_lowercase		:				Tuple      		=argparse.ArgumentParser()
		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(
		    "--tf_checkpoint_path",
		    default="",
		    type=str,
		    help="An optional path to a TensorFlow checkpoint path to be converted.",
		)
		parser.add_argument(
		    "--transfo_xl_config_file",
		    default="",
		    type=str,
		    help=(
		        "An optional config json file corresponding to the pre-trained BERT model. \n"
		        "This specifies the model architecture."
		    ),
		)
		parser.add_argument(
		    "--transfo_xl_dataset_file",
		    default="",
		    type=str,
		    help="An optional dataset file to be converted in a vocabulary.",
		)
		_lowercase		:				Dict      		=parser.parse_args()
		convert_transfo_xl_checkpoint_to_pytorch(
		    args.tf_checkpoint_path,
		    args.transfo_xl_config_file,
		    args.pytorch_dump_folder_path,
		    args.transfo_xl_dataset_file,
		)




