

import argparse
import dataclasses
import json
import logging
import os
import shutil
from typing import List, Optional

import datasets
from accelerate import Accelerator
from datasets import load_dataset
from finetuning import finetune
from tqdm.auto import tqdm

import transformers
from transformers import AutoConfig, set_seed
from transformers.trainer_utils import IntervalStrategy


A_       : Optional[int]         =			logging.getLogger(__name__)

A_       : Optional[int]         =			'pytorch_model.bin'



@dataclasses.dataclass
class     _a						:


					'''simple docstring'''






					UpperCAmelCase__:   str  							=       dataclasses.field(
					    metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models.'''}				)
					UpperCAmelCase__:   Optional[str]  							=       dataclasses.field(
					    default=__magic_name__		,   metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co.'''}		,   )



@dataclasses.dataclass
class     _a						:


					'''simple docstring'''






					UpperCAmelCase__:   str  							=       dataclasses.field(metadata={'''help''': '''A csv or a json file containing the training data.'''}				)
					UpperCAmelCase__:   str  							=       dataclasses.field(metadata={'''help''': '''A csv or a json file containing the data to predict on.'''}				)
					UpperCAmelCase__:   Optional[str]  							=       dataclasses.field(
					    default=__magic_name__		,   metadata={'''help''': '''A csv or a json file containing the validation data.'''}				)
					UpperCAmelCase__:   Optional[str]  							=       dataclasses.field(
					    default=__magic_name__		,   metadata={'''help''': '''The name of the task to train on.'''}		,   )
					UpperCAmelCase__:   Optional[List[str]]  							=       dataclasses.field(
					    default=__magic_name__		,   metadata={'''help''': '''The list of labels for the task.'''}				)



@dataclasses.dataclass
class     _a						:


					'''simple docstring'''






					UpperCAmelCase__:   str  							=       dataclasses.field(
					    metadata={'''help''': '''The output directory where the model predictions and checkpoints will be written.'''}				)
					UpperCAmelCase__:   Optional[str]  							=       dataclasses.field(
					    default='''accuracy'''		,   metadata={'''help''': '''The evaluation metric used for the task.'''}				)
					UpperCAmelCase__:   Optional[str]  							=       dataclasses.field(
					    default='''no'''		,   metadata={
					        '''help''': '''The evaluation strategy to adopt during training. Possible values are: ["no", "step", "epoch]'''
					    }		,   )
					UpperCAmelCase__:   Optional[int]  							=       dataclasses.field(
					    default=10		,   metadata={'''help''': '''Number of evaluation calls with no improvement after which training will be stopped.'''}		,   )
					UpperCAmelCase__:   Optional[float]  							=       dataclasses.field(
					    default=0.0		,   metadata={
					        '''help''': '''How much the specified evaluation metric must improve to satisfy early stopping conditions.'''
					    }		,   )
					UpperCAmelCase__:   Optional[bool]  							=       dataclasses.field(
					    default=__magic_name__		,   metadata={'''help''': '''Whether to filter the pseudo-labeled data based on the confidence score.'''}		,   )
					UpperCAmelCase__:   Optional[bool]  							=       dataclasses.field(
					    default=__magic_name__		,   metadata={'''help''': '''Whether to filter the pseudo-labeled data based on the validation performance.'''}		,   )
					UpperCAmelCase__:   Optional[bool]  							=       dataclasses.field(
					    default=__magic_name__		,   metadata={'''help''': '''Whether to fine-tune on labeled data after pseudo training.'''}		,   )
					UpperCAmelCase__:   Optional[float]  							=       dataclasses.field(
					    default=0.0		,   metadata={'''help''': '''Confidence threshold for pseudo-labeled data filtering.'''}		,   )
					UpperCAmelCase__:   Optional[int]  							=       dataclasses.field(
					    default=1_00		,   metadata={'''help''': '''Number of evaluation calls with no improvement after which training will be stopped.'''}		,   )
					UpperCAmelCase__:   Optional[int]  							=       dataclasses.field(
					    default=__magic_name__		,   metadata={'''help''': '''Random seed for initialization.'''}		,   )



def 						UpperCamelCase					(lowercase_: Tuple    ,      lowercase_: Optional[int]    ,      lowercase_: List[str]    ,      lowercase_: List[str]    ,      lowercase_: Dict    ,      lowercase_: int					)     ->      List[str]:
	A__							:    int        =		datasets.concatenate_datasets([infer_input, infer_output]    ,      axis=1					)

	if args.do_filter_by_confidence:
		A__							:    Optional[int]        =		dataset.filter(lambda lowercase_					: example["probability"] > args.confidence_threshold					)

	if args.do_filter_by_val_performance:
		assert eval_result >= 0.0 and eval_result <= 1.0
		A__							:    Any        =		int(eval_result * len(lowercase_					)					)
		print(lowercase_					)
		A__							:    List[str]        =		dataset.sort("""probability"""    ,      reverse=lowercase_					)
		A__							:    Optional[int]        =		dataset.select(range(lowercase_					)					)

	A__							:    Dict        =		dataset.remove_columns(["""label""", """probability"""]					)
	A__							:    Optional[Any]        =		dataset.rename_column("""prediction"""    ,      """label"""					)
	A__							:    List[Any]        =		dataset.map(lambda lowercase_					: {"label": idalabel[example["label"]]}					)
	A__							:    Union[str, Any]        =		dataset.shuffle(seed=args.seed					)

	A__							:    List[Any]        =		os.path.join(lowercase_    ,      f"""train_pseudo.{args.data_file_extension}"""					)
	if args.data_file_extension == "csv":
		dataset.to_csv(lowercase_    ,      index=lowercase_					)
	else:
		dataset.to_json(lowercase_					)


def 						UpperCamelCase					(lowercase_: int    ,      lowercase_: Union[str, Any]    ,      lowercase_: Optional[int]    ,      lowercase_: Optional[int]    ,      **lowercase_: Dict					)     ->      List[Any]:
	A__							:    List[Any]        =		Accelerator()
	# Make one log on every process with the configuration for debugging.
	logging.basicConfig(
	    format="""%(asctime)s - %(levelname)s - %(name)s - %(message)s"""    ,      datefmt="""%m/%d/%Y %H:%M:%S"""    ,      level=logging.INFO    ,      )
	logger.info(accelerator.state					)

	# Setup logging, we only want one process per machine to log things on the
	# screen. accelerator.is_local_main_process is only True for one process per
	# machine.
	logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR					)

	if accelerator.is_local_main_process:
		datasets.utils.logging.set_verbosity_warning()
		transformers.utils.logging.set_verbosity_info()
	else:
		datasets.utils.logging.set_verbosity_error()
		transformers.utils.logging.set_verbosity_error()

	A__							:    int        =		STModelArguments(model_name_or_path=lowercase_					)
	A__							:    List[Any]        =		STDataArguments(train_file=lowercase_    ,      infer_file=lowercase_					)
	A__							:    List[str]        =		STTrainingArguments(output_dir=lowercase_					)
	A__							:    List[Any]        =		argparse.Namespace()

	for arg_class in (model_args, data_args, training_args):
		for key, value in vars(lowercase_					).items():
			setattr(lowercase_    ,      lowercase_    ,      lowercase_					)

	for key, value in kwargs.items():
		if hasattr(lowercase_    ,      lowercase_					):
			setattr(lowercase_    ,      lowercase_    ,      lowercase_					)

    # Sanity checks
	A__							:    Any        =		{}
	A__							:    List[Any]        =		None

	# You need to provide the training data and the data to predict on
	assert args.train_file is not None
	assert args.infer_file is not None
	A__							:    Tuple        =		args.train_file
	A__							:    Tuple        =		args.infer_file

	if args.evaluation_strategy != IntervalStrategy.NO.value:
		assert args.eval_file is not None
		A__							:    Optional[Any]        =		args.eval_file

	for key in data_files:
		A__							:    Tuple        =		data_files[key].split("""."""					)[-1]
		assert extension in ["csv", "json"], f"""`{key}_file` should be a csv or a json file."""
		if args.data_file_extension is None:
			A__							:    Optional[int]        =		extension
		else:
			assert extension == args.data_file_extension, f"""`{key}_file` should be a {args.data_file_extension} file`."""

	assert (
	    args.eval_metric in datasets.list_metrics()
	), f"""{args.eval_metric} not in the list of supported metrics {datasets.list_metrics()}."""

	# If passed along, set the training seed now.
	if args.seed is not None:
		set_seed(args.seed					)

	logger.info("""Creating the initial data directory for self-training..."""					)
	A__							:    Tuple        =		f"""{args.output_dir}/self-train_iter-{{}}""".format
	A__							:    Union[str, Any]        =		data_dir_format(0					)

	if accelerator.is_main_process:
		if args.output_dir is not None:
			os.makedirs(args.output_dir    ,      exist_ok=lowercase_					)
			os.makedirs(lowercase_    ,      exist_ok=lowercase_					)
	accelerator.wait_for_everyone()

	A__							:    Optional[Any]        =		None
	A__							:    Union[str, Any]        =		None
	A__							:    Optional[int]        =		0
	A__							:    List[Any]        =		False
	# Show the progress bar
	A__							:    Any        =		tqdm(range(args.max_selftrain_iterations					)    ,      disable=not accelerator.is_local_main_process					)

	# Self-train
	for iteration in range(0    ,      int(args.max_selftrain_iterations					)					):
		A__							:    List[Any]        =		data_dir_format(lowercase_					)
		assert os.path.exists(lowercase_					)

		# Stage 1: initial fine-tuning for iteration = 0 or pseudo-training for
		# iteration > 0
		A__							:    Optional[int]        =		os.path.join(lowercase_    ,      """stage-1"""					)
		A__							:    str        =		{
		    """accelerator""": accelerator,
		    """model_name_or_path""": args.model_name_or_path,
		    """cache_dir""": args.cache_dir,
		    """do_train""": True,
		    """train_file""": data_files["""train"""] if iteration == 0 else data_files["""train_pseudo"""],
		    """do_eval""": True if args.eval_file is not None else False,
		    """eval_file""": data_files["""eval"""],
		    """do_predict""": True,
		    """infer_file""": data_files["""infer"""],
		    """task_name""": args.task_name,
		    """label_list""": args.label_list,
		    """output_dir""": current_output_dir,
		    """eval_metric""": args.eval_metric,
		    """evaluation_strategy""": args.evaluation_strategy,
		    """early_stopping_patience""": args.early_stopping_patience,
		    """early_stopping_threshold""": args.early_stopping_threshold,
		    """seed""": args.seed,
		}
		# Add additional training arguments
		for key, value in kwargs.items():
			if key not in arguments_dict and not hasattr(lowercase_    ,      lowercase_					):
				arguments_dict.update({key: value}					)

		A__							:    Dict        =		os.path.join(lowercase_    ,      """best-checkpoint"""    ,      lowercase_					)
		if os.path.exists(lowercase_					):
			logger.info(
			    """Found existing model checkpoint at %s. Skipping self-training: iteration: %d, stage: 1."""    ,      lowercase_    ,      lowercase_    ,      )
		else:
			logger.info("""***** Running self-training: iteration: %d, stage: 1 *****"""    ,      lowercase_					)
			finetune(**lowercase_					)
			accelerator.wait_for_everyone()
			assert os.path.exists(lowercase_					)
			logger.info("""Self-training job completed: iteration: %d, stage: 1."""    ,      lowercase_					)

		if iteration > 0 and args.finetune_on_labeled_data:
			# Stage 2 (optional): fine-tuning on the original labeled data
			A__							:    Tuple        =		os.path.join(lowercase_    ,      """best-checkpoint"""					)
			A__							:    Tuple        =		os.path.join(lowercase_    ,      """stage-2"""					)
			# Update arguments_dict
			A__							:    List[str]        =		model_path
			A__							:    List[Any]        =		data_files["""train"""]
			A__							:    Any        =		current_output_dir

			A__							:    List[str]        =		os.path.join(lowercase_    ,      """best-checkpoint"""    ,      lowercase_					)
			if os.path.exists(lowercase_					):
				logger.info(
				    """Found existing model checkpoint at %s. Skipping self-training: iteration: %d, stage: 2."""    ,      lowercase_    ,      lowercase_    ,      )
			else:
				logger.info("""***** Running self-training: iteration: %d, stage: 2 *****"""    ,      lowercase_					)
				finetune(**lowercase_					)
				accelerator.wait_for_everyone()
				assert os.path.exists(lowercase_					)
				logger.info("""Self-training job completed: iteration: %d, stage: 2."""    ,      lowercase_					)

		A__							:    Dict        =		iteration
		A__							:    Tuple        =		data_dir_format(iteration + 1					)

		A__							:    Optional[int]        =		AutoConfig.from_pretrained(os.path.join(lowercase_    ,      """best-checkpoint"""					)					)
		A__							:    Dict        =		config.idalabel
		A__							:    int        =		os.path.join(lowercase_    ,      """eval_results_best-checkpoint.json"""					)
		A__							:    List[Any]        =		os.path.join(lowercase_    ,      """test_results_best-checkpoint.json"""					)
		assert os.path.exists(lowercase_					)

		with open(lowercase_    ,      """r"""					) as f:
			A__							:    Dict        =		float(json.load(lowercase_					)[args.eval_metric]					)
		A__							:    int        =		os.path.join(lowercase_    ,      """infer_output_best-checkpoint.csv"""					)
		assert os.path.exists(lowercase_					)
		# Loading the dataset from local csv or json files.
		A__							:    List[Any]        =		load_dataset(args.data_file_extension    ,      data_files={"""data""": data_files["""infer"""]}					)["""data"""]
		A__							:    Dict        =		load_dataset("""csv"""    ,      data_files={"""data""": infer_output_file}					)["""data"""]

		if accelerator.is_main_process:
			os.makedirs(lowercase_    ,      exist_ok=lowercase_					)
			shutil.copy(lowercase_    ,      os.path.join(lowercase_    ,      f"""eval_results_iter-{iteration}.json"""					)					)
			if os.path.exists(lowercase_					):
				shutil.copy(lowercase_    ,      os.path.join(lowercase_    ,      f"""test_results_iter-{iteration}.json"""					)					)
			create_pseudo_labeled_data(lowercase_    ,      lowercase_    ,      lowercase_    ,      lowercase_    ,      lowercase_    ,      lowercase_					)
		accelerator.wait_for_everyone()

		A__							:    str        =		os.path.join(lowercase_    ,      f"""train_pseudo.{args.data_file_extension}"""					)

		if args.evaluation_strategy != IntervalStrategy.NO.value:
			A__							:    Tuple        =		eval_result

			if best_iteration is None:
				A__							:    Union[str, Any]        =		new_iteration
				A__							:    str        =		new_eval_result
			else:
				if new_eval_result - best_eval_result > args.early_stopping_threshold:
					A__							:    Dict        =		new_iteration
					A__							:    Union[str, Any]        =		new_eval_result
					A__							:    List[str]        =		0
				else:
					if new_eval_result == best_eval_result:
						A__							:    Union[str, Any]        =		new_iteration
						A__							:    Tuple        =		new_eval_result
					early_stopping_patience_counter += 1

				if early_stopping_patience_counter >= args.early_stopping_patience:
					A__							:    List[str]        =		True

		progress_bar.update(1					)

		if should_training_stop:
			break

	if best_iteration is not None:
		# Save the best iteration
		logger.info("""Best iteration: %d"""    ,      lowercase_					)
		logger.info("""Best evaluation result: %s = %f"""    ,      args.eval_metric    ,      lowercase_					)
		accelerator.wait_for_everyone()
		if accelerator.is_main_process:
			shutil.copy(
			    os.path.join(lowercase_    ,      f"""eval_results_iter-{iteration}.json"""					)    ,      os.path.join(lowercase_    ,      """eval_results_best-iteration.json"""					)    ,      )
	else:
		# Assume that the last iteration is the best
		logger.info("""Best iteration: %d"""    ,      args.max_selftrain_iterations - 1					)
		logger.info("""Best evaluation result: %s = %f"""    ,      args.eval_metric    ,      lowercase_					)
		accelerator.wait_for_everyone()
		if accelerator.is_main_process:
			shutil.copy(
			    os.path.join(lowercase_    ,      f"""eval_results_iter-{args.max_selftrain_iterations - 1}.json"""					)    ,      os.path.join(lowercase_    ,      """eval_results_best-iteration.json"""					)    ,      )






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					(lowercase_: str					)     ->      Dict:
	A__							:    int        =		int(lowercase_					)
	A__   ,   A__   ,   A__							:    Tuple        =		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					(lowercase_: str    ,      lowercase_: Optional[Any]    ,      lowercase_: Union[str, Any]    ,      lowercase_: Tuple    ,      lowercase_: Any=300					)     ->      Optional[int]:
	# 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					(lowercase_: Tuple					)     ->      Optional[int]:
	A__							:    Tuple        =		"""<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:
			A__							:    str        =		f"""{elt:.6f}""" if isinstance(lowercase_    ,      lowercase_					) else str(lowercase_					)
			html_code += f"""      <td>{elt}</td>\n"""
		html_code += "    </tr>\n"
	html_code += "  </tbody>\n</table><p>"
	return html_code



class     _a						:


					'''simple docstring'''






					UpperCAmelCase__:   str  							=       5
					UpperCAmelCase__:   int  							=       0.2






					def __init__( self    ,       A__    ,       A__ = None    ,       A__ = True    ,       A__ = None    ,       A__ = 300    ,       ):
						A__							:    Optional[int]        =		total
						A__							:    Tuple        =		"""""" if prefix is None else prefix
						A__							:    str        =		leave
						A__							:    str        =		parent
						A__							:    int        =		width
						A__							:    Dict        =		None
						A__							:    List[str]        =		None
						A__							:    Optional[int]        =		None






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






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






					def   __A			( self			):
						if self.parent is None and self.output is not None:
							self.output.update(disp.HTML(""""""			)			)



class     _a						(__magic_name__				):


					'''simple docstring'''






					def __init__( self    ,       A__    ,       A__=None			):
						super().__init__(A__			)
						A__							:    Optional[Any]        =		None if column_names is None else [column_names]
						A__							:    Optional[Any]        =		None






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






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






					def   __A			( self    ,       A__    ,       A__=None    ,       A__=300			):
						A__							:    Optional[Any]        =		NotebookProgressBar(A__    ,       prefix=A__    ,       parent=self    ,       width=A__			)
						return self.child_bar






					def   __A			( self			):
						A__							:    List[str]        =		None
						self.display()



class     _a						(__magic_name__				):


					'''simple docstring'''






					def __init__( self			):
						A__							:    int        =		None
						A__							:    List[str]        =		None
						A__							:    Union[str, Any]        =		False






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       **A__			):
						A__							:    List[str]        =		"""Epoch""" if args.evaluation_strategy == IntervalStrategy.EPOCH else """Step"""
						A__							:    Dict        =		0
						A__							:    Tuple        =		0
						A__							:    Optional[int]        =		[self.first_column] + ["""Training Loss"""]
						if args.evaluation_strategy != IntervalStrategy.NO:
							column_names.append("""Validation Loss"""			)
						A__							:    Union[str, Any]        =		NotebookTrainingTracker(state.max_steps    ,       A__			)






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       **A__			):
						A__							:    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    ,       )
						A__							:    str        =		False






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       A__=None    ,       **A__			):
						if not has_length(A__			):
							return
						if self.prediction_bar is None:
							if self.training_tracker is not None:
								A__							:    Union[str, Any]        =		self.training_tracker.add_child(len(A__			)			)
							else:
								A__							:    Tuple        =		NotebookProgressBar(len(A__			)			)
							self.prediction_bar.update(1			)
						else:
							self.prediction_bar.update(self.prediction_bar.value + 1			)






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       **A__			):
						if self.prediction_bar is not None:
							self.prediction_bar.close()
						A__							:    List[str]        =		None






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






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       A__=None    ,       **A__			):
						if self.training_tracker is not None:
							A__							:    Tuple        =		{"""Training Loss""": """No log""", """Validation Loss""": """No log"""}
							for log in reversed(state.log_history			):
								if "loss" in log:
									A__							:    Dict        =		log["""loss"""]
									break

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






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       **A__			):
						self.training_tracker.update(
						    state.global_step    ,       comment=F"""Epoch {int(state.epoch			)}/{state.num_train_epochs}"""    ,       force_update=A__			)
						A__							:    Optional[int]        =		None



