
"""simple docstring"""

from __future__ import annotations

import unittest

from transformers import MPNetConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
   import tensorflow as tf

   from transformers.models.mpnet.modeling_tf_mpnet import (
       TFMPNetForMaskedLM,
       TFMPNetForMultipleChoice,
       TFMPNetForQuestionAnswering,
       TFMPNetForSequenceClassification,
       TFMPNetForTokenClassification,
       TFMPNetModel,
   )
class    _UpperCAmelCase							:




      def __init__(							self      ,				lowercase_      ,				lowercase_=1_3      ,				lowercase_=7      ,				lowercase_=True      ,				lowercase_=True      ,				lowercase_=False      ,				lowercase_=True      ,				lowercase_=9_9      ,				lowercase_=6_4      ,				lowercase_=2      ,				lowercase_=4      ,				lowercase_=6_4      ,				lowercase_="gelu"      ,				lowercase_=0.1      ,				lowercase_=0.1      ,				lowercase_=5_1_2      ,				lowercase_=1_6      ,				lowercase_=2      ,				lowercase_=0.0_2      ,				lowercase_=3      ,				lowercase_=4      ,				lowercase_=None      ,				)  ->    Union[str, Any]:
          UpperCAmelCase  	=    parent
          UpperCAmelCase  	=    batch_size
          UpperCAmelCase  	=    seq_length
          UpperCAmelCase  	=    is_training
          UpperCAmelCase  	=    use_input_mask
          UpperCAmelCase  	=    use_token_type_ids
          UpperCAmelCase  	=    use_labels
          UpperCAmelCase  	=    vocab_size
          UpperCAmelCase  	=    hidden_size
          UpperCAmelCase  	=    num_hidden_layers
          UpperCAmelCase  	=    num_attention_heads
          UpperCAmelCase  	=    intermediate_size
          UpperCAmelCase  	=    hidden_act
          UpperCAmelCase  	=    hidden_dropout_prob
          UpperCAmelCase  	=    attention_probs_dropout_prob
          UpperCAmelCase  	=    max_position_embeddings
          UpperCAmelCase  	=    type_vocab_size
          UpperCAmelCase  	=    type_sequence_label_size
          UpperCAmelCase  	=    initializer_range
          UpperCAmelCase  	=    num_labels
          UpperCAmelCase  	=    num_choices
          UpperCAmelCase  	=    scope




      def 			a_					(							self							)  ->    Union[str, Any]:
          UpperCAmelCase  	=    ids_tensor([self.batch_size, self.seq_length]      ,				self.vocab_size							)

          UpperCAmelCase  	=    None
          if self.use_input_mask:
              UpperCAmelCase  	=    random_attention_mask([self.batch_size, self.seq_length]							)

          UpperCAmelCase  	=    None
          UpperCAmelCase  	=    None
          UpperCAmelCase  	=    None
          if self.use_labels:
              UpperCAmelCase  	=    ids_tensor([self.batch_size]      ,				self.type_sequence_label_size							)
              UpperCAmelCase  	=    ids_tensor([self.batch_size, self.seq_length]      ,				self.num_labels							)
              UpperCAmelCase  	=    ids_tensor([self.batch_size]      ,				self.num_choices							)

          UpperCAmelCase  	=    MPNetConfig(
              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      ,				initializer_range=self.initializer_range      ,				)
          return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels




      def 			a_					(							self      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_							)  ->    List[Any]:
          UpperCAmelCase  	=    TFMPNetModel(config=lowercase_							)
          UpperCAmelCase  	=    {'input_ids': input_ids, 'attention_mask': input_mask}
          UpperCAmelCase  	=    model(lowercase_							)
          UpperCAmelCase  	=    [input_ids, input_mask]
          UpperCAmelCase  	=    model(lowercase_							)
          self.parent.assertEqual(result.last_hidden_state.shape      ,				(self.batch_size, self.seq_length, self.hidden_size)							)




      def 			a_					(							self      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_							)  ->    List[str]:
          UpperCAmelCase  	=    TFMPNetForMaskedLM(config=lowercase_							)
          UpperCAmelCase  	=    {'input_ids': input_ids, 'attention_mask': input_mask}
          UpperCAmelCase  	=    model(lowercase_							)
          self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.seq_length, self.vocab_size)							)




      def 			a_					(							self      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_							)  ->    str:
          UpperCAmelCase  	=    TFMPNetForQuestionAnswering(config=lowercase_							)
          UpperCAmelCase  	=    {
              'input_ids': input_ids,
              'attention_mask': input_mask,
          }
          UpperCAmelCase  	=    model(lowercase_							)
          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 			a_					(							self      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_							)  ->    Union[str, Any]:
          UpperCAmelCase  	=    self.num_labels
          UpperCAmelCase  	=    TFMPNetForSequenceClassification(lowercase_							)
          UpperCAmelCase  	=    {'input_ids': input_ids, 'attention_mask': input_mask}
          UpperCAmelCase  	=    model(lowercase_							)
          self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.num_labels)							)




      def 			a_					(							self      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_							)  ->    Optional[Any]:
          UpperCAmelCase  	=    self.num_choices
          UpperCAmelCase  	=    TFMPNetForMultipleChoice(lowercase_							)
          UpperCAmelCase  	=    tf.tile(tf.expand_dims(lowercase_      ,				1							)      ,				(1, self.num_choices, 1)							)
          UpperCAmelCase  	=    tf.tile(tf.expand_dims(lowercase_      ,				1							)      ,				(1, self.num_choices, 1)							)
          UpperCAmelCase  	=    {
              'input_ids': multiple_choice_inputs_ids,
              'attention_mask': multiple_choice_input_mask,
          }
          UpperCAmelCase  	=    model(lowercase_							)
          self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.num_choices)							)




      def 			a_					(							self      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_							)  ->    Dict:
          UpperCAmelCase  	=    self.num_labels
          UpperCAmelCase  	=    TFMPNetForTokenClassification(lowercase_							)
          UpperCAmelCase  	=    {'input_ids': input_ids, 'attention_mask': input_mask}
          UpperCAmelCase  	=    model(lowercase_							)
          self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.seq_length, self.num_labels)							)




      def 			a_					(							self							)  ->    int:
          UpperCAmelCase  	=    self.prepare_config_and_inputs()
          ((UpperCAmelCase)	,							(UpperCAmelCase)	,							(UpperCAmelCase)	,							(UpperCAmelCase)	,							(UpperCAmelCase)	,							(UpperCAmelCase))  	=    config_and_inputs
          UpperCAmelCase  	=    {'input_ids': input_ids, 'attention_mask': input_mask}
          return config, inputs_dict

@require_tf
class    _UpperCAmelCase							(						SCREAMING_SNAKE_CASE_   ,	SCREAMING_SNAKE_CASE_   ,	unittest.TestCase							):
      __SCREAMING_SNAKE_CASE      :							Union[str, Any]												=		(
          (
              TFMPNetForMaskedLM,
              TFMPNetForMultipleChoice,
              TFMPNetForQuestionAnswering,
              TFMPNetForSequenceClassification,
              TFMPNetForTokenClassification,
              TFMPNetModel,
          )
          if is_tf_available()
          else ()
      )
      __SCREAMING_SNAKE_CASE      :							Union[str, Any]												=		(
          {
              "feature-extraction": TFMPNetModel,
              "fill-mask": TFMPNetForMaskedLM,
              "question-answering": TFMPNetForQuestionAnswering,
              "text-classification": TFMPNetForSequenceClassification,
              "token-classification": TFMPNetForTokenClassification,
              "zero-shot": TFMPNetForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      __SCREAMING_SNAKE_CASE      :							Dict												=		False
      __SCREAMING_SNAKE_CASE      :							List[Any]												=		False




      def 			a_					(							self							)  ->    Tuple:
          UpperCAmelCase  	=    TFMPNetModelTester(self							)
          UpperCAmelCase  	=    ConfigTester(self      ,				config_class=lowercase_      ,				hidden_size=3_7							)




      def 			a_					(							self							)  ->    Union[str, Any]:
          self.config_tester.run_common_tests()




      def 			a_					(							self							)  ->    Tuple:
          UpperCAmelCase  	=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mpnet_model(*lowercase_							)




      def 			a_					(							self							)  ->    Dict:
          UpperCAmelCase  	=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mpnet_for_masked_lm(*lowercase_							)




      def 			a_					(							self							)  ->    Dict:
          UpperCAmelCase  	=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mpnet_for_question_answering(*lowercase_							)




      def 			a_					(							self							)  ->    str:
          UpperCAmelCase  	=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mpnet_for_sequence_classification(*lowercase_							)




      def 			a_					(							self							)  ->    List[Any]:
          UpperCAmelCase  	=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mpnet_for_multiple_choice(*lowercase_							)




      def 			a_					(							self							)  ->    Union[str, Any]:
          UpperCAmelCase  	=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_mpnet_for_token_classification(*lowercase_							)




      @slow
      def 			a_					(							self							)  ->    Any:
          for model_name in ["microsoft/mpnet-base"]:
              UpperCAmelCase  	=    TFMPNetModel.from_pretrained(lowercase_							)
              self.assertIsNotNone(lowercase_							)

@require_tf
class    _UpperCAmelCase							(						unittest.TestCase							):




      @slow
      def 			a_					(							self							)  ->    Dict:
          UpperCAmelCase  	=    TFMPNetModel.from_pretrained('microsoft/mpnet-base'							)
          UpperCAmelCase  	=    tf.constant([[0, 1, 2, 3, 4, 5]]							)
          UpperCAmelCase  	=    model(lowercase_							)[0]

          UpperCAmelCase  	=    [1, 6, 7_6_8]
          self.assertEqual(output.shape      ,				lowercase_							)

          UpperCAmelCase  	=    tf.constant(
              [
                  [
                      [-0.1_0_6_7_1_7_2, 0.0_8_2_1_6_4_7_3, 0.0_0_2_4_5_4_3],
                      [-0.0_3_4_6_5_8_7_9, 0.8_3_5_4_1_1_8, -0.0_3_2_5_2_2_8_8],
                      [-0.0_6_5_6_9_4_7_6, -0.1_2_4_2_4_1_1_1, -0.0_4_9_4_4_3_6],
                  ]
              ]							)
          tf.debugging.assert_near(output[:, :3, :3]      ,				lowercase_      ,				atol=1E-4							)





"""simple docstring"""

from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve




def 			lowercase__  (    lowerCAmelCase :    np.ndarray       , lowerCAmelCase :    np.ndarray       , lowerCAmelCase :    np.ndarray       )       ->					np.ndarray:






    """simple docstring"""




    UpperCAmelCase  	=    np.stack((horizontal_flow, vertical_flow)       , 2       )

    # Create a grid of all pixel coordinates and subtract the flow to get the
    # target pixels coordinates
    UpperCAmelCase  	=    np.stack(
        np.meshgrid(np.arange(0       , image.shape[1]       )       , np.arange(0       , image.shape[0]       )       )       , 2       )
    UpperCAmelCase  	=    np.round(grid - flow       ).astype(np.intaa       )

    # Find the locations outside of the original image
    UpperCAmelCase  	=    (grid < 0) | (grid >= np.array([image.shape[1], image.shape[0]]       ))
    UpperCAmelCase  	=    0

    UpperCAmelCase  	=    image[grid[:, :, 1], grid[:, :, 0]]

    # Set pixels at invalid locations to 0
    UpperCAmelCase  	=    0

    return warped




def 			lowercase__  (    lowerCAmelCase :    np.ndarray       , lowerCAmelCase :    np.ndarray       , lowerCAmelCase :    SupportsIndex       , lowerCAmelCase :    float | None = None       , )       ->					tuple[np.ndarray, np.ndarray]:






    """simple docstring"""




    if alpha is None:
        UpperCAmelCase  	=    0.1

    # Initialize flow
    UpperCAmelCase  	=    np.zeros_like(lowerCAmelCase       )
    UpperCAmelCase  	=    np.zeros_like(lowerCAmelCase       )

    # Prepare kernels for the calculation of the derivatives and the average velocity
    UpperCAmelCase  	=    np.array([[-1, 1], [-1, 1]]       ) * 0.25
    UpperCAmelCase  	=    np.array([[-1, -1], [1, 1]]       ) * 0.25
    UpperCAmelCase  	=    np.array([[1, 1], [1, 1]]       ) * 0.25
    UpperCAmelCase  	=    np.array(
        [[1 / 12, 1 / 6, 1 / 12], [1 / 6, 0, 1 / 6], [1 / 12, 1 / 6, 1 / 12]]       )

    # Iteratively refine the flow
    for _ in range(lowerCAmelCase       ):
        UpperCAmelCase  	=    warp(lowerCAmelCase       , lowerCAmelCase       , lowerCAmelCase       )
        UpperCAmelCase  	=    convolve(lowerCAmelCase       , lowerCAmelCase       ) + convolve(lowerCAmelCase       , lowerCAmelCase       )
        UpperCAmelCase  	=    convolve(lowerCAmelCase       , lowerCAmelCase       ) + convolve(lowerCAmelCase       , lowerCAmelCase       )
        UpperCAmelCase  	=    convolve(lowerCAmelCase       , lowerCAmelCase       ) + convolve(lowerCAmelCase       , -kernel_t       )

        UpperCAmelCase  	=    convolve(lowerCAmelCase       , lowerCAmelCase       )
        UpperCAmelCase  	=    convolve(lowerCAmelCase       , lowerCAmelCase       )

        # This updates the flow as proposed in the paper (Step 12)
        UpperCAmelCase  	=    (
            derivative_x * avg_horizontal_velocity
            + derivative_y * avg_vertical_velocity
            + derivative_t
        )
        UpperCAmelCase  	=    update / (alpha**2 + derivative_x**2 + derivative_y**2)

        UpperCAmelCase  	=    avg_horizontal_velocity - derivative_x * update
        UpperCAmelCase  	=    avg_vertical_velocity - derivative_y * update

    return horizontal_flow, vertical_flow


if __name__ == "__main__":
   import doctest

   doctest.testmod()



