

'''simple docstring'''







from __future__ import annotations

import unittest

from transformers import 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 numpy
      import tensorflow as tf

      from transformers import (
          TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
          TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
          TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
          BertConfig,
          DPRConfig,
          TFDPRContextEncoder,
          TFDPRQuestionEncoder,
          TFDPRReader,
      )


class 					A__      :



    def __init__(					self		,				UpperCamelCase__		,				UpperCamelCase__=13		,				UpperCamelCase__=7		,				UpperCamelCase__=True		,				UpperCamelCase__=True		,				UpperCamelCase__=True		,				UpperCamelCase__=True		,				UpperCamelCase__=99		,				UpperCamelCase__=32		,				UpperCamelCase__=2		,				UpperCamelCase__=4		,				UpperCamelCase__=37		,				UpperCamelCase__="gelu"		,				UpperCamelCase__=0.1		,				UpperCamelCase__=0.1		,				UpperCamelCase__=512		,				UpperCamelCase__=16		,				UpperCamelCase__=2		,				UpperCamelCase__=0.02		,				UpperCamelCase__=3		,				UpperCamelCase__=4		,				UpperCamelCase__=None		,				UpperCamelCase__=0		,				)						->      Dict:





           '''simple docstring'''
           A_						=      parent
           A_						=      batch_size
           A_						=      seq_length
           A_						=      is_training
           A_						=      use_input_mask
           A_						=      use_token_type_ids
           A_						=      use_labels
           A_						=      vocab_size
           A_						=      hidden_size
           A_						=      num_hidden_layers
           A_						=      num_attention_heads
           A_						=      intermediate_size
           A_						=      hidden_act
           A_						=      hidden_dropout_prob
           A_						=      attention_probs_dropout_prob
           A_						=      max_position_embeddings
           A_						=      type_vocab_size
           A_						=      type_sequence_label_size
           A_						=      initializer_range
           A_						=      num_labels
           A_						=      num_choices
           A_						=      scope
           A_						=      projection_dim



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





           '''simple docstring'''
           A_						=      ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size       )

           A_						=      None
           if self.use_input_mask:
                  # follow test_modeling_tf_ctrl.py
                  A_						=      random_attention_mask([self.batch_size, self.seq_length]       )

           A_						=      None
           if self.use_token_type_ids:
                  A_						=      ids_tensor([self.batch_size, self.seq_length]		,				self.type_vocab_size       )

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

           A_						=      BertConfig(
               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		,				is_decoder=UpperCamelCase__		,				initializer_range=self.initializer_range		,				)
           A_						=      DPRConfig(projection_dim=self.projection_dim		,				**config.to_dict()       )

           return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__       )						->      Optional[Any]:





           '''simple docstring'''
           A_						=      TFDPRContextEncoder(config=UpperCamelCase__       )
           A_						=      model(UpperCamelCase__		,				attention_mask=UpperCamelCase__		,				token_type_ids=UpperCamelCase__       )
           A_						=      model(UpperCamelCase__		,				token_type_ids=UpperCamelCase__       )
           A_						=      model(UpperCamelCase__       )
           self.parent.assertEqual(result.pooler_output.shape		,				(self.batch_size, self.projection_dim or self.hidden_size)       )



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__       )						->      Tuple:





           '''simple docstring'''
           A_						=      TFDPRQuestionEncoder(config=UpperCamelCase__       )
           A_						=      model(UpperCamelCase__		,				attention_mask=UpperCamelCase__		,				token_type_ids=UpperCamelCase__       )
           A_						=      model(UpperCamelCase__		,				token_type_ids=UpperCamelCase__       )
           A_						=      model(UpperCamelCase__       )
           self.parent.assertEqual(result.pooler_output.shape		,				(self.batch_size, self.projection_dim or self.hidden_size)       )



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__       )						->      Optional[int]:





           '''simple docstring'''
           A_						=      TFDPRReader(config=UpperCamelCase__       )
           A_						=      model(UpperCamelCase__		,				attention_mask=UpperCamelCase__       )

           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)       )
           self.parent.assertEqual(result.relevance_logits.shape		,				(self.batch_size,)       )







    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           A_						=      self.prepare_config_and_inputs()
           (
               (
               A_
           )					,			(
               A_
           )					,			(
               A_
           )					,			(
               A_
           )					,			(
               A_
           )					,			(
               A_
           )					,			(
               A_
           )					,			
           )						=      config_and_inputs
           A_						=      {"""input_ids""": input_ids}
           return config, inputs_dict


@require_tf
class 					A__      (      _snake_case						,						_snake_case						,						unittest.TestCase			):
    lowercase						=		(
        (
            TFDPRContextEncoder,
            TFDPRQuestionEncoder,
            TFDPRReader,
        )
        if is_tf_available()
        else ()
    )
    lowercase						=		{"feature-extraction": TFDPRQuestionEncoder} if is_tf_available() else {}

    lowercase						=		False
    lowercase						=		False
    lowercase						=		False
    lowercase						=		False
    lowercase						=		False



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





           '''simple docstring'''
           A_						=      TFDPRModelTester(self       )
           A_						=      ConfigTester(self		,				config_class=UpperCamelCase__		,				hidden_size=37       )



    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           self.config_tester.run_common_tests()



    def 					snake_case_			(					self       )						->      Any:





           '''simple docstring'''
           A_						=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_context_encoder(*UpperCamelCase__       )



    def 					snake_case_			(					self       )						->      Optional[int]:





           '''simple docstring'''
           A_						=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_question_encoder(*UpperCamelCase__       )



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





           '''simple docstring'''
           A_						=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_reader(*UpperCamelCase__       )







    @slow
    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A_						=      TFDPRContextEncoder.from_pretrained(UpperCamelCase__       )
                  self.assertIsNotNone(UpperCamelCase__       )

           for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A_						=      TFDPRContextEncoder.from_pretrained(UpperCamelCase__       )
                  self.assertIsNotNone(UpperCamelCase__       )

           for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A_						=      TFDPRQuestionEncoder.from_pretrained(UpperCamelCase__       )
                  self.assertIsNotNone(UpperCamelCase__       )

           for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A_						=      TFDPRReader.from_pretrained(UpperCamelCase__       )
                  self.assertIsNotNone(UpperCamelCase__       )


@require_tf
class 					A__      (      unittest.TestCase			):



    @slow
    def 					snake_case_			(					self       )						->      Union[str, Any]:





           '''simple docstring'''
           A_						=      TFDPRQuestionEncoder.from_pretrained("""facebook/dpr-question_encoder-single-nq-base"""       )

           A_						=      tf.constant(
               [[101, 7592, 1010, 2003, 2026, 3899, 10140, 1029, 102]]       )  # [CLS] hello, is my dog cute? [SEP]
           A_						=      model(UpperCamelCase__       )[0]  # embedding shape = (1, 768)
           # compare the actual values for a slice.
           A_						=      tf.constant(
               [
                   [
                       0.03236253,
                       0.12753335,
                       0.16818509,
                       0.00279786,
                       0.3896933,
                       0.24264945,
                       0.2178971,
                       -0.02335227,
                       -0.08481959,
                       -0.14324117,
                   ]
               ]       )
           self.assertTrue(numpy.allclose(output[:, :10].numpy()		,				expected_slice.numpy()		,				atol=1e-4       )       )



'''simple docstring'''







import tensorflow as tf

from ...tf_utils import shape_list


class 					A__      (      tf.keras.layers.Layer			):



    def __init__(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__=1		,				UpperCamelCase__=False		,				**UpperCamelCase__       )						->      Union[str, Any]:





           '''simple docstring'''
           super().__init__(**UpperCamelCase__       )

           A_						=      vocab_size
           A_						=      d_embed
           A_						=      d_proj

           A_						=      cutoffs + [vocab_size]
           A_						=      [0] + self.cutoffs
           A_						=      div_val

           A_						=      self.cutoffs[0]
           A_						=      len(self.cutoffs       ) - 1
           A_						=      self.shortlist_size + self.n_clusters
           A_						=      keep_order

           A_						=      []
           A_						=      []



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





           '''simple docstring'''
           if self.n_clusters > 0:
                  A_						=      self.add_weight(
                      shape=(self.n_clusters, self.d_embed)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name="""cluster_weight"""       )
                  A_						=      self.add_weight(
                      shape=(self.n_clusters,)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name="""cluster_bias"""       )

           if self.div_val == 1:
                  for i in range(len(self.cutoffs       )       ):
                         if self.d_proj != self.d_embed:
                                A_						=      self.add_weight(
                                    shape=(self.d_embed, self.d_proj)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name=f'''out_projs_._{i}'''		,				)
                                self.out_projs.append(UpperCamelCase__       )
                         else:
                                self.out_projs.append(UpperCamelCase__       )
                         A_						=      self.add_weight(
                             shape=(self.vocab_size, self.d_embed)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name=f'''out_layers_._{i}_._weight'''		,				)
                         A_						=      self.add_weight(
                             shape=(self.vocab_size,)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name=f'''out_layers_._{i}_._bias'''		,				)
                         self.out_layers.append((weight, bias)       )
           else:
                  for i in range(len(self.cutoffs       )       ):
                         A_					,			A_						=      self.cutoff_ends[i], self.cutoff_ends[i + 1]
                         A_						=      self.d_embed // (self.div_val**i)

                         A_						=      self.add_weight(
                             shape=(d_emb_i, self.d_proj)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name=f'''out_projs_._{i}'''       )
                         self.out_projs.append(UpperCamelCase__       )
                         A_						=      self.add_weight(
                             shape=(r_idx - l_idx, d_emb_i)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name=f'''out_layers_._{i}_._weight'''		,				)
                         A_						=      self.add_weight(
                             shape=(r_idx - l_idx,)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name=f'''out_layers_._{i}_._bias'''		,				)
                         self.out_layers.append((weight, bias)       )
           super().build(UpperCamelCase__       )



    @staticmethod
    def 					snake_case_			(					UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__=None       )						->      List[Any]:





           '''simple docstring'''
           A_						=      x
           if proj is not None:
                  A_						=      tf.einsum("""ibd,ed->ibe"""		,				UpperCamelCase__		,				UpperCamelCase__       )
           return tf.einsum("""ibd,nd->ibn"""		,				UpperCamelCase__		,				UpperCamelCase__       ) + b



    @staticmethod
    def 					snake_case_			(					UpperCamelCase__		,				UpperCamelCase__       )						->      List[str]:





           '''simple docstring'''
           A_						=      shape_list(UpperCamelCase__       )
           A_						=      tf.range(lp_size[0]		,				dtype=target.dtype       )
           A_						=      tf.stack([r, target]		,				1       )
           return tf.gather_nd(UpperCamelCase__		,				UpperCamelCase__       )







    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__=True		,				UpperCamelCase__=False       )						->      Optional[int]:





           '''simple docstring'''
           A_						=      0
           if self.n_clusters == 0:
                  A_						=      self._logit(UpperCamelCase__		,				self.out_layers[0][0]		,				self.out_layers[0][1]		,				self.out_projs[0]       )
                  if target is not None:
                         A_						=      tf.nn.sparse_softmax_cross_entropy_with_logits(labels=UpperCamelCase__		,				logits=UpperCamelCase__       )
                  A_						=      tf.nn.log_softmax(UpperCamelCase__		,				axis=-1       )
           else:
                  A_						=      shape_list(UpperCamelCase__       )
                  A_						=      []
                  A_						=      tf.zeros(hidden_sizes[:2]       )
                  for i in range(len(self.cutoffs       )       ):
                         A_					,			A_						=      self.cutoff_ends[i], self.cutoff_ends[i + 1]
                         if target is not None:
                                A_						=      (target >= l_idx) & (target < r_idx)
                                A_						=      tf.where(UpperCamelCase__       )
                                A_						=      tf.boolean_mask(UpperCamelCase__		,				UpperCamelCase__       ) - l_idx

                         if self.div_val == 1:
                                A_						=      self.out_layers[0][0][l_idx:r_idx]
                                A_						=      self.out_layers[0][1][l_idx:r_idx]
                         else:
                                A_						=      self.out_layers[i][0]
                                A_						=      self.out_layers[i][1]

                         if i == 0:
                                A_						=      tf.concat([cur_W, self.cluster_weight]		,				0       )
                                A_						=      tf.concat([cur_b, self.cluster_bias]		,				0       )

                                A_						=      self._logit(UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				self.out_projs[0]       )
                                A_						=      tf.nn.log_softmax(UpperCamelCase__       )
                                out.append(head_logprob[..., : self.cutoffs[0]]       )
                                if target is not None:
                                       A_						=      tf.boolean_mask(UpperCamelCase__		,				UpperCamelCase__       )
                                       A_						=      self._gather_logprob(UpperCamelCase__		,				UpperCamelCase__       )
                         else:
                                A_						=      self._logit(UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				self.out_projs[i]       )
                                A_						=      tf.nn.log_softmax(UpperCamelCase__       )
                                A_						=      self.cutoffs[0] + i - 1  # No probability for the head cluster
                                A_						=      head_logprob[..., cluster_prob_idx, None] + tail_logprob
                                out.append(UpperCamelCase__       )
                                if target is not None:
                                       A_						=      tf.boolean_mask(UpperCamelCase__		,				UpperCamelCase__       )
                                       A_						=      tf.boolean_mask(UpperCamelCase__		,				UpperCamelCase__       )
                                       A_						=      self._gather_logprob(UpperCamelCase__		,				UpperCamelCase__       )
                                       cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
                         if target is not None:
                                loss += tf.scatter_nd(UpperCamelCase__		,				-cur_logprob		,				shape_list(UpperCamelCase__       )       )
                  A_						=      tf.concat(UpperCamelCase__		,				axis=-1       )

           if target is not None:
                  if return_mean:
                         A_						=      tf.reduce_mean(UpperCamelCase__       )
                  # Add the training-time loss value to the layer using `self.add_loss()`.
                  self.add_loss(UpperCamelCase__       )

                  # Log the loss as a metric (we could log arbitrary metrics,
                  # including different metrics for training and inference.
                  self.add_metric(UpperCamelCase__		,				name=self.name		,				aggregation="""mean""" if return_mean else """"""       )

           return out
