
from __future__ import annotations

import copy
import inspect
import unittest

import numpy as np

from transformers import is_tf_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow
from transformers.utils import cached_property

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


if is_tf_available():
      import tensorflow as tf

      from transformers import (
          TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST,
          TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
          TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
          TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
          TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
          LayoutLMvaConfig,
          TFLayoutLMvaForQuestionAnswering,
          TFLayoutLMvaForSequenceClassification,
          TFLayoutLMvaForTokenClassification,
          TFLayoutLMvaModel,
      )

if is_vision_available():
      from PIL import Image

      from transformers import LayoutLMvaImageProcessor






class   a_    :



       """simple docstring"""






       def __init__(						self							,    _lowerCamelCase							,    _lowerCamelCase=2							,    _lowerCamelCase=3							,    _lowerCamelCase=4							,    _lowerCamelCase=2							,    _lowerCamelCase=7							,    _lowerCamelCase=True							,    _lowerCamelCase=True							,    _lowerCamelCase=True							,    _lowerCamelCase=True							,    _lowerCamelCase=99							,    _lowerCamelCase=36							,    _lowerCamelCase=2							,    _lowerCamelCase=4							,    _lowerCamelCase=37							,    _lowerCamelCase="gelu"							,    _lowerCamelCase=0.1							,    _lowerCamelCase=0.1							,    _lowerCamelCase=512							,    _lowerCamelCase=16							,    _lowerCamelCase=2							,    _lowerCamelCase=0.0_2							,    _lowerCamelCase=6							,    _lowerCamelCase=6							,    _lowerCamelCase=3							,    _lowerCamelCase=4							,    _lowerCamelCase=None							,    _lowerCamelCase=1000							,    ) ->int:
           SCREAMING_SNAKE_CASE							:							Tuple								=      parent
           SCREAMING_SNAKE_CASE							:							int								=      batch_size
           SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      num_channels
           SCREAMING_SNAKE_CASE							:							int								=      image_size
           SCREAMING_SNAKE_CASE							:							Tuple								=      patch_size
           SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      is_training
           SCREAMING_SNAKE_CASE							:							List[Any]								=      use_input_mask
           SCREAMING_SNAKE_CASE							:							Tuple								=      use_token_type_ids
           SCREAMING_SNAKE_CASE							:							Optional[int]								=      use_labels
           SCREAMING_SNAKE_CASE							:							Any								=      vocab_size
           SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      hidden_size
           SCREAMING_SNAKE_CASE							:							Dict								=      num_hidden_layers
           SCREAMING_SNAKE_CASE							:							Dict								=      num_attention_heads
           SCREAMING_SNAKE_CASE							:							List[Any]								=      intermediate_size
           SCREAMING_SNAKE_CASE							:							Optional[Any]								=      hidden_act
           SCREAMING_SNAKE_CASE							:							Optional[Any]								=      hidden_dropout_prob
           SCREAMING_SNAKE_CASE							:							Dict								=      attention_probs_dropout_prob
           SCREAMING_SNAKE_CASE							:							Any								=      max_position_embeddings
           SCREAMING_SNAKE_CASE							:							int								=      type_vocab_size
           SCREAMING_SNAKE_CASE							:							Tuple								=      type_sequence_label_size
           SCREAMING_SNAKE_CASE							:							Optional[int]								=      initializer_range
           SCREAMING_SNAKE_CASE							:							List[Any]								=      coordinate_size
           SCREAMING_SNAKE_CASE							:							Dict								=      shape_size
           SCREAMING_SNAKE_CASE							:							Tuple								=      num_labels
           SCREAMING_SNAKE_CASE							:							str								=      num_choices
           SCREAMING_SNAKE_CASE							:							str								=      scope
           SCREAMING_SNAKE_CASE							:							Optional[Any]								=      range_bbox

           # LayoutLMv3's sequence length equals the number of text tokens + number of patches + 1 (we add 1 for the CLS token)
           SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      text_seq_length
           SCREAMING_SNAKE_CASE							:							Optional[Any]								=      (image_size // patch_size) ** 2 + 1
           SCREAMING_SNAKE_CASE							:							List[str]								=      self.text_seq_length + self.image_seq_length






       def      __lowerCAmelCase			(						self					) ->str:
           SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      ids_tensor([self.batch_size, self.text_seq_length]							,    self.vocab_size					)

           SCREAMING_SNAKE_CASE							:							Any								=      ids_tensor([self.batch_size, self.text_seq_length, 4]							,    self.range_bbox					)
           SCREAMING_SNAKE_CASE							:							Dict								=      bbox.numpy()
           # Ensure that bbox is legal
           for i in range(bbox.shape[0]					):
               for j in range(bbox.shape[1]					):
                   if bbox[i, j, 3] < bbox[i, j, 1]:
                       SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      bbox[i, j, 3]
                       SCREAMING_SNAKE_CASE							:							Optional[int]								=      bbox[i, j, 1]
                       SCREAMING_SNAKE_CASE							:							Dict								=      tmp_coordinate
                   if bbox[i, j, 2] < bbox[i, j, 0]:
                       SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      bbox[i, j, 2]
                       SCREAMING_SNAKE_CASE							:							Dict								=      bbox[i, j, 0]
                       SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      tmp_coordinate
           SCREAMING_SNAKE_CASE							:							Any								=      tf.constant(_a					)

           SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

           SCREAMING_SNAKE_CASE							:							List[Any]								=      None
           if self.use_input_mask:
               SCREAMING_SNAKE_CASE							:							List[str]								=      random_attention_mask([self.batch_size, self.text_seq_length]					)

           SCREAMING_SNAKE_CASE							:							List[Any]								=      None
           if self.use_token_type_ids:
               SCREAMING_SNAKE_CASE							:							str								=      ids_tensor([self.batch_size, self.text_seq_length]							,    self.type_vocab_size					)

           SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      None
           SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      None
           if self.use_labels:
               SCREAMING_SNAKE_CASE							:							int								=      ids_tensor([self.batch_size]							,    self.type_sequence_label_size					)
               SCREAMING_SNAKE_CASE							:							List[Any]								=      ids_tensor([self.batch_size, self.text_seq_length]							,    self.num_labels					)

           SCREAMING_SNAKE_CASE							:							Optional[int]								=      LayoutLMvaConfig(
               vocab_size=self.vocab_size							,    hidden_size=self.hidden_size							,    num_hidden_layers=self.num_hidden_layers							,    num_attention_heads=self.num_attention_heads							,    intermediate_size=self.intermediate_size							,    hidden_act=self.hidden_act							,    hidden_dropout_prob=self.hidden_dropout_prob							,    attention_probs_dropout_prob=self.attention_probs_dropout_prob							,    max_position_embeddings=self.max_position_embeddings							,    type_vocab_size=self.type_vocab_size							,    initializer_range=self.initializer_range							,    coordinate_size=self.coordinate_size							,    shape_size=self.shape_size							,    input_size=self.image_size							,    patch_size=self.patch_size							,    )

           return config, input_ids, bbox, pixel_values, token_type_ids, input_mask, sequence_labels, token_labels






       def      __lowerCAmelCase			(						self							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase					) ->Optional[int]:
           SCREAMING_SNAKE_CASE							:							str								=      TFLayoutLMvaModel(config=_a					)

           # text + image
           SCREAMING_SNAKE_CASE							:							List[str]								=      model(_a							,    pixel_values=_a							,    training=_a					)
           SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      model(
               _a							,    bbox=_a							,    pixel_values=_a							,    attention_mask=_a							,    token_type_ids=_a							,    training=_a							,    )
           SCREAMING_SNAKE_CASE							:							int								=      model(_a							,    bbox=_a							,    pixel_values=_a							,    training=_a					)

           self.parent.assertEqual(result.last_hidden_state.shape							,    (self.batch_size, self.seq_length, self.hidden_size)					)

           # text only
           SCREAMING_SNAKE_CASE							:							Dict								=      model(_a							,    training=_a					)
           self.parent.assertEqual(
               result.last_hidden_state.shape							,    (self.batch_size, self.text_seq_length, self.hidden_size)					)

           # image only
           SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      model({'''pixel_values''': pixel_values}							,    training=_a					)
           self.parent.assertEqual(
               result.last_hidden_state.shape							,    (self.batch_size, self.image_seq_length, self.hidden_size)					)






       def      __lowerCAmelCase			(						self							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase					) ->int:
           SCREAMING_SNAKE_CASE							:							str								=      self.num_labels
           SCREAMING_SNAKE_CASE							:							List[str]								=      TFLayoutLMvaForSequenceClassification(config=_a					)
           SCREAMING_SNAKE_CASE							:							int								=      model(
               _a							,    bbox=_a							,    pixel_values=_a							,    attention_mask=_a							,    token_type_ids=_a							,    labels=_a							,    training=_a							,    )
           self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.num_labels)					)






       def      __lowerCAmelCase			(						self							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase					) ->Any:
           SCREAMING_SNAKE_CASE							:							str								=      self.num_labels
           SCREAMING_SNAKE_CASE							:							Any								=      TFLayoutLMvaForTokenClassification(config=_a					)
           SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      model(
               _a							,    bbox=_a							,    pixel_values=_a							,    attention_mask=_a							,    token_type_ids=_a							,    labels=_a							,    training=_a							,    )
           self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.text_seq_length, self.num_labels)					)






       def      __lowerCAmelCase			(						self							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase					) ->Optional[int]:
           SCREAMING_SNAKE_CASE							:							int								=      2
           SCREAMING_SNAKE_CASE							:							Dict								=      TFLayoutLMvaForQuestionAnswering(config=_a					)
           SCREAMING_SNAKE_CASE							:							Tuple								=      model(
               _a							,    bbox=_a							,    pixel_values=_a							,    attention_mask=_a							,    token_type_ids=_a							,    start_positions=_a							,    end_positions=_a							,    training=_a							,    )
           self.parent.assertEqual(result.start_logits.shape							,    (self.batch_size, self.seq_length)					)
           self.parent.assertEqual(result.end_logits.shape							,    (self.batch_size, self.seq_length)					)






       def      __lowerCAmelCase			(						self					) ->Union[str, Any]:
           SCREAMING_SNAKE_CASE							:							int								=      self.prepare_config_and_inputs()
           (SCREAMING_SNAKE_CASE)							:							Any								=      config_and_inputs
           SCREAMING_SNAKE_CASE							:							List[Any]								=      {
               """input_ids""": input_ids,
               """bbox""": bbox,
               """pixel_values""": pixel_values,
               """token_type_ids""": token_type_ids,
               """attention_mask""": input_mask,
           }
           return config, inputs_dict






@require_tf
class   a_    (       UpperCamelCase__							,  UpperCamelCase__							,  unittest.TestCase				):



       """simple docstring"""
       __SCREAMING_SNAKE_CASE    : List[Any]		=   (
           (
               TFLayoutLMvaModel,
               TFLayoutLMvaForQuestionAnswering,
               TFLayoutLMvaForSequenceClassification,
               TFLayoutLMvaForTokenClassification,
           )
           if is_tf_available()
           else ()
       )
       __SCREAMING_SNAKE_CASE    : List[Any]		=   (
           {'document-question-answering': TFLayoutLMvaForQuestionAnswering, 'feature-extraction': TFLayoutLMvaModel}
           if is_tf_available()
           else {}
       )

       __SCREAMING_SNAKE_CASE    : List[str]		=   False
       __SCREAMING_SNAKE_CASE    : Optional[int]		=   False
       __SCREAMING_SNAKE_CASE    : int		=   False






       def      __lowerCAmelCase			(						self							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase					) ->Optional[int]:
           return True






       def      __lowerCAmelCase			(						self							,    _lowerCamelCase							,    _lowerCamelCase							,    _lowerCamelCase=False					) ->dict:
           SCREAMING_SNAKE_CASE							:							int								=      copy.deepcopy(_a					)

           if model_class in get_values(_a					):
               SCREAMING_SNAKE_CASE							:							List[Any]								=      {
                   k: tf.tile(tf.expand_dims(_a							,    1					)							,    (1, self.model_tester.num_choices) + (1,) * (v.ndim - 1)					)
                   if isinstance(_a							,    tf.Tensor					) and v.ndim > 0
                   else v
                   for k, v in inputs_dict.items()
               }

           if return_labels:
               if model_class in get_values(_a					):
                   SCREAMING_SNAKE_CASE							:							List[str]								=      tf.ones(self.model_tester.batch_size							,    dtype=tf.intaa					)
               elif model_class in get_values(_a					):
                   SCREAMING_SNAKE_CASE							:							Dict								=      tf.zeros(self.model_tester.batch_size							,    dtype=tf.intaa					)
                   SCREAMING_SNAKE_CASE							:							List[Any]								=      tf.zeros(self.model_tester.batch_size							,    dtype=tf.intaa					)
               elif model_class in get_values(_a					):
                   SCREAMING_SNAKE_CASE							:							str								=      tf.zeros(self.model_tester.batch_size							,    dtype=tf.intaa					)
               elif model_class in get_values(_a					):
                   SCREAMING_SNAKE_CASE							:							Tuple								=      tf.zeros(
                       (self.model_tester.batch_size, self.model_tester.text_seq_length)							,    dtype=tf.intaa					)

           return inputs_dict






       def      __lowerCAmelCase			(						self					) ->List[str]:
           SCREAMING_SNAKE_CASE							:							List[str]								=      TFLayoutLMvaModelTester(self					)
           SCREAMING_SNAKE_CASE							:							Optional[Any]								=      ConfigTester(self							,    config_class=_a							,    hidden_size=37					)






       def      __lowerCAmelCase			(						self					) ->Optional[Any]:
           self.config_tester.run_common_tests()






       def      __lowerCAmelCase			(						self					) ->List[str]:
           SCREAMING_SNAKE_CASE							:							Tuple								=      self.model_tester.prepare_config_and_inputs_for_common()
           for model_class in self.all_model_classes:
               SCREAMING_SNAKE_CASE							:							Dict								=      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
                   SCREAMING_SNAKE_CASE							:							int								=      self._prepare_for_class(inputs_dict.copy()							,    _a							,    return_labels=_a					)
                   SCREAMING_SNAKE_CASE							:							int								=      prepared_for_class[
                       sorted(prepared_for_class.keys() - inputs_dict.keys()							,    reverse=_a					)[0]
                   ]
                   SCREAMING_SNAKE_CASE							:							int								=      added_label.shape.as_list()[:1]

                   # Test that model correctly compute the loss with kwargs
                   SCREAMING_SNAKE_CASE							:							List[str]								=      self._prepare_for_class(inputs_dict.copy()							,    _a							,    return_labels=_a					)
                   SCREAMING_SNAKE_CASE							:							Union[str, Any]								=      prepared_for_class.pop('''input_ids'''					)

                   SCREAMING_SNAKE_CASE							:							Tuple								=      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
                   SCREAMING_SNAKE_CASE							:							Tuple								=      self._prepare_for_class(inputs_dict.copy()							,    _a							,    return_labels=_a					)
                   SCREAMING_SNAKE_CASE							:							Any								=      prepared_for_class.pop('''input_ids'''					)
                   if "labels" in prepared_for_class:
                       SCREAMING_SNAKE_CASE							:							int								=      prepared_for_class["""labels"""].numpy()
                       if len(labels.shape					) > 1 and labels.shape[1] != 1:
                           SCREAMING_SNAKE_CASE							:							Tuple								=      -100
                           SCREAMING_SNAKE_CASE							:							Dict								=      tf.convert_to_tensor(_a					)
                           SCREAMING_SNAKE_CASE							:							Any								=      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
                   SCREAMING_SNAKE_CASE							:							List[str]								=      self._prepare_for_class(inputs_dict.copy()							,    _a							,    return_labels=_a					)
                   SCREAMING_SNAKE_CASE							:							int								=      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
                   SCREAMING_SNAKE_CASE							:							List[str]								=      self._prepare_for_class(inputs_dict.copy()							,    _a							,    return_labels=_a					)

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

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

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

                   for index, value in sorted_tuple_index_mapping:
                       SCREAMING_SNAKE_CASE							:							Tuple								=      prepared_for_class[value]

                   SCREAMING_SNAKE_CASE							:							int								=      tuple(_a					)

                   # Send to model
                   SCREAMING_SNAKE_CASE							:							List[str]								=      model(tuple_input[:-1]					)[0]

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






       def      __lowerCAmelCase			(						self					) ->Optional[int]:
           (
               SCREAMING_SNAKE_CASE   
           )							:							Optional[int]								=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(_a							,    _a							,    _a							,    _a							,    _a							,    _a					)






       def      __lowerCAmelCase			(						self					) ->Any:
           (
               SCREAMING_SNAKE_CASE   
           )							:							Union[str, Any]								=      self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
               SCREAMING_SNAKE_CASE							:							Optional[Any]								=      type
               self.model_tester.create_and_check_model(_a							,    _a							,    _a							,    _a							,    _a							,    _a					)






       def      __lowerCAmelCase			(						self					) ->Tuple:
           (
               SCREAMING_SNAKE_CASE   
           )							:							List[Any]								=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(
               _a							,    _a							,    _a							,    _a							,    _a							,    _a							,    _a					)






       def      __lowerCAmelCase			(						self					) ->Tuple:
           (
               SCREAMING_SNAKE_CASE   
           )							:							List[Any]								=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(
               _a							,    _a							,    _a							,    _a							,    _a							,    _a							,    _a					)






       def      __lowerCAmelCase			(						self					) ->List[str]:
           (
               SCREAMING_SNAKE_CASE   
           )							:							Dict								=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(
               _a							,    _a							,    _a							,    _a							,    _a							,    _a							,    _a					)






       @slow
       def      __lowerCAmelCase			(						self					) ->str:
           for model_name in TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               SCREAMING_SNAKE_CASE							:							Tuple								=      TFLayoutLMvaModel.from_pretrained(_a					)
               self.assertIsNotNone(_a					)







def 							UpperCAmelCase_(  ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE							:							List[Any]								=      Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''					)
    return image






@require_tf
class   a_    (       unittest.TestCase				):



       """simple docstring"""






       @cached_property
       def      __lowerCAmelCase			(						self					) ->Tuple:
           return LayoutLMvaImageProcessor(apply_ocr=_a					) if is_vision_available() else None






       @slow
       def      __lowerCAmelCase			(						self					) ->Union[str, Any]:
           SCREAMING_SNAKE_CASE							:							Optional[int]								=      TFLayoutLMvaModel.from_pretrained('''microsoft/layoutlmv3-base'''					)

           SCREAMING_SNAKE_CASE							:							Tuple								=      self.default_image_processor
           SCREAMING_SNAKE_CASE							:							int								=      prepare_img()
           SCREAMING_SNAKE_CASE							:							List[Any]								=      image_processor(images=_a							,    return_tensors='''tf'''					).pixel_values

           SCREAMING_SNAKE_CASE							:							str								=      tf.constant([[1, 2]]					)
           SCREAMING_SNAKE_CASE							:							int								=      tf.expand_dims(tf.constant([[1, 2, 3, 4], [5, 6, 7, 8]]					)							,    axis=0					)

           # forward pass
           SCREAMING_SNAKE_CASE							:							Tuple								=      model(input_ids=_a							,    bbox=_a							,    pixel_values=_a							,    training=_a					)

           # verify the logits
           SCREAMING_SNAKE_CASE							:							int								=      (1, 199, 768)
           self.assertEqual(outputs.last_hidden_state.shape							,    _a					)

           SCREAMING_SNAKE_CASE							:							List[Any]								=      tf.constant(
               [[-0.0_5_2_9, 0.3_6_1_8, 0.1_6_3_2], [-0.1_5_8_7, -0.1_6_6_7, -0.0_4_0_0], [-0.1_5_5_7, -0.1_6_7_1, -0.0_5_0_5]]					)

           self.assertTrue(np.allclose(outputs.last_hidden_state[0, :3, :3]							,    _a							,    atol=1e-4					)					)






from __future__ import absolute_import, division, print_function, unicode_literals

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

from transformers import RobertaConfig
from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.roberta.modeling_roberta import (
    ROBERTA_INPUTS_DOCSTRING,
    ROBERTA_START_DOCSTRING,
    RobertaEmbeddings,
)

from .modeling_highway_bert import BertPreTrainedModel, DeeBertModel, HighwayException, entropy







@add_start_docstrings(
    "The RoBERTa Model transformer with early exiting (DeeRoBERTa). ",UpperCamelCase__,)
class 						lowercase							(   UpperCamelCase__		):
     _a							 =						RobertaConfig
     _a							 =						"roberta"

     def __init__(  self		,     _a				)						->							Optional[int]:
           super().__init__(_a				)

           _A							:     Union[str, Any]			=						RobertaEmbeddings(_a				)
           self.init_weights()







@add_start_docstrings(
    "RoBERTa Model (with early exiting - DeeRoBERTa) with a classifier on top,\n    also takes care of multi-layer training. ",UpperCamelCase__,)
class 						lowercase							(   UpperCamelCase__		):
     _a							 =						RobertaConfig
     _a							 =						"roberta"

     def __init__(  self		,     _a				)						->							str:
           super().__init__(_a				)
           _A							:     Any			=						config.num_labels
           _A							:     Dict			=						config.num_hidden_layers

           _A							:     List[str]			=						DeeRobertaModel(_a				)
           _A							:     int			=						nn.Dropout(config.hidden_dropout_prob				)
           _A							:     int			=						nn.Linear(config.hidden_size		,     self.config.num_labels				)



     @add_start_docstrings_to_model_forward(_a				)
     def 					a__	(  self		,     _a=None		,     _a=None		,     _a=None		,     _a=None		,     _a=None		,     _a=None		,     _a=None		,     _a=-1		,     _a=False		,     )						->							Any:
           _A							:     Optional[int]			=						self.num_layers
           try:
                 _A							:     List[str]			=						self.roberta(
                     _a		,     attention_mask=_a		,     token_type_ids=_a		,     position_ids=_a		,     head_mask=_a		,     inputs_embeds=_a		,     )

                 _A							:     List[str]			=						outputs[1]

                 _A							:     List[str]			=						self.dropout(_a				)
                 _A							:     Optional[Any]			=						self.classifier(_a				)
                 _A							:     List[Any]			=						(logits,) + outputs[2:]  # add hidden states and attention if they are here
           except HighwayException as e:
                 _A							:     List[Any]			=						e.message
                 _A							:     Optional[int]			=						e.exit_layer
                 _A							:     Optional[int]			=						outputs[0]

           if not self.training:
                 _A							:     int			=						entropy(_a				)
                 _A							:     int			=						[]
                 _A							:     int			=						[]
           if labels is not None:
                 if self.num_labels == 1:
                       #  We are doing regression
                       _A							:     Union[str, Any]			=						MSELoss()
                       _A							:     Tuple			=						loss_fct(logits.view(-1				)		,     labels.view(-1				)				)
                 else:
                       _A							:     List[Any]			=						CrossEntropyLoss()
                       _A							:     Dict			=						loss_fct(logits.view(-1		,     self.num_labels				)		,     labels.view(-1				)				)

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

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

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

