
'''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 		__UpperCAmelCase :


							def 					lowerCamelCase							(			self       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_    ):





								"""simple docstring"""






								return None


class 		__UpperCAmelCase :


							def 					lowerCamelCase							(			self       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_    ):





								"""simple docstring"""






								return None


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


							@require_tf
							@slow
							def 					lowerCamelCase							(			self    ):





								"""simple docstring"""






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


							@require_torch
							@slow
							def 					lowerCamelCase							(			self    ):





								"""simple docstring"""






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


							@require_torch
							@slow
							def 					lowerCamelCase							(			self    ):





								"""simple docstring"""






								from transformers import BertModel

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

								with TemporaryDirectory() as bert_save_dir:
									_snake_case  		=						BertModel(BertConfig(vocab_size=len(lowerCAmelCase__    )    )    )
									model.save_pretrained(lowerCAmelCase__    )
									self._test_export(lowerCAmelCase__       ,  'pt'       ,  12       ,  lowerCAmelCase__    )


							@require_tf
							@slow
							def 					lowerCamelCase							(			self    ):





								"""simple docstring"""






								for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
									_snake_case  		=						self._test_export(lowerCAmelCase__       ,  'tf'       ,  12       ,  **lowerCAmelCase__    )
									_snake_case  		=						quantize(Path(lowerCAmelCase__    )    )

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


							@require_torch
							@slow
							def 					lowerCamelCase							(			self    ):





								"""simple docstring"""






								for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
									_snake_case  		=						self._test_export(lowerCAmelCase__       ,  'pt'       ,  12       ,  **lowerCAmelCase__    )
									_snake_case  		=						quantize(lowerCAmelCase__    )

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


							def 					lowerCamelCase							(			self       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_=None       ,  **lowerCAmelCase_    ):





								"""simple docstring"""






								try:
									# Compute path
									with TemporaryDirectory() as tempdir:
										_snake_case  		=						Path(lowerCAmelCase__    ).joinpath('model.onnx'    )

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

									# Export
									convert(lowerCAmelCase__       ,  lowerCAmelCase__       ,  lowerCAmelCase__       ,  lowerCAmelCase__       ,  lowerCAmelCase__       ,  **lowerCAmelCase__    )

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


							@require_torch
							@require_tokenizers
							@slow
							def 					lowerCamelCase							(			self    ):





								"""simple docstring"""






								from transformers import BertModel

								_snake_case  		=						BertModel(BertConfig.from_pretrained('lysandre/tiny-bert-random'    )    )
								_snake_case  		=						BertTokenizerFast.from_pretrained('lysandre/tiny-bert-random'    )
								self._test_infer_dynamic_axis(lowerCAmelCase__       ,  lowerCAmelCase__       ,  'pt'    )


							@require_tf
							@require_tokenizers
							@slow
							def 					lowerCamelCase							(			self    ):





								"""simple docstring"""






								from transformers import TFBertModel

								_snake_case  		=						TFBertModel(BertConfig.from_pretrained('lysandre/tiny-bert-random'    )    )
								_snake_case  		=						BertTokenizerFast.from_pretrained('lysandre/tiny-bert-random'    )
								self._test_infer_dynamic_axis(lowerCAmelCase__       ,  lowerCAmelCase__       ,  'tf'    )


							def 					lowerCamelCase							(			self       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_    ):





								"""simple docstring"""






								_snake_case  		=						FeatureExtractionPipeline(lowerCAmelCase__       ,  lowerCAmelCase__    )

								_snake_case  		=						['input_ids', 'token_type_ids', 'attention_mask', 'output_0', 'output_1']
								_snake_case		,					_snake_case		,					_snake_case		,					_snake_case  		=						infer_shapes(lowerCAmelCase__       ,  lowerCAmelCase__    )

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

								# 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 					lowerCamelCase							(			self    ):





								"""simple docstring"""






								_snake_case  		=						['input_ids', 'attention_mask', 'token_type_ids']
								_snake_case  		=						{'input_ids': [1, 2, 3, 4], 'attention_mask': [0, 0, 0, 0], 'token_type_ids': [1, 1, 1, 1]}
								_snake_case		,					_snake_case  		=						ensure_valid_input(FuncContiguousArgs()       ,  lowerCAmelCase__       ,  lowerCAmelCase__    )

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

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

								# Parameter should be reordered according to their respective place in the function:
								# (input_ids, token_type_ids, attention_mask)
								self.assertEqual(lowerCAmelCase__       ,  (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  		=						ensure_valid_input(FuncNonContiguousArgs()       ,  lowerCAmelCase__       ,  lowerCAmelCase__    )

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

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




							def 					lowerCamelCase							(			self    ):





								"""simple docstring"""






								_snake_case  		=						generate_identified_filename(Path('/home/something/my_fake_model.onnx'    )       ,  '-test'    )
								self.assertEqual('/home/something/my_fake_model-test.onnx'       ,  generated.as_posix()    )



import os
import tempfile
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
						import torch
						from torch import nn

						from transformers import (
						    Adafactor,
						    AdamW,
						    get_constant_schedule,
						    get_constant_schedule_with_warmup,
						    get_cosine_schedule_with_warmup,
						    get_cosine_with_hard_restarts_schedule_with_warmup,
						    get_inverse_sqrt_schedule,
						    get_linear_schedule_with_warmup,
						    get_polynomial_decay_schedule_with_warmup,
						)

def  __snake_case							(	__magic_name__		,    __magic_name__=10						):







				'''simple docstring'''


				lowercase         =      []
				for _ in range(__magic_name__						):
								lrs.append(scheduler.get_lr()[0]						)
								scheduler.step()
				return lrs

def  __snake_case							(	__magic_name__		,    __magic_name__=10						):







				'''simple docstring'''


				lowercase         =      []
				for step in range(__magic_name__						):
								lrs.append(scheduler.get_lr()[0]						)
								scheduler.step()
								if step == num_steps // 2:
												with tempfile.TemporaryDirectory() as tmpdirname:
																lowercase         =      os.path.join(__magic_name__		,    "schedule.bin"						)
																torch.save(scheduler.state_dict()		,    __magic_name__						)

																lowercase         =      torch.load(__magic_name__						)
																scheduler.load_state_dict(__magic_name__						)
				return lrs





@require_torch
class   UpperCamelCase_      (      unittest.TestCase		):







						'''simple docstring'''



						def        SCREAMING_SNAKE_CASE(			self							:Optional[Any]		,    lowerCAmelCase__							:List[str]		,    lowerCAmelCase__							:Tuple		,    lowerCAmelCase__							:Dict     )  ->Tuple:
										self.assertEqual(len(lowerCAmelCase__     )		,    len(lowerCAmelCase__     )     )
										for a, b in zip(lowerCAmelCase__		,    lowerCAmelCase__     ):
														self.assertAlmostEqual(lowerCAmelCase__		,    lowerCAmelCase__		,    delta=lowerCAmelCase__     )





						def        SCREAMING_SNAKE_CASE(			self							:Optional[int]     )  ->Dict:
										lowercase         =      torch.tensor([0.1, -0.2, -0.1]		,    requires_grad=lowerCAmelCase__     )
										lowercase         =      torch.tensor([0.4, 0.2, -0.5]     )
										lowercase         =      nn.MSELoss()
										# No warmup, constant schedule, no gradient clipping
										lowercase         =      AdamW(params=[w]		,    lr=2E-1		,    weight_decay=0.0     )
										for _ in range(100     ):
														lowercase         =      criterion(lowerCAmelCase__		,    lowerCAmelCase__     )
														loss.backward()
														optimizer.step()
														w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
														w.grad.zero_()
										self.assertListAlmostEqual(w.tolist()		,    [0.4, 0.2, -0.5]		,    tol=1E-2     )





						def        SCREAMING_SNAKE_CASE(			self							:int     )  ->List[Any]:
										lowercase         =      torch.tensor([0.1, -0.2, -0.1]		,    requires_grad=lowerCAmelCase__     )
										lowercase         =      torch.tensor([0.4, 0.2, -0.5]     )
										lowercase         =      nn.MSELoss()
										# No warmup, constant schedule, no gradient clipping
										lowercase         =      Adafactor(
										    params=[w]		,    lr=1E-2		,    eps=(1E-3_0, 1E-3)		,    clip_threshold=1.0		,    decay_rate=-0.8		,    betaa=lowerCAmelCase__		,    weight_decay=0.0		,    relative_step=lowerCAmelCase__		,    scale_parameter=lowerCAmelCase__		,    warmup_init=lowerCAmelCase__		,    )
										for _ in range(1000     ):
														lowercase         =      criterion(lowerCAmelCase__		,    lowerCAmelCase__     )
														loss.backward()
														optimizer.step()
														w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
														w.grad.zero_()
										self.assertListAlmostEqual(w.tolist()		,    [0.4, 0.2, -0.5]		,    tol=1E-2     )





@require_torch
class   UpperCamelCase_      (      unittest.TestCase		):







						'''simple docstring'''



						UpperCamelCase      :  Any				   =					nn.Linear(50		,    50		) if is_torch_available() else None
						UpperCamelCase      :  Union[str, Any]				   =					AdamW(m.parameters()		,    lr=10.0		) if is_torch_available() else None
						UpperCamelCase      :  int				   =					10





						def        SCREAMING_SNAKE_CASE(			self							:Any		,    lowerCAmelCase__							:Optional[Any]		,    lowerCAmelCase__							:Optional[Any]		,    lowerCAmelCase__							:Any		,    lowerCAmelCase__							:Any=None     )  ->Dict:
										self.assertEqual(len(lowerCAmelCase__     )		,    len(lowerCAmelCase__     )     )
										for a, b in zip(lowerCAmelCase__		,    lowerCAmelCase__     ):
														self.assertAlmostEqual(lowerCAmelCase__		,    lowerCAmelCase__		,    delta=lowerCAmelCase__		,    msg=lowerCAmelCase__     )





						def        SCREAMING_SNAKE_CASE(			self							:Any     )  ->str:
										lowercase         =      {"num_warmup_steps": 2, "num_training_steps": 10}
										# schedulers doct format
										# function: (sched_args_dict, expected_learning_rates)
										lowercase         =      {
										    get_constant_schedule: ({}, [10.0] * self.num_steps),
										    get_constant_schedule_with_warmup: (
										        {"num_warmup_steps": 4},
										        [0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0],
										    ),
										    get_linear_schedule_with_warmup: (
										        {**common_kwargs},
										        [0.0, 5.0, 10.0, 8.75, 7.5, 6.25, 5.0, 3.75, 2.5, 1.25],
										    ),
										    get_cosine_schedule_with_warmup: (
										        {**common_kwargs},
										        [0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38],
										    ),
										    get_cosine_with_hard_restarts_schedule_with_warmup: (
										        {**common_kwargs, "num_cycles": 2},
										        [0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46],
										    ),
										    get_polynomial_decay_schedule_with_warmup: (
										        {**common_kwargs, "power": 2.0, "lr_end": 1E-7},
										        [0.0, 5.0, 10.0, 7.6_56, 5.6_25, 3.9_06, 2.5, 1.4_06, 0.6_25, 0.1_56],
										    ),
										    get_inverse_sqrt_schedule: (
										        {"num_warmup_steps": 2},
										        [0.0, 5.0, 10.0, 8.1_65, 7.0_71, 6.3_25, 5.7_74, 5.3_45, 5.0, 4.7_14],
										    ),
										}

										for scheduler_func, data in scheds.items():
														lowercase     ,	lowercase         =      data

														lowercase         =      scheduler_func(self.optimizer		,    **lowerCAmelCase__     )
														self.assertEqual(len([scheduler.get_lr()[0]]     )		,    1     )
														lowercase         =      unwrap_schedule(lowerCAmelCase__		,    self.num_steps     )
														self.assertListAlmostEqual(
														    lowerCAmelCase__		,    lowerCAmelCase__		,    tol=1E-2		,    msg=F'''failed for {scheduler_func} in normal scheduler'''		,    )

														lowercase         =      scheduler_func(self.optimizer		,    **lowerCAmelCase__     )
														if scheduler_func.__name__ != "get_constant_schedule":
																		LambdaScheduleWrapper.wrap_scheduler(lowerCAmelCase__     )  # wrap to test picklability of the schedule
														lowercase         =      unwrap_and_save_reload_schedule(lowerCAmelCase__		,    self.num_steps     )
														self.assertListEqual(lowerCAmelCase__		,    lowerCAmelCase__		,    msg=F'''failed for {scheduler_func} in save and reload'''     )





class   UpperCamelCase_      :







						'''simple docstring'''



						def __init__(			self							:List[str]		,    lowerCAmelCase__							:int     )  ->Optional[int]:
										lowercase         =      fn





						def __call__(			self							:Optional[Any]		,    *lowerCAmelCase__							:int		,    **lowerCAmelCase__							:List[str]     )  ->List[str]:
										return self.fn(*lowerCAmelCase__		,    **lowerCAmelCase__     )





						@classmethod
						def        SCREAMING_SNAKE_CASE(			self							:List[Any]		,    lowerCAmelCase__							:Optional[Any]     )  ->str:
										lowercase         =      list(map(self		,    scheduler.lr_lambdas     )     )
