
import logging

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


__lowerCamelCase				=			logging.getLogger(__name__)






class     UpperCAmelCase						(		A_					):



   def 							_SCREAMING_SNAKE_CASE  (self						:   int       ,		snake_case__						:   Tuple       ,		snake_case__						:   List[Any]       ,		snake_case__						:   Tuple=None       ,		snake_case__						:   Any=None					)						->      str:


          '''simple docstring'''


          snake_case			:  int     							= self.layer[current_layer](snake_case__       ,		snake_case__       ,		head_mask[current_layer]					)

          snake_case			:  Dict     							= layer_outputs[0]

          return hidden_states






@add_start_docstrings(
    "The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top."			,A_			,)
class     UpperCAmelCase						(		A_					):



   def __init__(self						:   int       ,		snake_case__						:   List[Any]					)						->      str:


          '''simple docstring'''


          super().__init__(snake_case__					)

          snake_case			:  List[Any]     							= BertEncoderWithPabee(snake_case__					)

          self.init_weights()
          snake_case			:  Optional[int]     							= 0
          snake_case			:  str     							= 0
          snake_case			:  str     							= 0

          snake_case			:  Optional[int]     							= 0



   def 							_SCREAMING_SNAKE_CASE  (self						:   Optional[Any]       ,		snake_case__						:   Dict					)						->      str:


          '''simple docstring'''


          snake_case			:  int     							= threshold



   def 							_SCREAMING_SNAKE_CASE  (self						:   List[Any]       ,		snake_case__						:   Optional[int]					)						->      Optional[int]:


          '''simple docstring'''


          snake_case			:  Union[str, Any]     							= patience



   def 							_SCREAMING_SNAKE_CASE  (self						:   str					)						->      Dict:


          '''simple docstring'''


          snake_case			:  List[str]     							= 0
          snake_case			:  str     							= 0



   def 							_SCREAMING_SNAKE_CASE  (self						:   Optional[Any]					)						->      str:


          '''simple docstring'''


          snake_case			:  Dict     							= self.inference_layers_num / self.inference_instances_num
          snake_case			:  Union[str, Any]     							= (
              f"""*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ="""
              f""" {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***"""
          )
          print(snake_case__					)



   @add_start_docstrings_to_model_forward(snake_case__					)
   def 							_SCREAMING_SNAKE_CASE  (self						:   Tuple       ,		snake_case__						:   str=None       ,		snake_case__						:   List[str]=None       ,		snake_case__						:   List[str]=None       ,		snake_case__						:   Union[str, Any]=None       ,		snake_case__						:   List[Any]=None       ,		snake_case__						:   str=None       ,		snake_case__						:   Union[str, Any]=None       ,		snake_case__						:   int=None       ,		snake_case__						:   List[Any]=None       ,		snake_case__						:   Dict=None       ,		snake_case__						:   Tuple=False       ,		)						->      Dict:


          '''simple docstring'''


          if input_ids is not None and inputs_embeds is not None:
                 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time"					)
          elif input_ids is not None:
                 snake_case			:  List[str]     							= input_ids.size()
          elif inputs_embeds is not None:
                 snake_case			:  Union[str, Any]     							= inputs_embeds.size()[:-1]
          else:
                 raise ValueError("You have to specify either input_ids or inputs_embeds"					)

          snake_case			:  Union[str, Any]     							= input_ids.device if input_ids is not None else inputs_embeds.device

          if attention_mask is None:
                 snake_case			:  Optional[Any]     							= torch.ones(snake_case__       ,		device=snake_case__					)
          if token_type_ids is None:
                 snake_case			:  int     							= torch.zeros(snake_case__       ,		dtype=torch.long       ,		device=snake_case__					)

          # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
          # ourselves in which case we just need to make it broadcastable to all heads.
          snake_case			:  torch.Tensor     							= self.get_extended_attention_mask(snake_case__       ,		snake_case__       ,		snake_case__					)

          # If a 2D ou 3D attention mask is provided for the cross-attention
          # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
          if self.config.is_decoder and encoder_hidden_states is not None:
                 snake_case							,   snake_case							,   snake_case			:  Dict     							= encoder_hidden_states.size()
                 snake_case			:  Tuple     							= (encoder_batch_size, encoder_sequence_length)
                 if encoder_attention_mask is None:
                        snake_case			:  int     							= torch.ones(snake_case__       ,		device=snake_case__					)
                 snake_case			:  List[str]     							= self.invert_attention_mask(snake_case__					)
          else:
                 snake_case			:  str     							= None

          # Prepare head mask if needed
          # 1.0 in head_mask indicate we keep the head
          # attention_probs has shape bsz x n_heads x N x N
          # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
          # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
          snake_case			:  str     							= self.get_head_mask(snake_case__       ,		self.config.num_hidden_layers					)

          snake_case			:  List[str]     							= self.embeddings(
              input_ids=snake_case__       ,		position_ids=snake_case__       ,		token_type_ids=snake_case__       ,		inputs_embeds=snake_case__					)
          snake_case			:  int     							= embedding_output

          if self.training:
                 snake_case			:  Tuple     							= []
                 for i in range(self.config.num_hidden_layers					):
                        snake_case			:  int     							= self.encoder.adaptive_forward(
                            snake_case__       ,		current_layer=snake_case__       ,		attention_mask=snake_case__       ,		head_mask=snake_case__					)

                        snake_case			:  Optional[int]     							= self.pooler(snake_case__					)
                        snake_case			:  Optional[int]     							= output_layers[i](output_dropout(snake_case__					)					)
                        res.append(snake_case__					)
          elif self.patience == 0:  # Use all layers for inference
                 snake_case			:  Dict     							= self.encoder(
                     snake_case__       ,		attention_mask=snake_case__       ,		head_mask=snake_case__       ,		encoder_hidden_states=snake_case__       ,		encoder_attention_mask=snake_case__       ,		)
                 snake_case			:  Optional[int]     							= self.pooler(encoder_outputs[0]					)
                 snake_case			:  List[Any]     							= [output_layers[self.config.num_hidden_layers - 1](snake_case__					)]
          else:
                 snake_case			:  List[str]     							= 0
                 snake_case			:  int     							= None
                 snake_case			:  Any     							= 0
                 for i in range(self.config.num_hidden_layers					):
                        calculated_layer_num += 1
                        snake_case			:  List[str]     							= self.encoder.adaptive_forward(
                            snake_case__       ,		current_layer=snake_case__       ,		attention_mask=snake_case__       ,		head_mask=snake_case__					)

                        snake_case			:  Any     							= self.pooler(snake_case__					)
                        snake_case			:  List[str]     							= output_layers[i](snake_case__					)
                        if regression:
                               snake_case			:  Tuple     							= logits.detach()
                               if patient_result is not None:
                                      snake_case			:  Optional[int]     							= patient_result.detach()
                               if (patient_result is not None) and torch.abs(patient_result - labels					) < self.regression_threshold:
                                      patient_counter += 1
                               else:
                                      snake_case			:  Optional[Any]     							= 0
                        else:
                               snake_case			:  List[Any]     							= logits.detach().argmax(dim=1					)
                               if patient_result is not None:
                                      snake_case			:  List[Any]     							= patient_result.detach().argmax(dim=1					)
                               if (patient_result is not None) and torch.all(labels.eq(snake_case__					)					):
                                      patient_counter += 1
                               else:
                                      snake_case			:  Any     							= 0

                        snake_case			:  int     							= logits
                        if patient_counter == self.patience:
                               break
                 snake_case			:  Optional[int]     							= [patient_result]
                 self.inference_layers_num += calculated_layer_num
                 self.inference_instances_num += 1

          return res






@add_start_docstrings(
    "Bert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of\n    the pooled output) e.g. for GLUE tasks. "			,A_			,)
class     UpperCAmelCase						(		A_					):



   def __init__(self						:   Optional[int]       ,		snake_case__						:   Dict					)						->      Any:


          '''simple docstring'''


          super().__init__(snake_case__					)
          snake_case			:  Any     							= config.num_labels

          snake_case			:  Dict     							= BertModelWithPabee(snake_case__					)
          snake_case			:  Dict     							= nn.Dropout(config.hidden_dropout_prob					)
          snake_case			:  Union[str, Any]     							= nn.ModuleList(
              [nn.Linear(config.hidden_size       ,		self.config.num_labels					) for _ in range(config.num_hidden_layers					)]					)

          self.init_weights()



   @add_start_docstrings_to_model_forward(snake_case__					)
   def 							_SCREAMING_SNAKE_CASE  (self						:   Optional[Any]       ,		snake_case__						:   List[Any]=None       ,		snake_case__						:   Optional[Any]=None       ,		snake_case__						:   Optional[int]=None       ,		snake_case__						:   Optional[int]=None       ,		snake_case__						:   Dict=None       ,		snake_case__						:   Dict=None       ,		snake_case__						:   Optional[int]=None       ,		)						->      List[str]:


          '''simple docstring'''


          snake_case			:  Optional[int]     							= self.bert(
              input_ids=snake_case__       ,		attention_mask=snake_case__       ,		token_type_ids=snake_case__       ,		position_ids=snake_case__       ,		head_mask=snake_case__       ,		inputs_embeds=snake_case__       ,		output_dropout=self.dropout       ,		output_layers=self.classifiers       ,		regression=self.num_labels == 1       ,		)

          snake_case			:  Any     							= (logits[-1],)

          if labels is not None:
                 snake_case			:  List[Any]     							= None
                 snake_case			:  List[Any]     							= 0
                 for ix, logits_item in enumerate(snake_case__					):
                        if self.num_labels == 1:
                               #  We are doing regression
                               snake_case			:  Optional[Any]     							= MSELoss()
                               snake_case			:  Tuple     							= loss_fct(logits_item.view(-1					)       ,		labels.view(-1					)					)
                        else:
                               snake_case			:  Optional[Any]     							= CrossEntropyLoss()
                               snake_case			:  str     							= loss_fct(logits_item.view(-1       ,		self.num_labels					)       ,		labels.view(-1					)					)
                        if total_loss is None:
                               snake_case			:  Tuple     							= loss
                        else:
                               total_loss += loss * (ix + 1)
                        total_weights += ix + 1
                 snake_case			:  Any     							= (total_loss / total_weights,) + outputs

          return outputs







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      UpperCamelCase   (       __lowerCamelCase	:					Namespace       ):
       return TrainCommand(__lowerCamelCase       )






class     UpperCAmelCase						(		A_					):



   @staticmethod
   def 							_SCREAMING_SNAKE_CASE  (snake_case__						:   ArgumentParser					)						->      int:


          '''simple docstring'''


          snake_case			:  Any     							= parser.add_parser("train"       ,		help="CLI tool to train a model on a task."					)

          train_parser.add_argument(
              "--train_data"       ,		type=snake_case__       ,		required=snake_case__       ,		help="path to train (and optionally evaluation) dataset as a csv with tab separated labels and sentences."       ,		)
          train_parser.add_argument(
              "--column_label"       ,		type=snake_case__       ,		default=0       ,		help="Column of the dataset csv file with example labels."					)
          train_parser.add_argument(
              "--column_text"       ,		type=snake_case__       ,		default=1       ,		help="Column of the dataset csv file with example texts."					)
          train_parser.add_argument(
              "--column_id"       ,		type=snake_case__       ,		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=snake_case__       ,		default=""       ,		help="path to validation dataset."					)
          train_parser.add_argument(
              "--validation_split"       ,		type=snake_case__       ,		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=snake_case__       ,		default="./"       ,		help="path to saved the trained model."					)

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



   def __init__(self						:   List[Any]       ,		snake_case__						:   Namespace					)						->      Tuple:


          '''simple docstring'''


          snake_case			:  Any     							= logging.get_logger("transformers-cli/training"					)

          snake_case			:  List[Any]     							= "tf" if is_tf_available() else "torch"

          os.makedirs(args.output       ,		exist_ok=snake_case__					)
          snake_case			:  Any     							= args.output

          snake_case			:  List[Any]     							= args.column_label
          snake_case			:  Tuple     							= args.column_text
          snake_case			:  str     							= args.column_id

          self.logger.info(f"""Loading {args.task} pipeline for {args.model}"""					)
          if args.task == "text_classification":
                 snake_case			:  Optional[int]     							= 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}"""					)
          snake_case			:  Tuple     							= 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       ,		)
          snake_case			:  Optional[Any]     							= None
          if args.validation_data:
                 self.logger.info(f"""Loading validation dataset from {args.validation_data}"""					)
                 snake_case			:  Dict     							= 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       ,		)

          snake_case			:  Union[str, Any]     							= args.validation_split
          snake_case			:  Optional[Any]     							= args.train_batch_size
          snake_case			:  List[str]     							= args.valid_batch_size
          snake_case			:  List[Any]     							= args.learning_rate
          snake_case			:  List[Any]     							= args.adam_epsilon



   def 							_SCREAMING_SNAKE_CASE  (self						:   int					)						->      List[str]:


          '''simple docstring'''


          if self.framework == "tf":
                 return self.run_tf()
          return self.run_torch()



   def 							_SCREAMING_SNAKE_CASE  (self						:   Tuple					)						->      Optional[Any]:


          '''simple docstring'''


          raise NotImplementedError



   def 							_SCREAMING_SNAKE_CASE  (self						:   Optional[Any]					)						->      Tuple:


          '''simple docstring'''


          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					)





