





'''simple docstring'''
import unittest

from transformers import ElectraConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        MODEL_FOR_PRETRAINING_MAPPING,
        ElectraForCausalLM,
        ElectraForMaskedLM,
        ElectraForMultipleChoice,
        ElectraForPreTraining,
        ElectraForQuestionAnswering,
        ElectraForSequenceClassification,
        ElectraForTokenClassification,
        ElectraModel,
    )
    from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST

class 					snake_case					:




 """simple docstring"""



 def __init__(	self	:							List[str]							,     __A	:							List[Any]							,     __A	:							Dict=1_3							,     __A	:							List[str]=7							,     __A	:							Tuple=True							,     __A	:							Dict=True							,     __A	:							List[str]=True							,     __A	:							Dict=True							,     __A	:							Any=9_9							,     __A	:							Optional[int]=3_2							,     __A	:							Union[str, Any]=5							,     __A	:							Optional[Any]=4							,     __A	:							List[Any]=3_7							,     __A	:							List[Any]="gelu"							,     __A	:							int=0.1							,     __A	:							Tuple=0.1							,     __A	:							Tuple=5_1_2							,     __A	:							List[Any]=1_6							,     __A	:							Any=2							,     __A	:							int=0.02							,     __A	:							int=3							,     __A	:							Tuple=4							,     __A	:							Tuple=None							,     ):
        __UpperCamelCase	=  parent
        __UpperCamelCase	=  batch_size
        __UpperCamelCase	=  seq_length
        __UpperCamelCase	=  is_training
        __UpperCamelCase	=  use_input_mask
        __UpperCamelCase	=  use_token_type_ids
        __UpperCamelCase	=  use_labels
        __UpperCamelCase	=  vocab_size
        __UpperCamelCase	=  hidden_size
        __UpperCamelCase	=  num_hidden_layers
        __UpperCamelCase	=  num_attention_heads
        __UpperCamelCase	=  intermediate_size
        __UpperCamelCase	=  hidden_act
        __UpperCamelCase	=  hidden_dropout_prob
        __UpperCamelCase	=  attention_probs_dropout_prob
        __UpperCamelCase	=  max_position_embeddings
        __UpperCamelCase	=  type_vocab_size
        __UpperCamelCase	=  type_sequence_label_size
        __UpperCamelCase	=  initializer_range
        __UpperCamelCase	=  num_labels
        __UpperCamelCase	=  num_choices
        __UpperCamelCase	=  scope






 def 	_lowerCamelCase (	self	:							Union[str, Any]     ):
        __UpperCamelCase	=  ids_tensor([self.batch_size, self.seq_length]							,     self.vocab_size     )

        __UpperCamelCase	=  None
        if self.use_input_mask:
               __UpperCamelCase	=  random_attention_mask([self.batch_size, self.seq_length]     )

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

        __UpperCamelCase	=  None
        __UpperCamelCase	=  None
        __UpperCamelCase	=  None
        if self.use_labels:
               __UpperCamelCase	=  ids_tensor([self.batch_size]							,     self.type_sequence_label_size     )
               __UpperCamelCase	=  ids_tensor([self.batch_size, self.seq_length]							,     self.num_labels     )
               __UpperCamelCase	=  ids_tensor([self.batch_size]							,     self.num_choices     )
               __UpperCamelCase	=  ids_tensor([self.batch_size, self.seq_length]							,     1     )

        __UpperCamelCase	=  self.get_config()

        return (
            config,
            input_ids,
            token_type_ids,
            input_mask,
            sequence_labels,
            token_labels,
            choice_labels,
            fake_token_labels,
        )






 def 	_lowerCamelCase (	self	:							int     ):
        return ElectraConfig(
            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=__A							,     initializer_range=self.initializer_range							,     )






 def 	_lowerCamelCase (	self	:							List[str]     ):
        (
            (
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						
        )	=  self.prepare_config_and_inputs()

        __UpperCamelCase	=  True
        __UpperCamelCase	=  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
        __UpperCamelCase	=  ids_tensor([self.batch_size, self.seq_length]							,     vocab_size=2     )

        return (
            config,
            input_ids,
            token_type_ids,
            input_mask,
            sequence_labels,
            token_labels,
            choice_labels,
            encoder_hidden_states,
            encoder_attention_mask,
        )






 def 	_lowerCamelCase (	self	:							List[str]							,     __A	:							Dict							,     __A	:							Dict							,     __A	:							int							,     __A	:							Optional[int]							,     __A	:							Optional[int]							,     __A	:							List[str]							,     __A	:							Any							,     __A	:							str							,     ):
        __UpperCamelCase	=  ElectraModel(config=__A     )
        model.to(__A     )
        model.eval()
        __UpperCamelCase	=  model(__A							,     attention_mask=__A							,     token_type_ids=__A     )
        __UpperCamelCase	=  model(__A							,     token_type_ids=__A     )
        __UpperCamelCase	=  model(__A     )
        self.parent.assertEqual(result.last_hidden_state.shape							,     (self.batch_size, self.seq_length, self.hidden_size)     )






 def 	_lowerCamelCase (	self	:							List[str]							,     __A	:							Optional[Any]							,     __A	:							List[Any]							,     __A	:							Any							,     __A	:							List[Any]							,     __A	:							Tuple							,     __A	:							Optional[int]							,     __A	:							Optional[int]							,     __A	:							Union[str, Any]							,     __A	:							str							,     ):
        __UpperCamelCase	=  True
        __UpperCamelCase	=  ElectraModel(__A     )
        model.to(__A     )
        model.eval()
        __UpperCamelCase	=  model(
            __A							,     attention_mask=__A							,     token_type_ids=__A							,     encoder_hidden_states=__A							,     encoder_attention_mask=__A							,     )
        __UpperCamelCase	=  model(
            __A							,     attention_mask=__A							,     token_type_ids=__A							,     encoder_hidden_states=__A							,     )
        __UpperCamelCase	=  model(__A							,     attention_mask=__A							,     token_type_ids=__A     )
        self.parent.assertEqual(result.last_hidden_state.shape							,     (self.batch_size, self.seq_length, self.hidden_size)     )






 def 	_lowerCamelCase (	self	:							int							,     __A	:							Optional[int]							,     __A	:							List[Any]							,     __A	:							Dict							,     __A	:							Optional[Any]							,     __A	:							int							,     __A	:							Dict							,     __A	:							List[Any]							,     __A	:							str							,     ):
        __UpperCamelCase	=  ElectraForMaskedLM(config=__A     )
        model.to(__A     )
        model.eval()
        __UpperCamelCase	=  model(__A							,     attention_mask=__A							,     token_type_ids=__A							,     labels=__A     )
        self.parent.assertEqual(result.logits.shape							,     (self.batch_size, self.seq_length, self.vocab_size)     )






 def 	_lowerCamelCase (	self	:							List[str]							,     __A	:							Optional[int]							,     __A	:							Union[str, Any]							,     __A	:							List[str]							,     __A	:							List[Any]							,     __A	:							List[Any]							,     __A	:							Tuple							,     __A	:							List[Any]							,     __A	:							Optional[Any]							,     __A	:							str							,     ):
        __UpperCamelCase	=  ElectraForCausalLM(config=__A     )
        model.to(__A     )
        model.eval()
        __UpperCamelCase	=  model(__A							,     attention_mask=__A							,     token_type_ids=__A							,     labels=__A     )
        self.parent.assertEqual(result.logits.shape							,     (self.batch_size, self.seq_length, self.vocab_size)     )






 def 	_lowerCamelCase (	self	:							Any							,     __A	:							Union[str, Any]							,     __A	:							Tuple							,     __A	:							str							,     __A	:							str							,     __A	:							List[Any]							,     __A	:							Optional[int]							,     __A	:							Tuple							,     __A	:							Optional[Any]							,     ):
        __UpperCamelCase	=  self.num_labels
        __UpperCamelCase	=  ElectraForTokenClassification(config=__A     )
        model.to(__A     )
        model.eval()
        __UpperCamelCase	=  model(__A							,     attention_mask=__A							,     token_type_ids=__A							,     labels=__A     )
        self.parent.assertEqual(result.logits.shape							,     (self.batch_size, self.seq_length, self.num_labels)     )






 def 	_lowerCamelCase (	self	:							Tuple							,     __A	:							Optional[Any]							,     __A	:							Optional[int]							,     __A	:							Tuple							,     __A	:							List[str]							,     __A	:							List[str]							,     __A	:							List[Any]							,     __A	:							Tuple							,     __A	:							int							,     ):
        __UpperCamelCase	=  self.num_labels
        __UpperCamelCase	=  ElectraForPreTraining(config=__A     )
        model.to(__A     )
        model.eval()
        __UpperCamelCase	=  model(__A							,     attention_mask=__A							,     token_type_ids=__A							,     labels=__A     )
        self.parent.assertEqual(result.logits.shape							,     (self.batch_size, self.seq_length)     )






 def 	_lowerCamelCase (	self	:							int							,     __A	:							int							,     __A	:							Union[str, Any]							,     __A	:							List[Any]							,     __A	:							Optional[int]							,     __A	:							int							,     __A	:							Union[str, Any]							,     __A	:							int							,     __A	:							str							,     ):
        __UpperCamelCase	=  self.num_labels
        __UpperCamelCase	=  ElectraForSequenceClassification(__A     )
        model.to(__A     )
        model.eval()
        __UpperCamelCase	=  model(__A							,     attention_mask=__A							,     token_type_ids=__A							,     labels=__A     )
        self.parent.assertEqual(result.logits.shape							,     (self.batch_size, self.num_labels)     )






 def 	_lowerCamelCase (	self	:							Union[str, Any]							,     __A	:							Optional[int]							,     __A	:							Any							,     __A	:							Tuple							,     __A	:							Tuple							,     __A	:							Tuple							,     __A	:							str							,     __A	:							Optional[int]							,     __A	:							Tuple							,     ):
        __UpperCamelCase	=  ElectraForQuestionAnswering(config=__A     )
        model.to(__A     )
        model.eval()
        __UpperCamelCase	=  model(
            __A							,     attention_mask=__A							,     token_type_ids=__A							,     start_positions=__A							,     end_positions=__A							,     )
        self.parent.assertEqual(result.start_logits.shape							,     (self.batch_size, self.seq_length)     )
        self.parent.assertEqual(result.end_logits.shape							,     (self.batch_size, self.seq_length)     )






 def 	_lowerCamelCase (	self	:							Union[str, Any]							,     __A	:							Union[str, Any]							,     __A	:							Optional[int]							,     __A	:							Dict							,     __A	:							Any							,     __A	:							str							,     __A	:							List[Any]							,     __A	:							Any							,     __A	:							List[str]							,     ):
        __UpperCamelCase	=  self.num_choices
        __UpperCamelCase	=  ElectraForMultipleChoice(config=__A     )
        model.to(__A     )
        model.eval()
        __UpperCamelCase	=  input_ids.unsqueeze(1     ).expand(-1							,     self.num_choices							,     -1     ).contiguous()
        __UpperCamelCase	=  token_type_ids.unsqueeze(1     ).expand(-1							,     self.num_choices							,     -1     ).contiguous()
        __UpperCamelCase	=  input_mask.unsqueeze(1     ).expand(-1							,     self.num_choices							,     -1     ).contiguous()
        __UpperCamelCase	=  model(
            __A							,     attention_mask=__A							,     token_type_ids=__A							,     labels=__A							,     )
        self.parent.assertEqual(result.logits.shape							,     (self.batch_size, self.num_choices)     )






 def 	_lowerCamelCase (	self	:							List[Any]     ):
        __UpperCamelCase	=  self.prepare_config_and_inputs()
        (
            (
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						(
            __UpperCamelCase
        )	,						
        )	=  config_and_inputs
        __UpperCamelCase	=  {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
        return config, inputs_dict






@require_torch
class 					snake_case					(   __lowerCamelCase      ,	__lowerCamelCase      ,	unittest.TestCase	):




 """simple docstring"""



 SCREAMING_SNAKE_CASE_     : Optional[int]      		=(
     (
         ElectraModel,
         ElectraForPreTraining,
         ElectraForMaskedLM,
         ElectraForCausalLM,
         ElectraForMultipleChoice,
         ElectraForTokenClassification,
         ElectraForSequenceClassification,
         ElectraForQuestionAnswering,
     )
     if is_torch_available()
     else ()
 )
 SCREAMING_SNAKE_CASE_     : Optional[Any]      		=(
     {
         "feature-extraction": ElectraModel,
         "fill-mask": ElectraForMaskedLM,
         "question-answering": ElectraForQuestionAnswering,
         "text-classification": ElectraForSequenceClassification,
         "text-generation": ElectraForCausalLM,
         "token-classification": ElectraForTokenClassification,
         "zero-shot": ElectraForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 SCREAMING_SNAKE_CASE_     : Optional[Any]      		=True






 def 	_lowerCamelCase (	self	:							Union[str, Any]							,     __A	:							str							,     __A	:							Optional[int]							,     __A	:							int=False     ):
        __UpperCamelCase	=  super()._prepare_for_class(__A							,     __A							,     return_labels=__A     )

        if return_labels:
               if model_class in get_values(__A     ):
                      __UpperCamelCase	=  torch.zeros(
                          (self.model_tester.batch_size, self.model_tester.seq_length)							,     dtype=torch.long							,     device=__A     )
        return inputs_dict






 def 	_lowerCamelCase (	self	:							Dict     ):
        __UpperCamelCase	=  ElectraModelTester(self     )
        __UpperCamelCase	=  ConfigTester(self							,     config_class=__A							,     hidden_size=3_7     )






 def 	_lowerCamelCase (	self	:							Optional[Any]     ):
        self.config_tester.run_common_tests()






 def 	_lowerCamelCase (	self	:							int     ):
        __UpperCamelCase	=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_model(*__A     )






 def 	_lowerCamelCase (	self	:							Optional[Any]     ):
        __UpperCamelCase	=  self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_electra_model_as_decoder(*__A     )






 def 	_lowerCamelCase (	self	:							Optional[int]     ):
        __UpperCamelCase	=  self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
               __UpperCamelCase	=  type
               self.model_tester.create_and_check_electra_model(*__A     )






 def 	_lowerCamelCase (	self	:							Union[str, Any]     ):
        __UpperCamelCase	=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_for_masked_lm(*__A     )






 def 	_lowerCamelCase (	self	:							Union[str, Any]     ):
        __UpperCamelCase	=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_for_token_classification(*__A     )






 def 	_lowerCamelCase (	self	:							str     ):
        __UpperCamelCase	=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_for_pretraining(*__A     )






 def 	_lowerCamelCase (	self	:							Optional[Any]     ):
        __UpperCamelCase	=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_for_sequence_classification(*__A     )






 def 	_lowerCamelCase (	self	:							List[str]     ):
        __UpperCamelCase	=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_for_question_answering(*__A     )






 def 	_lowerCamelCase (	self	:							Union[str, Any]     ):
        __UpperCamelCase	=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_for_multiple_choice(*__A     )






 @slow
 def 	_lowerCamelCase (	self	:							Tuple     ):
        for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __UpperCamelCase	=  ElectraModel.from_pretrained(__A     )
               self.assertIsNotNone(__A     )






 def 	_lowerCamelCase (	self	:							Dict     ):
        __UpperCamelCase	=  self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_electra_for_causal_lm(*__A     )






@require_torch
class 					snake_case					(   unittest.TestCase	):




 """simple docstring"""



 @slow
 def 	_lowerCamelCase (	self	:							str     ):
        __UpperCamelCase	=  ElectraModel.from_pretrained('google/electra-small-discriminator'     )
        __UpperCamelCase	=  torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]     )
        __UpperCamelCase	=  torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]     )
        __UpperCamelCase	=  model(__A							,     attention_mask=__A     )[0]
        __UpperCamelCase	=  torch.Size((1, 1_1, 2_5_6)     )
        self.assertEqual(output.shape							,     __A     )
        __UpperCamelCase	=  torch.tensor(
            [[[0.4471, 0.6821, -0.3265], [0.4627, 0.5255, -0.3668], [0.4532, 0.3313, -0.4344]]]     )

        self.assertTrue(torch.allclose(output[:, 1:4, 1:4]							,     __A							,     atol=1e-4     )     )







'''simple docstring'''
import itertools
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa

import datasets
from datasets.table import table_cast


a__							:							Optional[Any]			=datasets.utils.logging.get_logger(__name__)

@dataclass
class 					snake_case					(   datasets.BuilderConfig	):




 """simple docstring"""



 SCREAMING_SNAKE_CASE_     : Optional[datasets.Features]      		=None






class 					snake_case					(   datasets.ArrowBasedBuilder	):




 """simple docstring"""



 SCREAMING_SNAKE_CASE_     : int      		=ArrowConfig






 def 	_lowerCamelCase (	self	:							List[str]     ):
        return datasets.DatasetInfo(features=self.config.features     )






 def 	_lowerCamelCase (	self	:							Union[str, Any]							,     __A	:							Dict     ):
        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}'''     )
        __UpperCamelCase	=  dl_manager.download_and_extract(self.config.data_files     )
        if isinstance(__A							,     (str, list, tuple)     ):
               __UpperCamelCase	=  data_files
               if isinstance(__A							,     __A     ):
                      __UpperCamelCase	=  [files]
               # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
               __UpperCamelCase	=  [dl_manager.iter_files(__A     ) for file in files]
               return [datasets.SplitGenerator(name=datasets.Split.TRAIN							,     gen_kwargs={'files': files}     )]
        __UpperCamelCase	=  []
        for split_name, files in data_files.items():
               if isinstance(__A							,     __A     ):
                      __UpperCamelCase	=  [files]
               # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
               __UpperCamelCase	=  [dl_manager.iter_files(__A     ) for file in files]
               # Infer features is they are stoed in the arrow schema
               if self.info.features is None:
                      for file in itertools.chain.from_iterable(__A     ):
                             with open(__A							,     'rb'     ) as f:
                                    __UpperCamelCase	=  datasets.Features.from_arrow_schema(pa.ipc.open_stream(__A     ).schema     )
                             break
               splits.append(datasets.SplitGenerator(name=__A							,     gen_kwargs={'files': files}     )     )
        return splits






 def 	_lowerCamelCase (	self	:							Any							,     __A	:							pa.Table     ):
        if self.info.features is not None:
               # more expensive cast to support nested features with keys in a different order
               # allows str <-> int/float or str to Audio for example
               __UpperCamelCase	=  table_cast(__A							,     self.info.features.arrow_schema     )
        return pa_table






 def 	_lowerCamelCase (	self	:							int							,     __A	:							Optional[Any]     ):
        for file_idx, file in enumerate(itertools.chain.from_iterable(__A     )     ):
               with open(__A							,     'rb'     ) as f:
                      try:
                             for batch_idx, record_batch in enumerate(pa.ipc.open_stream(__A     )     ):
                                    __UpperCamelCase	=  pa.Table.from_batches([record_batch]     )
                                    # 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 f'''{file_idx}_{batch_idx}''', self._cast_table(__A     )
                      except ValueError as e:
                             logger.error(f'''Failed to read file \'{file}\' with error {type(__A     )}: {e}'''     )
                             raise
