

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


lowerCamelCase 	=   datasets.utils.logging.get_logger(__name__)







@dataclass
class 		__magic_name__     (				datasets.BuilderConfig				):

   '''simple docstring'''




   lowerCamelCase__					:      Optional[datasets.Features]        =  None
   lowerCamelCase__					:      str        =  "utf-8"
   lowerCamelCase__					:      Optional[str]        =  None
   lowerCamelCase__					:      int        =  10 << 20  # 10MB
   lowerCamelCase__					:      bool        =  False
   lowerCamelCase__					:      str        =  "line"







class 		__magic_name__     (				datasets.ArrowBasedBuilder				):

   '''simple docstring'''




   lowerCamelCase__					:      Union[str, Any]        =  TextConfig



   def 					_UpperCAmelCase					(							self			)				->       Dict:






     """simple docstring"""






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



   def 					_UpperCAmelCase					(							self,       lowercase_			)				->       Any:






     """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 					_UpperCAmelCase					(							self,       lowercase_			)				->       pa.Table:






     """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 					_UpperCAmelCase					(							self,       lowercase_			)				->       Dict:






     """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_			)
