


import inspect
import re
from hashlib import shaaaa
from typing import Dict, List

from .arrow import arrow
from .audiofolder import audiofolder
from .csv import csv
from .imagefolder import imagefolder
from .json import json
from .pandas import pandas
from .parquet import parquet
from .sql import sql  # noqa F401
from .text import text






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE )->  str:




      """simple docstring"""

      snake_case_	=				[]
      for line in lines:
            snake_case_	=				re.sub(R'''#.*''' , '''''' , SCREAMING_SNAKE_CASE )  # remove comments
            if line:
                  filtered_lines.append(SCREAMING_SNAKE_CASE )
      snake_case_	=				'''\n'''.join(SCREAMING_SNAKE_CASE )

      # Make a hash from all this code
      snake_case_	=				full_str.encode('''utf-8''' )
      return shaaaa(SCREAMING_SNAKE_CASE ).hexdigest()


# get importable module names and hash for caching
UpperCAmelCase      =   {
    """csv""": (csv.__name__, _hash_python_lines(inspect.getsource(csv).splitlines())),
    """json""": (json.__name__, _hash_python_lines(inspect.getsource(json).splitlines())),
    """pandas""": (pandas.__name__, _hash_python_lines(inspect.getsource(pandas).splitlines())),
    """parquet""": (parquet.__name__, _hash_python_lines(inspect.getsource(parquet).splitlines())),
    """arrow""": (arrow.__name__, _hash_python_lines(inspect.getsource(arrow).splitlines())),
    """text""": (text.__name__, _hash_python_lines(inspect.getsource(text).splitlines())),
    """imagefolder""": (imagefolder.__name__, _hash_python_lines(inspect.getsource(imagefolder).splitlines())),
    """audiofolder""": (audiofolder.__name__, _hash_python_lines(inspect.getsource(audiofolder).splitlines())),
}

# Used to infer the module to use based on the data files extensions
UpperCAmelCase      =   {
    """.csv""": ("""csv""", {}),
    """.tsv""": ("""csv""", {"""sep""": """\t"""}),
    """.json""": ("""json""", {}),
    """.jsonl""": ("""json""", {}),
    """.parquet""": ("""parquet""", {}),
    """.arrow""": ("""arrow""", {}),
    """.txt""": ("""text""", {}),
}
_EXTENSION_TO_MODULE.update({ext: ("""imagefolder""", {}) for ext in imagefolder.ImageFolder.EXTENSIONS})
_EXTENSION_TO_MODULE.update({ext.upper(): ("""imagefolder""", {}) for ext in imagefolder.ImageFolder.EXTENSIONS})
_EXTENSION_TO_MODULE.update({ext: ("""audiofolder""", {}) for ext in audiofolder.AudioFolder.EXTENSIONS})
_EXTENSION_TO_MODULE.update({ext.upper(): ("""audiofolder""", {}) for ext in audiofolder.AudioFolder.EXTENSIONS})
UpperCAmelCase      =   {"""imagefolder""", """audiofolder"""}

# Used to filter data files based on extensions given a module name
UpperCAmelCase      =   {}
for _ext, (_module, _) in _EXTENSION_TO_MODULE.items():
   _MODULE_TO_EXTENSIONS.setdefault(_module, []).append(_ext)

_MODULE_TO_EXTENSIONS["imagefolder"].append(""".zip""")




_MODULE_TO_EXTENSIONS["audiofolder"].append(""".zip""")



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,
    BertEmbeddings,
    BertLayer,
    BertPooler,
    BertPreTrainedModel,
)






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE )->  List[Any]:




      """simple docstring"""

      snake_case_	=				torch.exp(SCREAMING_SNAKE_CASE )
      snake_case_	=				torch.sum(SCREAMING_SNAKE_CASE , dim=1 )  # sum of exp(x_i)
      snake_case_	=				torch.sum(x * exp_x , dim=1 )  # sum of x_i * exp(x_i)
      return torch.log(SCREAMING_SNAKE_CASE ) - B / A







class lowerCAmelCase_						(			nn.Module	):

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase ):
       super().__init__()
       snake_case_	=				config.output_attentions
       snake_case_	=				config.output_hidden_states
       snake_case_	=				nn.ModuleList([BertLayer(_UpperCAmelCase ) for _ in range(config.num_hidden_layers )] )
       snake_case_	=				nn.ModuleList([BertHighway(_UpperCAmelCase ) for _ in range(config.num_hidden_layers )] )

       snake_case_	=				[-1 for _ in range(config.num_hidden_layers )]



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase ):
       if (type(_UpperCAmelCase ) is float) or (type(_UpperCAmelCase ) is int):
             for i in range(len(self.early_exit_entropy ) ):
                   snake_case_	=				x
       else:
             snake_case_	=				x



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase ):
       snake_case_	=				pooler.state_dict()
       for highway in self.highway:
             for name, param in highway.pooler.state_dict().items():
                   param.copy_(loaded_model[name] )





 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				):
       snake_case_	=				()
       snake_case_	=				()
       snake_case_	=				()
       for i, layer_module in enumerate(self.layer ):
             if self.output_hidden_states:
                   snake_case_	=				all_hidden_states + (hidden_states,)

             snake_case_	=				layer_module(
                 _UpperCAmelCase		,				_UpperCAmelCase		,				head_mask[i]		,				_UpperCAmelCase		,				_UpperCAmelCase )
             snake_case_	=				layer_outputs[0]

             if self.output_attentions:
                   snake_case_	=				all_attentions + (layer_outputs[1],)

             snake_case_	=				(hidden_states,)
             if self.output_hidden_states:
                   snake_case_	=				current_outputs + (all_hidden_states,)
             if self.output_attentions:
                   snake_case_	=				current_outputs + (all_attentions,)

             snake_case_	=				self.highway[i](_UpperCAmelCase )
             # logits, pooled_output

             if not self.training:
                   snake_case_	=				highway_exit[0]
                   snake_case_	=				entropy(_UpperCAmelCase )
                   snake_case_	=				highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
                   snake_case_	=				all_highway_exits + (highway_exit,)

                   if highway_entropy < self.early_exit_entropy[i]:
                         snake_case_	=				(highway_logits,) + current_outputs[1:] + (all_highway_exits,)
                         raise HighwayException(_UpperCAmelCase		,				i + 1 )
             else:
                   snake_case_	=				all_highway_exits + (highway_exit,)

        # Add last layer
       if self.output_hidden_states:
             snake_case_	=				all_hidden_states + (hidden_states,)

       snake_case_	=				(hidden_states,)
       if self.output_hidden_states:
             snake_case_	=				outputs + (all_hidden_states,)
       if self.output_attentions:
             snake_case_	=				outputs + (all_attentions,)

       snake_case_	=				outputs + (all_highway_exits,)
       return outputs  # last-layer hidden state, (all hidden states), (all attentions), all highway exits







@add_start_docstrings(
    "The Bert Model transformer with early exiting (DeeBERT). " ,						lowerCamelCase__ ,						)
class lowerCAmelCase_						(			lowerCamelCase__	):

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase ):
       super().__init__(_UpperCAmelCase )
       snake_case_	=				config

       snake_case_	=				BertEmbeddings(_UpperCAmelCase )
       snake_case_	=				DeeBertEncoder(_UpperCAmelCase )
       snake_case_	=				BertPooler(_UpperCAmelCase )

       self.init_weights()



 def 						UpperCamelCase__					(     self ):
       self.encoder.init_highway_pooler(self.pooler )



 def 						UpperCamelCase__					(     self ):
       return self.embeddings.word_embeddings



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase ):
       snake_case_	=				value



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase ):
       for layer, heads in heads_to_prune.items():
             self.encoder.layer[layer].attention.prune_heads(_UpperCAmelCase )





 @add_start_docstrings_to_model_forward(_UpperCAmelCase )
 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				):
       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_	=				input_ids.size()
       elif inputs_embeds is not None:
             snake_case_	=				inputs_embeds.size()[:-1]
       else:
             raise ValueError('''You have to specify either input_ids or inputs_embeds''' )

       snake_case_	=				input_ids.device if input_ids is not None else inputs_embeds.device

       if attention_mask is None:
             snake_case_	=				torch.ones(_UpperCAmelCase		,				device=_UpperCAmelCase )
       if encoder_attention_mask is None:
             snake_case_	=				torch.ones(_UpperCAmelCase		,				device=_UpperCAmelCase )
       if token_type_ids is None:
             snake_case_	=				torch.zeros(_UpperCAmelCase		,				dtype=torch.long		,				device=_UpperCAmelCase )

       # 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_	=				self.get_extended_attention_mask(_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase )

       # 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 encoder_attention_mask.dim() == 3:
             snake_case_	=				encoder_attention_mask[:, None, :, :]
       if encoder_attention_mask.dim() == 2:
             snake_case_	=				encoder_attention_mask[:, None, None, :]

       snake_case_	=				encoder_extended_attention_mask.to(
           dtype=next(self.parameters() ).dtype )  # fp16 compatibility
       snake_case_	=				(1.0 - encoder_extended_attention_mask) * -10_000.0

       # 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_	=				self.get_head_mask(_UpperCAmelCase		,				self.config.num_hidden_layers )

       snake_case_	=				self.embeddings(
           input_ids=_UpperCAmelCase		,				position_ids=_UpperCAmelCase		,				token_type_ids=_UpperCAmelCase		,				inputs_embeds=_UpperCAmelCase )
       snake_case_	=				self.encoder(
           _UpperCAmelCase		,				attention_mask=_UpperCAmelCase		,				head_mask=_UpperCAmelCase		,				encoder_hidden_states=_UpperCAmelCase		,				encoder_attention_mask=_UpperCAmelCase		,				)
       snake_case_	=				encoder_outputs[0]
       snake_case_	=				self.pooler(_UpperCAmelCase )

       snake_case_	=				(
           sequence_output,
           pooled_output,
       ) + encoder_outputs[
           1:
       ]  # add hidden_states and attentions if they are here
       return outputs  # sequence_output, pooled_output, (hidden_states), (attentions), highway exits







class lowerCAmelCase_						(			lowerCamelCase__	):

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       snake_case_	=				message
       snake_case_	=				exit_layer  # start from 1!







class lowerCAmelCase_						(			nn.Module	):

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase ):
       super().__init__()
       snake_case_	=				BertPooler(_UpperCAmelCase )
       snake_case_	=				nn.Dropout(config.hidden_dropout_prob )
       snake_case_	=				nn.Linear(config.hidden_size		,				config.num_labels )





 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase ):
       # Pooler
       snake_case_	=				encoder_outputs[0]
       snake_case_	=				self.pooler(_UpperCAmelCase )
       # "return" pooler_output

       # BertModel
       snake_case_	=				(pooler_input, pooler_output) + encoder_outputs[1:]
       # "return" bmodel_output

       # Dropout and classification
       snake_case_	=				bmodel_output[1]

       snake_case_	=				self.dropout(_UpperCAmelCase )
       snake_case_	=				self.classifier(_UpperCAmelCase )

       return logits, pooled_output







@add_start_docstrings(
    "Bert Model (with early exiting - DeeBERT) with a classifier on top,\n    also takes care of multi-layer training. " ,						lowerCamelCase__ ,						)
class lowerCAmelCase_						(			lowerCamelCase__	):

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase ):
       super().__init__(_UpperCAmelCase )
       snake_case_	=				config.num_labels
       snake_case_	=				config.num_hidden_layers

       snake_case_	=				DeeBertModel(_UpperCAmelCase )
       snake_case_	=				nn.Dropout(config.hidden_dropout_prob )
       snake_case_	=				nn.Linear(config.hidden_size		,				self.config.num_labels )

       self.init_weights()





 @add_start_docstrings_to_model_forward(_UpperCAmelCase )
 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=None		,				_UpperCAmelCase=-1		,				_UpperCAmelCase=False		,				):
       snake_case_	=				self.num_layers
       try:
             snake_case_	=				self.bert(
                 _UpperCAmelCase		,				attention_mask=_UpperCAmelCase		,				token_type_ids=_UpperCAmelCase		,				position_ids=_UpperCAmelCase		,				head_mask=_UpperCAmelCase		,				inputs_embeds=_UpperCAmelCase		,				)
             # sequence_output, pooled_output, (hidden_states), (attentions), highway exits

             snake_case_	=				outputs[1]

             snake_case_	=				self.dropout(_UpperCAmelCase )
             snake_case_	=				self.classifier(_UpperCAmelCase )
             snake_case_	=				(logits,) + outputs[2:]  # add hidden states and attention if they are here
       except HighwayException as e:
             snake_case_	=				e.message
             snake_case_	=				e.exit_layer
             snake_case_	=				outputs[0]

       if not self.training:
             snake_case_	=				entropy(_UpperCAmelCase )
             snake_case_	=				[]
             snake_case_	=				[]
       if labels is not None:
             if self.num_labels == 1:
                   #  We are doing regression
                   snake_case_	=				MSELoss()
                   snake_case_	=				loss_fct(logits.view(-1 )		,				labels.view(-1 ) )
             else:
                   snake_case_	=				CrossEntropyLoss()
                   snake_case_	=				loss_fct(logits.view(-1		,				self.num_labels )		,				labels.view(-1 ) )

             # work with highway exits
             snake_case_	=				[]
             for highway_exit in outputs[-1]:
                   snake_case_	=				highway_exit[0]
                   if not self.training:
                         highway_logits_all.append(_UpperCAmelCase )
                         highway_entropy.append(highway_exit[2] )
                   if self.num_labels == 1:
                         #  We are doing regression
                         snake_case_	=				MSELoss()
                         snake_case_	=				loss_fct(highway_logits.view(-1 )		,				labels.view(-1 ) )
                   else:
                         snake_case_	=				CrossEntropyLoss()
                         snake_case_	=				loss_fct(highway_logits.view(-1		,				self.num_labels )		,				labels.view(-1 ) )
                   highway_losses.append(_UpperCAmelCase )

             if train_highway:
                   snake_case_	=				(sum(highway_losses[:-1] ),) + outputs
                   # exclude the final highway, of course
             else:
                   snake_case_	=				(loss,) + outputs
       if not self.training:
             snake_case_	=				outputs + ((original_entropy, highway_entropy), exit_layer)
             if output_layer >= 0:
                   snake_case_	=				(
                       (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
                   )  # use the highway of the last layer

       return outputs  # (loss), logits, (hidden_states), (attentions), (highway_exits)