



from __future__ import annotations

import unittest

from transformers import FunnelConfig, is_tf_available
from transformers.testing_utils import require_tf

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import tensorflow as tf

     from transformers import (
         TFFunnelBaseModel,
         TFFunnelForMaskedLM,
         TFFunnelForMultipleChoice,
         TFFunnelForPreTraining,
         TFFunnelForQuestionAnswering,
         TFFunnelForSequenceClassification,
         TFFunnelForTokenClassification,
         TFFunnelModel,
     )
class    __UpperCamelCase			:

 """simple docstring"""



 def __init__(	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE=1_3		,				SCREAMING_SNAKE_CASE=7		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=9_9		,				SCREAMING_SNAKE_CASE=[1, 1, 2]		,				SCREAMING_SNAKE_CASE=1		,				SCREAMING_SNAKE_CASE=3_2		,				SCREAMING_SNAKE_CASE=4		,				SCREAMING_SNAKE_CASE=8		,				SCREAMING_SNAKE_CASE=3_7		,				SCREAMING_SNAKE_CASE="gelu_new"		,				SCREAMING_SNAKE_CASE=0.1		,				SCREAMING_SNAKE_CASE=0.1		,				SCREAMING_SNAKE_CASE=0.0		,				SCREAMING_SNAKE_CASE=5_1_2		,				SCREAMING_SNAKE_CASE=3		,				SCREAMING_SNAKE_CASE=0.02		,				SCREAMING_SNAKE_CASE=3		,				SCREAMING_SNAKE_CASE=4		,				SCREAMING_SNAKE_CASE=None		,				SCREAMING_SNAKE_CASE=False		,				) ->			Dict:
  a__												= parent
  a__												= batch_size
  a__												= seq_length
  a__												= is_training
  a__												= use_input_mask
  a__												= use_token_type_ids
  a__												= use_labels
  a__												= vocab_size
  a__												= block_sizes
  a__												= num_decoder_layers
  a__												= d_model
  a__												= n_head
  a__												= d_head
  a__												= d_inner
  a__												= hidden_act
  a__												= hidden_dropout
  a__												= attention_dropout
  a__												= activation_dropout
  a__												= max_position_embeddings
  a__												= type_vocab_size
  a__												= 2
  a__												= num_labels
  a__												= num_choices
  a__												= scope
  a__												= initializer_std

  # Used in the tests to check the size of the first attention layer
  a__												= n_head
  # Used in the tests to check the size of the first hidden state
  a__												= self.d_model
  # Used in the tests to check the number of output hidden states/attentions
  a__												= sum(self.block_sizes      ) + (0 if base else self.num_decoder_layers)
  # FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
  # the last hidden state of the first block (which is the first hidden state of the decoder).
  if not base:
   a__												= self.num_hidden_layers + 2

 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size      )

  a__												= None
  if self.use_input_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__												= None
  a__												= None
  a__												= None
  if self.use_labels:
   a__												= ids_tensor([self.batch_size]		,				self.type_sequence_label_size      )
   a__												= ids_tensor([self.batch_size, self.seq_length]		,				self.num_labels      )
   a__												= ids_tensor([self.batch_size]		,				self.num_choices      )

  a__												= FunnelConfig(
      vocab_size=self.vocab_size		,				block_sizes=self.block_sizes		,				num_decoder_layers=self.num_decoder_layers		,				d_model=self.d_model		,				n_head=self.n_head		,				d_head=self.d_head		,				d_inner=self.d_inner		,				hidden_act=self.hidden_act		,				hidden_dropout=self.hidden_dropout		,				attention_dropout=self.attention_dropout		,				activation_dropout=self.activation_dropout		,				max_position_embeddings=self.max_position_embeddings		,				type_vocab_size=self.type_vocab_size		,				initializer_std=self.initializer_std		,				)

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

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				) ->			List[Any]:
  a__												= TFFunnelModel(config=SCREAMING_SNAKE_CASE      )
  a__												= {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
  a__												= model(SCREAMING_SNAKE_CASE      )

  a__												= [input_ids, input_mask]
  a__												= model(SCREAMING_SNAKE_CASE      )

  a__												= model(SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.last_hidden_state.shape		,				(self.batch_size, self.seq_length, self.d_model)      )

  a__												= False
  a__												= TFFunnelModel(config=SCREAMING_SNAKE_CASE      )
  a__												= model(SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.last_hidden_state.shape		,				(self.batch_size, self.seq_length, self.d_model)      )

  a__												= False
  a__												= TFFunnelModel(config=SCREAMING_SNAKE_CASE      )
  a__												= model(SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.last_hidden_state.shape		,				(self.batch_size, self.seq_length, self.d_model)      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				) ->			str:
  a__												= TFFunnelBaseModel(config=SCREAMING_SNAKE_CASE      )
  a__												= {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
  a__												= model(SCREAMING_SNAKE_CASE      )

  a__												= [input_ids, input_mask]
  a__												= model(SCREAMING_SNAKE_CASE      )

  a__												= model(SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.last_hidden_state.shape		,				(self.batch_size, 2, self.d_model)      )

  a__												= False
  a__												= TFFunnelBaseModel(config=SCREAMING_SNAKE_CASE      )
  a__												= model(SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.last_hidden_state.shape		,				(self.batch_size, 3, self.d_model)      )

  a__												= False
  a__												= TFFunnelBaseModel(config=SCREAMING_SNAKE_CASE      )
  a__												= model(SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.last_hidden_state.shape		,				(self.batch_size, 2, self.d_model)      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				) ->			Any:
  a__												= TFFunnelForPreTraining(config=SCREAMING_SNAKE_CASE      )
  a__												= {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
  a__												= model(SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length)      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				) ->			Dict:
  a__												= TFFunnelForMaskedLM(config=SCREAMING_SNAKE_CASE      )
  a__												= {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
  a__												= model(SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.vocab_size)      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				) ->			List[str]:
  a__												= self.num_labels
  a__												= TFFunnelForSequenceClassification(config=SCREAMING_SNAKE_CASE      )
  a__												= {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
  a__												= model(SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.num_labels)      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				) ->			int:
  a__												= self.num_choices
  a__												= TFFunnelForMultipleChoice(config=SCREAMING_SNAKE_CASE      )
  a__												= tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE		,				1      )		,				(1, self.num_choices, 1)      )
  a__												= tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE		,				1      )		,				(1, self.num_choices, 1)      )
  a__												= tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE		,				1      )		,				(1, self.num_choices, 1)      )
  a__												= {
      '''input_ids''': multiple_choice_inputs_ids,
      '''attention_mask''': multiple_choice_input_mask,
      '''token_type_ids''': multiple_choice_token_type_ids,
  }
  a__												= model(SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.num_choices)      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				) ->			Optional[Any]:
  a__												= self.num_labels
  a__												= TFFunnelForTokenClassification(config=SCREAMING_SNAKE_CASE      )
  a__												= {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
  a__												= model(SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.num_labels)      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				) ->			List[str]:
  a__												= TFFunnelForQuestionAnswering(config=SCREAMING_SNAKE_CASE      )
  a__												= {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
  a__												= model(SCREAMING_SNAKE_CASE      )
  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 						_UpperCAmelCase       (	self      ) ->			Dict:
  a__												= self.prepare_config_and_inputs()
  (
      (
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			
  )												= config_and_inputs
  a__												= {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
  return config, inputs_dict





@require_tf
class    __UpperCamelCase			(       _lowercase	,	_lowercase	,	unittest.TestCase				):

 """simple docstring"""



 _lowercase   :    List[Any]      						=					(
     (
         TFFunnelModel,
         TFFunnelForMaskedLM,
         TFFunnelForPreTraining,
         TFFunnelForQuestionAnswering,
         TFFunnelForTokenClassification,
     )
     if is_tf_available()
     else ()
 )
 _lowercase   :    List[Any]      						=					(
     {
         '''feature-extraction''': (TFFunnelBaseModel, TFFunnelModel),
         '''fill-mask''': TFFunnelForMaskedLM,
         '''question-answering''': TFFunnelForQuestionAnswering,
         '''text-classification''': TFFunnelForSequenceClassification,
         '''token-classification''': TFFunnelForTokenClassification,
         '''zero-shot''': TFFunnelForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 _lowercase   :    Union[str, Any]      						=					False
 _lowercase   :    Optional[Any]      						=					False

 def 						_UpperCAmelCase       (	self      ) ->			List[Any]:
  a__												= TFFunnelModelTester(self      )
  a__												= ConfigTester(self		,				config_class=SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Optional[Any]:
  self.config_tester.run_common_tests()

 def 						_UpperCAmelCase       (	self      ) ->			str:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Tuple:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_pretraining(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			int:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Any:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE      )



 def 						_UpperCAmelCase       (	self      ) ->			Union[str, Any]:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE      )





@require_tf
class    __UpperCamelCase			(       _lowercase	,	unittest.TestCase				):

 """simple docstring"""



 _lowercase   :    List[str]      						=					(
     (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
 )
 _lowercase   :    Optional[int]      						=					False
 _lowercase   :    Union[str, Any]      						=					False

 def 						_UpperCAmelCase       (	self      ) ->			Any:
  a__												= TFFunnelModelTester(self		,				base=SCREAMING_SNAKE_CASE      )
  a__												= ConfigTester(self		,				config_class=SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			List[Any]:
  self.config_tester.run_common_tests()

 def 						_UpperCAmelCase       (	self      ) ->			Tuple:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_base_model(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Optional[Any]:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE      )



 def 						_UpperCAmelCase       (	self      ) ->			Any:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE      )










import argparse
import json
import os

from tensorflow.core.protobuf.saved_model_pba import SavedModel


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_copies.py
a_      :      Optional[int]    =     '.'

# Internal TensorFlow ops that can be safely ignored (mostly specific to a saved model)
a_      :      List[str]    =     [
    'Assert',
    'AssignVariableOp',
    'EmptyTensorList',
    'MergeV2Checkpoints',
    'ReadVariableOp',
    'ResourceGather',
    'RestoreV2',
    'SaveV2',
    'ShardedFilename',
    'StatefulPartitionedCall',
    'StaticRegexFullMatch',
    'VarHandleOp',
]






def     __a					(		__UpperCAmelCase							,  __UpperCAmelCase							,  __UpperCAmelCase       ):
 a__												= SavedModel()
 a__												= []

 with open(os.path.join(__UpperCAmelCase							,  '''utils'''							,  '''tf_ops'''							,  '''onnx.json'''       )       ) as f:
  a__												= json.load(__UpperCAmelCase       )['''opsets''']

 for i in range(1							,  opset + 1       ):
  onnx_ops.extend(onnx_opsets[str(__UpperCAmelCase       )]       )

 with open(__UpperCAmelCase							,  '''rb'''       ) as f:
  saved_model.ParseFromString(f.read()       )

 a__												= set()

 # Iterate over every metagraph in case there is more than one (a saved model can contain multiple graphs)
 for meta_graph in saved_model.meta_graphs:
  # Add operations in the graph definition
  model_op_names.update(node.op for node in meta_graph.graph_def.node       )

  # Go through the functions in the graph definition
  for func in meta_graph.graph_def.library.function:
   # Add operations in each function
   model_op_names.update(node.op for node in func.node_def       )

    # Convert to list, sorted if you want
 a__												= sorted(__UpperCAmelCase       )
 a__												= []

 for op in model_op_names:
  if op not in onnx_ops and op not in INTERNAL_OPS:
   incompatible_ops.append(__UpperCAmelCase       )

 if strict and len(__UpperCAmelCase       ) > 0:
  raise Exception(f"Found the following incompatible ops for the opset {opset}:\n" + incompatible_ops       )
 elif len(__UpperCAmelCase       ) > 0:
  print(f"Found the following incompatible ops for the opset {opset}:"       )
  print(*__UpperCAmelCase							,  sep='''\n'''       )
 else:
  print(f"The saved model {saved_model_path} can properly be converted with ONNX."       )


if __name__ == "__main__":
     a_      :      Dict    =     argparse.ArgumentParser()
     parser.add_argument('--saved_model_path', help='Path of the saved model to check (the .pb file).')
     parser.add_argument(
         '--opset', default=12, type=int, help='The ONNX opset against which the model has to be tested.'
     )
     parser.add_argument(
         '--framework', choices=['onnx'], default='onnx', help='Frameworks against which to test the saved model.'
     )
     parser.add_argument(
         '--strict', action='store_true', help='Whether make the checking strict (raise errors) or not (raise warnings)'
     )
     a_      :      str    =     parser.parse_args()

     if args.framework == "onnx":
          onnx_compliancy(args.saved_model_path, args.strict, args.opset)





