



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

import pyarrow as pa

import datasets
from datasets.features.features import require_storage_cast
from datasets.table import table_cast


__snake_case		:Dict 						=	datasets.utils.logging.get_logger(__name__)

@dataclass
class        _A			(      datasets.BuilderConfig       ):
       UpperCamelCase__    :       Optional[datasets.Features]		   =		None
       UpperCamelCase__    :       str		   =		"utf-8"
       UpperCamelCase__    :       Optional[str]		   =		None
       UpperCamelCase__    :       int		   =		10 << 20  # 10MB
       UpperCamelCase__    :       bool		   =		False
       UpperCamelCase__    :       str		   =		"line"


class        _A			(      datasets.ArrowBasedBuilder       ):
       UpperCamelCase__    :       Union[str, Any]		   =		TextConfig



       def 						_lowerCamelCase     (						self				:	int):


           '''simple docstring'''

           return datasets.DatasetInfo(features=self.config.features)



       def 						_lowerCamelCase     (						self				:	List[Any]	,						__SCREAMING_SNAKE_CASE				:	List[str]):


           '''simple docstring'''

           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}')
           __a					=							dl_manager.download_and_extract(self.config.data_files)
           if isinstance(lowercase_	,						(str, list, tuple)):
               __a					=							data_files
               if isinstance(lowercase_	,						lowercase_):
                   __a					=							[files]
               __a					=							[dl_manager.iter_files(lowercase_) for file in files]
               return [datasets.SplitGenerator(name=datasets.Split.TRAIN	,						gen_kwargs={'''files''': files})]
           __a					=							[]
           for split_name, files in data_files.items():
               if isinstance(lowercase_	,						lowercase_):
                   __a					=							[files]
               __a					=							[dl_manager.iter_files(lowercase_) for file in files]
               splits.append(datasets.SplitGenerator(name=lowercase_	,						gen_kwargs={'''files''': files}))
           return splits



       def 						_lowerCamelCase     (						self				:	Tuple	,						__SCREAMING_SNAKE_CASE				:	Dict):


           '''simple docstring'''

           if self.config.features is not None:
               __a					=							self.config.features.arrow_schema
               if all(not require_storage_cast(lowercase_) for feature in self.config.features.values()):
                   # cheaper cast
                   __a					=							pa_table.cast(lowercase_)
               else:
                   # more expensive cast; allows str <-> int/float or str to Audio for example
                   __a					=							table_cast(lowercase_	,						lowercase_)
               return pa_table
           else:
               return pa_table.cast(pa.schema({'''text''': pa.string()}))





       def 						_lowerCamelCase     (						self				:	int	,						__SCREAMING_SNAKE_CASE				:	Optional[int]):


           '''simple docstring'''

           __a					=							list(self.config.features) if self.config.features is not None else ['''text''']
           for file_idx, file in enumerate(itertools.chain.from_iterable(lowercase_)):
               # open in text mode, by default translates universal newlines ("\n", "\r\n" and "\r") into "\n"
               with open(lowercase_	,						encoding=self.config.encoding	,						errors=self.config.errors) as f:
                   if self.config.sample_by == "line":
                       __a					=							0
                       while True:
                           __a					=							f.read(self.config.chunksize)
                           if not batch:
                               break
                           batch += f.readline()  # finish current line
                           # StringIO.readlines, by default splits only on "\n" (and keeps line breaks)
                           __a					=							StringIO(lowercase_).readlines()
                           if not self.config.keep_linebreaks:
                               __a					=							[line.rstrip('''\n''') for line in batch]
                           __a					=							pa.Table.from_arrays([pa.array(lowercase_)]	,						names=lowercase_)
                           # 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(lowercase_)
                           batch_idx += 1
                   elif self.config.sample_by == "paragraph":
                       __a					=							0
                       __a					=							''''''
                       while True:
                           __a					=							f.read(self.config.chunksize)
                           if not new_batch:
                               break
                           batch += new_batch
                           batch += f.readline()  # finish current line
                           __a					=							batch.split('''\n\n''')
                           __a					=							pa.Table.from_arrays(
                               [pa.array([example for example in batch[:-1] if example])]	,						names=lowercase_)
                           # 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(lowercase_)
                           batch_idx += 1
                           __a					=							batch[-1]
                       if batch:
                           __a					=							pa.Table.from_arrays([pa.array([batch])]	,						names=lowercase_)
                           yield (file_idx, batch_idx), self._cast_table(lowercase_)
                   elif self.config.sample_by == "document":
                       __a					=							f.read()
                       __a					=							pa.Table.from_arrays([pa.array([text])]	,						names=lowercase_)
                       yield file_idx, self._cast_table(lowercase_)








import unittest

import numpy as np

from transformers import ElectraConfig, 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():
       from transformers.models.electra.modeling_flax_electra import (
           FlaxElectraForCausalLM,
           FlaxElectraForMaskedLM,
           FlaxElectraForMultipleChoice,
           FlaxElectraForPreTraining,
           FlaxElectraForQuestionAnswering,
           FlaxElectraForSequenceClassification,
           FlaxElectraForTokenClassification,
           FlaxElectraModel,
       )







class 		__magic_name__     (				unittest.TestCase				):

   '''simple docstring'''




   def __init__(							self,       lowercase_,       lowercase_=13,       lowercase_=7,       lowercase_=True,       lowercase_=True,       lowercase_=True,       lowercase_=True,       lowercase_=99,       lowercase_=24,       lowercase_=32,       lowercase_=5,       lowercase_=4,       lowercase_=37,       lowercase_="gelu",       lowercase_=0.1,       lowercase_=0.1,       lowercase_=512,       lowercase_=16,       lowercase_=2,       lowercase_=0.02,       lowercase_=4,       )				->       int:






     """simple docstring"""






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



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






     """simple docstring"""






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

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

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

     a__    =ElectraConfig(
         vocab_size=self.vocab_size,       hidden_size=self.hidden_size,       embedding_size=self.embedding_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,       initializer_range=self.initializer_range,       )

     return config, input_ids, token_type_ids, attention_mask



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






     """simple docstring"""






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







@require_flax
class 		__magic_name__     (				lowerCamelCase__						,			unittest.TestCase				):

   '''simple docstring'''




   lowerCamelCase__					:      List[str]        =  True

   lowerCamelCase__					:      Tuple        =  (
       (
           FlaxElectraModel,
           FlaxElectraForCausalLM,
           FlaxElectraForMaskedLM,
           FlaxElectraForPreTraining,
           FlaxElectraForTokenClassification,
           FlaxElectraForQuestionAnswering,
           FlaxElectraForMultipleChoice,
           FlaxElectraForSequenceClassification,
       )
       if is_flax_available()
       else ()
   )



   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     a__    =FlaxElectraModelTester(self			)



   @slow
   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     for model_class_name in self.all_model_classes:
       if model_class_name == FlaxElectraForMaskedLM:
         a__    =model_class_name.from_pretrained('''google/electra-small-generator'''			)
       else:
         a__    =model_class_name.from_pretrained('''google/electra-small-discriminator'''			)
       a__    =model(np.ones((1, 1)			)			)
       self.assertIsNotNone(lowercase_			)
