

from __future__ import annotations

import random
import unittest

from transformers import TransfoXLConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers import (
           TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
           TFTransfoXLForSequenceClassification,
           TFTransfoXLLMHeadModel,
           TFTransfoXLModel,
       )



class       __snake_case  :

  def __init__(					self						,		_A						,		):
      SCREAMING_SNAKE_CASE_									=						parent
      SCREAMING_SNAKE_CASE_									=						13
      SCREAMING_SNAKE_CASE_									=						7
      SCREAMING_SNAKE_CASE_									=						30
      SCREAMING_SNAKE_CASE_									=						self.seq_length + self.mem_len
      SCREAMING_SNAKE_CASE_									=						15
      SCREAMING_SNAKE_CASE_									=						True
      SCREAMING_SNAKE_CASE_									=						True
      SCREAMING_SNAKE_CASE_									=						99
      SCREAMING_SNAKE_CASE_									=						[10, 50, 80]
      SCREAMING_SNAKE_CASE_									=						32
      SCREAMING_SNAKE_CASE_									=						32
      SCREAMING_SNAKE_CASE_									=						4
      SCREAMING_SNAKE_CASE_									=						8
      SCREAMING_SNAKE_CASE_									=						128
      SCREAMING_SNAKE_CASE_									=						2
      SCREAMING_SNAKE_CASE_									=						2
      SCREAMING_SNAKE_CASE_									=						None
      SCREAMING_SNAKE_CASE_									=						1
      SCREAMING_SNAKE_CASE_									=						0
      SCREAMING_SNAKE_CASE_									=						3
      SCREAMING_SNAKE_CASE_									=						self.vocab_size - 1
      SCREAMING_SNAKE_CASE_									=						0.0_1

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						ids_tensor([self.batch_size, self.seq_length]						,		self.vocab_size)
      SCREAMING_SNAKE_CASE_									=						ids_tensor([self.batch_size, self.seq_length]						,		self.vocab_size)

      SCREAMING_SNAKE_CASE_									=						None
      if self.use_labels:
          SCREAMING_SNAKE_CASE_									=						ids_tensor([self.batch_size, self.seq_length]						,		self.vocab_size)

      SCREAMING_SNAKE_CASE_									=						TransfoXLConfig(
          vocab_size=self.vocab_size						,		mem_len=self.mem_len						,		clamp_len=self.clamp_len						,		cutoffs=self.cutoffs						,		d_model=self.hidden_size						,		d_embed=self.d_embed						,		n_head=self.num_attention_heads						,		d_head=self.d_head						,		d_inner=self.d_inner						,		div_val=self.div_val						,		n_layer=self.num_hidden_layers						,		eos_token_id=self.eos_token_id						,		pad_token_id=self.vocab_size - 1						,		init_range=self.init_range						,		num_labels=self.num_labels						,		)

      return (config, input_ids_a, input_ids_a, lm_labels)

  def 						lowerCAmelCase__     (					self):
      random.seed(self.seed)
      tf.random.set_seed(self.seed)

  def 						lowerCAmelCase__     (					self						,		_A						,		_A						,		_A						,		_A):
      SCREAMING_SNAKE_CASE_									=						TFTransfoXLModel(_A)

      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						model(_A).to_tuple()

      SCREAMING_SNAKE_CASE_									=						{'input_ids': input_ids_a, 'mems': mems_a}

      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						model(_A).to_tuple()

      self.parent.assertEqual(hidden_states_a.shape						,		(self.batch_size, self.seq_length, self.hidden_size))
      self.parent.assertEqual(hidden_states_a.shape						,		(self.batch_size, self.seq_length, self.hidden_size))
      self.parent.assertListEqual(
          [mem.shape for mem in mems_a]						,		[(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers						,		)
      self.parent.assertListEqual(
          [mem.shape for mem in mems_a]						,		[(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers						,		)

  def 						lowerCAmelCase__     (					self						,		_A						,		_A						,		_A						,		_A):
      SCREAMING_SNAKE_CASE_									=						TFTransfoXLLMHeadModel(_A)

      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						model(_A).to_tuple()

      SCREAMING_SNAKE_CASE_									=						{'input_ids': input_ids_a, 'labels': lm_labels}
      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						model(_A).to_tuple()

      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						model([input_ids_a, mems_a]).to_tuple()

      SCREAMING_SNAKE_CASE_									=						{'input_ids': input_ids_a, 'mems': mems_a, 'labels': lm_labels}

      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						model(_A).to_tuple()

      self.parent.assertEqual(lm_logits_a.shape						,		(self.batch_size, self.seq_length, self.vocab_size))
      self.parent.assertListEqual(
          [mem.shape for mem in mems_a]						,		[(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers						,		)

      self.parent.assertEqual(lm_logits_a.shape						,		(self.batch_size, self.seq_length, self.vocab_size))
      self.parent.assertListEqual(
          [mem.shape for mem in mems_a]						,		[(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers						,		)

  def 						lowerCAmelCase__     (					self						,		_A						,		_A						,		_A						,		_A):
      SCREAMING_SNAKE_CASE_									=						TFTransfoXLForSequenceClassification(_A)
      SCREAMING_SNAKE_CASE_									=						model(_A)
      self.parent.assertEqual(result.logits.shape						,		(self.batch_size, self.num_labels))


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.prepare_config_and_inputs()
      ((SCREAMING_SNAKE_CASE_)							,			(SCREAMING_SNAKE_CASE_)							,			(SCREAMING_SNAKE_CASE_)							,			(SCREAMING_SNAKE_CASE_))									=						config_and_inputs
      SCREAMING_SNAKE_CASE_									=						{'input_ids': input_ids_a}
      return config, inputs_dict






@require_tf
class       __snake_case  (    lowerCAmelCase__     , lowerCAmelCase__     , unittest.TestCase       ):
  __lowerCAmelCase     :	List[str]						 =  (
      (TFTransfoXLModel, TFTransfoXLLMHeadModel, TFTransfoXLForSequenceClassification) if is_tf_available() else ()
  )
  __lowerCAmelCase     :	List[str]						 =  () if is_tf_available() else ()
  __lowerCAmelCase     :	List[Any]						 =  (
      {
          'feature-extraction': TFTransfoXLModel,
          'text-classification': TFTransfoXLForSequenceClassification,
          'text-generation': TFTransfoXLLMHeadModel,
          'zero-shot': TFTransfoXLForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  # TODO: add this test when TFTransfoXLLMHead has a linear output layer implemented
  __lowerCAmelCase     :	List[Any]						 =  False
  __lowerCAmelCase     :	List[str]						 =  False
  __lowerCAmelCase     :	Optional[int]						 =  False
  __lowerCAmelCase     :	Dict						 =  False

  def 						lowerCAmelCase__     (					self						,		_A						,		_A						,		_A						,		_A						,		_A):
      if pipeline_test_casse_name == "TextGenerationPipelineTests":
          # Get `ValueError: AttributeError: 'NoneType' object has no attribute 'new_ones'` or `AssertionError`.
          # `TransfoXLConfig` was never used in pipeline tests: cannot create a simple
          # tokenizer.
          return True

      return False

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						TFTransfoXLModelTester(self)
      SCREAMING_SNAKE_CASE_									=						ConfigTester(self						,		config_class=_A						,		d_embed=37)

  def 						lowerCAmelCase__     (					self):
      self.config_tester.run_common_tests()

  def 						lowerCAmelCase__     (					self):
      self.model_tester.set_seed()
      SCREAMING_SNAKE_CASE_									=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_transfo_xl_model(*_A)

  def 						lowerCAmelCase__     (					self):
      self.model_tester.set_seed()
      SCREAMING_SNAKE_CASE_									=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_transfo_xl_lm_head(*_A)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_transfo_xl_for_sequence_classification(*_A)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						self.model_tester.prepare_config_and_inputs_for_common()
      SCREAMING_SNAKE_CASE_									=						[TFTransfoXLForSequenceClassification]

      for model_class in self.all_model_classes:
          SCREAMING_SNAKE_CASE_									=						model_class(_A)
          assert isinstance(model.get_input_embeddings()						,		tf.keras.layers.Layer)
          if model_class in list_other_models_with_output_ebd:
              SCREAMING_SNAKE_CASE_									=						model.get_output_embeddings()
              assert isinstance(_A						,		tf.keras.layers.Layer)
              SCREAMING_SNAKE_CASE_									=						model.get_bias()
              assert name is None
          else:
              SCREAMING_SNAKE_CASE_									=						model.get_output_embeddings()
              assert x is None
              SCREAMING_SNAKE_CASE_									=						model.get_bias()
              assert name is None

  def 						lowerCAmelCase__     (					self):
      # TODO JP: Make TransfoXL XLA compliant
      pass

  @slow
  def 						lowerCAmelCase__     (					self):
      for model_name in TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          SCREAMING_SNAKE_CASE_									=						TFTransfoXLModel.from_pretrained(_A)
          self.assertIsNotNone(_A)


  @unittest.skip(reason='This model doesn\'t play well with fit() due to not returning a single loss.')
  def 						lowerCAmelCase__     (					self):
      pass






@require_tf
class       __snake_case  (    unittest.TestCase       ):

  @unittest.skip('Skip test until #12651 is resolved.')
  @slow
  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						TFTransfoXLLMHeadModel.from_pretrained('transfo-xl-wt103')
      # fmt: off
      SCREAMING_SNAKE_CASE_									=						tf.convert_to_tensor([[33,1297,2,1,1009,4,1109,11739,4762,358,5,25,245,22,1706,17,20098,5,3215,21,37,1110,3,13,1041,4,24,603,490,2,71477,20098,104447,2,20961,1,2604,4,1,329,3,6224,831,16002,2,8,603,78967,29546,23,803,20,25,416,5,8,232,4,277,6,1855,4601,3,29546,54,8,3609,5,57211,49,4,1,277,18,8,1755,15691,3,341,25,416,693,42573,71,17,401,94,31,17919,2,29546,7873,18,1,435,23,11011,755,5,5167,3,7983,98,84,2,29546,3267,8,3609,4,1,4865,1075,2,6087,71,6,346,8,5854,3,29546,824,1400,1868,2,19,160,2,311,8,5496,2,20920,17,25,15097,3,24,24,0]]						,		dtype=tf.intaa)  # noqa: E231
      # fmt: on
      #  In 1991 , the remains of Russian Tsar Nicholas II and his family
      #  ( except for Alexei and Maria ) are discovered .
      #  The voice of Nicholas's young son , Tsarevich Alexei Nikolaevich , narrates the
      #  remainder of the story . 1883 Western Siberia ,
      #  a young Grigori Rasputin is asked by his father and a group of men to perform magic .
      #  Rasputin has a vision and denounces one of the men as a horse thief . Although his
      #  father initially slaps him for making such an accusation , Rasputin watches as the
      #  man is chased outside and beaten . Twenty years later , Rasputin sees a vision of
      #  the Virgin Mary , prompting him to become a priest . Rasputin quickly becomes famous ,
      #  with people , even a bishop , begging for his blessing . <eod> </s> <eos>

      # fmt: off
      SCREAMING_SNAKE_CASE_									=						[33,1297,2,1,1009,4,1109,11739,4762,358,5,25,245,22,1706,17,20098,5,3215,21,37,1110,3,13,1041,4,24,603,490,2,71477,20098,104447,2,20961,1,2604,4,1,329,3,6224,831,16002,2,8,603,78967,29546,23,803,20,25,416,5,8,232,4,277,6,1855,4601,3,29546,54,8,3609,5,57211,49,4,1,277,18,8,1755,15691,3,341,25,416,693,42573,71,17,401,94,31,17919,2,29546,7873,18,1,435,23,11011,755,5,5167,3,7983,98,84,2,29546,3267,8,3609,4,1,4865,1075,2,6087,71,6,346,8,5854,3,29546,824,1400,1868,2,19,160,2,311,8,5496,2,20920,17,25,15097,3,24,24,0,33,1,1857,2,1,1009,4,1109,11739,4762,358,5,25,245,28,1110,3,13,1041,4,24,603,490,2,71477,20098,104447,2,20961,1,2604,4,1,329,3,0]  # noqa: E231
      # fmt: on
      #  In 1991, the remains of Russian Tsar Nicholas II and his family (
      #  except for Alexei and Maria ) are discovered. The voice of young son,
      #  Tsarevich Alexei Nikolaevich, narrates the remainder of the story.
      #  1883 Western Siberia, a young Grigori Rasputin is asked by his father
      #  and a group of men to perform magic. Rasputin has a vision and
      #  denounces one of the men as a horse thief. Although his father initially
      #  slaps him for making such an accusation, Rasputin watches as the man
      #  is chased outside and beaten. Twenty years later, Rasputin sees a vision
      #  of the Virgin Mary, prompting him to become a priest.
      #  Rasputin quickly becomes famous, with people, even a bishop, begging for
      #  his blessing. <unk> <unk> <eos> In the 1990s, the remains of Russian Tsar
      # Nicholas II and his family were discovered. The voice of <unk> young son,
      # Tsarevich Alexei Nikolaevich, narrates the remainder of the story.<eos>

      SCREAMING_SNAKE_CASE_									=						model.generate(_A						,		max_length=200						,		do_sample=_A)
      self.assertListEqual(output_ids[0].numpy().tolist()						,		_A)





import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


UpperCamelCase__  :      Optional[int]			=    datasets.utils.logging.get_logger(__name__)



@dataclass
class       __snake_case  (    datasets.BuilderConfig       ):
  __lowerCAmelCase     :	Optional[datasets.Features]						 =  None
  __lowerCAmelCase     :	str						 =  "utf-8"
  __lowerCAmelCase     :	Optional[str]						 =  None
  __lowerCAmelCase     :	Optional[str]						 =  None
  __lowerCAmelCase     :	bool						 =  True  # deprecated
  __lowerCAmelCase     :	Optional[int]						 =  None  # deprecated
  __lowerCAmelCase     :	int						 =  10 << 20  # 10MB
  __lowerCAmelCase     :	Optional[bool]						 =  None






class       __snake_case  (    datasets.ArrowBasedBuilder       ):
  __lowerCAmelCase     :	int						 =  JsonConfig

  def 						lowerCAmelCase__     (					self):
      if self.config.block_size is not None:
          logger.warning('The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead')
          SCREAMING_SNAKE_CASE_									=						self.config.block_size
      if self.config.use_threads is not True:
          logger.warning(
              'The JSON loader parameter `use_threads` is deprecated and doesn\'t have any effect anymore.')
      if self.config.newlines_in_values is not None:
          raise ValueError('The JSON loader parameter `newlines_in_values` is no longer supported')
      return datasets.DatasetInfo(features=self.config.features)

  def 						lowerCAmelCase__     (					self						,		_A):
      if not self.config.data_files:
          raise ValueError(f"""At least one data file must be specified, but got data_files={self.config.data_files}""")
      SCREAMING_SNAKE_CASE_									=						dl_manager.download_and_extract(self.config.data_files)
      if isinstance(_A						,		(str, list, tuple)):
          SCREAMING_SNAKE_CASE_									=						data_files
          if isinstance(_A						,		_A):
              SCREAMING_SNAKE_CASE_									=						[files]
          SCREAMING_SNAKE_CASE_									=						[dl_manager.iter_files(_A) for file in files]
          return [datasets.SplitGenerator(name=datasets.Split.TRAIN						,		gen_kwargs={'files': files})]
      SCREAMING_SNAKE_CASE_									=						[]
      for split_name, files in data_files.items():
          if isinstance(_A						,		_A):
              SCREAMING_SNAKE_CASE_									=						[files]
          SCREAMING_SNAKE_CASE_									=						[dl_manager.iter_files(_A) for file in files]
          splits.append(datasets.SplitGenerator(name=_A						,		gen_kwargs={'files': files}))
      return splits

  def 						lowerCAmelCase__     (					self						,		_A):
      if self.config.features is not None:
          # adding missing columns
          for column_name in set(self.config.features) - set(pa_table.column_names):
              SCREAMING_SNAKE_CASE_									=						self.config.features.arrow_schema.field(_A).type
              SCREAMING_SNAKE_CASE_									=						pa_table.append_column(_A						,		pa.array([None] * len(_A)						,		type=_A))
          # more expensive cast to support nested structures with keys in a different order
          # allows str <-> int/float or str to Audio for example
          SCREAMING_SNAKE_CASE_									=						table_cast(_A						,		self.config.features.arrow_schema)
      return pa_table


  def 						lowerCAmelCase__     (					self						,		_A):
      for file_idx, file in enumerate(itertools.chain.from_iterable(_A)):
          # If the file is one json object and if we need to look at the list of items in one specific field
          if self.config.field is not None:
              with open(_A						,		encoding=self.config.encoding						,		errors=self.config.encoding_errors) as f:
                  SCREAMING_SNAKE_CASE_									=						json.load(_A)

              # We keep only the field we are interested in
              SCREAMING_SNAKE_CASE_									=						dataset[self.config.field]

              # We accept two format: a list of dicts or a dict of lists
              if isinstance(_A						,		(list, tuple)):
                  SCREAMING_SNAKE_CASE_									=						set().union(*[row.keys() for row in dataset])
                  SCREAMING_SNAKE_CASE_									=						{col: [row.get(_A) for row in dataset] for col in keys}
              else:
                  SCREAMING_SNAKE_CASE_									=						dataset
              SCREAMING_SNAKE_CASE_									=						pa.Table.from_pydict(_A)
              yield file_idx, self._cast_table(_A)

          # If the file has one json object per line
          else:
              with open(_A						,		'rb') as f:
                  SCREAMING_SNAKE_CASE_									=						0
                  # Use block_size equal to the chunk size divided by 32 to leverage multithreading
                  # Set a default minimum value of 16kB if the chunk size is really small
                  SCREAMING_SNAKE_CASE_									=						max(self.config.chunksize // 32						,		16 << 10)
                  SCREAMING_SNAKE_CASE_									=						(
                      self.config.encoding_errors if self.config.encoding_errors is not None else 'strict'
                  )
                  while True:
                      SCREAMING_SNAKE_CASE_									=						f.read(self.config.chunksize)
                      if not batch:
                          break
                      # Finish current line
                      try:
                          batch += f.readline()
                      except (AttributeError, io.UnsupportedOperation):
                          batch += readline(_A)
                      # PyArrow only accepts utf-8 encoded bytes
                      if self.config.encoding != "utf-8":
                          SCREAMING_SNAKE_CASE_									=						batch.decode(self.config.encoding						,		errors=_A).encode('utf-8')
                      try:
                          while True:
                              try:
                                  SCREAMING_SNAKE_CASE_									=						paj.read_json(
                                      io.BytesIO(_A)						,		read_options=paj.ReadOptions(block_size=_A))
                                  break
                              except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
                                  if (
                                      isinstance(_A						,		pa.ArrowInvalid)
                                      and "straddling" not in str(_A)
                                      or block_size > len(_A)
                                  ):
                                      raise
                                  else:
                                      # Increase the block size in case it was too small.
                                      # The block size will be reset for the next file.
                                      logger.debug(
                                          f"""Batch of {len(_A)} bytes couldn't be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}.""")
                                      block_size *= 2
                      except pa.ArrowInvalid as e:
                          try:
                              with open(
                                  _A						,		encoding=self.config.encoding						,		errors=self.config.encoding_errors) as f:
                                  SCREAMING_SNAKE_CASE_									=						json.load(_A)
                          except json.JSONDecodeError:
                              logger.error(f"""Failed to read file '{file}' with error {type(_A)}: {e}""")
                              raise e
                          # If possible, parse the file as a list of json objects and exit the loop
                          if isinstance(_A						,		_A):  # list is the only sequence type supported in JSON
                              try:
                                  SCREAMING_SNAKE_CASE_									=						set().union(*[row.keys() for row in dataset])
                                  SCREAMING_SNAKE_CASE_									=						{col: [row.get(_A) for row in dataset] for col in keys}
                                  SCREAMING_SNAKE_CASE_									=						pa.Table.from_pydict(_A)
                              except (pa.ArrowInvalid, AttributeError) as e:
                                  logger.error(f"""Failed to read file '{file}' with error {type(_A)}: {e}""")
                                  raise ValueError(f"""Not able to read records in the JSON file at {file}.""") from None
                              yield file_idx, self._cast_table(_A)
                              break
                          else:
                              logger.error(f"""Failed to read file '{file}' with error {type(_A)}: {e}""")
                              raise ValueError(
                                  f"""Not able to read records in the JSON file at {file}. """
                                  f"""You should probably indicate the field of the JSON file containing your records. """
                                  f"""This JSON file contain the following fields: {str(list(dataset.keys()))}. """
                                  f"""Select the correct one and provide it as `field='XXX'` to the dataset loading method. """) from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                      yield (file_idx, batch_idx), self._cast_table(_A)
                      batch_idx += 1


