

import re
import time
from typing import Optional

import IPython.display as disp

from ..trainer_callback import TrainerCallback
from ..trainer_utils import IntervalStrategy, has_length




def 					_lowercase     (				a_				:					List[str]				)				->	str:




					'''simple docstring'''







					__magic_name__							=    int(a_				)
					__magic_name__,	__magic_name__,	__magic_name__							=    t // 3_6_0_0, (t // 6_0) % 6_0, t % 6_0
					return F'{h}:{m:02d}:{s:02d}' if h != 0 else F'{m:02d}:{s:02d}'




def 					_lowercase     (				a_				:					Dict					,a_				:					Tuple					,a_				:					Union[str, Any]					,a_				:					List[str]					,a_				:					int=3_0_0				)				->	Optional[int]:




					'''simple docstring'''







					return F'\n    <div>\n      {prefix}\n      <progress value=\'{value}\' max=\'{total}\' style=\'width:{width}px; height:20px; vertical-align: middle;\'></progress>\n      {label}\n    </div>\n    '




def 					_lowercase     (				a_				:					str				)				->	Dict:




					'''simple docstring'''







					__magic_name__							=    '<table border="1" class="dataframe">\n'
					html_code += """  <thead>\n <tr style="text-align: left;">\n"""
					for i in items[0]:
										html_code += F'      <th>{i}</th>\n'
					html_code += "    </tr>\n  </thead>\n  <tbody>\n"
					for line in items[1:]:
										html_code += "    <tr>\n"
										for elt in line:
															__magic_name__							=    F'{elt:.6f}' if isinstance(a_					,a_				) else str(a_				)
															html_code += F'      <td>{elt}</td>\n'
										html_code += "    </tr>\n"
					html_code += "  </tbody>\n</table><p>"
					return html_code

class 				__UpperCamelCase				:
	_lowercase  :   Any       				=  5
	_lowercase  :   Dict       				=  0.2







	def __init__(   self: Tuple  , __UpperCamelCase: int  , __UpperCamelCase: Optional[str] = None  , __UpperCamelCase: bool = True  , __UpperCamelCase: Optional["NotebookTrainingTracker"] = None  , __UpperCamelCase: int = 3_00  , ):

						'''simple docstring'''





						__magic_name__							=    total
						__magic_name__							=    '' if prefix is None else prefix
						__magic_name__							=    leave
						__magic_name__							=    parent
						__magic_name__							=    width
						__magic_name__							=    None
						__magic_name__							=    None
						__magic_name__							=    None







	def 	_SCREAMING_SNAKE_CASE							(   self: int  , __UpperCamelCase: int  , __UpperCamelCase: bool = False  , __UpperCamelCase: str = None      ):

						'''simple docstring'''





						__magic_name__							=    value
						if comment is not None:
											__magic_name__							=    comment
						if self.last_value is None:
											__magic_name__							=    __magic_name__							=    time.time()
											__magic_name__							=    __magic_name__							=    value
											__magic_name__							=    __magic_name__							=    None
											__magic_name__							=    self.warmup
											__magic_name__							=    1
											self.update_bar(__UpperCamelCase      )
						elif value <= self.last_value and not force_update:
											return
						elif force_update or self.first_calls > 0 or value >= min(self.last_value + self.wait_for  , self.total      ):
											if self.first_calls > 0:
																self.first_calls -= 1
											__magic_name__							=    time.time()
											__magic_name__							=    current_time - self.start_time
											# We could have value = self.start_value if the update is called twixe with the same start value.
											if value > self.start_value:
																__magic_name__							=    self.elapsed_time / (value - self.start_value)
											else:
																__magic_name__							=    None
											if value >= self.total:
																__magic_name__							=    self.total
																__magic_name__							=    None
																if not self.leave:
																					self.close()
											elif self.average_time_per_item is not None:
																__magic_name__							=    self.average_time_per_item * (self.total - value)
											self.update_bar(__UpperCamelCase      )
											__magic_name__							=    value
											__magic_name__							=    current_time
											if self.average_time_per_item is None:
																__magic_name__							=    1
											else:
																__magic_name__							=    max(int(self.update_every / self.average_time_per_item      )  , 1      )







	def 	_SCREAMING_SNAKE_CASE							(   self: List[str]  , __UpperCamelCase: Dict  , __UpperCamelCase: Optional[Any]=None      ):

						'''simple docstring'''





						__magic_name__							=    ' ' * (len(str(self.total      )      ) - len(str(__UpperCamelCase      )      )) + str(__UpperCamelCase      )
						if self.elapsed_time is None:
											__magic_name__							=    F'[{spaced_value}/{self.total} : < :'
						elif self.predicted_remaining is None:
											__magic_name__							=    F'[{spaced_value}/{self.total} {format_time(self.elapsed_time      )}'
						else:
											__magic_name__							=    (
											    F'[{spaced_value}/{self.total} {format_time(self.elapsed_time      )} <'
											    F' {format_time(self.predicted_remaining      )}'
											)
											self.label += F', {1/self.average_time_per_item:.2f} it/s'
						self.label += "]" if self.comment is None or len(self.comment      ) == 0 else F', {self.comment}]'
						self.display()







	def 	_SCREAMING_SNAKE_CASE							(   self: Dict      ):

						'''simple docstring'''





						__magic_name__							=    html_progress_bar(self.value  , self.total  , self.prefix  , self.label  , self.width      )
						if self.parent is not None:
											# If this is a child bar, the parent will take care of the display.
											self.parent.display()
											return
						if self.output is None:
											__magic_name__							=    disp.display(disp.HTML(self.html_code      )  , display_id=__UpperCamelCase      )
						else:
											self.output.update(disp.HTML(self.html_code      )      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Union[str, Any]      ):

						'''simple docstring'''





						if self.parent is None and self.output is not None:
											self.output.update(disp.HTML(''      )      )




class 				__UpperCamelCase				(					SCREAMING_SNAKE_CASE			):







	def __init__(   self: Tuple  , __UpperCamelCase: Dict  , __UpperCamelCase: Dict=None      ):

						'''simple docstring'''





						super().__init__(__UpperCamelCase      )
						__magic_name__							=    None if column_names is None else [column_names]
						__magic_name__							=    None







	def 	_SCREAMING_SNAKE_CASE							(   self: Any      ):

						'''simple docstring'''





						__magic_name__							=    html_progress_bar(self.value  , self.total  , self.prefix  , self.label  , self.width      )
						if self.inner_table is not None:
											self.html_code += text_to_html_table(self.inner_table      )
						if self.child_bar is not None:
											self.html_code += self.child_bar.html_code
						if self.output is None:
											__magic_name__							=    disp.display(disp.HTML(self.html_code      )  , display_id=__UpperCamelCase      )
						else:
											self.output.update(disp.HTML(self.html_code      )      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Optional[int]  , __UpperCamelCase: Optional[Any]      ):

						'''simple docstring'''





						if self.inner_table is None:
											__magic_name__							=    [list(values.keys()      ), list(values.values()      )]
						else:
											__magic_name__							=    self.inner_table[0]
											if len(self.inner_table      ) == 1:
																# We give a chance to update the column names at the first iteration
																for key in values.keys():
																					if key not in columns:
																										columns.append(__UpperCamelCase      )
																__magic_name__							=    columns
											self.inner_table.append([values[c] for c in columns]      )







	def 	_SCREAMING_SNAKE_CASE							(   self: List[str]  , __UpperCamelCase: Optional[int]  , __UpperCamelCase: List[Any]=None  , __UpperCamelCase: List[str]=3_00      ):

						'''simple docstring'''





						__magic_name__							=    NotebookProgressBar(__UpperCamelCase  , prefix=__UpperCamelCase  , parent=self  , width=__UpperCamelCase      )
						return self.child_bar







	def 	_SCREAMING_SNAKE_CASE							(   self: Union[str, Any]      ):

						'''simple docstring'''





						__magic_name__							=    None
						self.display()




class 				__UpperCamelCase				(					SCREAMING_SNAKE_CASE			):







	def __init__(   self: Optional[Any]      ):

						'''simple docstring'''





						__magic_name__							=    None
						__magic_name__							=    None
						__magic_name__							=    False







	def 	_SCREAMING_SNAKE_CASE							(   self: List[str]  , __UpperCamelCase: int  , __UpperCamelCase: Dict  , __UpperCamelCase: Optional[int]  , **__UpperCamelCase: List[str]      ):

						'''simple docstring'''





						__magic_name__							=    'Epoch' if args.evaluation_strategy == IntervalStrategy.EPOCH else 'Step'
						__magic_name__							=    0
						__magic_name__							=    0
						__magic_name__							=    [self.first_column] + ['Training Loss']
						if args.evaluation_strategy != IntervalStrategy.NO:
											column_names.append('Validation Loss'      )
						__magic_name__							=    NotebookTrainingTracker(state.max_steps  , __UpperCamelCase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: int  , __UpperCamelCase: int  , __UpperCamelCase: List[str]  , __UpperCamelCase: List[Any]  , **__UpperCamelCase: Tuple      ):

						'''simple docstring'''





						__magic_name__							=    int(state.epoch      ) if int(state.epoch      ) == state.epoch else F'{state.epoch:.2f}'
						self.training_tracker.update(
						    state.global_step + 1  , comment=F'Epoch {epoch}/{state.num_train_epochs}'  , force_update=self._force_next_update  , )
						__magic_name__							=    False







	def 	_SCREAMING_SNAKE_CASE							(   self: Union[str, Any]  , __UpperCamelCase: int  , __UpperCamelCase: Optional[int]  , __UpperCamelCase: int  , __UpperCamelCase: Tuple=None  , **__UpperCamelCase: Tuple      ):

						'''simple docstring'''





						if not has_length(__UpperCamelCase      ):
											return
						if self.prediction_bar is None:
											if self.training_tracker is not None:
																__magic_name__							=    self.training_tracker.add_child(len(__UpperCamelCase      )      )
											else:
																__magic_name__							=    NotebookProgressBar(len(__UpperCamelCase      )      )
											self.prediction_bar.update(1      )
						else:
											self.prediction_bar.update(self.prediction_bar.value + 1      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Union[str, Any]  , __UpperCamelCase: Optional[Any]  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: Optional[int]  , **__UpperCamelCase: Dict      ):

						'''simple docstring'''





						if self.prediction_bar is not None:
											self.prediction_bar.close()
						__magic_name__							=    None







	def 	_SCREAMING_SNAKE_CASE							(   self: Union[str, Any]  , __UpperCamelCase: Tuple  , __UpperCamelCase: str  , __UpperCamelCase: int  , __UpperCamelCase: List[Any]=None  , **__UpperCamelCase: Optional[Any]      ):

						'''simple docstring'''





						if args.evaluation_strategy == IntervalStrategy.NO and "loss" in logs:
											__magic_name__							=    {'Training Loss': logs['loss']}
											# First column is necessarily Step sine we're not in epoch eval strategy
											__magic_name__							=    state.global_step
											self.training_tracker.write_line(__UpperCamelCase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Optional[int]  , __UpperCamelCase: Tuple  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: Optional[int]  , __UpperCamelCase: Union[str, Any]=None  , **__UpperCamelCase: Optional[int]      ):

						'''simple docstring'''





						if self.training_tracker is not None:
											__magic_name__							=    {'Training Loss': 'No log', 'Validation Loss': 'No log'}
											for log in reversed(state.log_history      ):
																if "loss" in log:
																					__magic_name__							=    log['loss']
																					break

											if self.first_column == "Epoch":
																__magic_name__							=    int(state.epoch      )
											else:
																__magic_name__							=    state.global_step
											__magic_name__							=    'eval'
											for k in metrics:
																if k.endswith('_loss'      ):
																					__magic_name__							=    re.sub(r'\_loss$'  , ''  , __UpperCamelCase      )
											__magic_name__							=    metrics.pop('total_flos'  , __UpperCamelCase      )
											__magic_name__							=    metrics.pop('epoch'  , __UpperCamelCase      )
											__magic_name__							=    metrics.pop(F'{metric_key_prefix}_runtime'  , __UpperCamelCase      )
											__magic_name__							=    metrics.pop(F'{metric_key_prefix}_samples_per_second'  , __UpperCamelCase      )
											__magic_name__							=    metrics.pop(F'{metric_key_prefix}_steps_per_second'  , __UpperCamelCase      )
											__magic_name__							=    metrics.pop(F'{metric_key_prefix}_jit_compilation_time'  , __UpperCamelCase      )
											for k, v in metrics.items():
																if k == F'{metric_key_prefix}_loss':
																					__magic_name__							=    v
																else:
																					__magic_name__							=    k.split('_'      )
																					__magic_name__							=    ' '.join([part.capitalize() for part in splits[1:]]      )
																					__magic_name__							=    v
											self.training_tracker.write_line(__UpperCamelCase      )
											self.training_tracker.remove_child()
											__magic_name__							=    None
											# Evaluation takes a long time so we should force the next update.
											__magic_name__							=    True







	def 	_SCREAMING_SNAKE_CASE							(   self: Dict  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: int  , __UpperCamelCase: List[str]  , **__UpperCamelCase: Union[str, Any]      ):

						'''simple docstring'''





						self.training_tracker.update(
						    state.global_step  , comment=F'Epoch {int(state.epoch      )}/{state.num_train_epochs}'  , force_update=__UpperCamelCase      )
						__magic_name__							=    None







import unittest

from transformers import AutoConfig, AutoTokenizer, BertConfig, TensorType, is_flax_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, slow


if is_flax_available():
						import jax

						from transformers.models.auto.modeling_flax_auto import FlaxAutoModel
						from transformers.models.bert.modeling_flax_bert import FlaxBertModel
						from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel

@require_flax
class 				__UpperCamelCase				(					unittest.TestCase			):







	@slow
	def 	_SCREAMING_SNAKE_CASE							(   self: Any      ):

						'''simple docstring'''





						for model_name in ["bert-base-cased", "bert-large-uncased"]:
											with self.subTest(__UpperCamelCase      ):
																__magic_name__							=    AutoConfig.from_pretrained(__UpperCamelCase      )
																self.assertIsNotNone(__UpperCamelCase      )
																self.assertIsInstance(__UpperCamelCase  , __UpperCamelCase      )

																__magic_name__							=    FlaxAutoModel.from_pretrained(__UpperCamelCase      )
																self.assertIsNotNone(__UpperCamelCase      )
																self.assertIsInstance(__UpperCamelCase  , __UpperCamelCase      )







	@slow
	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						for model_name in ["roberta-base", "roberta-large"]:
											with self.subTest(__UpperCamelCase      ):
																__magic_name__							=    AutoConfig.from_pretrained(__UpperCamelCase      )
																self.assertIsNotNone(__UpperCamelCase      )
																self.assertIsInstance(__UpperCamelCase  , __UpperCamelCase      )

																__magic_name__							=    FlaxAutoModel.from_pretrained(__UpperCamelCase      )
																self.assertIsNotNone(__UpperCamelCase      )
																self.assertIsInstance(__UpperCamelCase  , __UpperCamelCase      )







	@slow
	def 	_SCREAMING_SNAKE_CASE							(   self: Dict      ):

						'''simple docstring'''





						for model_name in ["bert-base-cased", "bert-large-uncased"]:
											__magic_name__							=    AutoTokenizer.from_pretrained(__UpperCamelCase      )
											__magic_name__							=    FlaxBertModel.from_pretrained(__UpperCamelCase      )
											__magic_name__							=    tokenizer('Do you support jax jitted function?'  , return_tensors=TensorType.JAX      )

											@jax.jit
											def eval(**__UpperCamelCase: Optional[Any]      ):
																return model(**__UpperCamelCase      )

											eval(**__UpperCamelCase      ).block_until_ready()







	@slow
	def 	_SCREAMING_SNAKE_CASE							(   self: int      ):

						'''simple docstring'''





						for model_name in ["roberta-base", "roberta-large"]:
											__magic_name__							=    AutoTokenizer.from_pretrained(__UpperCamelCase      )
											__magic_name__							=    FlaxRobertaModel.from_pretrained(__UpperCamelCase      )
											__magic_name__							=    tokenizer('Do you support jax jitted function?'  , return_tensors=TensorType.JAX      )

											@jax.jit
											def eval(**__UpperCamelCase: Any      ):
																return model(**__UpperCamelCase      )

											eval(**__UpperCamelCase      ).block_until_ready()







	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						with self.assertRaisesRegex(
						    __UpperCamelCase  , 'bert-base is not a local folder and is not a valid model identifier'      ):
											__magic_name__							=    FlaxAutoModel.from_pretrained('bert-base'      )







	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





						with self.assertRaisesRegex(
						    __UpperCamelCase  , r'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'      ):
											__magic_name__							=    FlaxAutoModel.from_pretrained(__UpperCamelCase  , revision='aaaaaa'      )







	def 	_SCREAMING_SNAKE_CASE							(   self: str      ):

						'''simple docstring'''





						with self.assertRaisesRegex(
						    __UpperCamelCase  , 'hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack'  , ):
											__magic_name__							=    FlaxAutoModel.from_pretrained('hf-internal-testing/config-no-model'      )







	def 	_SCREAMING_SNAKE_CASE							(   self: int      ):

						'''simple docstring'''





						with self.assertRaisesRegex(__UpperCamelCase  , 'Use `from_pt=True` to load this model'      ):
											__magic_name__							=    FlaxAutoModel.from_pretrained('hf-internal-testing/tiny-bert-pt-only'      )




