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







						'''simple docstring'''



						def        SCREAMING_SNAKE_CASE(			self							:Union[str, Any]		,    lowerCAmelCase__							:List[Any]		,    lowerCAmelCase__							:List[str]		,    lowerCAmelCase__							:Union[str, Any]     )  ->Union[str, Any]:
										return None





class   UpperCamelCase_      :







						'''simple docstring'''



						def        SCREAMING_SNAKE_CASE(			self							:Optional[int]		,    lowerCAmelCase__							:Optional[int]		,    lowerCAmelCase__							:str		,    lowerCAmelCase__							:Optional[Any]		,    lowerCAmelCase__							:Optional[int]     )  ->List[Any]:
										return None





class   UpperCamelCase_      (      unittest.TestCase		):







						'''simple docstring'''



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





						@require_tf
						@slow
						def        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->List[Any]:




										for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
														self._test_export(lowercase_		,    "tf"		,    12		,    **lowercase_     )





						@require_torch
						@slow
						def        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->Optional[int]:




										for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
														self._test_export(lowercase_		,    "pt"		,    12		,    **lowercase_     )





						@require_torch
						@slow
						def        SCREAMING_SNAKE_CASE(			self							:Tuple     )  ->List[str]:
										from transformers import BertModel

										lowercase         =      ['[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()
														lowercase         =      BertTokenizerFast(vocab_file.name     )

										with TemporaryDirectory() as bert_save_dir:
														lowercase         =      BertModel(BertConfig(vocab_size=len(lowercase_     )     )     )
														model.save_pretrained(lowercase_     )
														self._test_export(lowercase_		,    "pt"		,    12		,    lowercase_     )





						@require_tf
						@slow
						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->List[str]:




										for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
														lowercase         =      self._test_export(lowercase_		,    "tf"		,    12		,    **lowercase_     )
														lowercase         =      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        SCREAMING_SNAKE_CASE(			self							:List[Any]     )  ->List[Any]:




										for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
														lowercase         =      self._test_export(lowercase_		,    "pt"		,    12		,    **lowercase_     )
														lowercase         =      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        SCREAMING_SNAKE_CASE(			self							:Dict		,    lowerCAmelCase__							:List[Any]		,    lowerCAmelCase__							:Union[str, Any]		,    lowerCAmelCase__							:str		,    lowerCAmelCase__							:List[str]=None		,    **lowerCAmelCase__							:List[str]     )  ->Optional[int]:




										try:
														# Compute path
														with TemporaryDirectory() as tempdir:
																		lowercase         =      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        SCREAMING_SNAKE_CASE(			self							:List[Any]     )  ->List[Any]:
										from transformers import BertModel

										lowercase         =      BertModel(BertConfig.from_pretrained("lysandre/tiny-bert-random"     )     )
										lowercase         =      BertTokenizerFast.from_pretrained("lysandre/tiny-bert-random"     )
										self._test_infer_dynamic_axis(lowercase_		,    lowercase_		,    "pt"     )





						@require_tf
						@require_tokenizers
						@slow
						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->Dict:
										from transformers import TFBertModel

										lowercase         =      TFBertModel(BertConfig.from_pretrained("lysandre/tiny-bert-random"     )     )
										lowercase         =      BertTokenizerFast.from_pretrained("lysandre/tiny-bert-random"     )
										self._test_infer_dynamic_axis(lowercase_		,    lowercase_		,    "tf"     )





						def        SCREAMING_SNAKE_CASE(			self							:List[Any]		,    lowerCAmelCase__							:Optional[int]		,    lowerCAmelCase__							:str		,    lowerCAmelCase__							:int     )  ->List[str]:
										lowercase         =      FeatureExtractionPipeline(lowercase_		,    lowercase_     )

										lowercase         =      ['input_ids', 'token_type_ids', 'attention_mask', 'output_0', 'output_1']
										lowercase         =      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        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->Optional[Any]:
										lowercase         =      ['input_ids', 'attention_mask', 'token_type_ids']
										lowercase         =      {'input_ids': [1, 2, 3, 4], 'attention_mask': [0, 0, 0, 0], 'token_type_ids': [1, 1, 1, 1]}
										lowercase         =      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)
										lowercase         =      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        SCREAMING_SNAKE_CASE(			self							:str     )  ->List[Any]:
										lowercase         =      generate_identified_filename(Path("/home/something/my_fake_model.onnx"     )		,    "-test"     )
										self.assertEqual("/home/something/my_fake_model-test.onnx"		,    generated.as_posix()     )




import shutil
import tempfile
import unittest
from unittest.mock import patch

from transformers import (
    DefaultFlowCallback,
    IntervalStrategy,
    PrinterCallback,
    ProgressCallback,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    is_torch_available,
)
from transformers.testing_utils import require_torch


if is_torch_available():
 from transformers.trainer import DEFAULT_CALLBACKS

 from .test_trainer import RegressionDataset, RegressionModelConfig, RegressionPreTrainedModel




class A    (  UpperCamelCase_	):






      def __init__(   self		:	str					)   ->		Union[str, Any]:
             """simple docstring"""




             _lowerCamelCase			:						Optional[Any]            =[]






      def 			lowerCamelCase			(   self		:	Dict     ,  lowercase_		:	List[Any]     ,  lowercase_		:	Optional[Any]     ,  lowercase_		:	Optional[int]     ,  **lowercase_		:	List[Any]					)   ->		List[Any]:
             """simple docstring"""




             self.events.append('on_init_end'					)






      def 			lowerCamelCase			(   self		:	int     ,  lowercase_		:	Tuple     ,  lowercase_		:	Any     ,  lowercase_		:	Optional[Any]     ,  **lowercase_		:	Any					)   ->		Union[str, Any]:
             """simple docstring"""




             self.events.append('on_train_begin'					)






      def 			lowerCamelCase			(   self		:	Optional[int]     ,  lowercase_		:	Dict     ,  lowercase_		:	Dict     ,  lowercase_		:	Tuple     ,  **lowercase_		:	str					)   ->		Optional[Any]:
             """simple docstring"""




             self.events.append('on_train_end'					)






      def 			lowerCamelCase			(   self		:	Optional[Any]     ,  lowercase_		:	str     ,  lowercase_		:	Dict     ,  lowercase_		:	int     ,  **lowercase_		:	List[str]					)   ->		Dict:
             """simple docstring"""




             self.events.append('on_epoch_begin'					)






      def 			lowerCamelCase			(   self		:	Dict     ,  lowercase_		:	Any     ,  lowercase_		:	Dict     ,  lowercase_		:	List[Any]     ,  **lowercase_		:	Dict					)   ->		Optional[Any]:
             """simple docstring"""




             self.events.append('on_epoch_end'					)






      def 			lowerCamelCase			(   self		:	List[str]     ,  lowercase_		:	List[Any]     ,  lowercase_		:	Optional[Any]     ,  lowercase_		:	Any     ,  **lowercase_		:	Dict					)   ->		Tuple:
             """simple docstring"""




             self.events.append('on_step_begin'					)






      def 			lowerCamelCase			(   self		:	List[str]     ,  lowercase_		:	Dict     ,  lowercase_		:	str     ,  lowercase_		:	Optional[Any]     ,  **lowercase_		:	Optional[int]					)   ->		Tuple:
             """simple docstring"""




             self.events.append('on_step_end'					)






      def 			lowerCamelCase			(   self		:	Tuple     ,  lowercase_		:	int     ,  lowercase_		:	List[Any]     ,  lowercase_		:	Union[str, Any]     ,  **lowercase_		:	List[Any]					)   ->		Any:
             """simple docstring"""




             self.events.append('on_evaluate'					)






      def 			lowerCamelCase			(   self		:	Tuple     ,  lowercase_		:	Union[str, Any]     ,  lowercase_		:	Dict     ,  lowercase_		:	List[str]     ,  **lowercase_		:	Dict					)   ->		Dict:
             """simple docstring"""




             self.events.append('on_predict'					)






      def 			lowerCamelCase			(   self		:	Union[str, Any]     ,  lowercase_		:	List[str]     ,  lowercase_		:	Optional[int]     ,  lowercase_		:	List[Any]     ,  **lowercase_		:	Optional[int]					)   ->		Any:
             """simple docstring"""




             self.events.append('on_save'					)






      def 			lowerCamelCase			(   self		:	str     ,  lowercase_		:	List[Any]     ,  lowercase_		:	Any     ,  lowercase_		:	Optional[int]     ,  **lowercase_		:	Optional[int]					)   ->		List[str]:
             """simple docstring"""




             self.events.append('on_log'					)






      def 			lowerCamelCase			(   self		:	Any     ,  lowercase_		:	List[Any]     ,  lowercase_		:	int     ,  lowercase_		:	List[Any]     ,  **lowercase_		:	str					)   ->		List[Any]:
             """simple docstring"""




             self.events.append('on_prediction_step'					)




@require_torch
class A    (  unittest.TestCase	):






      def 			lowerCamelCase			(   self		:	List[Any]					)   ->		Union[str, Any]:
             """simple docstring"""




             _lowerCamelCase			:						Dict            =tempfile.mkdtemp()






      def 			lowerCamelCase			(   self		:	Union[str, Any]					)   ->		int:
             """simple docstring"""




             shutil.rmtree(self.output_dir					)






      def 			lowerCamelCase			(   self		:	Optional[Any]     ,  lowercase_		:	Optional[int]=0     ,  lowercase_		:	int=0     ,  lowercase_		:	Dict=64     ,  lowercase_		:	List[str]=64     ,  lowercase_		:	int=None     ,  lowercase_		:	Tuple=False     ,  **lowercase_		:	Dict					)   ->		Optional[int]:
             """simple docstring"""




             _lowerCamelCase			:						int            =RegressionDataset(length=lowercase_					)
             _lowerCamelCase			:						Optional[Any]            =RegressionDataset(length=lowercase_					)
             _lowerCamelCase			:						str            =RegressionModelConfig(a=lowercase_     ,  b=lowercase_					)
             _lowerCamelCase			:						Optional[Any]            =RegressionPreTrainedModel(lowercase_					)

             _lowerCamelCase			:						str            =TrainingArguments(self.output_dir     ,  disable_tqdm=lowercase_     ,  report_to=[]     ,  **lowercase_					)
             return Trainer(
                 lowercase_     ,  lowercase_     ,  train_dataset=lowercase_     ,  eval_dataset=lowercase_     ,  callbacks=lowercase_     ,  )






      def 			lowerCamelCase			(   self		:	List[str]     ,  lowercase_		:	int     ,  lowercase_		:	int					)   ->		str:
             """simple docstring"""




             self.assertEqual(len(lowercase_					)     ,  len(lowercase_					)					)

             # Order doesn't matter
             _lowerCamelCase			:						Any            =sorted(lowercase_     ,  key=lambda lowercase_					: cb.__name__ if isinstance(lowercase_     ,  lowercase_					) else cb.__class__.__name__					)
             _lowerCamelCase			:						Any            =sorted(lowercase_     ,  key=lambda lowercase_					: cb.__name__ if isinstance(lowercase_     ,  lowercase_					) else cb.__class__.__name__					)

             for cba, cba in zip(lowercase_     ,  lowercase_					):
                    if isinstance(lowercase_     ,  lowercase_					) and isinstance(lowercase_     ,  lowercase_					):
                           self.assertEqual(lowercase_     ,  lowercase_					)
                    elif isinstance(lowercase_     ,  lowercase_					) and not isinstance(lowercase_     ,  lowercase_					):
                           self.assertEqual(lowercase_     ,  cba.__class__					)
                    elif not isinstance(lowercase_     ,  lowercase_					) and isinstance(lowercase_     ,  lowercase_					):
                           self.assertEqual(cba.__class__     ,  lowercase_					)
                    else:
                           self.assertEqual(lowercase_     ,  lowercase_					)






      def 			lowerCamelCase			(   self		:	Dict     ,  lowercase_		:	List[str]					)   ->		str:
             """simple docstring"""




             _lowerCamelCase			:						List[Any]            =['on_init_end', 'on_train_begin']
             _lowerCamelCase			:						Dict            =0
             _lowerCamelCase			:						str            =len(trainer.get_eval_dataloader()					)
             _lowerCamelCase			:						List[Any]            =['on_prediction_step'] * len(trainer.get_eval_dataloader()					) + ['on_log', 'on_evaluate']
             for _ in range(trainer.state.num_train_epochs					):
                    expected_events.append('on_epoch_begin'					)
                    for _ in range(lowercase_					):
                           step += 1
                           expected_events += ["on_step_begin", "on_step_end"]
                           if step % trainer.args.logging_steps == 0:
                                  expected_events.append('on_log'					)
                           if trainer.args.evaluation_strategy == IntervalStrategy.STEPS and step % trainer.args.eval_steps == 0:
                                  expected_events += evaluation_events.copy()
                           if step % trainer.args.save_steps == 0:
                                  expected_events.append('on_save'					)
                    expected_events.append('on_epoch_end'					)
                    if trainer.args.evaluation_strategy == IntervalStrategy.EPOCH:
                           expected_events += evaluation_events.copy()
             expected_events += ["on_log", "on_train_end"]
             return expected_events






      def 			lowerCamelCase			(   self		:	Optional[Any]					)   ->		List[str]:
             """simple docstring"""




             _lowerCamelCase			:						Any            =self.get_trainer()
             _lowerCamelCase			:						Optional[Any]            =DEFAULT_CALLBACKS.copy() + [ProgressCallback]
             self.check_callbacks_equality(trainer.callback_handler.callbacks     ,  lowercase_					)

             # Callbacks passed at init are added to the default callbacks
             _lowerCamelCase			:						Dict            =self.get_trainer(callbacks=[MyTestTrainerCallback]					)
             expected_callbacks.append(lowercase_					)
             self.check_callbacks_equality(trainer.callback_handler.callbacks     ,  lowercase_					)

             # TrainingArguments.disable_tqdm controls if use ProgressCallback or PrinterCallback
             _lowerCamelCase			:						List[str]            =self.get_trainer(disable_tqdm=lowercase_					)
             _lowerCamelCase			:						Tuple            =DEFAULT_CALLBACKS.copy() + [PrinterCallback]
             self.check_callbacks_equality(trainer.callback_handler.callbacks     ,  lowercase_					)






      def 			lowerCamelCase			(   self		:	str					)   ->		Optional[Any]:
             """simple docstring"""




             _lowerCamelCase			:						int            =DEFAULT_CALLBACKS.copy() + [ProgressCallback]
             _lowerCamelCase			:						List[str]            =self.get_trainer()

             # We can add, pop, or remove by class name
             trainer.remove_callback(lowercase_					)
             expected_callbacks.remove(lowercase_					)
             self.check_callbacks_equality(trainer.callback_handler.callbacks     ,  lowercase_					)

             _lowerCamelCase			:						Tuple            =self.get_trainer()
             _lowerCamelCase			:						Dict            =trainer.pop_callback(lowercase_					)
             self.assertEqual(cb.__class__     ,  lowercase_					)
             self.check_callbacks_equality(trainer.callback_handler.callbacks     ,  lowercase_					)

             trainer.add_callback(lowercase_					)
             expected_callbacks.insert(0     ,  lowercase_					)
             self.check_callbacks_equality(trainer.callback_handler.callbacks     ,  lowercase_					)

             # We can also add, pop, or remove by instance
             _lowerCamelCase			:						Optional[int]            =self.get_trainer()
             _lowerCamelCase			:						Any            =trainer.callback_handler.callbacks[0]
             trainer.remove_callback(lowercase_					)
             expected_callbacks.remove(lowercase_					)
             self.check_callbacks_equality(trainer.callback_handler.callbacks     ,  lowercase_					)

             _lowerCamelCase			:						Any            =self.get_trainer()
             _lowerCamelCase			:						Optional[int]            =trainer.callback_handler.callbacks[0]
             _lowerCamelCase			:						Dict            =trainer.pop_callback(lowercase_					)
             self.assertEqual(lowercase_     ,  lowercase_					)
             self.check_callbacks_equality(trainer.callback_handler.callbacks     ,  lowercase_					)

             trainer.add_callback(lowercase_					)
             expected_callbacks.insert(0     ,  lowercase_					)
             self.check_callbacks_equality(trainer.callback_handler.callbacks     ,  lowercase_					)






      def 			lowerCamelCase			(   self		:	Tuple					)   ->		Union[str, Any]:
             """simple docstring"""




             import warnings

             # XXX: for now ignore scatter_gather warnings in this test since it's not relevant to what's being tested
             warnings.simplefilter(action='ignore'     ,  category=lowercase_					)

             _lowerCamelCase			:						str            =self.get_trainer(callbacks=[MyTestTrainerCallback]					)
             trainer.train()
             _lowerCamelCase			:						List[Any]            =trainer.callback_handler.callbacks[-2].events
             self.assertEqual(lowercase_     ,  self.get_expected_events(lowercase_					)					)

             # Independent log/save/eval
             _lowerCamelCase			:						str            =self.get_trainer(callbacks=[MyTestTrainerCallback]     ,  logging_steps=5					)
             trainer.train()
             _lowerCamelCase			:						Union[str, Any]            =trainer.callback_handler.callbacks[-2].events
             self.assertEqual(lowercase_     ,  self.get_expected_events(lowercase_					)					)

             _lowerCamelCase			:						int            =self.get_trainer(callbacks=[MyTestTrainerCallback]     ,  save_steps=5					)
             trainer.train()
             _lowerCamelCase			:						Optional[Any]            =trainer.callback_handler.callbacks[-2].events
             self.assertEqual(lowercase_     ,  self.get_expected_events(lowercase_					)					)

             _lowerCamelCase			:						str            =self.get_trainer(callbacks=[MyTestTrainerCallback]     ,  eval_steps=5     ,  evaluation_strategy='steps'					)
             trainer.train()
             _lowerCamelCase			:						Any            =trainer.callback_handler.callbacks[-2].events
             self.assertEqual(lowercase_     ,  self.get_expected_events(lowercase_					)					)

             _lowerCamelCase			:						Optional[int]            =self.get_trainer(callbacks=[MyTestTrainerCallback]     ,  evaluation_strategy='epoch'					)
             trainer.train()
             _lowerCamelCase			:						Optional[int]            =trainer.callback_handler.callbacks[-2].events
             self.assertEqual(lowercase_     ,  self.get_expected_events(lowercase_					)					)

             # A bit of everything
             _lowerCamelCase			:						Optional[Any]            =self.get_trainer(
                 callbacks=[MyTestTrainerCallback]     ,  logging_steps=3     ,  save_steps=10     ,  eval_steps=5     ,  evaluation_strategy='steps'     ,  )
             trainer.train()
             _lowerCamelCase			:						List[str]            =trainer.callback_handler.callbacks[-2].events
             self.assertEqual(lowercase_     ,  self.get_expected_events(lowercase_					)					)

             # warning should be emitted for duplicated callbacks
             with patch('transformers.trainer_callback.logger.warning'					) as warn_mock:
                    _lowerCamelCase			:						Union[str, Any]            =self.get_trainer(
                        callbacks=[MyTestTrainerCallback, MyTestTrainerCallback]     ,  )
                    assert str(lowercase_					) in warn_mock.call_args[0][0]


