

"""simple docstring"""






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        _lowercase					(     __UpperCAmelCase     ):







     def 	lowerCAmelCase__    (      self ,						UpperCamelCase_ ,						UpperCamelCase_ ,						UpperCamelCase_=None ,						UpperCamelCase_=None       ):
      __magic_name__     =      self.layer[current_layer](UpperCamelCase_ ,						UpperCamelCase_ ,						head_mask[current_layer]       )

      __magic_name__     =      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.'''   , __UpperCAmelCase   , )
class        _lowercase					(     __UpperCAmelCase     ):







     def __init__(      self ,						UpperCamelCase_       ):
      super().__init__(UpperCamelCase_       )

      __magic_name__     =      BertEncoderWithPabee(UpperCamelCase_       )

      self.init_weights()
      __magic_name__     =      0
      __magic_name__     =      0
      __magic_name__     =      0

      __magic_name__     =      0







     def 	lowerCAmelCase__    (      self ,						UpperCamelCase_       ):
      __magic_name__     =      threshold







     def 	lowerCAmelCase__    (      self ,						UpperCamelCase_       ):
      __magic_name__     =      patience







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      0
      __magic_name__     =      0







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      self.inference_layers_num / self.inference_instances_num
      __magic_name__     =      (
          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(UpperCamelCase_       )







     @add_start_docstrings_to_model_forward(UpperCamelCase_       )
     def 	lowerCAmelCase__    (      self ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=False ,						):

      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:
       __magic_name__     =      input_ids.size()
      elif inputs_embeds is not None:
       __magic_name__     =      inputs_embeds.size()[:-1]
      else:
       raise ValueError('''You have to specify either input_ids or inputs_embeds'''       )

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

      if attention_mask is None:
       __magic_name__     =      torch.ones(UpperCamelCase_ ,						device=UpperCamelCase_       )
      if token_type_ids is None:
       __magic_name__     =      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.
      __magic_name__     =      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 self.config.is_decoder and encoder_hidden_states is not None:
       __magic_name__   ,						__magic_name__   ,						__magic_name__     =      encoder_hidden_states.size()
       __magic_name__     =      (encoder_batch_size, encoder_sequence_length)
       if encoder_attention_mask is None:
        __magic_name__     =      torch.ones(UpperCamelCase_ ,						device=UpperCamelCase_       )
       __magic_name__     =      self.invert_attention_mask(UpperCamelCase_       )
      else:
       __magic_name__     =      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]
      __magic_name__     =      self.get_head_mask(UpperCamelCase_ ,						self.config.num_hidden_layers       )

      __magic_name__     =      self.embeddings(
          input_ids=UpperCamelCase_ ,						position_ids=UpperCamelCase_ ,						token_type_ids=UpperCamelCase_ ,						inputs_embeds=UpperCamelCase_       )
      __magic_name__     =      embedding_output

      if self.training:
       __magic_name__     =      []
       for i in range(self.config.num_hidden_layers       ):
        __magic_name__     =      self.encoder.adaptive_forward(
            UpperCamelCase_ ,						current_layer=UpperCamelCase_ ,						attention_mask=UpperCamelCase_ ,						head_mask=UpperCamelCase_       )

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

        __magic_name__     =      self.pooler(UpperCamelCase_       )
        __magic_name__     =      output_layers[i](UpperCamelCase_       )
        if regression:
         __magic_name__     =      logits.detach()
         if patient_result is not None:
          __magic_name__     =      patient_result.detach()
         if (patient_result is not None) and torch.abs(patient_result - labels       ) < self.regression_threshold:
          patient_counter += 1
         else:
          __magic_name__     =      0
        else:
         __magic_name__     =      logits.detach().argmax(dim=1       )
         if patient_result is not None:
          __magic_name__     =      patient_result.detach().argmax(dim=1       )
         if (patient_result is not None) and torch.all(labels.eq(UpperCamelCase_       )       ):
          patient_counter += 1
         else:
          __magic_name__     =      0

        __magic_name__     =      logits
        if patient_counter == self.patience:
         break
       __magic_name__     =      [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
    the pooled output) e.g. for GLUE tasks. '''   , __UpperCAmelCase   , )
class        _lowercase					(     __UpperCAmelCase     ):







     def __init__(      self ,						UpperCamelCase_       ):
      super().__init__(UpperCamelCase_       )
      __magic_name__     =      config.num_labels

      __magic_name__     =      BertModelWithPabee(UpperCamelCase_       )
      __magic_name__     =      nn.Dropout(config.hidden_dropout_prob       )
      __magic_name__     =      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(UpperCamelCase_       )
     def 	lowerCAmelCase__    (      self ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						):
      __magic_name__     =      self.bert(
          input_ids=UpperCamelCase_ ,						attention_mask=UpperCamelCase_ ,						token_type_ids=UpperCamelCase_ ,						position_ids=UpperCamelCase_ ,						head_mask=UpperCamelCase_ ,						inputs_embeds=UpperCamelCase_ ,						output_dropout=self.dropout ,						output_layers=self.classifiers ,						regression=self.num_labels == 1 ,						)

      __magic_name__     =      (logits[-1],)

      if labels is not None:
       __magic_name__     =      None
       __magic_name__     =      0
       for ix, logits_item in enumerate(UpperCamelCase_       ):
        if self.num_labels == 1:
         #  We are doing regression
         __magic_name__     =      MSELoss()
         __magic_name__     =      loss_fct(logits_item.view(-1       ) ,						labels.view(-1       )       )
        else:
         __magic_name__     =      CrossEntropyLoss()
         __magic_name__     =      loss_fct(logits_item.view(-1 ,						self.num_labels       ) ,						labels.view(-1       )       )
        if total_loss is None:
         __magic_name__     =      loss
        else:
         total_loss += loss * (ix + 1)
        total_weights += ix + 1
       __magic_name__     =      (total_loss / total_weights,) + outputs

      return outputs




"""simple docstring"""






import os
from tempfile import TemporaryDirectory
from unittest import TestCase

import pytest
from absl.testing import parameterized

from datasets import config
from datasets.arrow_reader import HF_GCP_BASE_URL
from datasets.builder import DatasetBuilder
from datasets.dataset_dict import IterableDatasetDict
from datasets.iterable_dataset import IterableDataset
from datasets.load import dataset_module_factory, import_main_class
from datasets.utils.file_utils import cached_path


__lowerCamelCase										=		[
    {"dataset": "wikipedia", "config_name": "20220301.de"},
    {"dataset": "wikipedia", "config_name": "20220301.en"},
    {"dataset": "wikipedia", "config_name": "20220301.fr"},
    {"dataset": "wikipedia", "config_name": "20220301.frr"},
    {"dataset": "wikipedia", "config_name": "20220301.it"},
    {"dataset": "wikipedia", "config_name": "20220301.simple"},
    {"dataset": "snli", "config_name": "plain_text"},
    {"dataset": "eli5", "config_name": "LFQA_reddit"},
    {"dataset": "wiki40b", "config_name": "en"},
    {"dataset": "wiki_dpr", "config_name": "psgs_w100.nq.compressed"},
    {"dataset": "wiki_dpr", "config_name": "psgs_w100.nq.no_index"},
    {"dataset": "wiki_dpr", "config_name": "psgs_w100.multiset.no_index"},
    {"dataset": "natural_questions", "config_name": "default"},
]





def 	lowercase      (   __UpperCamelCase=True						) ->      List[Any]:
 if with_config:
  return [
      {
          "testcase_name": d["dataset"] + "/" + d["config_name"],
          "dataset": d["dataset"],
          "config_name": d["config_name"],
      }
      for d in DATASETS_ON_HF_GCP
  ]
 else:
  return [
      {"testcase_name": dataset, "dataset": dataset} for dataset in {d["dataset"] for d in DATASETS_ON_HF_GCP}
  ]
@parameterized.named_parameters(list_datasets_on_hf_gcp_parameters(with_config=__UpperCAmelCase     )     )
class        _lowercase					(     __UpperCAmelCase     ):
     _lowerCamelCase     =       None
     _lowerCamelCase     =       None







     def 	lowerCAmelCase__    (      self ,						UpperCamelCase_ ,						UpperCamelCase_       ):
      with TemporaryDirectory() as tmp_dir:
       __magic_name__     =      dataset_module_factory(UpperCamelCase_ ,						cache_dir=UpperCamelCase_       )

       __magic_name__     =      import_main_class(dataset_module.module_path ,						dataset=UpperCamelCase_       )

       __magic_name__     =      builder_cls(
           cache_dir=UpperCamelCase_ ,						config_name=UpperCamelCase_ ,						hash=dataset_module.hash ,						)

       __magic_name__     =      '''/'''.join(
           [
               HF_GCP_BASE_URL,
               builder_instance._relative_data_dir(with_hash=UpperCamelCase_       ).replace(os.sep ,						'''/'''       ),
               config.DATASET_INFO_FILENAME,
           ]       )
       __magic_name__     =      cached_path(UpperCamelCase_ ,						cache_dir=UpperCamelCase_       )
       self.assertTrue(os.path.exists(UpperCamelCase_       )       )





@pytest.mark.integration
def 	lowercase      (   __UpperCamelCase						) ->      Tuple:
 __magic_name__     =      tmp_path_factory.mktemp('''test_hf_gcp'''						) / '''test_wikipedia_simple'''
 __magic_name__     =      dataset_module_factory('''wikipedia'''      ,   cache_dir=__UpperCamelCase						)
 __magic_name__     =      import_main_class(dataset_module.module_path						)
 __magic_name__     =      builder_cls(
     cache_dir=__UpperCamelCase      ,   config_name='''20220301.frr'''      ,   hash=dataset_module.hash      ,   )
 # use the HF cloud storage, not the original download_and_prepare that uses apache-beam
 __magic_name__     =      None
 builder_instance.download_and_prepare()
 __magic_name__     =      builder_instance.as_dataset()
 assert ds





@pytest.mark.integration
def 	lowercase      (   __UpperCamelCase						) ->      Tuple:
 __magic_name__     =      dataset_module_factory('''wikipedia'''      ,   cache_dir=__UpperCamelCase						)
 __magic_name__     =      import_main_class(dataset_module.module_path      ,   dataset=__UpperCamelCase						)
 __magic_name__     =      builder_cls(
     cache_dir=__UpperCamelCase      ,   config_name='''20220301.frr'''      ,   hash=dataset_module.hash      ,   )
 __magic_name__     =      builder_instance.as_streaming_dataset()
 assert ds
 assert isinstance(__UpperCamelCase      ,   __UpperCamelCase						)
 assert "train" in ds
 assert isinstance(ds['''train''']      ,   __UpperCamelCase						)
 assert next(iter(ds['''train''']						)						)

