





'''simple docstring'''




import argparse
import os

import gluonnlp as nlp
import mxnet as mx
import numpy as np
import torch
from gluonnlp.base import get_home_dir
from gluonnlp.model.bert import BERTEncoder
from gluonnlp.model.utils import _load_vocab
from gluonnlp.vocab import Vocab
from packaging import version
from torch import nn

from transformers import BertConfig, BertForMaskedLM, BertModel, RobertaTokenizer
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


if version.parse(nlp.__version__) != version.parse('''0.8.3'''):
			raise Exception('''requires gluonnlp == 0.8.3''')

if version.parse(mx.__version__) != version.parse('''1.5.0'''):
			raise Exception('''requires mxnet == 1.5.0''')

logging.set_verbosity_info()
a__					:			int          =     logging.get_logger(__name__)

a__					:			Any          =     '''The Nymphenburg Palace is a beautiful palace in Munich!'''




def 						__lowerCamelCase			(	UpperCAmelCase_					,				UpperCAmelCase_							)   ->str:
	snake_case__	       =   {
	    'attention_cell': 'multi_head',
	    'num_layers': 4,
	    'units': 10_24,
	    'hidden_size': 7_68,
	    'max_length': 5_12,
	    'num_heads': 8,
	    'scaled': True,
	    'dropout': 0.1,
	    'use_residual': True,
	    'embed_size': 10_24,
	    'embed_dropout': 0.1,
	    'word_embed': None,
	    'layer_norm_eps': 1E-5,
	    'token_type_vocab_size': 2,
	}

	snake_case__	       =   bort_4_8_768_1024_hparams

	# Let's construct the original Bort model here
	# Taken from official BERT implementation, see:
	# https://github.com/alexa/bort/blob/master/bort/bort.py
	snake_case__	       =   BERTEncoder(
	    attention_cell=predefined_args['attention_cell']					,				num_layers=predefined_args['num_layers']					,				units=predefined_args['units']					,				hidden_size=predefined_args['hidden_size']					,				max_length=predefined_args['max_length']					,				num_heads=predefined_args['num_heads']					,				scaled=predefined_args['scaled']					,				dropout=predefined_args['dropout']					,				output_attention=UpperCAmelCase_					,				output_all_encodings=UpperCAmelCase_					,				use_residual=predefined_args['use_residual']					,				activation=predefined_args.get('activation'					,				'gelu'							)					,				layer_norm_eps=predefined_args.get('layer_norm_eps'					,				UpperCAmelCase_							)					,				)

	# Vocab information needs to be fetched first
	# It's the same as RoBERTa, so RobertaTokenizer can be used later
	snake_case__	       =   'openwebtext_ccnews_stories_books_cased'

	# Specify download folder to Gluonnlp's vocab
	snake_case__	       =   os.path.join(get_home_dir()					,				'models'							)
	snake_case__	       =   _load_vocab(UpperCAmelCase_					,				UpperCAmelCase_					,				UpperCAmelCase_					,				cls=UpperCAmelCase_							)

	snake_case__	       =   nlp.model.BERTModel(
	    UpperCAmelCase_					,				len(UpperCAmelCase_							)					,				units=predefined_args['units']					,				embed_size=predefined_args['embed_size']					,				embed_dropout=predefined_args['embed_dropout']					,				word_embed=predefined_args['word_embed']					,				use_pooler=UpperCAmelCase_					,				use_token_type_embed=UpperCAmelCase_					,				token_type_vocab_size=predefined_args['token_type_vocab_size']					,				use_classifier=UpperCAmelCase_					,				use_decoder=UpperCAmelCase_					,				)

	original_bort.load_parameters(UpperCAmelCase_					,				cast_dtype=UpperCAmelCase_					,				ignore_extra=UpperCAmelCase_							)
	snake_case__	       =   original_bort._collect_params_with_prefix()

	# Build our config 🤗
	snake_case__	       =   {
	    'architectures': ['BertForMaskedLM'],
	    'attention_probs_dropout_prob': predefined_args['dropout'],
	    'hidden_act': 'gelu',
	    'hidden_dropout_prob': predefined_args['dropout'],
	    'hidden_size': predefined_args['embed_size'],
	    'initializer_range': 0.02,
	    'intermediate_size': predefined_args['hidden_size'],
	    'layer_norm_eps': predefined_args['layer_norm_eps'],
	    'max_position_embeddings': predefined_args['max_length'],
	    'model_type': 'bort',
	    'num_attention_heads': predefined_args['num_heads'],
	    'num_hidden_layers': predefined_args['num_layers'],
	    'pad_token_id': 1,  # 2 = BERT, 1 = RoBERTa
	    'type_vocab_size': 1,  # 2 = BERT, 1 = RoBERTa
	    'vocab_size': len(UpperCAmelCase_							),
	}

	snake_case__	       =   BertConfig.from_dict(UpperCAmelCase_							)
	snake_case__	       =   BertForMaskedLM(UpperCAmelCase_							)
	hf_bort_model.eval()

	# Parameter mapping table (Gluonnlp to Transformers)
	# * denotes layer index
	#
	# | Gluon Parameter                                                | Transformers Parameter
	# | -------------------------------------------------------------- | ----------------------
	# | `encoder.layer_norm.beta`                                      | `bert.embeddings.LayerNorm.bias`
	# | `encoder.layer_norm.gamma`                                     | `bert.embeddings.LayerNorm.weight`
	# | `encoder.position_weight`                                      | `bert.embeddings.position_embeddings.weight`
	# | `word_embed.0.weight`                                          | `bert.embeddings.word_embeddings.weight`
	# | `encoder.transformer_cells.*.attention_cell.proj_key.bias`     | `bert.encoder.layer.*.attention.self.key.bias`
	# | `encoder.transformer_cells.*.attention_cell.proj_key.weight`   | `bert.encoder.layer.*.attention.self.key.weight`
	# | `encoder.transformer_cells.*.attention_cell.proj_query.bias`   | `bert.encoder.layer.*.attention.self.query.bias`
	# | `encoder.transformer_cells.*.attention_cell.proj_query.weight` | `bert.encoder.layer.*.attention.self.query.weight`
	# | `encoder.transformer_cells.*.attention_cell.proj_value.bias`   | `bert.encoder.layer.*.attention.self.value.bias`
	# | `encoder.transformer_cells.*.attention_cell.proj_value.weight` | `bert.encoder.layer.*.attention.self.value.weight`
	# | `encoder.transformer_cells.*.ffn.ffn_2.bias`                   | `bert.encoder.layer.*.attention.output.dense.bias`
	# | `encoder.transformer_cells.*.ffn.ffn_2.weight`                 | `bert.encoder.layer.*.attention.output.dense.weight`
	# | `encoder.transformer_cells.*.layer_norm.beta`                  | `bert.encoder.layer.*.attention.output.LayerNorm.bias`
	# | `encoder.transformer_cells.*.layer_norm.gamma`                 | `bert.encoder.layer.*.attention.output.LayerNorm.weight`
	# | `encoder.transformer_cells.*.ffn.ffn_1.bias`                   | `bert.encoder.layer.*.intermediate.dense.bias`
	# | `encoder.transformer_cells.*.ffn.ffn_1.weight`                 | `bert.encoder.layer.*.intermediate.dense.weight`
	# | `encoder.transformer_cells.*.ffn.layer_norm.beta`              | `bert.encoder.layer.*.output.LayerNorm.bias`
	# | `encoder.transformer_cells.*.ffn.layer_norm.gamma`             | `bert.encoder.layer.*.output.LayerNorm.weight`
	# | `encoder.transformer_cells.*.proj.bias`                        | `bert.encoder.layer.*.output.dense.bias`
	# | `encoder.transformer_cells.*.proj.weight`                      | `bert.encoder.layer.*.output.dense.weight`

	# Helper function to convert MXNET Arrays to PyTorch
	def to_torch(UpperCAmelCase_							) -> nn.Parameter:
		return nn.Parameter(torch.FloatTensor(mx_array.data().asnumpy()							)							)

	# Check param shapes and map new HF param back
	def check_and_map_params(UpperCAmelCase_					,				UpperCAmelCase_							):
		snake_case__	       =   hf_param.shape

		snake_case__	       =   to_torch(params[gluon_param]							)
		snake_case__	       =   gluon_param.shape

		assert (
		    shape_hf == shape_gluon
		), f'''The gluon parameter {gluon_param} has shape {shape_gluon}, but expects shape {shape_hf} for Transformers'''

		return gluon_param

	snake_case__	       =   check_and_map_params(
	    hf_bort_model.bert.embeddings.word_embeddings.weight					,				'word_embed.0.weight'							)
	snake_case__	       =   check_and_map_params(
	    hf_bort_model.bert.embeddings.position_embeddings.weight					,				'encoder.position_weight'							)
	snake_case__	       =   check_and_map_params(
	    hf_bort_model.bert.embeddings.LayerNorm.bias					,				'encoder.layer_norm.beta'							)
	snake_case__	       =   check_and_map_params(
	    hf_bort_model.bert.embeddings.LayerNorm.weight					,				'encoder.layer_norm.gamma'							)

	# Inspired by RoBERTa conversion script, we just zero them out (Bort does not use them)
	snake_case__	       =   torch.zeros_like(
	    hf_bort_model.bert.embeddings.token_type_embeddings.weight.data							)

	for i in range(hf_bort_config.num_hidden_layers							):
		snake_case__	       =   hf_bort_model.bert.encoder.layer[i]

		# self attention
		snake_case__	       =   layer.attention.self

		snake_case__	       =   check_and_map_params(
		    self_attn.key.bias.data					,				f'''encoder.transformer_cells.{i}.attention_cell.proj_key.bias'''							)

		snake_case__	       =   check_and_map_params(
		    self_attn.key.weight.data					,				f'''encoder.transformer_cells.{i}.attention_cell.proj_key.weight'''							)
		snake_case__	       =   check_and_map_params(
		    self_attn.query.bias.data					,				f'''encoder.transformer_cells.{i}.attention_cell.proj_query.bias'''							)
		snake_case__	       =   check_and_map_params(
		    self_attn.query.weight.data					,				f'''encoder.transformer_cells.{i}.attention_cell.proj_query.weight'''							)
		snake_case__	       =   check_and_map_params(
		    self_attn.value.bias.data					,				f'''encoder.transformer_cells.{i}.attention_cell.proj_value.bias'''							)
		snake_case__	       =   check_and_map_params(
		    self_attn.value.weight.data					,				f'''encoder.transformer_cells.{i}.attention_cell.proj_value.weight'''							)

		# self attention output
		snake_case__	       =   layer.attention.output

		snake_case__	       =   check_and_map_params(
		    self_output.dense.bias					,				f'''encoder.transformer_cells.{i}.proj.bias'''							)
		snake_case__	       =   check_and_map_params(
		    self_output.dense.weight					,				f'''encoder.transformer_cells.{i}.proj.weight'''							)
		snake_case__	       =   check_and_map_params(
		    self_output.LayerNorm.bias					,				f'''encoder.transformer_cells.{i}.layer_norm.beta'''							)
		snake_case__	       =   check_and_map_params(
		    self_output.LayerNorm.weight					,				f'''encoder.transformer_cells.{i}.layer_norm.gamma'''							)

		# intermediate
		snake_case__	       =   layer.intermediate

		snake_case__	       =   check_and_map_params(
		    intermediate.dense.bias					,				f'''encoder.transformer_cells.{i}.ffn.ffn_1.bias'''							)
		snake_case__	       =   check_and_map_params(
		    intermediate.dense.weight					,				f'''encoder.transformer_cells.{i}.ffn.ffn_1.weight'''							)

		# output
		snake_case__	       =   layer.output

		snake_case__	       =   check_and_map_params(
		    bert_output.dense.bias					,				f'''encoder.transformer_cells.{i}.ffn.ffn_2.bias'''							)
		snake_case__	       =   check_and_map_params(
		    bert_output.dense.weight					,				f'''encoder.transformer_cells.{i}.ffn.ffn_2.weight'''							)
		snake_case__	       =   check_and_map_params(
		    bert_output.LayerNorm.bias					,				f'''encoder.transformer_cells.{i}.ffn.layer_norm.beta'''							)
		snake_case__	       =   check_and_map_params(
		    bert_output.LayerNorm.weight					,				f'''encoder.transformer_cells.{i}.ffn.layer_norm.gamma'''							)

	# Save space and energy 🎄
	hf_bort_model.half()

	# Compare output of both models
	snake_case__	       =   RobertaTokenizer.from_pretrained('roberta-base'							)

	snake_case__	       =   tokenizer.encode_plus(UpperCAmelCase_							)['input_ids']

	# Get gluon output
	snake_case__	       =   mx.nd.array([input_ids]							)
	snake_case__	       =   original_bort(inputs=UpperCAmelCase_					,				token_types=[]							)

	# Get Transformer output (save and reload model again)
	hf_bort_model.save_pretrained(UpperCAmelCase_							)
	snake_case__	       =   BertModel.from_pretrained(UpperCAmelCase_							)
	hf_bort_model.eval()

	snake_case__	       =   tokenizer.encode_plus(UpperCAmelCase_					,				return_tensors='pt'							)
	snake_case__	       =   hf_bort_model(**UpperCAmelCase_							)[0]

	snake_case__	       =   output_gluon[0].asnumpy()
	snake_case__	       =   output_hf[0].detach().numpy()

	snake_case__	       =   np.max(np.abs(hf_layer - gluon_layer							)							).item()
	snake_case__	       =   np.allclose(UpperCAmelCase_					,				UpperCAmelCase_					,				atol=1E-3							)

	if success:
		print('✔️ Both model do output the same tensors'							)
	else:
		print('❌ Both model do **NOT** output the same tensors'							)
		print('Absolute difference is:'					,				UpperCAmelCase_							)


if __name__ == "__main__":
			a__					:			Optional[int]          =     argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    '''--bort_checkpoint_path''', default=None, type=str, required=True, help='''Path the official Bort params file.'''
			)
			parser.add_argument(
			    '''--pytorch_dump_folder_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
			)
			a__					:			str          =     parser.parse_args()
			convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)










'''simple docstring'''




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__					:			str          =     logging.getLogger(__name__)

a__					:			Tuple          =     '''pytorch_model.bin'''






@dataclasses.dataclass
class 				__snake_case     :
	__lowerCAmelCase       						=  dataclasses.field(
	    metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models.'''}							)
	__lowerCAmelCase       						=  dataclasses.field(
	    default=__magic_name__       ,       metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co.'''}       ,       )






@dataclasses.dataclass
class 				__snake_case     :
	__lowerCAmelCase       						=  dataclasses.field(metadata={'''help''': '''A csv or a json file containing the training data.'''}							)
	__lowerCAmelCase       						=  dataclasses.field(metadata={'''help''': '''A csv or a json file containing the data to predict on.'''}							)
	__lowerCAmelCase       						=  dataclasses.field(
	    default=__magic_name__       ,       metadata={'''help''': '''A csv or a json file containing the validation data.'''}							)
	__lowerCAmelCase       						=  dataclasses.field(
	    default=__magic_name__       ,       metadata={'''help''': '''The name of the task to train on.'''}       ,       )
	__lowerCAmelCase       						=  dataclasses.field(
	    default=__magic_name__       ,       metadata={'''help''': '''The list of labels for the task.'''}							)






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





def 						__lowerCamelCase			(	UpperCAmelCase_					,				UpperCAmelCase_					,				UpperCAmelCase_					,				UpperCAmelCase_					,				UpperCAmelCase_					,				UpperCAmelCase_							)   ->int:
	snake_case__	       =   datasets.concatenate_datasets([infer_input, infer_output]					,				axis=1							)

	if args.do_filter_by_confidence:
		snake_case__	       =   dataset.filter(lambda UpperCAmelCase_							: example["probability"] > args.confidence_threshold							)

	if args.do_filter_by_val_performance:
		assert eval_result >= 0.0 and eval_result <= 1.0
		snake_case__	       =   int(eval_result * len(UpperCAmelCase_							)							)
		print(UpperCAmelCase_							)
		snake_case__	       =   dataset.sort('probability'					,				reverse=UpperCAmelCase_							)
		snake_case__	       =   dataset.select(range(UpperCAmelCase_							)							)

	snake_case__	       =   dataset.remove_columns(['label', 'probability']							)
	snake_case__	       =   dataset.rename_column('prediction'					,				'label'							)
	snake_case__	       =   dataset.map(lambda UpperCAmelCase_							: {"label": idalabel[example["label"]]}							)
	snake_case__	       =   dataset.shuffle(seed=args.seed							)

	snake_case__	       =   os.path.join(UpperCAmelCase_					,				f'''train_pseudo.{args.data_file_extension}'''							)
	if args.data_file_extension == "csv":
		dataset.to_csv(UpperCAmelCase_					,				index=UpperCAmelCase_							)
	else:
		dataset.to_json(UpperCAmelCase_							)







def 						__lowerCamelCase			(	UpperCAmelCase_					,				UpperCAmelCase_					,				UpperCAmelCase_					,				UpperCAmelCase_					,				**UpperCAmelCase_							)   ->Tuple:
	snake_case__	       =   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()

	snake_case__	       =   STModelArguments(model_name_or_path=UpperCAmelCase_							)
	snake_case__	       =   STDataArguments(train_file=UpperCAmelCase_					,				infer_file=UpperCAmelCase_							)
	snake_case__	       =   STTrainingArguments(output_dir=UpperCAmelCase_							)
	snake_case__	       =   argparse.Namespace()

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

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

    # Sanity checks
	snake_case__	       =   {}
	snake_case__	       =   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
	snake_case__	       =   args.train_file
	snake_case__	       =   args.infer_file

	if args.evaluation_strategy != IntervalStrategy.NO.value:
		assert args.eval_file is not None
		snake_case__	       =   args.eval_file

	for key in data_files:
		snake_case__	       =   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:
			snake_case__	       =   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...'							)
	snake_case__	       =   f'''{args.output_dir}/self-train_iter-{{}}'''.format
	snake_case__	       =   data_dir_format(0							)

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

	snake_case__	       =   None
	snake_case__	       =   None
	snake_case__	       =   0
	snake_case__	       =   False
	# Show the progress bar
	snake_case__	       =   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							)							):
		snake_case__	       =   data_dir_format(UpperCAmelCase_							)
		assert os.path.exists(UpperCAmelCase_							)

		# Stage 1: initial fine-tuning for iteration = 0 or pseudo-training for
		# iteration > 0
		snake_case__	       =   os.path.join(UpperCAmelCase_					,				'stage-1'							)
		snake_case__	       =   {
		    '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(UpperCAmelCase_					,				UpperCAmelCase_							):
				arguments_dict.update({key: value}							)

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

		if iteration > 0 and args.finetune_on_labeled_data:
			# Stage 2 (optional): fine-tuning on the original labeled data
			snake_case__	       =   os.path.join(UpperCAmelCase_					,				'best-checkpoint'							)
			snake_case__	       =   os.path.join(UpperCAmelCase_					,				'stage-2'							)
			# Update arguments_dict
			snake_case__	       =   model_path
			snake_case__	       =   data_files['train']
			snake_case__	       =   current_output_dir

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

		snake_case__	       =   iteration
		snake_case__	       =   data_dir_format(iteration + 1							)

		snake_case__	       =   AutoConfig.from_pretrained(os.path.join(UpperCAmelCase_					,				'best-checkpoint'							)							)
		snake_case__	       =   config.idalabel
		snake_case__	       =   os.path.join(UpperCAmelCase_					,				'eval_results_best-checkpoint.json'							)
		snake_case__	       =   os.path.join(UpperCAmelCase_					,				'test_results_best-checkpoint.json'							)
		assert os.path.exists(UpperCAmelCase_							)

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

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

		snake_case__	       =   os.path.join(UpperCAmelCase_					,				f'''train_pseudo.{args.data_file_extension}'''							)

		if args.evaluation_strategy != IntervalStrategy.NO.value:
			snake_case__	       =   eval_result

			if best_iteration is None:
				snake_case__	       =   new_iteration
				snake_case__	       =   new_eval_result
			else:
				if new_eval_result - best_eval_result > args.early_stopping_threshold:
					snake_case__	       =   new_iteration
					snake_case__	       =   new_eval_result
					snake_case__	       =   0
				else:
					if new_eval_result == best_eval_result:
						snake_case__	       =   new_iteration
						snake_case__	       =   new_eval_result
					early_stopping_patience_counter += 1

				if early_stopping_patience_counter >= args.early_stopping_patience:
					snake_case__	       =   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'					,				UpperCAmelCase_							)
		logger.info('Best evaluation result: %s = %f'					,				args.eval_metric					,				UpperCAmelCase_							)
		accelerator.wait_for_everyone()
		if accelerator.is_main_process:
			shutil.copy(
			    os.path.join(UpperCAmelCase_					,				f'''eval_results_iter-{iteration}.json'''							)					,				os.path.join(UpperCAmelCase_					,				'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					,				UpperCAmelCase_							)
		accelerator.wait_for_everyone()
		if accelerator.is_main_process:
			shutil.copy(
			    os.path.join(UpperCAmelCase_					,				f'''eval_results_iter-{args.max_selftrain_iterations - 1}.json'''							)					,				os.path.join(UpperCAmelCase_					,				'eval_results_best-iteration.json'							)					,				)



