


import os
import sys
import tempfile

import torch

from .state import AcceleratorState
from .utils import PrecisionType, PrepareForLaunch, is_mps_available, patch_environment



def  _a			(						__UpperCamelCase  ,  __UpperCamelCase=()  ,  __UpperCamelCase=None  ,  __UpperCamelCase="no"  ,  __UpperCamelCase="29500" ):
     a_    :     Union[str, Any]										=    False
     a_    :     Tuple										=    False
     if any(key.startswith("""KAGGLE""" ) for key in os.environ.keys() ):
          a_    :     Union[str, Any]										=    True
     elif "IPython" in sys.modules:
          a_    :     Optional[int]										=    """google.colab""" in str(sys.modules["""IPython"""].get_ipython() )

     try:
          a_    :     Optional[Any]										=    PrecisionType(mixed_precision.lower() )
     except ValueError:
          raise ValueError(
              F'''Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}.''' )

     if (in_colab or in_kaggle) and (os.environ.get("""TPU_NAME"""  ,  __UpperCamelCase ) is not None):
          # TPU launch
          import torch_xla.distributed.xla_multiprocessing as xmp

          if len(AcceleratorState._shared_state ) > 0:
               raise ValueError(
                   """To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside """
                   """your training function. Restart your notebook and make sure no cells initializes an """
                   """`Accelerator`.""" )
          if num_processes is None:
               a_    :     Any										=    8

          a_    :     Tuple										=    PrepareForLaunch(__UpperCamelCase  ,  distributed_type="""TPU""" )
          print(F'''Launching a training on {num_processes} TPU cores.''' )
          xmp.spawn(__UpperCamelCase  ,  args=__UpperCamelCase  ,  nprocs=__UpperCamelCase  ,  start_method="""fork""" )
     elif in_colab:
          # No need for a distributed launch otherwise as it's either CPU or one GPU.
          if torch.cuda.is_available():
               print("""Launching training on one GPU.""" )
          else:
               print("""Launching training on one CPU.""" )
          function(*__UpperCamelCase )
     else:
          if num_processes is None:
               raise ValueError(
                   """You have to specify the number of GPUs you would like to use, add `num_processes=...` to your call.""" )

          if num_processes > 1:
               # Multi-GPU launch
               from torch.multiprocessing import start_processes
               from torch.multiprocessing.spawn import ProcessRaisedException

               if len(AcceleratorState._shared_state ) > 0:
                    raise ValueError(
                        """To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized """
                        """inside your training function. Restart your notebook and make sure no cells initializes an """
                        """`Accelerator`.""" )

               if torch.cuda.is_initialized():
                    raise ValueError(
                        """To launch a multi-GPU training from your notebook, you need to avoid running any instruction """
                        """using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA """
                        """function.""" )

               # torch.distributed will expect a few environment variable to be here. We set the ones common to each
               # process here (the other ones will be set be the launcher).
               with patch_environment(
                   world_size=__UpperCamelCase  ,  master_addr="""127.0.01"""  ,  master_port=__UpperCamelCase  ,  mixed_precision=__UpperCamelCase ):
                    a_    :     str										=    PrepareForLaunch(__UpperCamelCase  ,  distributed_type="""MULTI_GPU""" )
                    print(F'''Launching training on {num_processes} GPUs.''' )
                    try:
                         start_processes(__UpperCamelCase  ,  args=__UpperCamelCase  ,  nprocs=__UpperCamelCase  ,  start_method="""fork""" )
                    except ProcessRaisedException as e:
                         if "Cannot re-initialize CUDA in forked subprocess" in e.args[0]:
                              raise RuntimeError(
                                  """CUDA has been initialized before the `notebook_launcher` could create a forked subprocess. """
                                  """This likely stems from an outside import causing issues once the `notebook_launcher()` is called. """
                                  """Please review your imports and test them when running the `notebook_launcher()` to identify """
                                  """which one is problematic.""" ) from e

          else:
               # No need for a distributed launch otherwise as it's either CPU, GPU or MPS.
               if is_mps_available():
                    a_    :     str										=    """1"""
                    print("""Launching training on MPS.""" )
               elif torch.cuda.is_available():
                    print("""Launching training on one GPU.""" )
               else:
                    print("""Launching training on CPU.""" )
               function(*__UpperCamelCase )



def  _a			(						__UpperCamelCase  ,  __UpperCamelCase=()  ,  __UpperCamelCase=2 ):
     from torch.multiprocessing import start_processes

     with tempfile.NamedTemporaryFile() as tmp_file:
          # torch.distributed will expect a few environment variable to be here. We set the ones common to each
          # process here (the other ones will be set be the launcher).
          with patch_environment(
              world_size=__UpperCamelCase  ,  master_addr="""127.0.01"""  ,  master_port="""29500"""  ,  accelerate_mixed_precision="""no"""  ,  accelerate_debug_rdv_file=tmp_file.name  ,  accelerate_use_cpu="""yes"""  ,  ):
               a_    :     List[Any]										=    PrepareForLaunch(__UpperCamelCase  ,  debug=__UpperCamelCase )
               start_processes(__UpperCamelCase  ,  args=__UpperCamelCase  ,  nprocs=__UpperCamelCase  ,  start_method="""fork""" )







import os
from argparse import ArgumentParser, Namespace

from ..data import SingleSentenceClassificationProcessor as Processor
from ..pipelines import TextClassificationPipeline
from ..utils import is_tf_available, is_torch_available, logging
from . import BaseTransformersCLICommand


if not is_tf_available() and not is_torch_available():
       raise RuntimeError('''At least one of PyTorch or TensorFlow 2.0+ should be installed to use CLI training''')

# TF training parameters
__lowerCamelCase        =     False
__lowerCamelCase        =     False



def  _a			(						__UpperCamelCase ):
     return TrainCommand(__UpperCamelCase )


class   a__	(						lowerCAmelCase_    ):

       @staticmethod
       def 						UpperCAmelCase(   lowerCamelCase_      :     ArgumentParser    ):
            a_    :     List[Any]										=    parser.add_parser("""train"""      ,							help="""CLI tool to train a model on a task."""    )

            train_parser.add_argument(
                """--train_data"""      ,							type=lowerCamelCase_      ,							required=lowerCamelCase_      ,							help="""path to train (and optionally evaluation) dataset as a csv with tab separated labels and sentences."""      ,							)
            train_parser.add_argument(
                """--column_label"""      ,							type=lowerCamelCase_      ,							default=0      ,							help="""Column of the dataset csv file with example labels."""    )
            train_parser.add_argument(
                """--column_text"""      ,							type=lowerCamelCase_      ,							default=1      ,							help="""Column of the dataset csv file with example texts."""    )
            train_parser.add_argument(
                """--column_id"""      ,							type=lowerCamelCase_      ,							default=2      ,							help="""Column of the dataset csv file with example ids."""    )
            train_parser.add_argument(
                """--skip_first_row"""      ,							action="""store_true"""      ,							help="""Skip the first row of the csv file (headers)."""    )

            train_parser.add_argument("""--validation_data"""      ,							type=lowerCamelCase_      ,							default=""""""      ,							help="""path to validation dataset."""    )
            train_parser.add_argument(
                """--validation_split"""      ,							type=lowerCamelCase_      ,							default=0.1      ,							help="""if validation dataset is not provided, fraction of train dataset to use as validation dataset."""      ,							)

            train_parser.add_argument("""--output"""      ,							type=lowerCamelCase_      ,							default="""./"""      ,							help="""path to saved the trained model."""    )

            train_parser.add_argument(
                """--task"""      ,							type=lowerCamelCase_      ,							default="""text_classification"""      ,							help="""Task to train the model on."""    )
            train_parser.add_argument(
                """--model"""      ,							type=lowerCamelCase_      ,							default="""bert-base-uncased"""      ,							help="""Model's name or path to stored model."""    )
            train_parser.add_argument("""--train_batch_size"""      ,							type=lowerCamelCase_      ,							default=3_2      ,							help="""Batch size for training."""    )
            train_parser.add_argument("""--valid_batch_size"""      ,							type=lowerCamelCase_      ,							default=6_4      ,							help="""Batch size for validation."""    )
            train_parser.add_argument("""--learning_rate"""      ,							type=lowerCamelCase_      ,							default=3E-5      ,							help="""Learning rate."""    )
            train_parser.add_argument("""--adam_epsilon"""      ,							type=lowerCamelCase_      ,							default=1E-08      ,							help="""Epsilon for Adam optimizer."""    )
            train_parser.set_defaults(func=lowerCamelCase_    )

       def __init__(   self      :     int      ,							lowerCamelCase_      :     Namespace    ):
            a_    :     Optional[int]										=    logging.get_logger("""transformers-cli/training"""    )

            a_    :     List[str]										=    """tf""" if is_tf_available() else """torch"""

            os.makedirs(args.output      ,							exist_ok=lowerCamelCase_    )
            a_    :     List[str]										=    args.output

            a_    :     Optional[int]										=    args.column_label
            a_    :     Tuple										=    args.column_text
            a_    :     Optional[int]										=    args.column_id

            self.logger.info(F'''Loading {args.task} pipeline for {args.model}'''    )
            if args.task == "text_classification":
                 a_    :     str										=    TextClassificationPipeline.from_pretrained(args.model    )
            elif args.task == "token_classification":
                 raise NotImplementedError
            elif args.task == "question_answering":
                 raise NotImplementedError

            self.logger.info(F'''Loading dataset from {args.train_data}'''    )
            a_    :     int										=    Processor.create_from_csv(
                args.train_data      ,							column_label=args.column_label      ,							column_text=args.column_text      ,							column_id=args.column_id      ,							skip_first_row=args.skip_first_row      ,							)
            a_    :     Union[str, Any]										=    None
            if args.validation_data:
                 self.logger.info(F'''Loading validation dataset from {args.validation_data}'''    )
                 a_    :     Tuple										=    Processor.create_from_csv(
                     args.validation_data      ,							column_label=args.column_label      ,							column_text=args.column_text      ,							column_id=args.column_id      ,							skip_first_row=args.skip_first_row      ,							)

            a_    :     str										=    args.validation_split
            a_    :     Tuple										=    args.train_batch_size
            a_    :     List[Any]										=    args.valid_batch_size
            a_    :     str										=    args.learning_rate
            a_    :     str										=    args.adam_epsilon

       def 						UpperCAmelCase(   self      :     List[str]    ):
            if self.framework == "tf":
                 return self.run_tf()
            return self.run_torch()

       def 						UpperCAmelCase(   self      :     Any    ):
            raise NotImplementedError

       def 						UpperCAmelCase(   self      :     List[str]    ):
            self.pipeline.fit(
                self.train_dataset      ,							validation_data=self.valid_dataset      ,							validation_split=self.validation_split      ,							learning_rate=self.learning_rate      ,							adam_epsilon=self.adam_epsilon      ,							train_batch_size=self.train_batch_size      ,							valid_batch_size=self.valid_batch_size      ,							)

            # Save trained pipeline
            self.pipeline.save_pretrained(self.output    )



