





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     __snake_case       :





       '''simple docstring'''







       def 							UpperCAmelCase__ (  self :    str							,				A :    List[str]							,				A :    Union[str, Any]							,				A :    Tuple		):
         return None







class     __snake_case       :





       '''simple docstring'''







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







class     __snake_case       (							unittest.TestCase			):





       '''simple docstring'''







       lowerCAmelCase__      		=     [
           # (model_name, model_kwargs)
           ("""bert-base-cased""", {}),
           ("""gpt2""", {"""use_cache""": False}),  # We don't support exporting GPT2 past keys anymore
       ]




       @require_tf
       @slow
       def 							UpperCAmelCase__ (  self :    Tuple		):
         for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
           self._test_export(A							,				"""tf"""							,				12							,				**A		)




       @require_torch
       @slow
       def 							UpperCAmelCase__ (  self :    Union[str, Any]		):
         for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
           self._test_export(A							,				"""pt"""							,				12							,				**A		)




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

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

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




       @require_tf
       @slow
       def 							UpperCAmelCase__ (  self :    List[Any]		):
         for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
           __snake_case:						str   =						self._test_export(A							,				"""tf"""							,				12							,				**A		)
           __snake_case:						Any   =						quantize(Path(A		)		)

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




       @require_torch
       @slow
       def 							UpperCAmelCase__ (  self :    int		):
         for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
           __snake_case:						str   =						self._test_export(A							,				"""pt"""							,				12							,				**A		)
           __snake_case:						Dict   =						quantize(A		)

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




       def 							UpperCAmelCase__ (  self :    Union[str, Any]							,				A :    List[Any]							,				A :    str							,				A :    Union[str, Any]							,				A :    List[Any]=None							,				**A :    Any		):
         try:
           # Compute path
           with TemporaryDirectory() as tempdir:
             __snake_case:						Any   =						Path(A		).joinpath("""model.onnx"""		)

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

           # Export
           convert(A							,				A							,				A							,				A							,				A							,				**A		)

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




       @require_torch
       @require_tokenizers
       @slow
       def 							UpperCAmelCase__ (  self :    Optional[Any]		):
         from transformers import BertModel

         __snake_case:						Optional[Any]   =						BertModel(BertConfig.from_pretrained("""lysandre/tiny-bert-random"""		)		)
         __snake_case:						Union[str, Any]   =						BertTokenizerFast.from_pretrained("""lysandre/tiny-bert-random"""		)
         self._test_infer_dynamic_axis(A							,				A							,				"""pt"""		)




       @require_tf
       @require_tokenizers
       @slow
       def 							UpperCAmelCase__ (  self :    List[Any]		):
         from transformers import TFBertModel

         __snake_case:						Union[str, Any]   =						TFBertModel(BertConfig.from_pretrained("""lysandre/tiny-bert-random"""		)		)
         __snake_case:						Optional[int]   =						BertTokenizerFast.from_pretrained("""lysandre/tiny-bert-random"""		)
         self._test_infer_dynamic_axis(A							,				A							,				"""tf"""		)




       def 							UpperCAmelCase__ (  self :    Dict							,				A :    Tuple							,				A :    List[str]							,				A :    List[Any]		):
         __snake_case:						Tuple   =						FeatureExtractionPipeline(A							,				A		)

         __snake_case:						Tuple   =						["""input_ids""", """token_type_ids""", """attention_mask""", """output_0""", """output_1"""]
         __snake_case  ,   __snake_case  ,   __snake_case  ,   __snake_case:						Any   =						infer_shapes(A							,				A		)

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

         # 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 							UpperCAmelCase__ (  self :    str		):
         __snake_case:						List[Any]   =						["""input_ids""", """attention_mask""", """token_type_ids"""]
         __snake_case:						Tuple   =						{"""input_ids""": [1, 2, 3, 4], """attention_mask""": [0, 0, 0, 0], """token_type_ids""": [1, 1, 1, 1]}
         __snake_case  ,   __snake_case:						str   =						ensure_valid_input(FuncContiguousArgs()							,				A							,				A		)

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

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

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

         # Generated args are interleaved with another args (for instance parameter "past" in GPT2)
         __snake_case  ,   __snake_case:						int   =						ensure_valid_input(FuncNonContiguousArgs()							,				A							,				A		)

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

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






       def 							UpperCAmelCase__ (  self :    List[Any]		):
         __snake_case:						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()		)













from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__UpperCAmelCase       :   Optional[int]				       =							logging.get_logger(__name__)

__UpperCAmelCase       :   Optional[int]				       =							{
    "kssteven/ibert-roberta-base": "https://huggingface.co/kssteven/ibert-roberta-base/resolve/main/config.json",
    "kssteven/ibert-roberta-large": "https://huggingface.co/kssteven/ibert-roberta-large/resolve/main/config.json",
    "kssteven/ibert-roberta-large-mnli": (
        "https://huggingface.co/kssteven/ibert-roberta-large-mnli/resolve/main/config.json"
    ),
}







class     __snake_case       (							__lowerCamelCase			):





       '''simple docstring'''







       lowerCAmelCase__      		=     """ibert"""




       def __init__(  self :    Dict							,				A :    Union[str, Any]=30_522							,				A :    List[Any]=768							,				A :    List[Any]=12							,				A :    Optional[int]=12							,				A :    Optional[Any]=3_072							,				A :    int="gelu"							,				A :    str=0.1							,				A :    List[Any]=0.1							,				A :    Optional[Any]=512							,				A :    int=2							,				A :    Union[str, Any]=0.02							,				A :    List[str]=1E-12							,				A :    Optional[int]=1							,				A :    Optional[int]=0							,				A :    List[str]=2							,				A :    str="absolute"							,				A :    Any=False							,				A :    Optional[Any]="none"							,				**A :    Any							,				):
         super().__init__(pad_token_id=A							,				bos_token_id=A							,				eos_token_id=A							,				**A		)

         __snake_case:						List[Any]   =						vocab_size
         __snake_case:						Optional[Any]   =						hidden_size
         __snake_case:						List[str]   =						num_hidden_layers
         __snake_case:						Tuple   =						num_attention_heads
         __snake_case:						List[str]   =						hidden_act
         __snake_case:						Optional[Any]   =						intermediate_size
         __snake_case:						Tuple   =						hidden_dropout_prob
         __snake_case:						List[str]   =						attention_probs_dropout_prob
         __snake_case:						Any   =						max_position_embeddings
         __snake_case:						int   =						type_vocab_size
         __snake_case:						List[str]   =						initializer_range
         __snake_case:						List[Any]   =						layer_norm_eps
         __snake_case:						Optional[int]   =						position_embedding_type
         __snake_case:						str   =						quant_mode
         __snake_case:						Optional[int]   =						force_dequant







class     __snake_case       (							__lowerCamelCase			):





       '''simple docstring'''







       @property
       def 							UpperCAmelCase__ (  self :    Tuple		):
         if self.task == "multiple-choice":
           __snake_case:						List[Any]   =						{0: """batch""", 1: """choice""", 2: """sequence"""}
         else:
           __snake_case:						List[str]   =						{0: """batch""", 1: """sequence"""}
         return OrderedDict(
             [
                 ("""input_ids""", dynamic_axis),
                 ("""attention_mask""", dynamic_axis),
             ]		)






