


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 							UpperCamelCase_( lowerCamelCase_					)       ->		int:
    _lowercase			:				List[str] 						=						int(lowerCamelCase_					)
    _lowercase						,							_lowercase						,							_lowercase			:				Optional[Any] 						=						t // 3600, (t // 60) % 60, t % 60
    return F'''{h}:{m:02d}:{s:02d}''' if h != 0 else F'''{m:02d}:{s:02d}'''



def 							UpperCamelCase_( lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_=300					)       ->		Dict:
    # docstyle-ignore
    return F'''
    <div>
      {prefix}
      <progress value=\'{value}\' max=\'{total}\' style=\'width:{width}px; height:20px; vertical-align: middle;\'></progress>
      {label}
    </div>
    '''



def 							UpperCamelCase_( lowerCamelCase_					)       ->		Any:
    _lowercase			:				int 						=						'<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			:				Any 						=						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 _lowerCamelCase:
       lowercase_					:      str       					=     5
       lowercase_					:      str       					=     0.2
       def __init__(   self, lowerCamelCase, lowerCamelCase = None, lowerCamelCase = True, lowerCamelCase = None, lowerCamelCase = 3_00, )			->  Optional[Any]:




           """simple docstring"""

           _lowercase			:				Optional[int] 						=						total
           _lowercase			:				Optional[int] 						=						'' if prefix is None else prefix
           _lowercase			:				Tuple 						=						leave
           _lowercase			:				str 						=						parent
           _lowercase			:				str 						=						width
           _lowercase			:				List[Any] 						=						None
           _lowercase			:				List[str] 						=						None
           _lowercase			:				Tuple 						=						None
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase = False, lowerCamelCase = None)			->  Dict:




           """simple docstring"""

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




           """simple docstring"""

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




           """simple docstring"""

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




       def      UpperCamelCase	(   self)			->  Optional[Any]:




           """simple docstring"""

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






class _lowerCamelCase(		_a  ):
       def __init__(   self, lowerCamelCase, lowerCamelCase=None)			->  int:




           """simple docstring"""

           super().__init__(lowerCamelCase)
           _lowercase			:				Optional[Any] 						=						None if column_names is None else [column_names]
           _lowercase			:				Any 						=						None
       def      UpperCamelCase	(   self)			->  List[Any]:




           """simple docstring"""

           _lowercase			:				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, lowerCamelCase)			->  Dict:




           """simple docstring"""

           if self.inner_table is None:
               _lowercase			:				Dict 						=						[list(values.keys()), list(values.values())]
           else:
               _lowercase			:				Tuple 						=						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			:				str 						=						columns
               self.inner_table.append([values[c] for c in columns])
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase=None, lowerCamelCase=3_00)			->  Union[str, Any]:




           """simple docstring"""

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




       def      UpperCamelCase	(   self)			->  Union[str, Any]:




           """simple docstring"""

           _lowercase			:				Optional[Any] 						=						None
           self.display()






class _lowerCamelCase(		_a  ):
       def __init__(   self)			->  List[Any]:




           """simple docstring"""

           _lowercase			:				Union[str, Any] 						=						None
           _lowercase			:				Dict 						=						None
           _lowercase			:				Dict 						=						False
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, **lowerCamelCase)			->  Dict:




           """simple docstring"""

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




           """simple docstring"""

           _lowercase			:				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			:				str 						=						False
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase=None, **lowerCamelCase)			->  Any:




           """simple docstring"""

           if not has_length(lowerCamelCase):
               return
           if self.prediction_bar is None:
               if self.training_tracker is not None:
                   _lowercase			:				Optional[int] 						=						self.training_tracker.add_child(len(lowerCamelCase))
               else:
                   _lowercase			:				Optional[int] 						=						NotebookProgressBar(len(lowerCamelCase))
               self.prediction_bar.update(1)
           else:
               self.prediction_bar.update(self.prediction_bar.value + 1)
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, **lowerCamelCase)			->  Optional[int]:




           """simple docstring"""

           if self.prediction_bar is not None:
               self.prediction_bar.close()
           _lowercase			:				Any 						=						None
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase=None, **lowerCamelCase)			->  List[Any]:




           """simple docstring"""

           if args.evaluation_strategy == IntervalStrategy.NO and "loss" in logs:
               _lowercase			:				Dict 						=						{'Training Loss': logs['loss']}
               # First column is necessarily Step sine we're not in epoch eval strategy
               _lowercase			:				List[Any] 						=						state.global_step
               self.training_tracker.write_line(lowerCamelCase)
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase=None, **lowerCamelCase)			->  List[str]:




           """simple docstring"""

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

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




       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, **lowerCamelCase)			->  Union[str, Any]:




           """simple docstring"""

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










'''simple docstring'''
import argparse
import os

import torch

from transformers.utils import WEIGHTS_NAME


UpperCamelCase__			=      ['small', 'medium', 'large']

UpperCamelCase__			=      'lm_head.decoder.weight'
UpperCamelCase__			=      'lm_head.weight'




def 		__SCREAMING_SNAKE_CASE       (	_UpperCamelCase    ,		_UpperCamelCase				):


   """simple docstring"""




   lowercase_						:     Optional[Any]							     =  torch.load(_UpperCamelCase				)
   lowercase_						:     List[Any]							     =  d.pop(_UpperCamelCase				)
   os.makedirs(_UpperCamelCase    ,		exist_ok=_UpperCamelCase				)
   torch.save(_UpperCamelCase    ,		os.path.join(_UpperCamelCase    ,		_UpperCamelCase				)				)


if __name__ == "__main__":
 UpperCamelCase__			=      argparse.ArgumentParser()
 parser.add_argument('--dialogpt_path', default='.', type=str)
 UpperCamelCase__			=      parser.parse_args()
 for MODEL in DIALOGPT_MODELS:
  UpperCamelCase__			=      os.path.join(args.dialogpt_path, f"""{MODEL}_ft.pkl""")
  UpperCamelCase__			=      f"""./DialoGPT-{MODEL}"""
  convert_dialogpt_checkpoint(
      checkpoint_path,
      pytorch_dump_folder_path,
  )




