code: 
  
  
  
  
  
  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.albert.modeling_albert import (
      ALBERT_INPUTS_DOCSTRING,
      ALBERT_START_DOCSTRING,
      AlbertModel,
      AlbertPreTrainedModel,
      AlbertTransformer,
  )
  
  
  UpperCAmelCase			:			Tuple	  = logging.getLogger(__name__)
  
  
  
  
  
  
  
  class        __lowerCAmelCase (					UpperCamelCase_):
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None							)     ->		Tuple:
  
  
  
          '''simple docstring'''
  
  
  
  
          if current_layer == 0:
            a__			:      Dict            =self.embedding_hidden_mapping_in(__UpperCamelCase							)
          else:
            a__			:      List[Any]            =hidden_states[0]
  
          a__			:      str            =int(self.config.num_hidden_layers / self.config.num_hidden_groups							)
  
          # Index of the hidden group
          a__			:      Any            =int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)							)
  
          a__			:      Union[str, Any]            =self.albert_layer_groups[group_idx](
              __UpperCamelCase	,   __UpperCamelCase	,   head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]	,   )
          a__			:      str            =layer_group_output[0]
  
          return (hidden_states,)
  
  
  
  
  
  
  
  @add_start_docstrings(
      """The bare ALBERT Model transformer with PABEE outputting raw hidden-states without any specific head on top."""		,       UpperCamelCase_		,       )
  class        __lowerCAmelCase (					UpperCamelCase_):
  
  
  
  
  
  
  
        def __init__(					self	,   lowerCAmelCase__							)     ->		Tuple:
  
  
  
          '''simple docstring'''
  
  
  
  
          super().__init__(__UpperCamelCase							)
  
          a__			:      List[Any]            =AlbertTransformerWithPabee(__UpperCamelCase							)
  
          self.init_weights()
          a__			:      Union[str, Any]            =0
          a__			:      Optional[Any]            =0
          a__			:      Optional[Any]            =0
  
          a__			:      List[str]            =0
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__							)     ->		Union[str, Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Dict            =threshold
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__							)     ->		Optional[int]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      int            =patience
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		Dict:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Dict            =0
          a__			:      Any            =0
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		str:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Optional[Any]            =self.inference_layers_num / self.inference_instances_num
          a__			:      Dict            =(
              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      _lowercase						(					self	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None	,   lowerCAmelCase__=False	,   )     ->		Optional[int]:
  
  
  
          '''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:
            a__			:      str            =input_ids.size()
          elif inputs_embeds is not None:
            a__			:      Optional[Any]            =inputs_embeds.size()[:-1]
          else:
            raise ValueError("You have to specify either input_ids or inputs_embeds"							)
  
          a__			:      int            =input_ids.device if input_ids is not None else inputs_embeds.device
  
          if attention_mask is None:
            a__			:      Any            =torch.ones(__UpperCamelCase	,   device=__UpperCamelCase							)
          if token_type_ids is None:
            a__			:      Optional[int]            =torch.zeros(__UpperCamelCase	,   dtype=torch.long	,   device=__UpperCamelCase							)
  
          a__			:      Tuple            =attention_mask.unsqueeze(1							).unsqueeze(2							)
          a__			:      Union[str, Any]            =extended_attention_mask.to(dtype=self.dtype							)  # fp16 compatibility
          a__			:      List[str]            =(1.0 - extended_attention_mask) * -1_0_0_0_0.0
          a__			:      Optional[Any]            =self.get_head_mask(__UpperCamelCase	,   self.config.num_hidden_layers							)
  
          a__			:      int            =self.embeddings(
              __UpperCamelCase	,   position_ids=__UpperCamelCase	,   token_type_ids=__UpperCamelCase	,   inputs_embeds=__UpperCamelCase							)
          a__			:      int            =embedding_output
  
          if self.training:
            a__			:      Optional[Any]            =[]
            for i in range(self.config.num_hidden_layers							):
              a__			:      str            =self.encoder.adaptive_forward(
                  __UpperCamelCase	,   current_layer=__UpperCamelCase	,   attention_mask=__UpperCamelCase	,   head_mask=__UpperCamelCase	,   )
  
              a__			:      Tuple            =self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]							)							)
              a__			:      Dict            =output_layers[i](output_dropout(__UpperCamelCase							)							)
              res.append(__UpperCamelCase							)
          elif self.patience == 0:  # Use all layers for inference
            a__			:      List[Any]            =self.encoder(__UpperCamelCase	,   __UpperCamelCase	,   head_mask=__UpperCamelCase							)
            a__			:      Tuple            =self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]							)							)
            a__			:      Any            =[output_layers[self.config.num_hidden_layers - 1](__UpperCamelCase							)]
          else:
            a__			:      List[str]            =0
            a__			:      int            =None
            a__			:      Dict            =0
            for i in range(self.config.num_hidden_layers							):
              calculated_layer_num += 1
              a__			:      Optional[Any]            =self.encoder.adaptive_forward(
                  __UpperCamelCase	,   current_layer=__UpperCamelCase	,   attention_mask=__UpperCamelCase	,   head_mask=__UpperCamelCase	,   )
  
              a__			:      Optional[Any]            =self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]							)							)
              a__			:      Optional[int]            =output_layers[i](__UpperCamelCase							)
              if regression:
                a__			:      Dict            =logits.detach()
                if patient_result is not None:
                  a__			:      int            =patient_result.detach()
                if (patient_result is not None) and torch.abs(patient_result - labels							) < self.regression_threshold:
                  patient_counter += 1
                else:
                  a__			:      List[str]            =0
              else:
                a__			:      Optional[Any]            =logits.detach().argmax(dim=1							)
                if patient_result is not None:
                  a__			:      Tuple            =patient_result.detach().argmax(dim=1							)
                if (patient_result is not None) and torch.all(labels.eq(__UpperCamelCase							)							):
                  patient_counter += 1
                else:
                  a__			:      List[Any]            =0
  
              a__			:      Union[str, Any]            =logits
              if patient_counter == self.patience:
                break
            a__			:      Optional[int]            =[patient_result]
            self.inference_layers_num += calculated_layer_num
            self.inference_instances_num += 1
  
          return res
  
  
  
  
  
  
  
  @add_start_docstrings(
      """Albert 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        __lowerCAmelCase (					UpperCamelCase_):
  
  
  
  
  
  
  
        def __init__(					self	,   lowerCAmelCase__							)     ->		List[str]:
  
  
  
          '''simple docstring'''
  
  
  
  
          super().__init__(__UpperCamelCase							)
          a__			:      Union[str, Any]            =config.num_labels
  
          a__			:      Optional[Any]            =AlbertModelWithPabee(__UpperCamelCase							)
          a__			:      int            =nn.Dropout(config.classifier_dropout_prob							)
          a__			:      Dict            =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      _lowercase						(					self	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None	,   lowerCAmelCase__=None	,   )     ->		int:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Optional[int]            =self.albert(
              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	,   )
  
          a__			:      Tuple            =(logits[-1],)
  
          if labels is not None:
            a__			:      List[Any]            =None
            a__			:      Dict            =0
            for ix, logits_item in enumerate(__UpperCamelCase							):
              if self.num_labels == 1:
                #  We are doing regression
                a__			:      int            =MSELoss()
                a__			:      Union[str, Any]            =loss_fct(logits_item.view(-1							)	,   labels.view(-1							)							)
              else:
                a__			:      List[Any]            =CrossEntropyLoss()
                a__			:      Any            =loss_fct(logits_item.view(-1	,   self.num_labels							)	,   labels.view(-1							)							)
              if total_loss is None:
                a__			:      Optional[Any]            =loss
              else:
                total_loss += loss * (ix + 1)
              total_weights += ix + 1
            a__			:      Optional[int]            =(total_loss / total_weights,) + outputs
  
          return outputs
  
  
  
  
  
  
  
style_context: 
  
  
  
  
  
  import os
  import tempfile
  from pathlib import Path
  from unittest import TestCase
  
  import pyarrow as pa
  import pytest
  
  from datasets.arrow_dataset import Dataset
  from datasets.arrow_reader import ArrowReader, BaseReader, FileInstructions, ReadInstruction, make_file_instructions
  from datasets.info import DatasetInfo
  from datasets.splits import NamedSplit, Split, SplitDict, SplitInfo
  
  from .utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases
  
  
  
  
  
  
  
  class        __lowerCAmelCase (					UpperCamelCase__):
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__	,   lowerCAmelCase__=False							)     ->		List[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__      ,							a__      ,							a__			:      Optional[int]            =(
              filename_skip_take["filename"],
              filename_skip_take["skip"] if "skip" in filename_skip_take else None,
              filename_skip_take["take"] if "take" in filename_skip_take else None,
          )
          open(os.path.join(lowerCAmelCase__							)	,   "wb"							).close()
          a__			:      Optional[int]            =pa.Table.from_pydict({"filename": [Path(lowerCAmelCase__							).name] * 1_0_0}							)
          if take == -1:
            a__			:      int            =len(lowerCAmelCase__							) - skip
          if skip is not None and take is not None:
            a__			:      Optional[Any]            =pa_table.slice(lowerCAmelCase__	,   lowerCAmelCase__							)
          return pa_table
  
  
  
  
  
  
  
  class        __lowerCAmelCase (					UpperCamelCase__):
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		str:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      List[str]            ="my_name"
          a__			:      List[Any]            =SplitInfo(name="train"	,   num_examples=1_0_0							)
          a__			:      List[str]            =SplitInfo(name="test"	,   num_examples=1_0_0							)
          a__			:      List[Any]            =[train_info, test_info]
          a__			:      Union[str, Any]            =SplitDict()
          split_dict.add(lowerCAmelCase__							)
          split_dict.add(lowerCAmelCase__							)
          a__			:      Dict            =DatasetInfo(splits=lowerCAmelCase__							)
  
          with tempfile.TemporaryDirectory() as tmp_dir:
            a__			:      Optional[Any]            =ReaderTest(lowerCAmelCase__	,   lowerCAmelCase__							)
  
            a__			:      Any            ="test[:33%]"
            a__			:      int            =Dataset(**reader.read(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							)							)
            self.assertEqual(dset["filename"][0]	,   F'''{name}-test'''							)
            self.assertEqual(dset.num_rows	,   3_3							)
            self.assertEqual(dset.num_columns	,   1							)
  
            a__			:      Tuple            =["train", "test[:33%]"]
            a__			:      str            =[Split.TRAIN, ReadInstruction.from_spec("test[:33%]"							)]
            for instructions in [instructionsa, instructionsa]:
              a__			:      Optional[int]            =[reader.read(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							) for instr in instructions]
              a__      ,							a__			:      Optional[int]            =(Dataset(**lowerCAmelCase__							) for dataset_kwargs in datasets_kwargs)
              self.assertEqual(train_dset["filename"][0]	,   F'''{name}-train'''							)
              self.assertEqual(train_dset.num_rows	,   1_0_0							)
              self.assertEqual(train_dset.num_columns	,   1							)
              self.assertIsInstance(train_dset.split	,   lowerCAmelCase__							)
              self.assertEqual(str(train_dset.split							)	,   "train"							)
              self.assertEqual(test_dset["filename"][0]	,   F'''{name}-test'''							)
              self.assertEqual(test_dset.num_rows	,   3_3							)
              self.assertEqual(test_dset.num_columns	,   1							)
              self.assertIsInstance(test_dset.split	,   lowerCAmelCase__							)
              self.assertEqual(str(test_dset.split							)	,   "test[:33%]"							)
              del train_dset, test_dset
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		Optional[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Any            ="my_name"
          a__			:      List[Any]            =SplitInfo(name="train"	,   num_examples=1_0_0_0	,   shard_lengths=[1_0_0] * 1_0							)
          a__			:      Optional[Any]            =[train_info]
          a__			:      Any            =SplitDict()
          split_dict.add(lowerCAmelCase__							)
          a__			:      Any            =DatasetInfo(splits=lowerCAmelCase__							)
  
          with tempfile.TemporaryDirectory() as tmp_dir:
            a__			:      Optional[int]            =ReaderTest(lowerCAmelCase__	,   lowerCAmelCase__							)
  
            a__			:      Union[str, Any]            ="train[:33%]"
            a__			:      Union[str, Any]            =Dataset(**reader.read(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							)							)
            self.assertEqual(dset["filename"][0]	,   F'''{name}-train-00000-of-00010'''							)
            self.assertEqual(dset["filename"][-1]	,   F'''{name}-train-00003-of-00010'''							)
            self.assertEqual(dset.num_rows	,   3_3_0							)
            self.assertEqual(dset.num_columns	,   1							)
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		List[str]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Dict            =SplitInfo(name="train"	,   num_examples=1_0_0							)
          a__			:      str            =SplitInfo(name="test"	,   num_examples=1_0_0							)
          a__			:      int            =SplitDict()
          split_dict.add(lowerCAmelCase__							)
          split_dict.add(lowerCAmelCase__							)
          a__			:      int            =DatasetInfo(splits=lowerCAmelCase__							)
  
          with tempfile.TemporaryDirectory() as tmp_dir:
            a__			:      Tuple            =ReaderTest(lowerCAmelCase__	,   lowerCAmelCase__							)
  
            a__			:      Tuple            =[
                {"filename": os.path.join(lowerCAmelCase__	,   "train"							)},
                {"filename": os.path.join(lowerCAmelCase__	,   "test"							), "skip": 1_0, "take": 1_0},
            ]
            a__			:      Optional[Any]            =Dataset(**reader.read_files(lowerCAmelCase__	,   original_instructions="train+test[10:20]"							)							)
            self.assertEqual(dset.num_rows	,   1_1_0							)
            self.assertEqual(dset.num_columns	,   1							)
            del dset
  @pytest.mark.parametrize("in_memory" , [False, True]		)
  def   _A (	SCREAMING_SNAKE_CASE  :     Optional[Any] , SCREAMING_SNAKE_CASE  :     Optional[Any] , SCREAMING_SNAKE_CASE  :     Optional[Any]		):
  
  
  
  
  
    """simple docstring"""
  
  
  
    a__			:      int            =arrow_file
    with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
      a__			:      str            =ArrowReader.read_table(SCREAMING_SNAKE_CASE , in_memory=SCREAMING_SNAKE_CASE		)
    assert table.shape == dataset.data.shape
    assert set(table.column_names		) == set(dataset.data.column_names		)
    assert dict(table.to_pydict()		) == dict(dataset.data.to_pydict()		)  # to_pydict returns OrderedDict
  @pytest.mark.parametrize("in_memory" , [False, True]		)
  def   _A (	SCREAMING_SNAKE_CASE  :     Tuple , SCREAMING_SNAKE_CASE  :     List[Any] , SCREAMING_SNAKE_CASE  :     str		):
  
  
  
  
  
    """simple docstring"""
  
  
  
    a__			:      Optional[int]            =arrow_file
    a__			:      Tuple            =ArrowReader("" , SCREAMING_SNAKE_CASE		)
    with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
      a__			:      Union[str, Any]            =reader.read_files([{"filename": filename}] , in_memory=SCREAMING_SNAKE_CASE		)
    assert dataset_kwargs.keys() == {"arrow_table", "info", "split"}
    a__			:      Dict            =dataset_kwargs["arrow_table"]
    assert table.shape == dataset.data.shape
    assert set(table.column_names		) == set(dataset.data.column_names		)
    assert dict(table.to_pydict()		) == dict(dataset.data.to_pydict()		)  # to_pydict returns OrderedDict
  def   _A (	):
  
  
  
  
  
    """simple docstring"""
  
  
  
    assert ReadInstruction("train" , to=10 , unit="abs"		).to_spec() == "train[:10]"
    assert ReadInstruction("train" , from_=-80 , to=10 , unit="%"		).to_spec() == "train[-80%:10%]"
  
    a__			:      List[Any]            ="train+test"
    assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE		).to_spec() == spec_train_test
  
    a__			:      Union[str, Any]            ="train[2:10]"
    assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE		).to_spec() == spec_train_abs
  
    a__			:      List[Any]            ="train[15%:-20%]"
    assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE		).to_spec() == spec_train_pct
  
    a__			:      Dict            ="train[:10%](closest)"
    assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE		).to_spec() == "train[:10%]"
  
    a__			:      Tuple            ="train[:10%](pct1_dropremainder)"
    assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE		).to_spec() == spec_train_pct_rounding
  
    a__			:      List[Any]            ="train[:10%](pct1_dropremainder)+test[-10%:](pct1_dropremainder)"
    assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE		).to_spec() == spec_train_test_pct_rounding
  def   _A (	):
  
  
  
  
  
    """simple docstring"""
  
  
  
    a__			:      Optional[Any]            ="dummy"
    a__			:      Dict            =[SplitInfo(name="train" , num_examples=100		)]
    a__			:      Optional[Any]            ="train[:33%]"
    a__			:      int            ="arrow"
    a__			:      Any            ="prefix"
  
    a__			:      Tuple            =make_file_instructions(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		)
    assert isinstance(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		)
    assert file_instructions.num_examples == 33
    assert file_instructions.file_instructions == [
        {"filename": os.path.join(SCREAMING_SNAKE_CASE , f'''{name}-train.arrow'''		), "skip": 0, "take": 33}
    ]
  
    a__			:      List[str]            =[SplitInfo(name="train" , num_examples=100 , shard_lengths=[10] * 10		)]
    a__			:      Union[str, Any]            =make_file_instructions(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		)
    assert isinstance(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		)
    assert file_instructions.num_examples == 33
    assert file_instructions.file_instructions == [
        {"filename": os.path.join(SCREAMING_SNAKE_CASE , f'''{name}-train-00000-of-00010.arrow'''		), "skip": 0, "take": -1},
        {"filename": os.path.join(SCREAMING_SNAKE_CASE , f'''{name}-train-00001-of-00010.arrow'''		), "skip": 0, "take": -1},
        {"filename": os.path.join(SCREAMING_SNAKE_CASE , f'''{name}-train-00002-of-00010.arrow'''		), "skip": 0, "take": -1},
        {"filename": os.path.join(SCREAMING_SNAKE_CASE , f'''{name}-train-00003-of-00010.arrow'''		), "skip": 0, "take": 3},
    ]
  
  
  
  
  
  @pytest.mark.parametrize("name, expected_exception" , [(None, TypeError), ("", ValueError)]		)
  def   _A (	SCREAMING_SNAKE_CASE  :     Tuple , SCREAMING_SNAKE_CASE  :     Tuple		):
  
  
  
  
  
    """simple docstring"""
  
  
  
    a__			:      Dict            =[SplitInfo(name="train" , num_examples=100		)]
    a__			:      Any            ="train"
    a__			:      Optional[Any]            ="arrow"
    a__			:      List[Any]            ="prefix_path"
    with pytest.raises(SCREAMING_SNAKE_CASE		):
      a__			:      Optional[Any]            =make_file_instructions(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		)
  
  
  
  
  
  
  