



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		(  __lowerCamelCase	:    int	)  ->						Any:
     __lowerCAmelCase	   =int(_lowerCAmelCase	)
     __lowerCAmelCase   ,	__lowerCAmelCase   ,	__lowerCAmelCase	   =t // 3600, (t // 60) % 60, t % 60
     return f"""{h}:{m:02d}:{s:02d}""" if h != 0 else f"""{m:02d}:{s:02d}"""






def 					__lowerCAmelCase		(  __lowerCamelCase	:    str				,   __lowerCamelCase	:    Any				,   __lowerCamelCase	:    List[Any]				,   __lowerCamelCase	:    List[str]				,   __lowerCamelCase	:    Union[str, Any]=300	)  ->						Any:
     # docstyle-ignore
     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 					__lowerCAmelCase		(  __lowerCamelCase	:    List[str]	)  ->						List[str]:
     __lowerCAmelCase	   ="""<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:
               __lowerCAmelCase	   =f"""{elt:.6f}""" if isinstance(_lowerCAmelCase				,   _lowerCAmelCase	) else str(_lowerCAmelCase	)
               html_code += f"""      <td>{elt}</td>\n"""
          html_code += "    </tr>\n"
     html_code += "  </tbody>\n</table><p>"
     return html_code






class 						__a :
 SCREAMING_SNAKE_CASE       =       5
 SCREAMING_SNAKE_CASE       =       0.2



 def __init__(						self							:    str				,	snake_case_							:    int				,	snake_case_							:    Optional[str] = None				,	snake_case_							:    bool = True				,	snake_case_							:    Optional["NotebookTrainingTracker"] = None				,	snake_case_							:    int = 3_00				,	)->							str:
      __lowerCAmelCase	   =total
      __lowerCAmelCase	   ="""""" if prefix is None else prefix
      __lowerCAmelCase	   =leave
      __lowerCAmelCase	   =parent
      __lowerCAmelCase	   =width
      __lowerCAmelCase	   =None
      __lowerCAmelCase	   =None
      __lowerCAmelCase	   =None



 def 		UpperCamelCase							(						self							:    List[Any]				,	snake_case_							:    int				,	snake_case_							:    bool = False				,	snake_case_							:    str = None)->							Union[str, Any]:
      __lowerCAmelCase	   =value
      if comment is not None:
           __lowerCAmelCase	   =comment
      if self.last_value is None:
           __lowerCAmelCase	   =__lowerCAmelCase	   =time.time()
           __lowerCAmelCase	   =__lowerCAmelCase	   =value
           __lowerCAmelCase	   =__lowerCAmelCase	   =None
           __lowerCAmelCase	   =self.warmup
           __lowerCAmelCase	   =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
           __lowerCAmelCase	   =time.time()
           __lowerCAmelCase	   =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:
                __lowerCAmelCase	   =self.elapsed_time / (value - self.start_value)
           else:
                __lowerCAmelCase	   =None
           if value >= self.total:
                __lowerCAmelCase	   =self.total
                __lowerCAmelCase	   =None
                if not self.leave:
                     self.close()
           elif self.average_time_per_item is not None:
                __lowerCAmelCase	   =self.average_time_per_item * (self.total - value)
           self.update_bar(_lowerCamelCase)
           __lowerCAmelCase	   =value
           __lowerCAmelCase	   =current_time
           if self.average_time_per_item is None:
                __lowerCAmelCase	   =1
           else:
                __lowerCAmelCase	   =max(int(self.update_every / self.average_time_per_item)				,	1)



 def 		UpperCamelCase							(						self							:    int				,	snake_case_							:    Any				,	snake_case_							:    List[str]=None)->							List[Any]:
      __lowerCAmelCase	   =""" """ * (len(str(self.total)) - len(str(_lowerCamelCase))) + str(_lowerCamelCase)
      if self.elapsed_time is None:
           __lowerCAmelCase	   =F"""[{spaced_value}/{self.total} : < :"""
      elif self.predicted_remaining is None:
           __lowerCAmelCase	   =F"""[{spaced_value}/{self.total} {format_time(self.elapsed_time)}"""
      else:
           __lowerCAmelCase	   =(
               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							:    List[Any])->							str:
      __lowerCAmelCase	   =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:
           __lowerCAmelCase	   =disp.display(disp.HTML(self.html_code)				,	display_id=_lowerCamelCase)
      else:
           self.output.update(disp.HTML(self.html_code))



 def 		UpperCamelCase							(						self							:    List[str])->							Tuple:
      if self.parent is None and self.output is not None:
           self.output.update(disp.HTML(""""""))






class 						__a (     __UpperCAmelCase				):



 def __init__(						self							:    Dict				,	snake_case_							:    Union[str, Any]				,	snake_case_							:    Optional[int]=None)->							Tuple:
      super().__init__(_lowerCamelCase)
      __lowerCAmelCase	   =None if column_names is None else [column_names]
      __lowerCAmelCase	   =None



 def 		UpperCamelCase							(						self							:    Optional[int])->							Union[str, Any]:
      __lowerCAmelCase	   =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:
           __lowerCAmelCase	   =disp.display(disp.HTML(self.html_code)				,	display_id=_lowerCamelCase)
      else:
           self.output.update(disp.HTML(self.html_code))



 def 		UpperCamelCase							(						self							:    List[Any]				,	snake_case_							:    Any)->							List[Any]:
      if self.inner_table is None:
           __lowerCAmelCase	   =[list(values.keys()), list(values.values())]
      else:
           __lowerCAmelCase	   =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)
                __lowerCAmelCase	   =columns
           self.inner_table.append([values[c] for c in columns])



 def 		UpperCamelCase							(						self							:    List[Any]				,	snake_case_							:    Optional[Any]				,	snake_case_							:    Union[str, Any]=None				,	snake_case_							:    Any=3_00)->							List[Any]:
      __lowerCAmelCase	   =NotebookProgressBar(_lowerCamelCase				,	prefix=_lowerCamelCase				,	parent=self				,	width=_lowerCamelCase)
      return self.child_bar



 def 		UpperCamelCase							(						self							:    Union[str, Any])->							List[Any]:
      __lowerCAmelCase	   =None
      self.display()






class 						__a (     __UpperCAmelCase				):



 def __init__(						self							:    Tuple)->							List[Any]:
      __lowerCAmelCase	   =None
      __lowerCAmelCase	   =None
      __lowerCAmelCase	   =False



 def 		UpperCamelCase							(						self							:    int				,	snake_case_							:    Union[str, Any]				,	snake_case_							:    Union[str, Any]				,	snake_case_							:    Optional[int]				,	**snake_case_							:    List[Any])->							Union[str, Any]:
      __lowerCAmelCase	   ="""Epoch""" if args.evaluation_strategy == IntervalStrategy.EPOCH else """Step"""
      __lowerCAmelCase	   =0
      __lowerCAmelCase	   =0
      __lowerCAmelCase	   =[self.first_column] + ["""Training Loss"""]
      if args.evaluation_strategy != IntervalStrategy.NO:
           column_names.append("""Validation Loss""")
      __lowerCAmelCase	   =NotebookTrainingTracker(state.max_steps				,	_lowerCamelCase)



 def 		UpperCamelCase							(						self							:    str				,	snake_case_							:    str				,	snake_case_							:    Tuple				,	snake_case_							:    str				,	**snake_case_							:    Optional[int])->							Any:
      __lowerCAmelCase	   =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				,	)
      __lowerCAmelCase	   =False



 def 		UpperCamelCase							(						self							:    Optional[Any]				,	snake_case_							:    Optional[Any]				,	snake_case_							:    List[str]				,	snake_case_							:    int				,	snake_case_							:    Dict=None				,	**snake_case_							:    Optional[Any])->							Dict:
      if not has_length(_lowerCamelCase):
           return
      if self.prediction_bar is None:
           if self.training_tracker is not None:
                __lowerCAmelCase	   =self.training_tracker.add_child(len(_lowerCamelCase))
           else:
                __lowerCAmelCase	   =NotebookProgressBar(len(_lowerCamelCase))
           self.prediction_bar.update(1)
      else:
           self.prediction_bar.update(self.prediction_bar.value + 1)



 def 		UpperCamelCase							(						self							:    int				,	snake_case_							:    Optional[int]				,	snake_case_							:    Optional[Any]				,	snake_case_							:    int				,	**snake_case_							:    Union[str, Any])->							str:
      if self.prediction_bar is not None:
           self.prediction_bar.close()
      __lowerCAmelCase	   =None



 def 		UpperCamelCase							(						self							:    Any				,	snake_case_							:    Union[str, Any]				,	snake_case_							:    Union[str, Any]				,	snake_case_							:    Union[str, Any]				,	snake_case_							:    Dict=None				,	**snake_case_							:    str)->							int:
      # Only for when there is no evaluation
      if args.evaluation_strategy == IntervalStrategy.NO and "loss" in logs:
           __lowerCAmelCase	   ={"""Training Loss""": logs["""loss"""]}
           # First column is necessarily Step sine we're not in epoch eval strategy
           __lowerCAmelCase	   =state.global_step
           self.training_tracker.write_line(_lowerCamelCase)



 def 		UpperCamelCase							(						self							:    str				,	snake_case_							:    List[Any]				,	snake_case_							:    Union[str, Any]				,	snake_case_							:    Any				,	snake_case_							:    Dict=None				,	**snake_case_							:    List[Any])->							List[Any]:
      if self.training_tracker is not None:
           __lowerCAmelCase	   ={"""Training Loss""": """No log""", """Validation Loss""": """No log"""}
           for log in reversed(state.log_history):
                if "loss" in log:
                     __lowerCAmelCase	   =log["""loss"""]
                     break

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



 def 		UpperCamelCase							(						self							:    int				,	snake_case_							:    Union[str, Any]				,	snake_case_							:    int				,	snake_case_							:    Any				,	**snake_case_							:    List[str])->							str:
      self.training_tracker.update(
          state.global_step				,	comment=F"""Epoch {int(state.epoch)}/{state.num_train_epochs}"""				,	force_update=_lowerCamelCase)
      __lowerCAmelCase	   =None









from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase_	       =			logging.get_logger(__name__)

lowercase_	       =			{
    '''google/realm-cc-news-pretrained-embedder''': (
        '''https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/config.json'''
    ),
    '''google/realm-cc-news-pretrained-encoder''': (
        '''https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/config.json'''
    ),
    '''google/realm-cc-news-pretrained-scorer''': (
        '''https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/config.json'''
    ),
    '''google/realm-cc-news-pretrained-openqa''': (
        '''https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/config.json'''
    ),
    '''google/realm-orqa-nq-openqa''': '''https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/config.json''',
    '''google/realm-orqa-nq-reader''': '''https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/config.json''',
    '''google/realm-orqa-wq-openqa''': '''https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/config.json''',
    '''google/realm-orqa-wq-reader''': '''https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/config.json''',
    # See all REALM models at https://huggingface.co/models?filter=realm
}






class 						__a (     SCREAMING_SNAKE_CASE				):
	SCREAMING_SNAKE_CASE       =       "realm"



	def __init__(						self							:    Tuple				,	snake_case_							:    Tuple=3_05_22				,	snake_case_							:    List[Any]=7_68				,	snake_case_							:    Any=1_28				,	snake_case_							:    Union[str, Any]=12				,	snake_case_							:    List[str]=12				,	snake_case_							:    Any=8				,	snake_case_							:    Any=30_72				,	snake_case_							:    int="gelu_new"				,	snake_case_							:    Optional[int]=0.1				,	snake_case_							:    Tuple=0.1				,	snake_case_							:    List[Any]=5_12				,	snake_case_							:    List[Any]=2				,	snake_case_							:    Any=0.0_2				,	snake_case_							:    List[str]=1e-12				,	snake_case_							:    Optional[Any]=2_56				,	snake_case_							:    Dict=10				,	snake_case_							:    str=1e-3				,	snake_case_							:    List[str]=5				,	snake_case_							:    Optional[Any]=3_20				,	snake_case_							:    Optional[Any]=13_35_37_18				,	snake_case_							:    Tuple=50_00				,	snake_case_							:    Union[str, Any]=1				,	snake_case_							:    Union[str, Any]=0				,	snake_case_							:    List[str]=2				,	**snake_case_							:    List[str]				,	)->							str:
						super().__init__(pad_token_id=snake_case_				,	bos_token_id=snake_case_				,	eos_token_id=snake_case_				,	**snake_case_)

						# Common config
						__lowerCAmelCase	   =vocab_size
						__lowerCAmelCase	   =max_position_embeddings
						__lowerCAmelCase	   =hidden_size
						__lowerCAmelCase	   =retriever_proj_size
						__lowerCAmelCase	   =num_hidden_layers
						__lowerCAmelCase	   =num_attention_heads
						__lowerCAmelCase	   =num_candidates
						__lowerCAmelCase	   =intermediate_size
						__lowerCAmelCase	   =hidden_act
						__lowerCAmelCase	   =hidden_dropout_prob
						__lowerCAmelCase	   =attention_probs_dropout_prob
						__lowerCAmelCase	   =initializer_range
						__lowerCAmelCase	   =type_vocab_size
						__lowerCAmelCase	   =layer_norm_eps

						# Reader config
						__lowerCAmelCase	   =span_hidden_size
						__lowerCAmelCase	   =max_span_width
						__lowerCAmelCase	   =reader_layer_norm_eps
						__lowerCAmelCase	   =reader_beam_size
						__lowerCAmelCase	   =reader_seq_len

						# Retrieval config
						__lowerCAmelCase	   =num_block_records
						__lowerCAmelCase	   =searcher_beam_size




