






"""simple docstring"""




from __future__ import annotations

import inspect
import unittest

from transformers import MobileViTConfig
from transformers.file_utils import is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import numpy as np
     import tensorflow as tf

     from transformers import TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation, TFMobileViTModel
     from transformers.models.mobilevit.modeling_tf_mobilevit import TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import MobileViTImageProcessor




class  lowercase_     (_UpperCAmelCase  ):






      def       lowerCamelCase__     (     self						)						->Optional[Any]:







        '''simple docstring'''


        _a					      =						self.config_class(**self.inputs_dict						)
        self.parent.assertTrue(hasattr(a_  ,  "hidden_sizes"						)						)
        self.parent.assertTrue(hasattr(a_  ,  "neck_hidden_sizes"						)						)
        self.parent.assertTrue(hasattr(a_  ,  "num_attention_heads"						)						)







class  lowercase_     :






      def __init__(     self  ,  a_  ,  a_=1_3  ,  a_=3_2  ,  a_=2  ,  a_=3  ,  a_=6_4_0  ,  a_=4  ,  a_="silu"  ,  a_=3  ,  a_=3_2  ,  a_=0.1  ,  a_=0.1  ,  a_=0.1  ,  a_=0.02  ,  a_=True  ,  a_=True  ,  a_=1_0  ,  a_=None  ,  )						->Optional[Any]:







        '''simple docstring'''


        _a					      =						parent
        _a					      =						batch_size
        _a					      =						image_size
        _a					      =						patch_size
        _a					      =						num_channels
        _a					      =						last_hidden_size
        _a					      =						num_attention_heads
        _a					      =						hidden_act
        _a					      =						conv_kernel_size
        _a					      =						output_stride
        _a					      =						hidden_dropout_prob
        _a					      =						attention_probs_dropout_prob
        _a					      =						classifier_dropout_prob
        _a					      =						use_labels
        _a					      =						is_training
        _a					      =						num_labels
        _a					      =						initializer_range
        _a					      =						scope






      def       lowerCamelCase__     (     self						)						->List[str]:







        '''simple docstring'''


        _a					      =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

        _a					      =						None
        _a					      =						None
        if self.use_labels:
          _a					      =						ids_tensor([self.batch_size]  ,  self.num_labels						)
          _a					      =						ids_tensor([self.batch_size, self.image_size, self.image_size]  ,  self.num_labels						)

        _a					      =						self.get_config()

        return config, pixel_values, labels, pixel_labels






      def       lowerCamelCase__     (     self						)						->Tuple:







        '''simple docstring'''


        return MobileViTConfig(
            image_size=self.image_size  ,  patch_size=self.patch_size  ,  num_channels=self.num_channels  ,  num_attention_heads=self.num_attention_heads  ,  hidden_act=self.hidden_act  ,  conv_kernel_size=self.conv_kernel_size  ,  output_stride=self.output_stride  ,  hidden_dropout_prob=self.hidden_dropout_prob  ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,  classifier_dropout_prob=self.classifier_dropout_prob  ,  initializer_range=self.initializer_range  ,  )






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_  ,  a_						)						->str:







        '''simple docstring'''


        _a					      =						TFMobileViTModel(config=a_						)
        _a					      =						model(a_  ,  training=a_						)
        _a					      =						_a					      =						self.image_size // self.output_stride
        self.parent.assertEqual(
            result.last_hidden_state.shape  ,  (self.batch_size, self.last_hidden_size, expected_height, expected_width)						)






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_  ,  a_						)						->Optional[Any]:







        '''simple docstring'''


        _a					      =						self.num_labels
        _a					      =						TFMobileViTForImageClassification(a_						)
        _a					      =						model(a_  ,  labels=a_  ,  training=a_						)
        self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels)						)






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_  ,  a_						)						->Optional[Any]:







        '''simple docstring'''


        _a					      =						self.num_labels
        _a					      =						TFMobileViTForSemanticSegmentation(a_						)
        _a					      =						_a					      =						self.image_size // self.output_stride

        _a					      =						model(a_  ,  training=a_						)
        self.parent.assertEqual(
            result.logits.shape  ,  (self.batch_size, self.num_labels, expected_height, expected_width)						)

        _a					      =						model(a_  ,  labels=a_  ,  training=a_						)
        self.parent.assertEqual(
            result.logits.shape  ,  (self.batch_size, self.num_labels, expected_height, expected_width)						)






      def       lowerCamelCase__     (     self						)						->Any:







        '''simple docstring'''


        _a					      =						self.prepare_config_and_inputs()
        _a      ,						_a      ,						_a      ,						_a					      =						config_and_inputs
        _a					      =						{"pixel_values": pixel_values}
        return config, inputs_dict







@require_tf
class  lowercase_     (_UpperCAmelCase,     _UpperCAmelCase,     unittest.TestCase  ):
      A__						:		Tuple							      =					(
          (TFMobileViTModel, TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation)
          if is_tf_available()
          else ()
      )
      A__						:		int							      =					(
          {'''feature-extraction''': TFMobileViTModel, '''image-classification''': TFMobileViTForImageClassification}
          if is_tf_available()
          else {}
      )

      A__						:		Dict							      =					False
      A__						:		Any							      =					False
      A__						:		Union[str, Any]							      =					False
      A__						:		str							      =					False
      A__						:		Any							      =					False






      def       lowerCamelCase__     (     self						)						->List[Any]:







        '''simple docstring'''


        _a					      =						TFMobileViTModelTester(self						)
        _a					      =						TFMobileViTConfigTester(self  ,  config_class=a_  ,  has_text_modality=a_						)






      def       lowerCamelCase__     (     self						)						->Any:







        '''simple docstring'''


        self.config_tester.run_common_tests()






      @unittest.skip(reason="MobileViT does not use inputs_embeds"						)
      def       lowerCamelCase__     (     self						)						->Tuple:







        '''simple docstring'''


        pass






      @unittest.skip(reason="MobileViT does not support input and output embeddings"						)
      def       lowerCamelCase__     (     self						)						->Tuple:







        '''simple docstring'''


        pass






      @unittest.skip(reason="MobileViT does not output attentions"						)
      def       lowerCamelCase__     (     self						)						->Any:







        '''simple docstring'''


        pass






      def       lowerCamelCase__     (     self						)						->List[Any]:







        '''simple docstring'''


        _a      ,						_a					      =						self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          _a					      =						model_class(a_						)
          _a					      =						inspect.signature(model.call						)
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          _a					      =						[*signature.parameters.keys()]

          _a					      =						["pixel_values"]
          self.assertListEqual(arg_names[:1]  ,  a_						)






      def       lowerCamelCase__     (     self						)						->Union[str, Any]:







        '''simple docstring'''


        _a					      =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*a_						)






      def       lowerCamelCase__     (     self						)						->Tuple:







        '''simple docstring'''


        def check_hidden_states_output(a_  ,  a_  ,  a_						):
          _a					      =						model_class(a_						)

          _a					      =						model(**self._prepare_for_class(a_  ,  a_						)						)

          _a					      =						outputs.hidden_states

          _a					      =						5
          self.assertEqual(len(a_						)  ,  a_						)

          # MobileViT's feature maps are of shape (batch_size, num_channels, height, width)
          # with the width and height being successively divided by 2.
          _a					      =						2
          for i in range(len(a_						)						):
            self.assertListEqual(
                list(hidden_states[i].shape[-2:]						)  ,  [self.model_tester.image_size // divisor, self.model_tester.image_size // divisor]  ,  )
            divisor *= 2

          self.assertEqual(self.model_tester.output_stride  ,  divisor // 2						)

        _a      ,						_a					      =						self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          _a					      =						True
          check_hidden_states_output(a_  ,  a_  ,  a_						)

          # check that output_hidden_states also work using config
          del inputs_dict["output_hidden_states"]
          _a					      =						True

          check_hidden_states_output(a_  ,  a_  ,  a_						)






      def       lowerCamelCase__     (     self						)						->str:







        '''simple docstring'''


        _a					      =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*a_						)






      def       lowerCamelCase__     (     self						)						->Dict:







        '''simple docstring'''


        _a					      =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_semantic_segmentation(*a_						)






      @unittest.skipIf(
          not is_tf_available() or len(tf.config.list_physical_devices("GPU"						)						) == 0  ,  reason="TF does not support backprop for grouped convolutions on CPU."  ,  )
      def       lowerCamelCase__     (     self						)						->Optional[Any]:







        '''simple docstring'''


        super().test_dataset_conversion()






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_=2E-1  ,  a_=2E-1						)						->Optional[int]:







        '''simple docstring'''


        self.assertTrue(np.allclose(a_  ,  a_  ,  atol=a_  ,  rtol=a_						)						)






      @unittest.skipIf(
          not is_tf_available() or len(tf.config.list_physical_devices("GPU"						)						) == 0  ,  reason="TF does not support backprop for grouped convolutions on CPU."  ,  )
      @slow
      def       lowerCamelCase__     (     self						)						->List[Any]:







        '''simple docstring'''


        _a      ,						_a					      =						self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          # Since `TFMobileViTModel` cannot operate with the default `fit()` method.
          if model_class.__name__ != "TFMobileViTModel":
            _a					      =						model_class(a_						)
            if getattr(a_  ,  "hf_compute_loss"  ,  a_						):
              super().test_keras_fit()






      def       lowerCamelCase__     (     self						)						->Optional[Any]:







        '''simple docstring'''


        _a      ,						_a					      =						self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
          # set an ignore index to correctly test the masked loss used in
          # `TFMobileViTForSemanticSegmentation`.
          if model_class.__name__ != "TFMobileViTForSemanticSegmentation":
            _a					      =						5

          _a					      =						model_class(a_						)
          if getattr(a_  ,  "hf_compute_loss"  ,  a_						):
            # The number of elements in the loss should be the same as the number of elements in the label
            _a					      =						self._prepare_for_class(inputs_dict.copy()  ,  a_  ,  return_labels=a_						)
            _a					      =						prepared_for_class[
                sorted(prepared_for_class.keys() - inputs_dict.keys()  ,  reverse=a_						)[0]
            ]
            _a					      =						added_label.shape.as_list()[:1]

            # Test that model correctly compute the loss with kwargs
            _a					      =						self._prepare_for_class(inputs_dict.copy()  ,  a_  ,  return_labels=a_						)
            _a					      =						{"input_ids", "pixel_values", "input_features"}
            _a					      =						possible_input_names.intersection(set(a_						)						).pop()
            _a					      =						prepared_for_class.pop(a_						)

            _a					      =						model(a_  ,  **a_						)[0]
            self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]						)

            # Test that model correctly compute the loss when we mask some positions
            _a					      =						self._prepare_for_class(inputs_dict.copy()  ,  a_  ,  return_labels=a_						)
            _a					      =						{"input_ids", "pixel_values", "input_features"}
            _a					      =						possible_input_names.intersection(set(a_						)						).pop()
            _a					      =						prepared_for_class.pop(a_						)
            if "labels" in prepared_for_class:
              _a					      =						prepared_for_class["labels"].numpy()
              if len(labels.shape						) > 1 and labels.shape[1] != 1:
                # labels[0] = -100
                _a					      =						tf.convert_to_tensor(a_						)
                _a					      =						model(a_  ,  **a_						)[0]
                self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]						)
                self.assertTrue(not np.any(np.isnan(loss.numpy()						)						)						)

                # Test that model correctly compute the loss with a dict
            _a					      =						self._prepare_for_class(inputs_dict.copy()  ,  a_  ,  return_labels=a_						)
            _a					      =						model(a_						)[0]
            self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]						)

            # Test that model correctly compute the loss with a tuple
            _a					      =						self._prepare_for_class(inputs_dict.copy()  ,  a_  ,  return_labels=a_						)

            # Get keys that were added with the _prepare_for_class function
            _a					      =						prepared_for_class.keys() - inputs_dict.keys()
            _a					      =						inspect.signature(model.call						).parameters
            _a					      =						list(signature.keys()						)

            # Create a dictionary holding the location of the tensors in the tuple
            _a					      =						{0: input_name}
            for label_key in label_keys:
              _a					      =						signature_names.index(a_						)
              _a					      =						label_key
            _a					      =						sorted(tuple_index_mapping.items()						)
            # Initialize a list with their default values, update the values and convert to a tuple
            _a					      =						[]

            for name in signature_names:
              if name != "kwargs":
                list_input.append(signature[name].default						)

            for index, value in sorted_tuple_index_mapping:
              _a					      =						prepared_for_class[value]

            _a					      =						tuple(a_						)

            # Send to model
            _a					      =						model(tuple_input[:-1]						)[0]

            self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]						)






      @slow
      def       lowerCamelCase__     (     self						)						->Any:







        '''simple docstring'''


        for model_name in TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          _a					      =						TFMobileViTModel.from_pretrained(a_						)
          self.assertIsNotNone(a_						)







def 							lowerCAmelCase			(   )							->				int:







  '''simple docstring'''







  _a					      =						Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"       )
  return image




@require_tf
class  lowercase_     (unittest.TestCase  ):






      @slow
      def       lowerCamelCase__     (     self						)						->List[str]:







        '''simple docstring'''


        _a					      =						TFMobileViTForImageClassification.from_pretrained("apple/mobilevit-xx-small"						)

        _a					      =						MobileViTImageProcessor.from_pretrained("apple/mobilevit-xx-small"						)
        _a					      =						prepare_img()
        _a					      =						image_processor(images=a_  ,  return_tensors="tf"						)

        # forward pass
        _a					      =						model(**a_  ,  training=a_						)

        # verify the logits
        _a					      =						tf.TensorShape((1, 1_0_0_0)						)
        self.assertEqual(outputs.logits.shape  ,  a_						)

        _a					      =						tf.constant([-1.9_364, -1.2_327, -0.4_653]						)

        tf.debugging.assert_near(outputs.logits[0, :3]  ,  a_  ,  atol=1E-4  ,  rtol=1E-04						)






      @slow
      def       lowerCamelCase__     (     self						)						->int:







        '''simple docstring'''


        _a					      =						TFMobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-xx-small"						)

        _a					      =						MobileViTImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-xx-small"						)

        _a					      =						prepare_img()
        _a					      =						image_processor(images=a_  ,  return_tensors="tf"						)

        # forward pass
        _a					      =						model(inputs.pixel_values  ,  training=a_						)
        _a					      =						outputs.logits

        # verify the logits
        _a					      =						tf.TensorShape((1, 2_1, 3_2, 3_2)						)
        self.assertEqual(logits.shape  ,  a_						)

        _a					      =						tf.constant(
            [
                [[6.9_713, 6.9_786, 7.2_422], [7.2_893, 7.2_825, 7.4_446], [7.6_580, 7.8_797, 7.9_420]],
                [[-10.6_869, -10.3_250, -10.3_471], [-10.4_228, -9.9_868, -9.7_132], [-11.0_405, -11.0_221, -10.7_318]],
                [[-3.3_089, -2.8_539, -2.6_740], [-3.2_706, -2.5_621, -2.5_108], [-3.2_534, -2.6_615, -2.6_651]],
            ]						)

        tf.debugging.assert_near(logits[0, :3, :3, :3]  ,  a_  ,  rtol=1E-4  ,  atol=1E-4						)










"""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.albert.modeling_albert import (
    ALBERT_INPUTS_DOCSTRING,
    ALBERT_START_DOCSTRING,
    AlbertModel,
    AlbertPreTrainedModel,
    AlbertTransformer,
)


UpperCamelCase											=     logging.getLogger(__name__)




class  lowercase_     (_UpperCAmelCase  ):






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_=None  ,  a_=None						)						->Optional[int]:







        '''simple docstring'''


        if current_layer == 0:
          _a					      =						self.embedding_hidden_mapping_in(a_						)
        else:
          _a					      =						hidden_states[0]

        _a					      =						int(self.config.num_hidden_layers / self.config.num_hidden_groups						)

        # Index of the hidden group
        _a					      =						int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)						)

        _a					      =						self.albert_layer_groups[group_idx](
            a_  ,  a_  ,  head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]  ,  )
        _a					      =						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  lowercase_     (_UpperCAmelCase  ):






      def __init__(     self  ,  a_						)						->Dict:







        '''simple docstring'''


        super().__init__(a_						)

        _a					      =						AlbertTransformerWithPabee(a_						)

        self.init_weights()
        _a					      =						0
        _a					      =						0
        _a					      =						0

        _a					      =						0






      def       lowerCamelCase__     (     self  ,  a_						)						->Optional[int]:







        '''simple docstring'''


        _a					      =						threshold






      def       lowerCamelCase__     (     self  ,  a_						)						->Union[str, Any]:







        '''simple docstring'''


        _a					      =						patience






      def       lowerCamelCase__     (     self						)						->int:







        '''simple docstring'''


        _a					      =						0
        _a					      =						0






      def       lowerCamelCase__     (     self						)						->Tuple:







        '''simple docstring'''


        _a					      =						self.inference_layers_num / self.inference_instances_num
        _a					      =						(
            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(a_						)






      @add_start_docstrings_to_model_forward(a_						)
      def       lowerCamelCase__     (     self  ,  a_=None  ,  a_=None  ,  a_=None  ,  a_=None  ,  a_=None  ,  a_=None  ,  a_=None  ,  a_=None  ,  a_=False  ,  )						->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					      =						input_ids.size()
        elif inputs_embeds is not None:
          _a					      =						inputs_embeds.size()[:-1]
        else:
          raise ValueError("You have to specify either input_ids or inputs_embeds"						)

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

        if attention_mask is None:
          _a					      =						torch.ones(a_  ,  device=a_						)
        if token_type_ids is None:
          _a					      =						torch.zeros(a_  ,  dtype=torch.long  ,  device=a_						)

        _a					      =						attention_mask.unsqueeze(1						).unsqueeze(2						)
        _a					      =						extended_attention_mask.to(dtype=self.dtype						)  # fp16 compatibility
        _a					      =						(1.0 - extended_attention_mask) * -10_000.0
        _a					      =						self.get_head_mask(a_  ,  self.config.num_hidden_layers						)

        _a					      =						self.embeddings(
            a_  ,  position_ids=a_  ,  token_type_ids=a_  ,  inputs_embeds=a_						)
        _a					      =						embedding_output

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

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

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

            _a					      =						logits
            if patient_counter == self.patience:
              break
          _a					      =						[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  lowercase_     (_UpperCAmelCase  ):






      def __init__(     self  ,  a_						)						->Dict:







        '''simple docstring'''


        super().__init__(a_						)
        _a					      =						config.num_labels

        _a					      =						AlbertModelWithPabee(a_						)
        _a					      =						nn.Dropout(config.classifier_dropout_prob						)
        _a					      =						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(a_						)
      def       lowerCamelCase__     (     self  ,  a_=None  ,  a_=None  ,  a_=None  ,  a_=None  ,  a_=None  ,  a_=None  ,  a_=None  ,  )						->List[Any]:







        '''simple docstring'''


        _a					      =						self.albert(
            input_ids=a_  ,  attention_mask=a_  ,  token_type_ids=a_  ,  position_ids=a_  ,  head_mask=a_  ,  inputs_embeds=a_  ,  output_dropout=self.dropout  ,  output_layers=self.classifiers  ,  regression=self.num_labels == 1  ,  )

        _a					      =						(logits[-1],)

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

        return outputs


