


'''simple docstring'''






import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

from transformers import AutoConfig, TFAutoModel, is_tensorflow_text_available, is_tf_available
from transformers.models.bert.tokenization_bert import BertTokenizer
from transformers.testing_utils import require_tensorflow_text, require_tf, slow


if is_tf_available():
		import tensorflow as tf

if is_tensorflow_text_available():
		from transformers.models.bert import TFBertTokenizer


__UpperCAmelCase       =							['''bert-base-uncased''', '''bert-base-cased''']
__UpperCAmelCase       =							'''hf-internal-testing/tiny-bert-tf-only'''

if is_tf_available():




		class 		a__    (							tf.keras.Model  ):





									'''simple docstring'''






									def __init__(		self ,   lowerCamelCase_       )	->							List[str]:
																super().__init__()
																lowerCAmelCase__	       =     tokenizer
																lowerCAmelCase__	       =     AutoConfig.from_pretrained(snake_case_       )
																lowerCAmelCase__	       =     TFAutoModel.from_config(snake_case_       )







									def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_       )	->							Union[str, Any]:
																lowerCAmelCase__	       =     self.tokenizer(snake_case_       )
																lowerCAmelCase__	       =     self.bert(**snake_case_       )
																return out["pooler_output"]




@require_tf
@require_tensorflow_text
class 		a__    (							unittest.TestCase  ):





							'''simple docstring'''






							def   __SCREAMING_SNAKE_CASE							(		self       )	->							List[str]:
														super().setUp()

														lowerCAmelCase__	       =     [
														    BertTokenizer.from_pretrained(snake_case_       ) for checkpoint in (TOKENIZER_CHECKPOINTS * 2)
														]  # repeat for when fast_bert_tokenizer=false
														lowerCAmelCase__	       =     [TFBertTokenizer.from_pretrained(snake_case_       ) for checkpoint in TOKENIZER_CHECKPOINTS] + [
														    TFBertTokenizer.from_pretrained(snake_case_ ,   use_fast_bert_tokenizer=snake_case_       )
														    for checkpoint in TOKENIZER_CHECKPOINTS
														]
														assert len(self.tokenizers       ) == len(self.tf_tokenizers       )

														lowerCAmelCase__	       =     [
														    """This is a straightforward English test sentence.""",
														    """This one has some weird characters\rto\nsee\r\nif  those\u00E9break things.""",
														    """Now we're going to add some Chinese: 一 二 三 一二三""",
														    """And some much more rare Chinese: 齉 堃 齉堃""",
														    """Je vais aussi écrire en français pour tester les accents""",
														    """Classical Irish also has some unusual characters, so in they go: Gaelaċ, ꝼ""",
														]
														lowerCAmelCase__	       =     list(zip(self.test_sentences ,   self.test_sentences[::-1]       )       )







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Union[str, Any]:
														for tokenizer, tf_tokenizer in zip(self.tokenizers ,   self.tf_tokenizers       ):
																					for test_inputs in (self.test_sentences, self.paired_sentences):
																												lowerCAmelCase__	       =     tokenizer(snake_case_ ,   return_tensors='''tf''' ,   padding='''longest'''       )
																												lowerCAmelCase__	       =     tf_tokenizer(snake_case_       )

																												for key in python_outputs.keys():
																																			self.assertTrue(tf.reduce_all(python_outputs[key].shape == tf_outputs[key].shape       )       )
																																			self.assertTrue(tf.reduce_all(tf.cast(python_outputs[key] ,   tf.intaa       ) == tf_outputs[key]       )       )







							@slow
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Dict:
														for tf_tokenizer in self.tf_tokenizers:
																					lowerCAmelCase__	       =     tf_tokenizer(self.paired_sentences       )
																					lowerCAmelCase__	       =     tf_tokenizer(
																					    text=[sentence[0] for sentence in self.paired_sentences] ,   text_pair=[sentence[1] for sentence in self.paired_sentences] ,   )
																					for key in merged_outputs.keys():
																												self.assertTrue(tf.reduce_all(tf.cast(merged_outputs[key] ,   tf.intaa       ) == separated_outputs[key]       )       )







							@slow
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Any:
														for tf_tokenizer in self.tf_tokenizers:
																					lowerCAmelCase__	       =     tf.function(snake_case_       )
																					for test_inputs in (self.test_sentences, self.paired_sentences):
																												lowerCAmelCase__	       =     tf.constant(snake_case_       )
																												lowerCAmelCase__	       =     compiled_tokenizer(snake_case_       )
																												lowerCAmelCase__	       =     tf_tokenizer(snake_case_       )

																												for key in eager_outputs.keys():
																																			self.assertTrue(tf.reduce_all(eager_outputs[key] == compiled_outputs[key]       )       )







							@slow
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[Any]:
														for tf_tokenizer in self.tf_tokenizers:
																					lowerCAmelCase__	       =     ModelToSave(tokenizer=snake_case_       )
																					lowerCAmelCase__	       =     tf.convert_to_tensor(self.test_sentences       )
																					lowerCAmelCase__	       =     model(snake_case_       )  # Build model with some sample inputs
																					with TemporaryDirectory() as tempdir:
																												lowerCAmelCase__	       =     Path(snake_case_       ) / """saved.model"""
																												model.save(snake_case_       )
																												lowerCAmelCase__	       =     tf.keras.models.load_model(snake_case_       )
																					lowerCAmelCase__	       =     loaded_model(snake_case_       )
																					# We may see small differences because the loaded model is compiled, so we need an epsilon for the test
																					self.assertLessEqual(tf.reduce_max(tf.abs(out - loaded_output       )       ) ,   1e-5       )



'''simple docstring'''






import math
from enum import Enum
from typing import Optional, Union

from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR

from .utils import logging


__UpperCAmelCase       =							logging.get_logger(__name__)




class 		a__    (							a__  ):





							'''simple docstring'''






							lowercase__      :       int     						=				"linear"
							lowercase__      :       Any     						=				"cosine"
							lowercase__      :       Optional[int]     						=				"cosine_with_restarts"
							lowercase__      :       Optional[Any]     						=				"polynomial"
							lowercase__      :       Tuple     						=				"constant"
							lowercase__      :       Optional[int]     						=				"constant_with_warmup"
							lowercase__      :       Optional[int]     						=				"piecewise_constant"


def  _snake_case       (					A			,   A = -1	)						->   Any:
							return LambdaLR(A			,   lambda A	: 1			,   last_epoch=A	)


def  _snake_case       (					A			,   A			,   A = -1	)						->   Optional[Any]:

							def lr_lambda(A	):
														if current_step < num_warmup_steps:
																					return float(A	) / float(max(1.0			,   A	)	)
														return 1.0

							return LambdaLR(A			,   A			,   last_epoch=A	)


def  _snake_case       (					A			,   A			,   A = -1	)						->   Union[str, Any]:
							lowerCAmelCase__	       =     {}
							lowerCAmelCase__	       =     step_rules.split(''','''	)
							for rule_str in rule_list[:-1]:
														lowerCAmelCase__       , lowerCAmelCase__	       =     rule_str.split(''':'''	)
														lowerCAmelCase__	       =     int(A	)
														lowerCAmelCase__	       =     float(A	)
														lowerCAmelCase__	       =     value
							lowerCAmelCase__	       =     float(rule_list[-1]	)

							def create_rules_function(A			,   A	):
														def rule_func(A	) -> float:
																					lowerCAmelCase__	       =     sorted(rules_dict.keys()	)
																					for i, sorted_step in enumerate(A	):
																												if steps < sorted_step:
																																			return rules_dict[sorted_steps[i]]
																					return last_lr_multiple

														return rule_func

							lowerCAmelCase__	       =     create_rules_function(A			,   A	)

							return LambdaLR(A			,   A			,   last_epoch=A	)


def  _snake_case       (					A			,   A			,   A			,   A=-1	)						->   Optional[int]:

							def lr_lambda(A	):
														if current_step < num_warmup_steps:
																					return float(A	) / float(max(1			,   A	)	)
														return max(
														    0.0			,   float(num_training_steps - current_step	) / float(max(1			,   num_training_steps - num_warmup_steps	)	)	)

							return LambdaLR(A			,   A			,   A	)


def  _snake_case       (					A			,   A			,   A			,   A = 0.5			,   A = -1	)						->   List[str]:

							def lr_lambda(A	):
														if current_step < num_warmup_steps:
																					return float(A	) / float(max(1			,   A	)	)
														lowerCAmelCase__	       =     float(current_step - num_warmup_steps	) / float(max(1			,   num_training_steps - num_warmup_steps	)	)
														return max(0.0			,   0.5 * (1.0 + math.cos(math.pi * float(A	) * 2.0 * progress	))	)

							return LambdaLR(A			,   A			,   A	)


def  _snake_case       (					A			,   A			,   A			,   A = 1			,   A = -1	)						->   Union[str, Any]:

							def lr_lambda(A	):
														if current_step < num_warmup_steps:
																					return float(A	) / float(max(1			,   A	)	)
														lowerCAmelCase__	       =     float(current_step - num_warmup_steps	) / float(max(1			,   num_training_steps - num_warmup_steps	)	)
														if progress >= 1.0:
																					return 0.0
														return max(0.0			,   0.5 * (1.0 + math.cos(math.pi * ((float(A	) * progress) % 1.0)	))	)

							return LambdaLR(A			,   A			,   A	)


def  _snake_case       (					A			,   A			,   A			,   A=1E-7			,   A=1.0			,   A=-1	)						->   Union[str, Any]:
							lowerCAmelCase__	       =     optimizer.defaults['''lr''']
							if not (lr_init > lr_end):
														raise ValueError(F"""lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})"""	)

							def lr_lambda(A	):
														if current_step < num_warmup_steps:
																					return float(A	) / float(max(1			,   A	)	)
														elif current_step > num_training_steps:
																					return lr_end / lr_init  # as LambdaLR multiplies by lr_init
														else:
																					lowerCAmelCase__	       =     lr_init - lr_end
																					lowerCAmelCase__	       =     num_training_steps - num_warmup_steps
																					lowerCAmelCase__	       =     1 - (current_step - num_warmup_steps) / decay_steps
																					lowerCAmelCase__	       =     lr_range * pct_remaining**power + lr_end
																					return decay / lr_init  # as LambdaLR multiplies by lr_init

							return LambdaLR(A			,   A			,   A	)


__UpperCAmelCase       =							{
    SchedulerType.LINEAR: get_linear_schedule_with_warmup,
    SchedulerType.COSINE: get_cosine_schedule_with_warmup,
    SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,
    SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,
    SchedulerType.CONSTANT: get_constant_schedule,
    SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,
    SchedulerType.PIECEWISE_CONSTANT: get_piecewise_constant_schedule,
}






def  _snake_case       (					A			,   A			,   A = None			,   A = None			,   A = None			,   A = 1			,   A = 1.0			,   A = -1			,   )						->   int:
							lowerCAmelCase__	       =     SchedulerType(A	)
							lowerCAmelCase__	       =     TYPE_TO_SCHEDULER_FUNCTION[name]
							if name == SchedulerType.CONSTANT:
														return schedule_func(A			,   last_epoch=A	)

							if name == SchedulerType.PIECEWISE_CONSTANT:
														return schedule_func(A			,   step_rules=A			,   last_epoch=A	)

							# All other schedulers require `num_warmup_steps`
							if num_warmup_steps is None:
														raise ValueError(F"""{name} requires `num_warmup_steps`, please provide that argument."""	)

							if name == SchedulerType.CONSTANT_WITH_WARMUP:
														return schedule_func(A			,   num_warmup_steps=A			,   last_epoch=A	)

							# All other schedulers require `num_training_steps`
							if num_training_steps is None:
														raise ValueError(F"""{name} requires `num_training_steps`, please provide that argument."""	)

							if name == SchedulerType.COSINE_WITH_RESTARTS:
														return schedule_func(
														    A			,   num_warmup_steps=A			,   num_training_steps=A			,   num_cycles=A			,   last_epoch=A			,   )

							if name == SchedulerType.POLYNOMIAL:
														return schedule_func(
														    A			,   num_warmup_steps=A			,   num_training_steps=A			,   power=A			,   last_epoch=A			,   )

							return schedule_func(
							    A			,   num_warmup_steps=A			,   num_training_steps=A			,   last_epoch=A	)