


import os
from tempfile import TemporaryDirectory
from unittest import TestCase

import pytest
from absl.testing import parameterized

from datasets import config
from datasets.arrow_reader import HF_GCP_BASE_URL
from datasets.builder import DatasetBuilder
from datasets.dataset_dict import IterableDatasetDict
from datasets.iterable_dataset import IterableDataset
from datasets.load import dataset_module_factory, import_main_class
from datasets.utils.file_utils import cached_path


__lowerCAmelCase      :					Dict    							=[
    {'dataset': 'wikipedia', 'config_name': '20220301.de'},
    {'dataset': 'wikipedia', 'config_name': '20220301.en'},
    {'dataset': 'wikipedia', 'config_name': '20220301.fr'},
    {'dataset': 'wikipedia', 'config_name': '20220301.frr'},
    {'dataset': 'wikipedia', 'config_name': '20220301.it'},
    {'dataset': 'wikipedia', 'config_name': '20220301.simple'},
    {'dataset': 'snli', 'config_name': 'plain_text'},
    {'dataset': 'eli5', 'config_name': 'LFQA_reddit'},
    {'dataset': 'wiki40b', 'config_name': 'en'},
    {'dataset': 'wiki_dpr', 'config_name': 'psgs_w100.nq.compressed'},
    {'dataset': 'wiki_dpr', 'config_name': 'psgs_w100.nq.no_index'},
    {'dataset': 'wiki_dpr', 'config_name': 'psgs_w100.multiset.no_index'},
    {'dataset': 'natural_questions', 'config_name': 'default'},
]


def 			_UpperCamelCase				( lowercase__=True ):
						if with_config:
												return [
												    {
												        "testcase_name": d["dataset"] + "/" + d["config_name"],
												        "dataset": d["dataset"],
												        "config_name": d["config_name"],
												    }
												    for d in DATASETS_ON_HF_GCP
												]
						else:
												return [
												    {"testcase_name": dataset, "dataset": dataset} for dataset in {d["dataset"] for d in DATASETS_ON_HF_GCP}
												]
@parameterized.named_parameters(list_datasets_on_hf_gcp_parameters(with_config=__a			)			)
class 			_lowercase (     __a			):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__							:     List[str]         =						None
				SCREAMING_SNAKE_CASE__							:     Union[str, Any]         =						None


				def 						__magic_name__(			self     :List[Any]     ,     lowerCAmelCase__     :int     ,     lowerCAmelCase__     :int      )						->     Optional[Any]:



										with TemporaryDirectory() as tmp_dir:
																__SCREAMING_SNAKE_CASE		:	int						     =						dataset_module_factory(_lowercase     ,     cache_dir=_lowercase      )

																__SCREAMING_SNAKE_CASE		:	Dict						     =						import_main_class(dataset_module.module_path     ,     dataset=_lowercase      )

																__SCREAMING_SNAKE_CASE		:	DatasetBuilder						     =						builder_cls(
																    cache_dir=_lowercase     ,     config_name=_lowercase     ,     hash=dataset_module.hash     ,     )

																__SCREAMING_SNAKE_CASE		:	Tuple						     =						'''/'''.join(
																    [
																        HF_GCP_BASE_URL,
																        builder_instance._relative_data_dir(with_hash=_lowercase      ).replace(os.sep     ,     '''/'''      ),
																        config.DATASET_INFO_FILENAME,
																    ]      )
																__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						cached_path(_lowercase     ,     cache_dir=_lowercase      )
																self.assertTrue(os.path.exists(_lowercase      )      )




@pytest.mark.integration
def 			_UpperCamelCase				( lowercase__ ):
						__SCREAMING_SNAKE_CASE		:	List[str]						     =						tmp_path_factory.mktemp('''test_hf_gcp''' ) / '''test_wikipedia_simple'''
						__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						dataset_module_factory('''wikipedia'''       ,			cache_dir=lowercase__ )
						__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						import_main_class(dataset_module.module_path )
						__SCREAMING_SNAKE_CASE		:	DatasetBuilder						     =						builder_cls(
						    cache_dir=lowercase__       ,			config_name='''20220301.frr'''       ,			hash=dataset_module.hash       ,			)
						# use the HF cloud storage, not the original download_and_prepare that uses apache-beam
						__SCREAMING_SNAKE_CASE		:	int						     =						None
						builder_instance.download_and_prepare()
						__SCREAMING_SNAKE_CASE		:	Dict						     =						builder_instance.as_dataset()
						assert ds


@pytest.mark.integration
def 			_UpperCamelCase				( lowercase__ ):
						__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						dataset_module_factory('''wikipedia'''       ,			cache_dir=lowercase__ )
						__SCREAMING_SNAKE_CASE		:	Any						     =						import_main_class(dataset_module.module_path       ,			dataset=lowercase__ )
						__SCREAMING_SNAKE_CASE		:	DatasetBuilder						     =						builder_cls(
						    cache_dir=lowercase__       ,			config_name='''20220301.frr'''       ,			hash=dataset_module.hash       ,			)
						__SCREAMING_SNAKE_CASE		:	Any						     =						builder_instance.as_streaming_dataset()
						assert ds
						assert isinstance(lowercase__       ,			lowercase__ )
						assert "train" in ds
						assert isinstance(ds['''train''']       ,			lowercase__ )
						assert next(iter(ds['''train'''] ) )

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




def    lowerCAmelCase_  (		lowercase:				Tuple	,       lowercase:				Union[str, Any]	,       lowercase:				Any	,       lowercase:				List[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:
   _UpperCamelCase:  Optional[int]    						=			pickle.load(lowercase	,       encoding='''latin1'''		)
  # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
  _UpperCamelCase:  Tuple    						=			pytorch_dump_folder_path + '''/''' + VOCAB_FILES_NAMES['''pretrained_vocab_file''']
  print(F"""Save vocabulary to {pytorch_vocab_dump_path}"""		)
  _UpperCamelCase:  List[str]    						=			corpus.vocab.__dict__
  torch.save(lowercase	,       lowercase		)

  _UpperCamelCase:  str    						=			corpus.__dict__
  corpus_dict_no_vocab.pop('''vocab'''	,       lowercase		)
  _UpperCamelCase:  str    						=			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
  _UpperCamelCase:  Optional[Any]    						=			os.path.abspath(lowercase		)
  _UpperCamelCase:  List[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 == "":
   _UpperCamelCase:  Tuple    						=			TransfoXLConfig()
  else:
   _UpperCamelCase:  Union[str, Any]    						=			TransfoXLConfig.from_json_file(lowercase		)
  print(F"""Building PyTorch model from configuration: {config}"""		)
  _UpperCamelCase:  Optional[int]    						=			TransfoXLLMHeadModel(lowercase		)

  _UpperCamelCase:  int    						=			load_tf_weights_in_transfo_xl(lowercase	,       lowercase	,       lowercase		)
  # Save pytorch-model
  _UpperCamelCase:  List[str]    						=			os.path.join(lowercase	,       lowercase		)
  _UpperCamelCase:  Union[str, Any]    						=			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__":
  UpperCAmelCase_           =				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.''',
  )
  UpperCAmelCase_           =				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,
  )