






"""simple docstring"""
from collections import OrderedDict
from typing import Any, List, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


_lowerCAmelCase       =				logging.get_logger(__name__)


_lowerCAmelCase       =				{
    "Salesforce/codegen-350M-nl": "https://huggingface.co/Salesforce/codegen-350M-nl/resolve/main/config.json",
    "Salesforce/codegen-350M-multi": "https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/config.json",
    "Salesforce/codegen-350M-mono": "https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/config.json",
    "Salesforce/codegen-2B-nl": "https://huggingface.co/Salesforce/codegen-2B-nl/resolve/main/config.json",
    "Salesforce/codegen-2B-multi": "https://huggingface.co/Salesforce/codegen-2B-multi/resolve/main/config.json",
    "Salesforce/codegen-2B-mono": "https://huggingface.co/Salesforce/codegen-2B-mono/resolve/main/config.json",
    "Salesforce/codegen-6B-nl": "https://huggingface.co/Salesforce/codegen-6B-nl/resolve/main/config.json",
    "Salesforce/codegen-6B-multi": "https://huggingface.co/Salesforce/codegen-6B-multi/resolve/main/config.json",
    "Salesforce/codegen-6B-mono": "https://huggingface.co/Salesforce/codegen-6B-mono/resolve/main/config.json",
    "Salesforce/codegen-16B-nl": "https://huggingface.co/Salesforce/codegen-16B-nl/resolve/main/config.json",
    "Salesforce/codegen-16B-multi": "https://huggingface.co/Salesforce/codegen-16B-multi/resolve/main/config.json",
    "Salesforce/codegen-16B-mono": "https://huggingface.co/Salesforce/codegen-16B-mono/resolve/main/config.json",
}



class   SCREAMING_SNAKE_CASE	(					_SCREAMING_SNAKE_CASE    ):





	"""simple docstring"""




	_A							:					Optional[int] 			=					"""codegen"""
	_A							:					str 			=					{
	    """max_position_embeddings""": """n_positions""",
	    """hidden_size""": """n_embd""",
	    """num_attention_heads""": """n_head""",
	    """num_hidden_layers""": """n_layer""",
	}





	def __init__(self  ,  lowerCAmelCase_=50400  ,  lowerCAmelCase_=2048  ,  lowerCAmelCase_=2048  ,  lowerCAmelCase_=4096  ,  lowerCAmelCase_=28  ,  lowerCAmelCase_=16  ,  lowerCAmelCase_=64  ,  lowerCAmelCase_=None  ,  lowerCAmelCase_="gelu_new"  ,  lowerCAmelCase_=0.0  ,  lowerCAmelCase_=0.0  ,  lowerCAmelCase_=0.0  ,  lowerCAmelCase_=1e-5  ,  lowerCAmelCase_=0.02  ,  lowerCAmelCase_=True  ,  lowerCAmelCase_=50256  ,  lowerCAmelCase_=50256  ,  lowerCAmelCase_=False  ,  **lowerCAmelCase_  ,  ):
		A_		:    List[Any]						     = vocab_size
		A_		:    str						     = n_ctx
		A_		:    Tuple						     = n_positions
		A_		:    List[str]						     = n_embd
		A_		:    List[str]						     = n_layer
		A_		:    Any						     = n_head
		A_		:    Any						     = n_inner
		A_		:    Union[str, Any]						     = rotary_dim
		A_		:    str						     = activation_function
		A_		:    Optional[Any]						     = resid_pdrop
		A_		:    List[Any]						     = embd_pdrop
		A_		:    Optional[int]						     = attn_pdrop
		A_		:    Any						     = layer_norm_epsilon
		A_		:    Tuple						     = initializer_range
		A_		:    List[str]						     = use_cache

		A_		:    Any						     = bos_token_id
		A_		:    Any						     = eos_token_id

		super().__init__(
		    bos_token_id=lowerCAmelCase_  ,  eos_token_id=lowerCAmelCase_  ,  tie_word_embeddings=lowerCAmelCase_  ,  **lowerCAmelCase_				)






class   SCREAMING_SNAKE_CASE	(					_SCREAMING_SNAKE_CASE    ):





	"""simple docstring"""




	def __init__(self  ,  lowerCAmelCase_  ,  lowerCAmelCase_ = "default"  ,  lowerCAmelCase_ = None  ,  lowerCAmelCase_ = False  ,  ):
		super().__init__(lowerCAmelCase_  ,  task=lowerCAmelCase_  ,  patching_specs=lowerCAmelCase_  ,  use_past=lowerCAmelCase_				)
		if not getattr(self._config  ,  """pad_token_id"""  ,  lowerCAmelCase_				):
			# TODO: how to do that better?
			A_		:    Optional[Any]						     = 0





	@property
	def  lowerCamelCase(self				):
		A_		:    List[Any]						     = OrderedDict({"""input_ids""": {0: """batch""", 1: """sequence"""}}				)
		if self.use_past:
			self.fill_with_past_key_values_(lowerCAmelCase_  ,  direction="""inputs"""				)
			A_		:    Tuple						     = {0: """batch""", 1: """past_sequence + sequence"""}
		else:
			A_		:    Optional[int]						     = {0: """batch""", 1: """sequence"""}

		return common_inputs





	@property
	def  lowerCamelCase(self				):
		return self._config.n_layer





	@property
	def  lowerCamelCase(self				):
		return self._config.n_head





	def  lowerCamelCase(self  ,  lowerCAmelCase_  ,  lowerCAmelCase_ = -1  ,  lowerCAmelCase_ = -1  ,  lowerCAmelCase_ = False  ,  lowerCAmelCase_ = None  ,  ):
		A_		:    Tuple						     = super(lowerCAmelCase_  ,  self				).generate_dummy_inputs(
		    lowerCAmelCase_  ,  batch_size=lowerCAmelCase_  ,  seq_length=lowerCAmelCase_  ,  is_pair=lowerCAmelCase_  ,  framework=lowerCAmelCase_				)

		# We need to order the input in the way they appears in the forward()
		A_		:    str						     = OrderedDict({"""input_ids""": common_inputs["""input_ids"""]}				)

		# Need to add the past_keys
		if self.use_past:
			if not is_torch_available():
				raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""				)
			else:
				import torch

				A_     ,  A_		:    Union[str, Any]						     = common_inputs["""input_ids"""].shape
				# Not using the same length for past_key_values
				A_		:    List[Any]						     = seqlen + 2
				A_		:    Any						     = (
				    batch,
				    self.num_attention_heads,
				    past_key_values_length,
				    self._config.hidden_size // self.num_attention_heads,
				)
				A_		:    List[str]						     = [
				    (torch.zeros(lowerCAmelCase_				), torch.zeros(lowerCAmelCase_				)) for _ in range(self.num_layers				)
				]

		A_		:    int						     = common_inputs["""attention_mask"""]
		if self.use_past:
			A_		:    Tuple						     = ordered_inputs["""attention_mask"""].dtype
			A_		:    Optional[Any]						     = torch.cat(
			    [ordered_inputs["""attention_mask"""], torch.ones(lowerCAmelCase_  ,  lowerCAmelCase_  ,  dtype=lowerCAmelCase_				)]  ,  dim=1				)

		return ordered_inputs





	@property
	def  lowerCamelCase(self				):
		return 13









"""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


_lowerCAmelCase       =				logging.getLogger(__name__)

_lowerCAmelCase       =				"pytorch_model.bin"



@dataclasses.dataclass
class   SCREAMING_SNAKE_CASE	:





	"""simple docstring"""




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






@dataclasses.dataclass
class   SCREAMING_SNAKE_CASE	:





	"""simple docstring"""




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






@dataclasses.dataclass
class   SCREAMING_SNAKE_CASE	:





	"""simple docstring"""




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






def       __UpperCamelCase						(		snake_case__ , snake_case__ , snake_case__ , snake_case__ , snake_case__ , snake_case__       ):
	A_		:    Tuple						     = datasets.concatenate_datasets([infer_input, infer_output] , axis=1       )

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

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

	A_		:    int						     = dataset.remove_columns(["""label""", """probability"""]       )
	A_		:    Any						     = dataset.rename_column("""prediction""" , """label"""       )
	A_		:    str						     = dataset.map(lambda snake_case__       : {"label": idalabel[example["label"]]}       )
	A_		:    Optional[Any]						     = dataset.shuffle(seed=args.seed       )

	A_		:    Dict						     = os.path.join(snake_case__ , F"""train_pseudo.{args.data_file_extension}"""       )
	if args.data_file_extension == "csv":
		dataset.to_csv(snake_case__ , index=snake_case__       )
	else:
		dataset.to_json(snake_case__       )






def       __UpperCamelCase						(		snake_case__ , snake_case__ , snake_case__ , snake_case__ , **snake_case__       ):
	A_		:    Tuple						     = 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_		:    Optional[int]						     = STModelArguments(model_name_or_path=snake_case__       )
	A_		:    int						     = STDataArguments(train_file=snake_case__ , infer_file=snake_case__       )
	A_		:    Dict						     = STTrainingArguments(output_dir=snake_case__       )
	A_		:    Any						     = argparse.Namespace()

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

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

    # 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_		:    Any						     = args.train_file
	A_		:    Optional[Any]						     = 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_		:    Union[str, Any]						     = 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_		:    str						     = 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_		:    Optional[Any]						     = data_dir_format(0       )

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

	A_		:    Dict						     = None
	A_		:    Optional[int]						     = None
	A_		:    List[Any]						     = 0
	A_		:    List[Any]						     = False
	# Show the progress bar
	A_		:    Optional[int]						     = 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_		:    int						     = data_dir_format(snake_case__       )
		assert os.path.exists(snake_case__       )

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

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

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

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

		A_		:    Optional[Any]						     = iteration
		A_		:    Optional[Any]						     = data_dir_format(iteration + 1       )

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

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

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

		A_		:    str						     = os.path.join(snake_case__ , F"""train_pseudo.{args.data_file_extension}"""       )

		if args.evaluation_strategy != IntervalStrategy.NO.value:
			A_		:    List[str]						     = eval_result

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

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

