





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







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase    ,						_lowercase						):
       return None





class      lowercase__							:







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase    ,						_lowercase    ,						_lowercase						):
       return None





class      lowercase__							(							unittest.TestCase  ):
  __UpperCamelCase	     =  [
      # (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						):
       for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
            self._test_export(_lowercase    ,						"""tf"""    ,						12    ,						**_lowercase						)







  @require_torch
  @slow
  def 	UpperCAmelCase__		(	self						):
       for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
            self._test_export(_lowercase    ,						"""pt"""    ,						12    ,						**_lowercase						)







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

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

       with TemporaryDirectory() as bert_save_dir:
            lowerCAmelCase_					:       List[str]    =		BertModel(BertConfig(vocab_size=len(_lowercase						)						)						)
            model.save_pretrained(_lowercase						)
            self._test_export(_lowercase    ,						"""pt"""    ,						12    ,						_lowercase						)







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

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







  @require_torch
  @slow
  def 	UpperCAmelCase__		(	self						):
       for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
            lowerCAmelCase_					:       List[Any]    =		self._test_export(_lowercase    ,						"""pt"""    ,						12    ,						**_lowercase						)
            lowerCAmelCase_					:       Optional[int]    =		quantize(_lowercase						)

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







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase    ,						_lowercase    ,						_lowercase=None    ,						**_lowercase						):
       try:
            # Compute path
            with TemporaryDirectory() as tempdir:
                 lowerCAmelCase_					:       Optional[int]    =		Path(_lowercase						).joinpath("""model.onnx"""						)

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

            # Export
            convert(_lowercase    ,						_lowercase    ,						_lowercase    ,						_lowercase    ,						_lowercase    ,						**_lowercase						)

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







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

       lowerCAmelCase_					:       List[Any]    =		BertModel(BertConfig.from_pretrained("""lysandre/tiny-bert-random"""						)						)
       lowerCAmelCase_					:       str    =		BertTokenizerFast.from_pretrained("""lysandre/tiny-bert-random"""						)
       self._test_infer_dynamic_axis(_lowercase    ,						_lowercase    ,						"""pt"""						)







  @require_tf
  @require_tokenizers
  @slow
  def 	UpperCAmelCase__		(	self						):
       from transformers import TFBertModel

       lowerCAmelCase_					:       str    =		TFBertModel(BertConfig.from_pretrained("""lysandre/tiny-bert-random"""						)						)
       lowerCAmelCase_					:       Any    =		BertTokenizerFast.from_pretrained("""lysandre/tiny-bert-random"""						)
       self._test_infer_dynamic_axis(_lowercase    ,						_lowercase    ,						"""tf"""						)







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase    ,						_lowercase						):
       lowerCAmelCase_					:       List[Any]    =		FeatureExtractionPipeline(_lowercase    ,						_lowercase						)

       lowerCAmelCase_					:       Union[str, Any]    =		["""input_ids""", """token_type_ids""", """attention_mask""", """output_0""", """output_1"""]
       lowerCAmelCase_					:       List[Any]    =		infer_shapes(_lowercase    ,						_lowercase						)

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

       # 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						):
       lowerCAmelCase_					:       Optional[int]    =		["""input_ids""", """attention_mask""", """token_type_ids"""]
       lowerCAmelCase_					:       List[str]    =		{"""input_ids""": [1, 2, 3, 4], """attention_mask""": [0, 0, 0, 0], """token_type_ids""": [1, 1, 1, 1]}
       lowerCAmelCase_					:       Optional[int]    =		ensure_valid_input(FuncContiguousArgs()    ,						_lowercase    ,						_lowercase						)

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

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

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

       # Generated args are interleaved with another args (for instance parameter "past" in GPT2)
       lowerCAmelCase_					:       str    =		ensure_valid_input(FuncNonContiguousArgs()    ,						_lowercase    ,						_lowercase						)

       # Should have exactly the one arg (all before the one not provided "some_other_args")
       self.assertEqual(len(_lowercase						)    ,						1						)
       self.assertEqual(len(_lowercase						)    ,						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						):
       lowerCAmelCase_					:       Tuple    =		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 Any, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...feature_extraction_utils import FeatureExtractionMixin
from ...onnx import OnnxConfig
from ...onnx.utils import compute_effective_axis_dimension
from ...tokenization_utils_base import PreTrainedTokenizerBase
from ...utils import TensorType, logging


UpperCAmelCase_					:   Optional[Any] 							=						logging.get_logger(__name__)

UpperCAmelCase_					:   Any 							=						{
    """deepmind/language-perceiver""": """https://huggingface.co/deepmind/language-perceiver/resolve/main/config.json""",
    # See all Perceiver models at https://huggingface.co/models?filter=perceiver
}





class      lowercase__    (							__A  ):
  __UpperCamelCase	     =  """perceiver"""







  def __init__(	self    ,						_lowercase=256    ,						_lowercase=1_280    ,						_lowercase=768    ,						_lowercase=1    ,						_lowercase=26    ,						_lowercase=8    ,						_lowercase=8    ,						_lowercase=None    ,						_lowercase=None    ,						_lowercase="kv"    ,						_lowercase=1    ,						_lowercase=1    ,						_lowercase="gelu"    ,						_lowercase=0.1    ,						_lowercase=0.02    ,						_lowercase=1e-12    ,						_lowercase=True    ,						_lowercase=262    ,						_lowercase=2_048    ,						_lowercase=56    ,						_lowercase=[368, 496]    ,						_lowercase=16    ,						_lowercase=1_920    ,						_lowercase=16    ,						_lowercase=[1, 16, 224, 224]    ,						**_lowercase    ,						):
       super().__init__(**_lowercase						)

       lowerCAmelCase_					:       Optional[int]    =		num_latents
       lowerCAmelCase_					:       List[str]    =		d_latents
       lowerCAmelCase_					:       int    =		d_model
       lowerCAmelCase_					:       Dict    =		num_blocks
       lowerCAmelCase_					:       Union[str, Any]    =		num_self_attends_per_block
       lowerCAmelCase_					:       List[str]    =		num_self_attention_heads
       lowerCAmelCase_					:       List[str]    =		num_cross_attention_heads
       lowerCAmelCase_					:       List[Any]    =		qk_channels
       lowerCAmelCase_					:       Optional[Any]    =		v_channels
       lowerCAmelCase_					:       Optional[Any]    =		cross_attention_shape_for_attention
       lowerCAmelCase_					:       Optional[int]    =		self_attention_widening_factor
       lowerCAmelCase_					:       int    =		cross_attention_widening_factor
       lowerCAmelCase_					:       Union[str, Any]    =		hidden_act
       lowerCAmelCase_					:       int    =		attention_probs_dropout_prob
       lowerCAmelCase_					:       int    =		initializer_range
       lowerCAmelCase_					:       Optional[Any]    =		layer_norm_eps
       lowerCAmelCase_					:       Optional[Any]    =		use_query_residual
       # masked language modeling attributes
       lowerCAmelCase_					:       Tuple    =		vocab_size
       lowerCAmelCase_					:       List[Any]    =		max_position_embeddings
       # image classification attributes
       lowerCAmelCase_					:       List[Any]    =		image_size
       # flow attributes
       lowerCAmelCase_					:       Dict    =		train_size
       # multimodal autoencoding attributes
       lowerCAmelCase_					:       Optional[Any]    =		num_frames
       lowerCAmelCase_					:       int    =		audio_samples_per_frame
       lowerCAmelCase_					:       Any    =		samples_per_patch
       lowerCAmelCase_					:       Any    =		output_shape





class      lowercase__    (							__A  ):







  @property
  def 	UpperCAmelCase__		(	self						):
       if self.task == "multiple-choice":
            lowerCAmelCase_					:       Optional[Any]    =		{0: """batch""", 1: """choice""", 2: """sequence"""}
       else:
            lowerCAmelCase_					:       Optional[int]    =		{0: """batch""", 1: """sequence"""}
       return OrderedDict(
           [
               ("""inputs""", dynamic_axis),
               ("""attention_mask""", dynamic_axis),
           ]						)







  @property
  def 	UpperCAmelCase__		(	self						):
       return 1e-4







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase = -1    ,						_lowercase = -1    ,						_lowercase = -1    ,						_lowercase = False    ,						_lowercase = None    ,						_lowercase = 3    ,						_lowercase = 40    ,						_lowercase = 40    ,						):
       # copied from `transformers.onnx.config.OnnxConfig` and slightly altered/simplified

       if isinstance(_lowercase    ,						_lowercase						):
            # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
            lowerCAmelCase_					:       Union[str, Any]    =		compute_effective_axis_dimension(
                _lowercase    ,						fixed_dimension=OnnxConfig.default_fixed_batch    ,						num_token_to_add=0						)
            # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
            lowerCAmelCase_					:       int    =		preprocessor.num_special_tokens_to_add(_lowercase						)
            lowerCAmelCase_					:       int    =		compute_effective_axis_dimension(
                _lowercase    ,						fixed_dimension=OnnxConfig.default_fixed_sequence    ,						num_token_to_add=_lowercase						)
            # Generate dummy inputs according to compute batch and sequence
            lowerCAmelCase_					:       Optional[int]    =		[""" """.join(["""a"""]						) * seq_length] * batch_size
            lowerCAmelCase_					:       Optional[Any]    =		dict(preprocessor(_lowercase    ,						return_tensors=_lowercase						)						)
            lowerCAmelCase_					:       Optional[Any]    =		inputs.pop("""input_ids"""						)
            return inputs
       elif isinstance(_lowercase    ,						_lowercase						) and preprocessor.model_input_names[0] == "pixel_values":
            # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
            lowerCAmelCase_					:       Any    =		compute_effective_axis_dimension(_lowercase    ,						fixed_dimension=OnnxConfig.default_fixed_batch						)
            lowerCAmelCase_					:       str    =		self._generate_dummy_images(_lowercase    ,						_lowercase    ,						_lowercase    ,						_lowercase						)
            lowerCAmelCase_					:       List[str]    =		dict(preprocessor(images=_lowercase    ,						return_tensors=_lowercase						)						)
            lowerCAmelCase_					:       Tuple    =		inputs.pop("""pixel_values"""						)
            return inputs
       else:
            raise ValueError(
                """Unable to generate dummy inputs for the model. Please provide a tokenizer or a preprocessor."""						)
