"""simple docstring"""







import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import (
						    MODEL_FOR_PRETRAINING_MAPPING,
						    MegatronBertForCausalLM,
						    MegatronBertForMaskedLM,
						    MegatronBertForMultipleChoice,
						    MegatronBertForNextSentencePrediction,
						    MegatronBertForPreTraining,
						    MegatronBertForQuestionAnswering,
						    MegatronBertForSequenceClassification,
						    MegatronBertForTokenClassification,
						    MegatronBertModel,
						)



class  _lowerCamelCase			:







				def __init__( self    :       List[Any]			,					lowerCamelCase_    :       List[Any]			,					lowerCamelCase_    :       Optional[Any]=1_3			,					lowerCamelCase_    :       Optional[int]=7			,					lowerCamelCase_    :       Tuple=True			,					lowerCamelCase_    :       Union[str, Any]=True			,					lowerCamelCase_    :       List[Any]=True			,					lowerCamelCase_    :       str=True			,					lowerCamelCase_    :       Tuple=9_9			,					lowerCamelCase_    :       Union[str, Any]=6_4			,					lowerCamelCase_    :       Tuple=3_2			,					lowerCamelCase_    :       Optional[int]=5			,					lowerCamelCase_    :       int=4			,					lowerCamelCase_    :       Optional[int]=3_7			,					lowerCamelCase_    :       str="gelu"			,					lowerCamelCase_    :       List[str]=0.1			,					lowerCamelCase_    :       str=0.1			,					lowerCamelCase_    :       List[Any]=5_1_2			,					lowerCamelCase_    :       str=1_6			,					lowerCamelCase_    :       int=2			,					lowerCamelCase_    :       List[str]=0.02			,					lowerCamelCase_    :       List[Any]=3			,					lowerCamelCase_    :       Optional[int]=4			,					lowerCamelCase_    :       Tuple=None			,					):

					"""simple docstring"""






					_lowercase :    str							=							parent
					_lowercase :    Optional[Any]							=							batch_size
					_lowercase :    Dict							=							seq_length
					_lowercase :    int							=							is_training
					_lowercase :    Tuple							=							use_input_mask
					_lowercase :    Union[str, Any]							=							use_token_type_ids
					_lowercase :    int							=							use_labels
					_lowercase :    Union[str, Any]							=							vocab_size
					_lowercase :    Dict							=							hidden_size
					_lowercase :    int							=							embedding_size
					_lowercase :    Optional[int]							=							num_hidden_layers
					_lowercase :    List[Any]							=							num_attention_heads
					_lowercase :    Dict							=							intermediate_size
					_lowercase :    List[str]							=							hidden_act
					_lowercase :    List[str]							=							hidden_dropout_prob
					_lowercase :    List[str]							=							attention_probs_dropout_prob
					_lowercase :    List[str]							=							max_position_embeddings
					_lowercase :    List[str]							=							type_vocab_size
					_lowercase :    List[str]							=							type_sequence_label_size
					_lowercase :    List[str]							=							initializer_range
					_lowercase :    Optional[Any]							=							num_labels
					_lowercase :    str							=							num_choices
					_lowercase :    List[Any]							=							scope







				def 							__UpperCAmelCase					( self    :       int						):

					"""simple docstring"""






					_lowercase :    Any							=							ids_tensor([self.batch_size, self.seq_length]			,					self.vocab_size						)

					_lowercase :    Optional[Any]							=							None
					if self.use_input_mask:
						_lowercase :    Dict							=							random_attention_mask([self.batch_size, self.seq_length]						)

					_lowercase :    int							=							None
					if self.use_token_type_ids:
						_lowercase :    List[Any]							=							ids_tensor([self.batch_size, self.seq_length]			,					self.type_vocab_size						)

					_lowercase :    Optional[Any]							=							None
					_lowercase :    Any							=							None
					_lowercase :    Optional[int]							=							None
					if self.use_labels:
						_lowercase :    int							=							ids_tensor([self.batch_size]			,					self.type_sequence_label_size						)
						_lowercase :    List[Any]							=							ids_tensor([self.batch_size, self.seq_length]			,					self.num_labels						)
						_lowercase :    Dict							=							ids_tensor([self.batch_size]			,					self.num_choices						)

					_lowercase :    Optional[Any]							=							self.get_config()

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







				def 							__UpperCAmelCase					( self    :       int						):

					"""simple docstring"""






					return MegatronBertConfig(
					    vocab_size=self.vocab_size			,					hidden_size=self.hidden_size			,					num_hidden_layers=self.num_hidden_layers			,					num_attention_heads=self.num_attention_heads			,					intermediate_size=self.intermediate_size			,					embedding_size=self.embedding_size			,					hidden_act=self.hidden_act			,					hidden_dropout_prob=self.hidden_dropout_prob			,					attention_probs_dropout_prob=self.attention_probs_dropout_prob			,					max_position_embeddings=self.max_position_embeddings			,					type_vocab_size=self.type_vocab_size			,					is_decoder=__A			,					initializer_range=self.initializer_range			,					)







				def 							__UpperCAmelCase					( self    :       List[str]			,					lowerCamelCase_    :       Optional[int]			,					lowerCamelCase_    :       List[Any]			,					lowerCamelCase_    :       List[Any]			,					lowerCamelCase_    :       Any			,					lowerCamelCase_    :       List[Any]			,					lowerCamelCase_    :       Any			,					lowerCamelCase_    :       Union[str, Any]						):

					"""simple docstring"""






					_lowercase :    Optional[Any]							=							MegatronBertModel(config=__A						)
					model.to(__A						)
					model.eval()
					_lowercase :    Union[str, Any]							=							model(__A			,					attention_mask=__A			,					token_type_ids=__A						)
					_lowercase :    Any							=							model(__A			,					token_type_ids=__A						)
					_lowercase :    Dict							=							model(__A						)

					self.parent.assertEqual(result.last_hidden_state.shape			,					(self.batch_size, self.seq_length, self.hidden_size)						)
					self.parent.assertEqual(result.pooler_output.shape			,					(self.batch_size, self.hidden_size)						)







				def 							__UpperCAmelCase					( self    :       Any			,					lowerCamelCase_    :       Union[str, Any]			,					lowerCamelCase_    :       List[str]			,					lowerCamelCase_    :       Union[str, Any]			,					lowerCamelCase_    :       Any			,					lowerCamelCase_    :       Tuple			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       int						):

					"""simple docstring"""






					_lowercase :    Tuple							=							MegatronBertForMaskedLM(config=__A						)
					model.to(__A						)
					model.eval()
					_lowercase :    Optional[int]							=							model(__A			,					attention_mask=__A			,					token_type_ids=__A			,					labels=__A						)
					self.parent.assertEqual(result.logits.shape			,					(self.batch_size, self.seq_length, self.vocab_size)						)







				def 							__UpperCAmelCase					( self    :       str			,					lowerCamelCase_    :       Tuple			,					lowerCamelCase_    :       List[Any]			,					lowerCamelCase_    :       Any			,					lowerCamelCase_    :       Optional[Any]			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       Union[str, Any]			,					lowerCamelCase_    :       Union[str, Any]						):

					"""simple docstring"""






					_lowercase :    Any							=							MegatronBertForCausalLM(config=__A						)
					model.to(__A						)
					model.eval()
					_lowercase :    Optional[Any]							=							model(__A			,					attention_mask=__A			,					token_type_ids=__A			,					labels=__A						)
					self.parent.assertEqual(result.logits.shape			,					(self.batch_size, self.seq_length, self.vocab_size)						)







				def 							__UpperCAmelCase					( self    :       Optional[int]			,					lowerCamelCase_    :       Optional[Any]			,					lowerCamelCase_    :       Dict			,					lowerCamelCase_    :       List[str]			,					lowerCamelCase_    :       List[str]			,					lowerCamelCase_    :       List[Any]			,					lowerCamelCase_    :       Optional[int]			,					lowerCamelCase_    :       str						):

					"""simple docstring"""






					_lowercase :    str							=							MegatronBertForNextSentencePrediction(config=__A						)
					model.to(__A						)
					model.eval()
					_lowercase :    Optional[int]							=							model(
					    __A			,					attention_mask=__A			,					token_type_ids=__A			,					labels=__A			,					)
					self.parent.assertEqual(result.logits.shape			,					(self.batch_size, 2)						)







				def 							__UpperCAmelCase					( self    :       Optional[int]			,					lowerCamelCase_    :       Optional[Any]			,					lowerCamelCase_    :       str			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       List[str]			,					lowerCamelCase_    :       Optional[Any]			,					lowerCamelCase_    :       List[Any]			,					lowerCamelCase_    :       List[str]						):

					"""simple docstring"""






					_lowercase :    Dict							=							MegatronBertForPreTraining(config=__A						)
					model.to(__A						)
					model.eval()
					_lowercase :    Union[str, Any]							=							model(
					    __A			,					attention_mask=__A			,					token_type_ids=__A			,					labels=__A			,					next_sentence_label=__A			,					)
					self.parent.assertEqual(result.prediction_logits.shape			,					(self.batch_size, self.seq_length, self.vocab_size)						)
					self.parent.assertEqual(result.seq_relationship_logits.shape			,					(self.batch_size, 2)						)







				def 							__UpperCAmelCase					( self    :       Dict			,					lowerCamelCase_    :       List[Any]			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       Any			,					lowerCamelCase_    :       Optional[int]			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       Union[str, Any]			,					lowerCamelCase_    :       int						):

					"""simple docstring"""






					_lowercase :    Union[str, Any]							=							MegatronBertForQuestionAnswering(config=__A						)
					model.to(__A						)
					model.eval()
					_lowercase :    Dict							=							model(
					    __A			,					attention_mask=__A			,					token_type_ids=__A			,					start_positions=__A			,					end_positions=__A			,					)
					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    :       str			,					lowerCamelCase_    :       Union[str, Any]			,					lowerCamelCase_    :       Any			,					lowerCamelCase_    :       Optional[Any]			,					lowerCamelCase_    :       Union[str, Any]			,					lowerCamelCase_    :       str			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       Union[str, Any]						):

					"""simple docstring"""






					_lowercase :    Dict							=							self.num_labels
					_lowercase :    List[str]							=							MegatronBertForSequenceClassification(__A						)
					model.to(__A						)
					model.eval()
					_lowercase :    str							=							model(__A			,					attention_mask=__A			,					token_type_ids=__A			,					labels=__A						)
					self.parent.assertEqual(result.logits.shape			,					(self.batch_size, self.num_labels)						)







				def 							__UpperCAmelCase					( self    :       List[Any]			,					lowerCamelCase_    :       Tuple			,					lowerCamelCase_    :       Union[str, Any]			,					lowerCamelCase_    :       Optional[Any]			,					lowerCamelCase_    :       List[Any]			,					lowerCamelCase_    :       Union[str, Any]			,					lowerCamelCase_    :       List[str]			,					lowerCamelCase_    :       List[str]						):

					"""simple docstring"""






					_lowercase :    List[Any]							=							self.num_labels
					_lowercase :    Optional[int]							=							MegatronBertForTokenClassification(config=__A						)
					model.to(__A						)
					model.eval()
					_lowercase :    Optional[int]							=							model(__A			,					attention_mask=__A			,					token_type_ids=__A			,					labels=__A						)
					self.parent.assertEqual(result.logits.shape			,					(self.batch_size, self.seq_length, self.num_labels)						)







				def 							__UpperCAmelCase					( self    :       Dict			,					lowerCamelCase_    :       str			,					lowerCamelCase_    :       List[str]			,					lowerCamelCase_    :       Dict			,					lowerCamelCase_    :       Union[str, Any]			,					lowerCamelCase_    :       List[str]			,					lowerCamelCase_    :       List[str]			,					lowerCamelCase_    :       Dict						):

					"""simple docstring"""






					_lowercase :    int							=							self.num_choices
					_lowercase :    str							=							MegatronBertForMultipleChoice(config=__A						)
					model.to(__A						)
					model.eval()
					_lowercase :    Optional[int]							=							input_ids.unsqueeze(1						).expand(-1			,					self.num_choices			,					-1						).contiguous()
					_lowercase :    List[Any]							=							token_type_ids.unsqueeze(1						).expand(-1			,					self.num_choices			,					-1						).contiguous()
					_lowercase :    Tuple							=							input_mask.unsqueeze(1						).expand(-1			,					self.num_choices			,					-1						).contiguous()
					_lowercase :    Optional[Any]							=							model(
					    __A			,					attention_mask=__A			,					token_type_ids=__A			,					labels=__A			,					)
					self.parent.assertEqual(result.logits.shape			,					(self.batch_size, self.num_choices)						)







				def 							__UpperCAmelCase					( self    :       Dict						):

					"""simple docstring"""






					_lowercase :    Union[str, Any]							=							self.prepare_config_and_inputs()
					(
					    (
					    _lowercase						
					)				,      (
					    _lowercase						
					)				,      (
					    _lowercase						
					)				,      (
					    _lowercase						
					)				,      (
					    _lowercase						
					)				,      (
					    _lowercase						
					)				,      (
					    _lowercase						
					)				,      						
					) :    int							=							config_and_inputs
					_lowercase :    int							=							{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
					return config, inputs_dict



@require_torch
class  _lowerCamelCase			(UpperCamelCase__						,							UpperCamelCase__						,							unittest.TestCase     ):
				_snake_case     =  (
				    (
				        MegatronBertModel,
				        MegatronBertForMaskedLM,
				        MegatronBertForCausalLM,
				        MegatronBertForMultipleChoice,
				        MegatronBertForNextSentencePrediction,
				        MegatronBertForPreTraining,
				        MegatronBertForQuestionAnswering,
				        MegatronBertForSequenceClassification,
				        MegatronBertForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				_snake_case     =  (
				    {
				        "feature-extraction": MegatronBertModel,
				        "fill-mask": MegatronBertForMaskedLM,
				        "question-answering": MegatronBertForQuestionAnswering,
				        "text-classification": MegatronBertForSequenceClassification,
				        "text-generation": MegatronBertForCausalLM,
				        "token-classification": MegatronBertForTokenClassification,
				        "zero-shot": MegatronBertForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				_snake_case     =  True
				# test_resize_embeddings = False
				_snake_case     =  False







				def 							__UpperCAmelCase					( self    :       Optional[Any]			,					lowerCamelCase_    :       Union[str, Any]			,					lowerCamelCase_    :       List[Any]			,					lowerCamelCase_    :       Optional[int]=False						):

					"""simple docstring"""






					_lowercase :    List[str]							=							super()._prepare_for_class(__A			,					__A			,					return_labels=__A						)

					if return_labels:
						if model_class in get_values(__A						):
							_lowercase :    Any							=							torch.zeros(
							    (self.model_tester.batch_size, self.model_tester.seq_length)			,					dtype=torch.long			,					device=__A						)
							_lowercase :    List[str]							=							torch.zeros(
							    self.model_tester.batch_size			,					dtype=torch.long			,					device=__A						)
					return inputs_dict







				def 							__UpperCAmelCase					( self    :       Optional[int]						):

					"""simple docstring"""






					_lowercase :    Tuple							=							MegatronBertModelTester(self						)
					_lowercase :    Any							=							ConfigTester(self			,					config_class=__A			,					hidden_size=3_7						)







				def 							__UpperCAmelCase					( self    :       Optional[Any]						):

					"""simple docstring"""






					self.config_tester.run_common_tests()







				def 							__UpperCAmelCase					( self    :       Tuple						):

					"""simple docstring"""






					_lowercase :    Union[str, Any]							=							self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_model(*__A						)







				def 							__UpperCAmelCase					( self    :       List[Any]						):

					"""simple docstring"""






					_lowercase :    Optional[int]							=							self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__A						)







				def 							__UpperCAmelCase					( self    :       Optional[Any]						):

					"""simple docstring"""






					_lowercase :    Any							=							self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__A						)







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

					"""simple docstring"""






					_lowercase :    List[str]							=							self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__A						)







				def 							__UpperCAmelCase					( self    :       List[Any]						):

					"""simple docstring"""






					_lowercase :    Optional[int]							=							self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_for_pretraining(*__A						)







				def 							__UpperCAmelCase					( self    :       Optional[Any]						):

					"""simple docstring"""






					_lowercase :    int							=							self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_for_question_answering(*__A						)







				def 							__UpperCAmelCase					( self    :       Optional[int]						):

					"""simple docstring"""






					_lowercase :    List[Any]							=							self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__A						)







				def 							__UpperCAmelCase					( self    :       Dict						):

					"""simple docstring"""






					_lowercase :    Dict							=							self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_for_token_classification(*__A						)



def      __lowerCAmelCase(		__UpperCAmelCase			):






	"""simple docstring"""





	return torch.tensor(
	    UpperCamelCase__     ,dtype=torch.long     ,device=UpperCamelCase__     ,)


SCREAMING_SNAKE_CASE									=						1E-4







@require_torch
@require_sentencepiece
@require_tokenizers
class  _lowerCamelCase			(unittest.TestCase     ):







				@slow
				@unittest.skip('Model is not available.'						)
				def 							__UpperCAmelCase					( self    :       str						):

					"""simple docstring"""






					_lowercase :    Optional[int]							=							'nvidia/megatron-bert-uncased-345m'
					if "MYDIR" in os.environ:
						_lowercase :    Dict							=							os.path.join(os.environ['MYDIR']			,					__A						)
					_lowercase :    Optional[Any]							=							MegatronBertModel.from_pretrained(__A						)
					model.to(__A						)
					model.half()
					_lowercase :    Union[str, Any]							=							_long_tensor([[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]]						)
					with torch.no_grad():
						_lowercase :    Optional[Any]							=							model(__A						)[0]
					_lowercase :    Optional[int]							=							torch.Size((1, 9, 1_0_2_4)						)
					self.assertEqual(output.shape			,					__A						)

					_lowercase :    str							=							[-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
					for ii in range(3						):
						for jj in range(3						):
							_lowercase :    int							=							output[0, ii, jj]
							_lowercase :    Any							=							expected[3 * ii + jj]
							_lowercase :    Union[str, Any]							=							'ii={} jj={} a={} b={}'.format(__A			,					__A			,					__A			,					__A						)
							self.assertTrue(math.isclose(__A			,					__A			,					rel_tol=__A			,					abs_tol=__A						)			,					msg=__A						)




"""simple docstring"""







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      __lowerCAmelCase(		__UpperCAmelCase			):






	"""simple docstring"""





	_lowercase :    Optional[int]							=							int(__UpperCAmelCase			)
	_lowercase				,      _lowercase				,      _lowercase :    Union[str, Any]							=							t // 3_600, (t // 60) % 60, t % 60
	return F'''{h}:{m:02d}:{s:02d}''' if h != 0 else F'''{m:02d}:{s:02d}'''


def      __lowerCAmelCase(		__UpperCAmelCase     ,__UpperCAmelCase     ,__UpperCAmelCase     ,__UpperCAmelCase     ,__UpperCAmelCase=300			):






	"""simple docstring"""





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


def      __lowerCAmelCase(		__UpperCAmelCase			):






	"""simple docstring"""





	_lowercase :    List[Any]							=							'<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:
			_lowercase :    Optional[Any]							=							F'''{elt:.6f}''' if isinstance(__UpperCAmelCase     ,__UpperCAmelCase			) else str(__UpperCAmelCase			)
			html_code += F'''      <td>{elt}</td>\n'''
		html_code += "    </tr>\n"
	html_code += "  </tbody>\n</table><p>"
	return html_code



class  _lowerCamelCase			:
				_snake_case     =  5
				_snake_case     =  0.2







				def __init__( self    :       List[str]			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       Optional[str] = None			,					lowerCamelCase_    :       bool = True			,					lowerCamelCase_    :       Optional["NotebookTrainingTracker"] = None			,					lowerCamelCase_    :       int = 3_0_0			,					):

					"""simple docstring"""






					_lowercase :    int							=							total
					_lowercase :    Dict							=							'' if prefix is None else prefix
					_lowercase :    List[Any]							=							leave
					_lowercase :    Dict							=							parent
					_lowercase :    str							=							width
					_lowercase :    Any							=							None
					_lowercase :    Optional[int]							=							None
					_lowercase :    Union[str, Any]							=							None







				def 							__UpperCAmelCase					( self    :       Dict			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       bool = False			,					lowerCamelCase_    :       str = None						):

					"""simple docstring"""






					_lowercase :    Dict							=							value
					if comment is not None:
						_lowercase :    List[str]							=							comment
					if self.last_value is None:
						_lowercase :    List[Any]							=							time.time()
						_lowercase :    str							=							value
						_lowercase :    Any							=							None
						_lowercase :    List[Any]							=							self.warmup
						_lowercase :    Union[str, Any]							=							1
						self.update_bar(lowerCamelCase_						)
					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
						_lowercase :    Tuple							=							time.time()
						_lowercase :    List[str]							=							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:
							_lowercase :    List[Any]							=							self.elapsed_time / (value - self.start_value)
						else:
							_lowercase :    List[Any]							=							None
						if value >= self.total:
							_lowercase :    Optional[Any]							=							self.total
							_lowercase :    str							=							None
							if not self.leave:
								self.close()
						elif self.average_time_per_item is not None:
							_lowercase :    int							=							self.average_time_per_item * (self.total - value)
						self.update_bar(lowerCamelCase_						)
						_lowercase :    Optional[int]							=							value
						_lowercase :    Dict							=							current_time
						if self.average_time_per_item is None:
							_lowercase :    Union[str, Any]							=							1
						else:
							_lowercase :    Optional[Any]							=							max(int(self.update_every / self.average_time_per_item						)			,					1						)







				def 							__UpperCAmelCase					( self    :       str			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       Optional[Any]=None						):

					"""simple docstring"""






					_lowercase :    str							=							' ' * (len(str(self.total						)						) - len(str(lowerCamelCase_						)						)) + str(lowerCamelCase_						)
					if self.elapsed_time is None:
						_lowercase :    Dict							=							F'''[{spaced_value}/{self.total} : < :'''
					elif self.predicted_remaining is None:
						_lowercase :    Optional[int]							=							F'''[{spaced_value}/{self.total} {format_time(self.elapsed_time						)}'''
					else:
						_lowercase :    int							=							(
						    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 							__UpperCAmelCase					( self    :       int						):

					"""simple docstring"""






					_lowercase :    int							=							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:
						_lowercase :    List[str]							=							disp.display(disp.HTML(self.html_code						)			,					display_id=lowerCamelCase_						)
					else:
						self.output.update(disp.HTML(self.html_code						)						)







				def 							__UpperCAmelCase					( self    :       Any						):

					"""simple docstring"""






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



class  _lowerCamelCase			(__lowerCamelCase     ):







				def __init__( self    :       Optional[int]			,					lowerCamelCase_    :       Any			,					lowerCamelCase_    :       Tuple=None						):

					"""simple docstring"""






					super().__init__(lowerCamelCase_						)
					_lowercase :    Tuple							=							None if column_names is None else [column_names]
					_lowercase :    Optional[int]							=							None







				def 							__UpperCAmelCase					( self    :       Optional[Any]						):

					"""simple docstring"""






					_lowercase :    Union[str, Any]							=							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:
						_lowercase :    Dict							=							disp.display(disp.HTML(self.html_code						)			,					display_id=lowerCamelCase_						)
					else:
						self.output.update(disp.HTML(self.html_code						)						)







				def 							__UpperCAmelCase					( self    :       Optional[Any]			,					lowerCamelCase_    :       Optional[int]						):

					"""simple docstring"""






					if self.inner_table is None:
						_lowercase :    Optional[int]							=							[list(values.keys()						), list(values.values()						)]
					else:
						_lowercase :    str							=							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(lowerCamelCase_						)
							_lowercase :    List[Any]							=							columns
						self.inner_table.append([values[c] for c in columns]						)







				def 							__UpperCAmelCase					( self    :       List[Any]			,					lowerCamelCase_    :       List[Any]			,					lowerCamelCase_    :       Union[str, Any]=None			,					lowerCamelCase_    :       int=3_0_0						):

					"""simple docstring"""






					_lowercase :    List[str]							=							NotebookProgressBar(lowerCamelCase_			,					prefix=lowerCamelCase_			,					parent=self			,					width=lowerCamelCase_						)
					return self.child_bar







				def 							__UpperCAmelCase					( self    :       Any						):

					"""simple docstring"""






					_lowercase :    Any							=							None
					self.display()



class  _lowerCamelCase			(__lowerCamelCase     ):







				def __init__( self    :       Optional[Any]						):

					"""simple docstring"""






					_lowercase :    int							=							None
					_lowercase :    Any							=							None
					_lowercase :    Optional[int]							=							False







				def 							__UpperCAmelCase					( self    :       Any			,					lowerCamelCase_    :       Optional[Any]			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       List[Any]			,					**lowerCamelCase_    :       Union[str, Any]						):

					"""simple docstring"""






					_lowercase :    Any							=							'Epoch' if args.evaluation_strategy == IntervalStrategy.EPOCH else 'Step'
					_lowercase :    Dict							=							0
					_lowercase :    Optional[Any]							=							0
					_lowercase :    Union[str, Any]							=							[self.first_column] + ['Training Loss']
					if args.evaluation_strategy != IntervalStrategy.NO:
						column_names.append('Validation Loss'						)
					_lowercase :    Tuple							=							NotebookTrainingTracker(state.max_steps			,					lowerCamelCase_						)







				def 							__UpperCAmelCase					( self    :       Union[str, Any]			,					lowerCamelCase_    :       Any			,					lowerCamelCase_    :       Optional[Any]			,					lowerCamelCase_    :       Union[str, Any]			,					**lowerCamelCase_    :       Dict						):

					"""simple docstring"""






					_lowercase :    Optional[Any]							=							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			,					)
					_lowercase :    Tuple							=							False







				def 							__UpperCAmelCase					( self    :       int			,					lowerCamelCase_    :       Dict			,					lowerCamelCase_    :       Optional[int]			,					lowerCamelCase_    :       Optional[int]			,					lowerCamelCase_    :       Union[str, Any]=None			,					**lowerCamelCase_    :       str						):

					"""simple docstring"""






					if not has_length(lowerCamelCase_						):
						return
					if self.prediction_bar is None:
						if self.training_tracker is not None:
							_lowercase :    Tuple							=							self.training_tracker.add_child(len(lowerCamelCase_						)						)
						else:
							_lowercase :    Optional[Any]							=							NotebookProgressBar(len(lowerCamelCase_						)						)
						self.prediction_bar.update(1						)
					else:
						self.prediction_bar.update(self.prediction_bar.value + 1						)







				def 							__UpperCAmelCase					( self    :       Tuple			,					lowerCamelCase_    :       Dict			,					lowerCamelCase_    :       Optional[Any]			,					lowerCamelCase_    :       Dict			,					**lowerCamelCase_    :       List[str]						):

					"""simple docstring"""






					if self.prediction_bar is not None:
						self.prediction_bar.close()
					_lowercase :    Union[str, Any]							=							None







				def 							__UpperCAmelCase					( self    :       List[str]			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       str			,					lowerCamelCase_    :       int=None			,					**lowerCamelCase_    :       List[Any]						):

					"""simple docstring"""






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







				def 							__UpperCAmelCase					( self    :       int			,					lowerCamelCase_    :       List[Any]			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       List[str]			,					lowerCamelCase_    :       str=None			,					**lowerCamelCase_    :       List[Any]						):

					"""simple docstring"""






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

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







				def 							__UpperCAmelCase					( self    :       Optional[int]			,					lowerCamelCase_    :       Tuple			,					lowerCamelCase_    :       Optional[Any]			,					lowerCamelCase_    :       Union[str, Any]			,					**lowerCamelCase_    :       Any						):

					"""simple docstring"""






					self.training_tracker.update(
					    state.global_step			,					comment=F'''Epoch {int(state.epoch						)}/{state.num_train_epochs}'''			,					force_update=lowerCamelCase_						)
					_lowercase :    List[Any]							=							None



