



'''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  			=						logging.getLogger(__name__)

__a  			=						"pytorch_model.bin"







@dataclasses.dataclass
class        UpperCAmelCase_			:
    """simple docstring"""







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







@dataclasses.dataclass
class        UpperCAmelCase_			:
    """simple docstring"""







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







@dataclasses.dataclass
class        UpperCAmelCase_			:
    """simple docstring"""







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





def 		__snake_case(							_lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase						)						->					List[Any]:
      snake_case__    :       List[Any]       		= datasets.concatenate_datasets([infer_input, infer_output]	,  axis=1						)

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

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

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

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





def 		__snake_case(							_lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  **_lowerCAmelCase						)						->					List[Any]:
      snake_case__    :       List[str]       		= 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__    :       Tuple       		= STModelArguments(model_name_or_path=_lowerCAmelCase						)
      snake_case__    :       Optional[Any]       		= STDataArguments(train_file=_lowerCAmelCase	,  infer_file=_lowerCAmelCase						)
      snake_case__    :       List[Any]       		= STTrainingArguments(output_dir=_lowerCAmelCase						)
      snake_case__    :       int       		= argparse.Namespace()

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

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

    # Sanity checks
      snake_case__    :       Union[str, Any]       		= {}
      snake_case__    :       str       		= 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__    :       Dict       		= args.train_file
      snake_case__    :       int       		= args.infer_file

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

      for key in data_files:
            snake_case__    :       List[str]       		= 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__    :       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..."""						)
      snake_case__    :       List[Any]       		= f"{args.output_dir}/self-train_iter-{{}}".format
      snake_case__    :       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=_lowerCAmelCase						)
                  os.makedirs(_lowerCAmelCase	,  exist_ok=_lowerCAmelCase						)
      accelerator.wait_for_everyone()

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

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

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

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

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

            snake_case__    :       Dict       		= iteration
            snake_case__    :       List[Any]       		= data_dir_format(iteration + 1						)

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

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

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

            snake_case__    :       Dict       		= os.path.join(_lowerCAmelCase	,  f"train_pseudo.{args.data_file_extension}"						)

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

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

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





'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__a  			=						{
    "configuration_bigbird_pegasus": [
        "BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "BigBirdPegasusConfig",
        "BigBirdPegasusOnnxConfig",
    ],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __a  			=						[
           "BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST",
           "BigBirdPegasusForCausalLM",
           "BigBirdPegasusForConditionalGeneration",
           "BigBirdPegasusForQuestionAnswering",
           "BigBirdPegasusForSequenceClassification",
           "BigBirdPegasusModel",
           "BigBirdPegasusPreTrainedModel",
       ]


if TYPE_CHECKING:
       from .configuration_bigbird_pegasus import (
           BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP,
           BigBirdPegasusConfig,
           BigBirdPegasusOnnxConfig,
       )

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_bigbird_pegasus import (
                  BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST,
                  BigBirdPegasusForCausalLM,
                  BigBirdPegasusForConditionalGeneration,
                  BigBirdPegasusForQuestionAnswering,
                  BigBirdPegasusForSequenceClassification,
                  BigBirdPegasusModel,
                  BigBirdPegasusPreTrainedModel,
              )


else:
       import sys

       __a  			=						_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
