


'''simple docstring'''







import unittest
from pathlib import Path
from tempfile import NamedTemporaryFile, TemporaryDirectory

from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline
from transformers.convert_graph_to_onnx import (
    convert,
    ensure_valid_input,
    generate_identified_filename,
    infer_shapes,
    quantize,
)
from transformers.testing_utils import require_tf, require_tokenizers, require_torch, slow


class 			__lowerCamelCase			:
    """simple docstring"""


    def        A   (  self      :							Optional[Any]				,							SCREAMING_SNAKE_CASE      :							Optional[Any]				,							SCREAMING_SNAKE_CASE      :							List[str]				,							SCREAMING_SNAKE_CASE      :							Any):
           return None




class 			__lowerCamelCase			:
    """simple docstring"""


    def        A   (  self      :							Union[str, Any]				,							SCREAMING_SNAKE_CASE      :							str				,							SCREAMING_SNAKE_CASE      :							Any				,							SCREAMING_SNAKE_CASE      :							Any				,							SCREAMING_SNAKE_CASE      :							str):
           return None




class 			__lowerCamelCase			(	unittest.TestCase	):
    """simple docstring"""
    a													=   [
        # (model_name, model_kwargs)
        ("bert-base-cased", {}),
        ("gpt2", {"use_cache": False}),  # We don't support exporting GPT2 past keys anymore
    ]


    @require_tf
    @slow
    def        A   (  self      :							int):
           for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
                  self._test_export(SCREAMING_SNAKE_CASE				,							'tf'				,							12				,							**SCREAMING_SNAKE_CASE)


    @require_torch
    @slow
    def        A   (  self      :							Tuple):
           for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
                  self._test_export(SCREAMING_SNAKE_CASE				,							'pt'				,							12				,							**SCREAMING_SNAKE_CASE)


    @require_torch
    @slow
    def        A   (  self      :							Any):
           from transformers import BertModel

           _A					:     Any     = ['[UNK]', '[SEP]', '[CLS]', '[PAD]', '[MASK]', 'some', 'other', 'words']
           with NamedTemporaryFile(mode='w+t') as vocab_file:
                  vocab_file.write('\n'.join(SCREAMING_SNAKE_CASE))
                  vocab_file.flush()
                  _A					:     str     = BertTokenizerFast(vocab_file.name)

           with TemporaryDirectory() as bert_save_dir:
                  _A					:     Optional[int]     = BertModel(BertConfig(vocab_size=len(SCREAMING_SNAKE_CASE)))
                  model.save_pretrained(SCREAMING_SNAKE_CASE)
                  self._test_export(SCREAMING_SNAKE_CASE				,							'pt'				,							12				,							SCREAMING_SNAKE_CASE)


    @require_tf
    @slow
    def        A   (  self      :							Tuple):
           for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
                  _A					:     int     = self._test_export(SCREAMING_SNAKE_CASE				,							'tf'				,							12				,							**SCREAMING_SNAKE_CASE)
                  _A					:     List[Any]     = quantize(Path(SCREAMING_SNAKE_CASE))

                  # Ensure the actual quantized model is not bigger than the original one
                  if quantized_path.stat().st_size >= Path(SCREAMING_SNAKE_CASE).stat().st_size:
                         self.fail('Quantized model is bigger than initial ONNX model')


    @require_torch
    @slow
    def        A   (  self      :							List[Any]):
           for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
                  _A					:     int     = self._test_export(SCREAMING_SNAKE_CASE				,							'pt'				,							12				,							**SCREAMING_SNAKE_CASE)
                  _A					:     List[Any]     = quantize(SCREAMING_SNAKE_CASE)

                  # Ensure the actual quantized model is not bigger than the original one
                  if quantized_path.stat().st_size >= Path(SCREAMING_SNAKE_CASE).stat().st_size:
                         self.fail('Quantized model is bigger than initial ONNX model')


    def        A   (  self      :							str				,							SCREAMING_SNAKE_CASE      :							Dict				,							SCREAMING_SNAKE_CASE      :							Union[str, Any]				,							SCREAMING_SNAKE_CASE      :							List[str]				,							SCREAMING_SNAKE_CASE      :							Optional[Any]=None				,							**SCREAMING_SNAKE_CASE      :							Any):
           try:
                  # Compute path
                  with TemporaryDirectory() as tempdir:
                         _A					:     List[str]     = Path(SCREAMING_SNAKE_CASE).joinpath('model.onnx')

                  # Remove folder if exists
                  if path.parent.exists():
                         path.parent.rmdir()

                  # Export
                  convert(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							**SCREAMING_SNAKE_CASE)

                  return path
           except Exception as e:
                  self.fail(SCREAMING_SNAKE_CASE)


    @require_torch
    @require_tokenizers
    @slow
    def        A   (  self      :							List[str]):
           from transformers import BertModel

           _A					:     int     = BertModel(BertConfig.from_pretrained('lysandre/tiny-bert-random'))
           _A					:     Optional[int]     = BertTokenizerFast.from_pretrained('lysandre/tiny-bert-random')
           self._test_infer_dynamic_axis(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							'pt')


    @require_tf
    @require_tokenizers
    @slow
    def        A   (  self      :							Optional[int]):
           from transformers import TFBertModel

           _A					:     Union[str, Any]     = TFBertModel(BertConfig.from_pretrained('lysandre/tiny-bert-random'))
           _A					:     Any     = BertTokenizerFast.from_pretrained('lysandre/tiny-bert-random')
           self._test_infer_dynamic_axis(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE				,							'tf')


    def        A   (  self      :							str				,							SCREAMING_SNAKE_CASE      :							Union[str, Any]				,							SCREAMING_SNAKE_CASE      :							Optional[int]				,							SCREAMING_SNAKE_CASE      :							Optional[Any]):
           _A					:     int     = FeatureExtractionPipeline(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)

           _A					:     Union[str, Any]     = ['input_ids', 'token_type_ids', 'attention_mask', 'output_0', 'output_1']
           _A  ,			_A  ,			_A  ,			_A					:     int     = infer_shapes(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)

           # Assert all variables are present
           self.assertEqual(len(SCREAMING_SNAKE_CASE)				,							len(SCREAMING_SNAKE_CASE))
           self.assertTrue(all(var_name in shapes for var_name in variable_names))
           self.assertSequenceEqual(variable_names[:3]				,							SCREAMING_SNAKE_CASE)
           self.assertSequenceEqual(variable_names[3:]				,							SCREAMING_SNAKE_CASE)

           # Assert inputs are {0: batch, 1: sequence}
           for var_name in ["input_ids", "token_type_ids", "attention_mask"]:
                  self.assertDictEqual(shapes[var_name]				,							{0: 'batch', 1: 'sequence'})

           # Assert outputs are {0: batch, 1: sequence} and {0: batch}
           self.assertDictEqual(shapes['output_0']				,							{0: 'batch', 1: 'sequence'})
           self.assertDictEqual(shapes['output_1']				,							{0: 'batch'})


    def        A   (  self      :							Dict):
           _A					:     int     = ['input_ids', 'attention_mask', 'token_type_ids']
           _A					:     int     = {'input_ids': [1, 2, 3, 4], 'attention_mask': [0, 0, 0, 0], 'token_type_ids': [1, 1, 1, 1]}
           _A  ,			_A					:     int     = ensure_valid_input(FuncContiguousArgs()				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)

           # Should have exactly the same number of args (all are valid)
           self.assertEqual(len(SCREAMING_SNAKE_CASE)				,							3)

           # Should have exactly the same input names
           self.assertEqual(set(SCREAMING_SNAKE_CASE)				,							set(SCREAMING_SNAKE_CASE))

           # Parameter should be reordered according to their respective place in the function:
           # (input_ids, token_type_ids, attention_mask)
           self.assertEqual(SCREAMING_SNAKE_CASE				,							(tokens['input_ids'], tokens['token_type_ids'], tokens['attention_mask']))

           # Generated args are interleaved with another args (for instance parameter "past" in GPT2)
           _A  ,			_A					:     Optional[Any]     = ensure_valid_input(FuncNonContiguousArgs()				,							SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE)

           # Should have exactly the one arg (all before the one not provided "some_other_args")
           self.assertEqual(len(SCREAMING_SNAKE_CASE)				,							1)
           self.assertEqual(len(SCREAMING_SNAKE_CASE)				,							1)

           # Should have only "input_ids"
           self.assertEqual(inputs_args[0]				,							tokens['input_ids'])
           self.assertEqual(ordered_input_names[0]				,							'input_ids')


    def        A   (  self      :							List[Any]):
           _A					:     List[str]     = generate_identified_filename(Path('/home/something/my_fake_model.onnx')				,							'-test')
           self.assertEqual('/home/something/my_fake_model-test.onnx'				,							generated.as_posix())










'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


A				:					List[str]    						=			{
    '''configuration_electra''': ['''ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ElectraConfig''', '''ElectraOnnxConfig'''],
    '''tokenization_electra''': ['''ElectraTokenizer'''],
}

try:
   if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   A				:					int    						=			['''ElectraTokenizerFast''']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   A				:					List[str]    						=			[
       '''ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''ElectraForCausalLM''',
       '''ElectraForMaskedLM''',
       '''ElectraForMultipleChoice''',
       '''ElectraForPreTraining''',
       '''ElectraForQuestionAnswering''',
       '''ElectraForSequenceClassification''',
       '''ElectraForTokenClassification''',
       '''ElectraModel''',
       '''ElectraPreTrainedModel''',
       '''load_tf_weights_in_electra''',
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   A				:					Dict    						=			[
       '''TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''TFElectraForMaskedLM''',
       '''TFElectraForMultipleChoice''',
       '''TFElectraForPreTraining''',
       '''TFElectraForQuestionAnswering''',
       '''TFElectraForSequenceClassification''',
       '''TFElectraForTokenClassification''',
       '''TFElectraModel''',
       '''TFElectraPreTrainedModel''',
   ]

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   A				:					List[Any]    						=			[
       '''FlaxElectraForCausalLM''',
       '''FlaxElectraForMaskedLM''',
       '''FlaxElectraForMultipleChoice''',
       '''FlaxElectraForPreTraining''',
       '''FlaxElectraForQuestionAnswering''',
       '''FlaxElectraForSequenceClassification''',
       '''FlaxElectraForTokenClassification''',
       '''FlaxElectraModel''',
       '''FlaxElectraPreTrainedModel''',
   ]


if TYPE_CHECKING:
   from .configuration_electra import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP, ElectraConfig, ElectraOnnxConfig
   from .tokenization_electra import ElectraTokenizer

   try:
      if not is_tokenizers_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_electra_fast import ElectraTokenizerFast

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_electra import (
          ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
          ElectraForCausalLM,
          ElectraForMaskedLM,
          ElectraForMultipleChoice,
          ElectraForPreTraining,
          ElectraForQuestionAnswering,
          ElectraForSequenceClassification,
          ElectraForTokenClassification,
          ElectraModel,
          ElectraPreTrainedModel,
          load_tf_weights_in_electra,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_electra import (
          TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFElectraForMaskedLM,
          TFElectraForMultipleChoice,
          TFElectraForPreTraining,
          TFElectraForQuestionAnswering,
          TFElectraForSequenceClassification,
          TFElectraForTokenClassification,
          TFElectraModel,
          TFElectraPreTrainedModel,
      )

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_flax_electra import (
          FlaxElectraForCausalLM,
          FlaxElectraForMaskedLM,
          FlaxElectraForMultipleChoice,
          FlaxElectraForPreTraining,
          FlaxElectraForQuestionAnswering,
          FlaxElectraForSequenceClassification,
          FlaxElectraForTokenClassification,
          FlaxElectraModel,
          FlaxElectraPreTrainedModel,
      )

else:
   import sys

   A				:					Tuple    						=			_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






