





'''simple docstring'''



import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
	import jax

	from transformers.models.big_bird.modeling_flax_big_bird import (
	    FlaxBigBirdForCausalLM,
	    FlaxBigBirdForMaskedLM,
	    FlaxBigBirdForMultipleChoice,
	    FlaxBigBirdForPreTraining,
	    FlaxBigBirdForQuestionAnswering,
	    FlaxBigBirdForSequenceClassification,
	    FlaxBigBirdForTokenClassification,
	    FlaxBigBirdModel,
	)






class 							__lowercase       (      unittest.TestCase						):







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase     ,_lowerCamelCase=2     ,_lowerCamelCase=56     ,_lowerCamelCase=True     ,_lowerCamelCase=True     ,_lowerCamelCase=True     ,_lowerCamelCase=True     ,_lowerCamelCase=99     ,_lowerCamelCase=32     ,_lowerCamelCase=2     ,_lowerCamelCase=2     ,_lowerCamelCase=7     ,_lowerCamelCase="gelu_new"     ,_lowerCamelCase=0.1     ,_lowerCamelCase=0.1     ,_lowerCamelCase=512     ,_lowerCamelCase=16     ,_lowerCamelCase=2     ,_lowerCamelCase=0.0_2     ,_lowerCamelCase=4     ,_lowerCamelCase="block_sparse"     ,_lowerCamelCase=True     ,_lowerCamelCase=False     ,_lowerCamelCase=2     ,_lowerCamelCase=3     ,)							->							Tuple:







				'''simple docstring'''





				__lowercase										=							parent
				__lowercase										=							batch_size
				__lowercase										=							seq_length
				__lowercase										=							is_training
				__lowercase										=							use_attention_mask
				__lowercase										=							use_token_type_ids
				__lowercase										=							use_labels
				__lowercase										=							vocab_size
				__lowercase										=							hidden_size
				__lowercase										=							num_hidden_layers
				__lowercase										=							num_attention_heads
				__lowercase										=							intermediate_size
				__lowercase										=							hidden_act
				__lowercase										=							hidden_dropout_prob
				__lowercase										=							attention_probs_dropout_prob
				__lowercase										=							max_position_embeddings
				__lowercase										=							type_vocab_size
				__lowercase										=							type_sequence_label_size
				__lowercase										=							initializer_range
				__lowercase										=							num_choices

				__lowercase										=							rescale_embeddings
				__lowercase										=							attention_type
				__lowercase										=							use_bias
				__lowercase										=							block_size
				__lowercase										=							num_random_blocks


	def 						_UpperCAmelCase     (self )							->							List[str]:







				'''simple docstring'''





				__lowercase										=							ids_tensor([self.batch_size, self.seq_length]     ,self.vocab_size )

				__lowercase										=							None
				if self.use_attention_mask:
							__lowercase										=							random_attention_mask([self.batch_size, self.seq_length] )

				__lowercase										=							None
				if self.use_token_type_ids:
							__lowercase										=							ids_tensor([self.batch_size, self.seq_length]     ,self.type_vocab_size )

				__lowercase										=							BigBirdConfig(
				    vocab_size=self.vocab_size     ,hidden_size=self.hidden_size     ,num_hidden_layers=self.num_hidden_layers     ,num_attention_heads=self.num_attention_heads     ,intermediate_size=self.intermediate_size     ,hidden_act=self.hidden_act     ,hidden_dropout_prob=self.hidden_dropout_prob     ,attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,max_position_embeddings=self.max_position_embeddings     ,type_vocab_size=self.type_vocab_size     ,is_decoder=_lowerCamelCase     ,initializer_range=self.initializer_range     ,attention_type=self.attention_type     ,block_size=self.block_size     ,num_random_blocks=self.num_random_blocks     ,use_bias=self.use_bias     ,rescale_embeddings=self.rescale_embeddings     ,)

				return config, input_ids, token_type_ids, attention_mask


	def 						_UpperCAmelCase     (self )							->							List[str]:







				'''simple docstring'''





				__lowercase										=							self.prepare_config_and_inputs()
				__lowercase						,		__lowercase						,		__lowercase						,		__lowercase										=							config_and_inputs
				__lowercase										=							{
				    '''input_ids''': input_ids,
				    '''token_type_ids''': token_type_ids,
				    '''attention_mask''': attention_mask,
				}
				return config, inputs_dict






@require_flax
class 							__lowercase       (      lowerCAmelCase__ ,      unittest.TestCase						):







	'''simple docstring'''

	a    :				List[str]											=       (
	    (
	        FlaxBigBirdForCausalLM,
	        FlaxBigBirdModel,
	        FlaxBigBirdForPreTraining,
	        FlaxBigBirdForMaskedLM,
	        FlaxBigBirdForMultipleChoice,
	        FlaxBigBirdForQuestionAnswering,
	        FlaxBigBirdForSequenceClassification,
	        FlaxBigBirdForTokenClassification,
	    )
	    if is_flax_available()
	    else ()
	)

	a    :				int											=       False
	a    :				Union[str, Any]											=       False


	def 						_UpperCAmelCase     (self )							->							Tuple:







				'''simple docstring'''





				__lowercase										=							FlaxBigBirdModelTester(self )


	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def 						_UpperCAmelCase     (self )							->							Optional[int]:







				'''simple docstring'''





				super().test_from_pretrained_save_pretrained()


	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def 						_UpperCAmelCase     (self )							->							str:







				'''simple docstring'''





				super().test_from_pretrained_with_no_automatic_init()


	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def 						_UpperCAmelCase     (self )							->							Optional[Any]:







				'''simple docstring'''





				super().test_no_automatic_init()


	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def 						_UpperCAmelCase     (self )							->							Dict:







				'''simple docstring'''





				super().test_hidden_states_output()


	@slow
	def 						_UpperCAmelCase     (self )							->							int:







				'''simple docstring'''





				for model_class_name in self.all_model_classes:
							__lowercase										=							model_class_name.from_pretrained('''google/bigbird-roberta-base''' )
							self.assertIsNotNone(_lowerCamelCase )


	def 						_UpperCAmelCase     (self )							->							List[str]:







				'''simple docstring'''





				if self.test_attn_probs:
							super().test_attention_outputs()


	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def 						_UpperCAmelCase     (self )							->							Optional[Any]:







				'''simple docstring'''





				__lowercase						,		__lowercase										=							self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
							with self.subTest(model_class.__name__ ):
										__lowercase										=							self._prepare_for_class(_lowerCamelCase     ,_lowerCamelCase )
										__lowercase										=							model_class(_lowerCamelCase )

										@jax.jit
										def model_jitted(_lowerCamelCase     ,_lowerCamelCase=None     ,**_lowerCamelCase ):
													return model(input_ids=_lowerCamelCase     ,attention_mask=_lowerCamelCase     ,**_lowerCamelCase )

										with self.subTest('''JIT Enabled''' ):
													__lowercase										=							model_jitted(**_lowerCamelCase ).to_tuple()

										with self.subTest('''JIT Disabled''' ):
													with jax.disable_jit():
																__lowercase										=							model_jitted(**_lowerCamelCase ).to_tuple()

										self.assertEqual(len(_lowerCamelCase )     ,len(_lowerCamelCase ) )
										for jitted_output, output in zip(_lowerCamelCase     ,_lowerCamelCase ):
													self.assertEqual(jitted_output.shape     ,output.shape )


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase=1E-5     ,_lowerCamelCase="outputs"     ,_lowerCamelCase=None )							->							Optional[int]:







				'''simple docstring'''





				if name.startswith('''outputs.attentions''' ):
							return
				else:
							super().check_pt_flax_outputs(_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase )












'''simple docstring'''



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

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


_SCREAMING_SNAKE_CASE					=    '''<<<<<<< This should probably be modified because it mentions: '''

_SCREAMING_SNAKE_CASE					=    '''=======
>>>>>>>
'''

_SCREAMING_SNAKE_CASE					=    [
    '''TextEncoderConfig''',
    '''ByteTextEncoder''',
    '''SubwordTextEncoder''',
    '''encoder_config''',
    '''maybe_build_from_corpus''',
    '''manual_dir''',
]

_SCREAMING_SNAKE_CASE					=    [
    # (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    _lowerCAmelCase						(				lowerCamelCase_			:  Namespace				):
			return ConvertCommand(args.tfds_path		,       args.datasets_directory				)






class 							__lowercase       (      lowerCAmelCase__						):







	'''simple docstring'''

	@staticmethod
	def 						_UpperCAmelCase     (_lowerCamelCase )							->							Any:







				'''simple docstring'''





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


	def __init__(self     ,_lowerCamelCase     ,_lowerCamelCase     ,*_lowerCamelCase )							->							List[Any]:







				'''simple docstring'''





				__lowercase										=							get_logger('''datasets-cli/converting''' )

				__lowercase										=							tfds_path
				__lowercase										=							datasets_directory


	def 						_UpperCAmelCase     (self )							->							str:







				'''simple docstring'''





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

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

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

				__lowercase										=							[]
				__lowercase										=							[]
				__lowercase										=							{}

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

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

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

							with open(_lowerCamelCase     ,encoding='''utf-8''' ) as f:
										__lowercase										=							f.readlines()

							__lowercase										=							[]
							__lowercase										=							False
							__lowercase										=							False
							__lowercase										=							[]
							for line in lines:
										__lowercase										=							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:
													__lowercase										=							'''import datasets\n'''
										elif "import tensorflow" in out_line:
													# order is important here
													__lowercase										=							''''''
													continue
										elif "from absl import logging" in out_line:
													__lowercase										=							'''from datasets import logging\n'''
										elif "getLogger" in out_line:
													__lowercase										=							out_line.replace('''getLogger'''     ,'''get_logger''' )
										elif any(expression in out_line for expression in TO_HIGHLIGHT ):
													__lowercase										=							True
													__lowercase										=							list(filter(lambda _lowerCamelCase : e in out_line     ,_lowerCamelCase ) )
													out_lines.append(HIGHLIGHT_MESSAGE_PRE + str(_lowerCamelCase ) + '''\n''' )
													out_lines.append(_lowerCamelCase )
													out_lines.append(_lowerCamelCase )
													continue
										else:
													for pattern, replacement in TO_CONVERT:
																__lowercase										=							re.sub(_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase )

                # Take care of saving utilities (to later move them together with main script)
										if "tensorflow_datasets" in out_line:
													__lowercase										=							re.match(R'''from\stensorflow_datasets.*import\s([^\.\r\n]+)'''     ,_lowerCamelCase )
													tfds_imports.extend(imp.strip() for imp in match.group(1 ).split(''',''' ) )
													__lowercase										=							'''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:
													__lowercase										=							True
										out_lines.append(_lowerCamelCase )

							if is_builder or "wmt" in f_name:
										# We create a new directory for each dataset
										__lowercase										=							f_name.replace('''.py'''     ,'''''' )
										__lowercase										=							os.path.join(_lowerCamelCase     ,_lowerCamelCase )
										__lowercase										=							os.path.join(_lowerCamelCase     ,_lowerCamelCase )
										os.makedirs(_lowerCamelCase     ,exist_ok=_lowerCamelCase )
										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(_lowerCamelCase )

							if needs_manual_update:
										with_manual_update.append(_lowerCamelCase )

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

				for utils_file in utils_files:
							try:
										__lowercase										=							os.path.basename(_lowerCamelCase )
										__lowercase										=							imports_to_builder_map[f_name.replace('''.py'''     ,'''''' )]
										self._logger.info(f"Moving {dest_folder} to {utils_file}" )
										shutil.copy(_lowerCamelCase     ,_lowerCamelCase )
							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'." )





