


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 __UpperCamelCase   (			__lowerCamelCase		: int      ,  __lowerCamelCase		: Tuple				)      ->     int:







     '''simple docstring'''






     assert isinstance(__lowerCamelCase      ,  __lowerCamelCase				)
     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 __UpperCamelCase   (			__lowerCamelCase		: Union[str, Any]      ,  __lowerCamelCase		: Any      ,  __lowerCamelCase		: Tuple				)      ->     Union[str, Any]:







     '''simple docstring'''






     _a						=		tmp_path / "cache"
     _a						=		{"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						=		JsonDatasetReader(__lowerCamelCase      ,  cache_dir=__lowerCamelCase      ,  keep_in_memory=__lowerCamelCase				).read()
     _check_json_dataset(__lowerCamelCase      ,  __lowerCamelCase				)






@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 __UpperCamelCase   (			__lowerCamelCase		: Union[str, Any]      ,  __lowerCamelCase		: Optional[Any]      ,  __lowerCamelCase		: Optional[int]				)      ->     Optional[Any]:







     '''simple docstring'''






     _a						=		tmp_path / "cache"
     _a						=		{"col_1": "string", "col_2": "int64", "col_3": "float64"}
     _a						=		features.copy() if features else default_expected_features
     _a						=		(
         Features({feature: Value(__lowerCamelCase				) for feature, dtype in features.items()}				) if features is not None else None
     )
     _a						=		JsonDatasetReader(__lowerCamelCase      ,  features=__lowerCamelCase      ,  cache_dir=__lowerCamelCase				).read()
     _check_json_dataset(__lowerCamelCase      ,  __lowerCamelCase				)






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







     '''simple docstring'''






     _a						=		tmp_path / "cache"
     _a						=		{"col_3": "float64", "col_1": "string", "col_2": "int64"}
     _a						=		features.copy() if features else default_expected_features
     _a						=		(
         Features({feature: Value(__lowerCamelCase				) for feature, dtype in features.items()}				) if features is not None else None
     )
     _a						=		JsonDatasetReader(__lowerCamelCase      ,  features=__lowerCamelCase      ,  cache_dir=__lowerCamelCase				).read()
     assert isinstance(__lowerCamelCase      ,  __lowerCamelCase				)
     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 __UpperCamelCase   (			__lowerCamelCase		: Any      ,  __lowerCamelCase		: Optional[Any]				)      ->     int:







     '''simple docstring'''






     _a						=		{"col_2": "int64", "col_3": "float64", "col_1": "string"}
     _a						=		features.copy()
     _a						=		(
         Features({feature: Value(__lowerCamelCase				) for feature, dtype in features.items()}				) if features is not None else None
     )
     _a						=		tmp_path / "cache"
     _a						=		JsonDatasetReader(__lowerCamelCase      ,  features=__lowerCamelCase      ,  cache_dir=__lowerCamelCase				).read()
     assert isinstance(__lowerCamelCase      ,  __lowerCamelCase				)
     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 __UpperCamelCase   (			__lowerCamelCase		: Tuple      ,  __lowerCamelCase		: List[Any]      ,  __lowerCamelCase		: Tuple				)      ->     Optional[int]:







     '''simple docstring'''






     _a						=		tmp_path / "cache"
     _a						=		{"col_1": "string", "col_2": "int64", "col_3": "float64"}
     _a						=		JsonDatasetReader(__lowerCamelCase      ,  cache_dir=__lowerCamelCase      ,  split=__lowerCamelCase				).read()
     _check_json_dataset(__lowerCamelCase      ,  __lowerCamelCase				)
     assert dataset.split == split if split else "train"






@pytest.mark.parametrize("path_type"      ,  [str, list]				)
def __UpperCamelCase   (			__lowerCamelCase		: List[Any]      ,  __lowerCamelCase		: str      ,  __lowerCamelCase		: Optional[Any]				)      ->     Union[str, Any]:







     '''simple docstring'''






     if issubclass(__lowerCamelCase      ,  __lowerCamelCase				):
          _a						=		jsonl_path
     elif issubclass(__lowerCamelCase      ,  __lowerCamelCase				):
          _a						=		[jsonl_path]
     _a						=		tmp_path / "cache"
     _a						=		{"col_1": "string", "col_2": "int64", "col_3": "float64"}
     _a						=		JsonDatasetReader(__lowerCamelCase      ,  cache_dir=__lowerCamelCase				).read()
     _check_json_dataset(__lowerCamelCase      ,  __lowerCamelCase				)






def __UpperCamelCase   (			__lowerCamelCase		: str      ,  __lowerCamelCase		: Optional[Any]      ,  __lowerCamelCase		: Dict=("train",)				)      ->     Tuple:







     '''simple docstring'''






     assert isinstance(__lowerCamelCase      ,  __lowerCamelCase				)
     for split in splits:
          _a						=		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 __UpperCamelCase   (			__lowerCamelCase		: int      ,  __lowerCamelCase		: int      ,  __lowerCamelCase		: str				)      ->     List[Any]:







     '''simple docstring'''






     _a						=		tmp_path / "cache"
     _a						=		{"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						=		JsonDatasetReader({"train": jsonl_path}      ,  cache_dir=__lowerCamelCase      ,  keep_in_memory=__lowerCamelCase				).read()
     _check_json_datasetdict(__lowerCamelCase      ,  __lowerCamelCase				)






@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 __UpperCamelCase   (			__lowerCamelCase		: Optional[Any]      ,  __lowerCamelCase		: int      ,  __lowerCamelCase		: Union[str, Any]				)      ->     Dict:







     '''simple docstring'''






     _a						=		tmp_path / "cache"
     _a						=		{"col_1": "string", "col_2": "int64", "col_3": "float64"}
     _a						=		features.copy() if features else default_expected_features
     _a						=		(
         Features({feature: Value(__lowerCamelCase				) for feature, dtype in features.items()}				) if features is not None else None
     )
     _a						=		JsonDatasetReader({"train": jsonl_path}      ,  features=__lowerCamelCase      ,  cache_dir=__lowerCamelCase				).read()
     _check_json_datasetdict(__lowerCamelCase      ,  __lowerCamelCase				)






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







     '''simple docstring'''






     if split:
          _a						=		{split: jsonl_path}
     else:
          _a						=		"train"
          _a						=		{"train": jsonl_path, "test": jsonl_path}
     _a						=		tmp_path / "cache"
     _a						=		{"col_1": "string", "col_2": "int64", "col_3": "float64"}
     _a						=		JsonDatasetReader(__lowerCamelCase      ,  cache_dir=__lowerCamelCase				).read()
     _check_json_datasetdict(__lowerCamelCase      ,  __lowerCamelCase      ,  splits=list(path.keys()				)				)
     assert all(dataset[split].split == split for split in path.keys()				)






def __UpperCamelCase   (			__lowerCamelCase		: Dict				)      ->     Tuple:







     '''simple docstring'''






     return json.load(__lowerCamelCase				)






def __UpperCamelCase   (			__lowerCamelCase		: Tuple				)      ->     Union[str, Any]:







     '''simple docstring'''






     return [json.loads(__lowerCamelCase				) for line in buffer]






class 						__SCREAMING_SNAKE_CASE					:






     @pytest.mark.parametrize("lines, load_json_function"     ,							[(True, load_json_lines), (False, load_json)]		)
     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    Optional[Any]:
          with io.BytesIO() as buffer:
               JsonDatasetWriter(__UpperCamelCase     ,							__UpperCamelCase     ,							lines=__UpperCamelCase		).write()
               buffer.seek(0		)
               _a						=		load_json_function(__UpperCamelCase		)
          assert isinstance(__UpperCamelCase     ,							__UpperCamelCase		)
          assert isinstance(exported_content[0]     ,							__UpperCamelCase		)
          assert len(__UpperCamelCase		) == 10






     @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  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    List[str]:
          with io.BytesIO() as buffer:
               JsonDatasetWriter(__UpperCamelCase     ,							__UpperCamelCase     ,							lines=__UpperCamelCase     ,							orient=__UpperCamelCase		).write()
               buffer.seek(0		)
               _a						=		load_json(__UpperCamelCase		)
          assert isinstance(__UpperCamelCase     ,							__UpperCamelCase		)
          if keys:
               if container is dict:
                    assert exported_content.keys() == keys
               else:
                    assert exported_content[0].keys() == keys
          else:
               assert not hasattr(__UpperCamelCase     ,							"keys"		) and not hasattr(exported_content[0]     ,							"keys"		)
          if len_at:
               assert len(exported_content[len_at]		) == 10
          else:
               assert len(__UpperCamelCase		) == 10






     @pytest.mark.parametrize("lines, load_json_function"     ,							[(True, load_json_lines), (False, load_json)]		)
     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    Optional[Any]:
          with io.BytesIO() as buffer:
               JsonDatasetWriter(__UpperCamelCase     ,							__UpperCamelCase     ,							lines=__UpperCamelCase     ,							num_proc=2		).write()
               buffer.seek(0		)
               _a						=		load_json_function(__UpperCamelCase		)
          assert isinstance(__UpperCamelCase     ,							__UpperCamelCase		)
          assert isinstance(exported_content[0]     ,							__UpperCamelCase		)
          assert len(__UpperCamelCase		) == 10






     @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  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    Tuple:
          with io.BytesIO() as buffer:
               JsonDatasetWriter(__UpperCamelCase     ,							__UpperCamelCase     ,							lines=__UpperCamelCase     ,							orient=__UpperCamelCase     ,							num_proc=2		).write()
               buffer.seek(0		)
               _a						=		load_json(__UpperCamelCase		)
          assert isinstance(__UpperCamelCase     ,							__UpperCamelCase		)
          if keys:
               if container is dict:
                    assert exported_content.keys() == keys
               else:
                    assert exported_content[0].keys() == keys
          else:
               assert not hasattr(__UpperCamelCase     ,							"keys"		) and not hasattr(exported_content[0]     ,							"keys"		)
          if len_at:
               assert len(exported_content[len_at]		) == 10
          else:
               assert len(__UpperCamelCase		) == 10






     def  a_					(			self     ,							__UpperCamelCase		)       ->    List[Any]:
          with pytest.raises(__UpperCamelCase		):
               with io.BytesIO() as buffer:
                    JsonDatasetWriter(__UpperCamelCase     ,							__UpperCamelCase     ,							num_proc=0		)






     @pytest.mark.parametrize("compression, extension"     ,							[("gzip", "gz"), ("bz2", "bz2"), ("xz", "xz")]		)
     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    Optional[int]:
          _a						=		tmp_path_factory.mktemp("data"		) / f"test.json.{extension}"
          _a						=		str(shared_datadir / f"test_file.json.{extension}"		)
          JsonDatasetWriter(__UpperCamelCase     ,							__UpperCamelCase     ,							compression=__UpperCamelCase		).write()

          with fsspec.open(__UpperCamelCase     ,							"rb"     ,							compression="infer"		) as f:
               _a						=		f.read()
          with fsspec.open(__UpperCamelCase     ,							"rb"     ,							compression="infer"		) as f:
               _a						=		f.read()
          assert exported_content == original_content





'''simple docstring'''

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__				     =		data_utils.TransfoXLTokenizer
lowercase__				     =		data_utils.TransfoXLCorpus
lowercase__				     =		data_utils
lowercase__				     =		data_utils






def __UpperCamelCase   (			__lowerCamelCase		: str      ,  __lowerCamelCase		: Tuple      ,  __lowerCamelCase		: Optional[int]      ,  __lowerCamelCase		: Tuple				)      ->     Dict:







     '''simple docstring'''






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

          _a						=		corpus.__dict__
          corpus_dict_no_vocab.pop("vocab"      ,  __lowerCamelCase				)
          _a						=		pytorch_dump_folder_path + "/" + CORPUS_NAME
          print(F"Save dataset to {pytorch_dataset_dump_path}"				)
          torch.save(__lowerCamelCase      ,  __lowerCamelCase				)

     if tf_checkpoint_path:
          # Convert a pre-trained TensorFlow model
          _a						=		os.path.abspath(__lowerCamelCase				)
          _a						=		os.path.abspath(__lowerCamelCase				)

          print(F"Converting Transformer XL checkpoint from {tf_path} with config at {config_path}."				)
          # Initialise PyTorch model
          if transfo_xl_config_file == "":
               _a						=		TransfoXLConfig()
          else:
               _a						=		TransfoXLConfig.from_json_file(__lowerCamelCase				)
          print(F"Building PyTorch model from configuration: {config}"				)
          _a						=		TransfoXLLMHeadModel(__lowerCamelCase				)

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


if __name__ == "__main__":
   lowercase__				     =		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__				     =		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,
   )

