

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 					__lowercase  :





							"""simple docstring"""

							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=5_12    ,						__UpperCAmelCase=16    ,						__UpperCAmelCase=2    ,						__UpperCAmelCase=0.0_2    ,						__UpperCAmelCase=3    ,						__UpperCAmelCase=4    ,						__UpperCAmelCase=None    ,						__UpperCAmelCase=0    ,						)		->       Tuple:
														A      :					int      		=					parent
														A      :					int      		=					batch_size
														A      :					Tuple      		=					seq_length
														A      :					Union[str, Any]      		=					is_training
														A      :					List[str]      		=					use_input_mask
														A      :					Dict      		=					use_token_type_ids
														A      :					List[Any]      		=					use_labels
														A      :					Dict      		=					vocab_size
														A      :					str      		=					hidden_size
														A      :					List[str]      		=					num_hidden_layers
														A      :					List[Any]      		=					num_attention_heads
														A      :					Optional[Any]      		=					intermediate_size
														A      :					Optional[Any]      		=					hidden_act
														A      :					int      		=					hidden_dropout_prob
														A      :					List[Any]      		=					attention_probs_dropout_prob
														A      :					Optional[Any]      		=					max_position_embeddings
														A      :					int      		=					type_vocab_size
														A      :					Dict      		=					type_sequence_label_size
														A      :					Dict      		=					initializer_range
														A      :					Union[str, Any]      		=					num_labels
														A      :					List[Any]      		=					num_choices
														A      :					Union[str, Any]      		=					scope
														A      :					int      		=					projection_dim




							def       snake_case			(					self						)		->       List[str]:
														A      :					Any      		=					ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size						)

														A      :					Optional[int]      		=					None
														if self.use_input_mask:
																					# follow test_modeling_tf_ctrl.py
																					A      :					Optional[Any]      		=					random_attention_mask([self.batch_size, self.seq_length]						)

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

														A      :					Dict      		=					None
														A      :					str      		=					None
														A      :					Dict      		=					None
														if self.use_labels:
																					A      :					Optional[int]      		=					ids_tensor([self.batch_size]    ,						self.type_sequence_label_size						)
																					A      :					Tuple      		=					ids_tensor([self.batch_size, self.seq_length]    ,						self.num_labels						)
																					A      :					int      		=					ids_tensor([self.batch_size]    ,						self.num_choices						)

														A      :					Optional[int]      		=					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      :					Union[str, Any]      		=					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						)		->       int:
														A      :					List[str]      		=					TFDPRContextEncoder(config=__UpperCAmelCase						)
														A      :					str      		=					model(__UpperCAmelCase    ,						attention_mask=__UpperCAmelCase    ,						token_type_ids=__UpperCAmelCase						)
														A      :					Dict      		=					model(__UpperCAmelCase    ,						token_type_ids=__UpperCAmelCase						)
														A      :					Dict      		=					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[Any]:
														A      :					Optional[int]      		=					TFDPRQuestionEncoder(config=__UpperCAmelCase						)
														A      :					int      		=					model(__UpperCAmelCase    ,						attention_mask=__UpperCAmelCase    ,						token_type_ids=__UpperCAmelCase						)
														A      :					Optional[int]      		=					model(__UpperCAmelCase    ,						token_type_ids=__UpperCAmelCase						)
														A      :					Optional[Any]      		=					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						)		->       List[Any]:
														A      :					List[Any]      		=					TFDPRReader(config=__UpperCAmelCase						)
														A      :					Union[str, Any]      		=					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						)		->       Dict:
														A      :					Optional[int]      		=					self.prepare_config_and_inputs()
														(
														    A			
														)      :					List[Any]      		=					config_and_inputs
														A      :					str      		=					{"input_ids": input_ids}
														return config, inputs_dict







@require_tf
class 					__lowercase  (						_UpperCamelCase		,       _UpperCamelCase		,       unittest.TestCase				):





							"""simple docstring"""

							UpperCAmelCase_						:		List[str]								=	(
							    (
							        TFDPRContextEncoder,
							        TFDPRQuestionEncoder,
							        TFDPRReader,
							    )
							    if is_tf_available()
							    else ()
							)
							UpperCAmelCase_						:		Optional[Any]								=	{'feature-extraction': TFDPRQuestionEncoder} if is_tf_available() else {}

							UpperCAmelCase_						:		List[str]								=	False
							UpperCAmelCase_						:		Optional[int]								=	False
							UpperCAmelCase_						:		Optional[int]								=	False
							UpperCAmelCase_						:		int								=	False
							UpperCAmelCase_						:		Optional[Any]								=	False




							def       snake_case			(					self						)		->       Union[str, Any]:
														A      :					Any      		=					TFDPRModelTester(self						)
														A      :					str      		=					ConfigTester(self    ,						config_class=__UpperCAmelCase    ,						hidden_size=37						)




							def       snake_case			(					self						)		->       List[str]:
														self.config_tester.run_common_tests()




							def       snake_case			(					self						)		->       Union[str, Any]:
														A      :					int      		=					self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_dpr_context_encoder(*__UpperCAmelCase						)




							def       snake_case			(					self						)		->       List[Any]:
														A      :					str      		=					self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_dpr_question_encoder(*__UpperCAmelCase						)




							def       snake_case			(					self						)		->       List[str]:
														A      :					List[Any]      		=					self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_dpr_reader(*__UpperCAmelCase						)






							@slow
							def       snake_case			(					self						)		->       Tuple:
														for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					A      :					List[Any]      		=					TFDPRContextEncoder.from_pretrained(__UpperCAmelCase						)
																					self.assertIsNotNone(__UpperCAmelCase						)

														for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					A      :					Optional[int]      		=					TFDPRContextEncoder.from_pretrained(__UpperCAmelCase						)
																					self.assertIsNotNone(__UpperCAmelCase						)

														for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					A      :					Any      		=					TFDPRQuestionEncoder.from_pretrained(__UpperCAmelCase						)
																					self.assertIsNotNone(__UpperCAmelCase						)

														for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					A      :					Union[str, Any]      		=					TFDPRReader.from_pretrained(__UpperCAmelCase						)
																					self.assertIsNotNone(__UpperCAmelCase						)







@require_tf
class 					__lowercase  (						unittest.TestCase				):





							"""simple docstring"""

							@slow
							def       snake_case			(					self						)		->       Any:
														A      :					Dict      		=					TFDPRQuestionEncoder.from_pretrained('''facebook/dpr-question_encoder-single-nq-base'''						)

														A      :					Union[str, Any]      		=					tf.constant(
														    [[1_01, 75_92, 10_10, 20_03, 20_26, 38_99, 1_01_40, 10_29, 1_02]]						)  # [CLS] hello, is my dog cute? [SEP]
														A      :					Optional[Any]      		=					model(__UpperCAmelCase						)[0]  # embedding shape = (1, 768)
														# compare the actual values for a slice.
														A      :					Union[str, Any]      		=					tf.constant(
														    [
														        [
														            0.0_3_2_3_6_2_5_3,
														            0.1_2_7_5_3_3_3_5,
														            0.1_6_8_1_8_5_0_9,
														            0.0_0_2_7_9_7_8_6,
														            0.3_8_9_6_9_3_3,
														            0.2_4_2_6_4_9_4_5,
														            0.2_1_7_8_9_7_1,
														            -0.0_2_3_3_5_2_2_7,
														            -0.0_8_4_8_1_9_5_9,
														            -0.1_4_3_2_4_1_1_7,
														        ]
														    ]						)
														self.assertTrue(numpy.allclose(output[:, :10].numpy()    ,						expected_slice.numpy()    ,						atol=1E-4						)						)











"""simple docstring"""




from manim import *






class       UpperCAmelCase_    (							_UpperCamelCase		):



      def 		snake_case_   (	self	:      Tuple		):
           _UpperCAmelCase     :      str														=       Rectangle(height=0.5	,      width=0.5		)
           _UpperCAmelCase     :      List[Any]														=       Rectangle(height=0.25	,      width=0.25		)
           _UpperCAmelCase     :      Tuple														=       Rectangle(height=0.46	,      width=0.46		).set_stroke(width=0		)

           _UpperCAmelCase     :      List[Any]														=       [mem.copy() for i in range(6		)]
           _UpperCAmelCase     :      List[str]														=       [mem.copy() for i in range(6		)]
           _UpperCAmelCase     :      int														=       VGroup(*A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      Dict														=       VGroup(*A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      int														=       VGroup(A	,      A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      str														=       Text("CPU"	,      font_size=2_4		)
           _UpperCAmelCase     :      Union[str, Any]														=       Group(A	,      A		).arrange(A	,      buff=0.5	,      aligned_edge=A		)
           cpu.move_to([-2.5, -0.5, 0]		)
           self.add(A		)

           _UpperCAmelCase     :      Union[str, Any]														=       [mem.copy() for i in range(4		)]
           _UpperCAmelCase     :      Optional[int]														=       VGroup(*A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      Any														=       Text("GPU"	,      font_size=2_4		)
           _UpperCAmelCase     :      Dict														=       Group(A	,      A		).arrange(A	,      buff=0.5	,      aligned_edge=A		)
           gpu.move_to([-1, -1, 0]		)
           self.add(A		)

           _UpperCAmelCase     :      Tuple														=       [mem.copy() for i in range(6		)]
           _UpperCAmelCase     :      str														=       VGroup(*A		).arrange(A	,      buff=0		)

           _UpperCAmelCase     :      Any														=       Text("Model"	,      font_size=2_4		)
           _UpperCAmelCase     :      Any														=       Group(A	,      A		).arrange(A	,      buff=0.5	,      aligned_edge=A		)
           model.move_to([3, -1.0, 0]		)
           self.add(A		)

           _UpperCAmelCase     :      List[str]														=       []
           _UpperCAmelCase     :      Optional[int]														=       []
           _UpperCAmelCase     :      str														=       []

           for i, rect in enumerate(A		):
                rect.set_stroke(A		)

                _UpperCAmelCase     :      Dict														=       Rectangle(height=0.46 / 4	,      width=0.46 / 3		).set_stroke(width=0.0		).set_fill(A	,      opacity=0.7		)

                if i == 0:
                     cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT		)	,      buff=0.02	,      direction=A		)
                     cpu_target.set_x(cpu_target.get_x() + 0.1		)
                elif i == 3:
                     cpu_target.next_to(model_cpu_arr[0]	,      direction=A	,      buff=0.0		)
                else:
                     cpu_target.next_to(model_cpu_arr[i - 1]	,      direction=A	,      buff=0.0		)
                self.add(A		)
                model_cpu_arr.append(A		)

           self.add(*A	,      *A	,      *A		)

           _UpperCAmelCase     :      Dict														=       [mem.copy() for i in range(6		)]
           _UpperCAmelCase     :      List[str]														=       VGroup(*A		).arrange(A	,      buff=0		)

           _UpperCAmelCase     :      List[str]														=       Text("Loaded Checkpoint"	,      font_size=2_4		)
           _UpperCAmelCase     :      Optional[int]														=       Group(A	,      A		).arrange(A	,      buff=0.5	,      aligned_edge=A		)
           checkpoint.move_to([3, 0.5, 0]		)

           self.add(A		)

           _UpperCAmelCase     :      Union[str, Any]														=       []
           _UpperCAmelCase     :      Union[str, Any]														=       []

           for i, rect in enumerate(A		):
                _UpperCAmelCase     :      str														=       fill.copy().set_fill(A	,      opacity=0.7		)
                target.move_to(A		)
                ckpt_arr.append(A		)

                _UpperCAmelCase     :      Dict														=       target.copy()
                if i < 5:
                     cpu_target.move_to(cpu_left_col_base[i + 1]		)
                else:
                     cpu_target.move_to(cpu_right_col_base[i - 5]		)
                ckpt_cpu_arr.append(A		)
           self.add(*A	,      *A		)

           _UpperCAmelCase     :      List[str]														=       Square(side_length=2.2		)
           key.move_to([-5, 2, 0]		)

           _UpperCAmelCase     :      List[Any]														=       MarkupText(
               f'<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'	,      font_size=1_8	,      )

           key_text.move_to([-5, 2.4, 0]		)

           self.add(A	,      A		)

           _UpperCAmelCase     :      Tuple														=       MarkupText(
               f'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'	,      font_size=1_8	,      )

           blue_text.next_to(A	,      DOWN * 2.4	,      aligned_edge=key_text.get_left()		)
           self.add(A		)

           _UpperCAmelCase     :      Any														=       MarkupText(
               f'Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device.'	,      font_size=2_4	,      )
           step_a.move_to([2, 2, 0]		)

           _UpperCAmelCase     :      int														=       [meta_mem.copy() for i in range(6		)]
           _UpperCAmelCase     :      Optional[int]														=       [meta_mem.copy() for i in range(6		)]
           _UpperCAmelCase     :      List[str]														=       VGroup(*A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      Optional[int]														=       VGroup(*A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      List[Any]														=       VGroup(A	,      A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      Union[str, Any]														=       Text("Disk"	,      font_size=2_4		)
           _UpperCAmelCase     :      int														=       Group(A	,      A		).arrange(A	,      buff=0.5	,      aligned_edge=A		)
           disk.move_to([-4.0, -1.25, 0]		)
           self.play(Write(A	,      run_time=3		)	,      Write(A	,      run_time=1		)	,      Create(A	,      run_time=1		)		)

           _UpperCAmelCase     :      List[Any]														=       []
           for i, rect in enumerate(A		):
                _UpperCAmelCase     :      Any														=       rect.copy()
                target.generate_target()
                target.target.move_to(disk_left_col_base[i]		).scale(0.5		)
                animations.append(MoveToTarget(A	,      run_time=1.5		)		)
           self.play(*A		)

           self.play(FadeOut(A		)		)

           _UpperCAmelCase     :      int														=       MarkupText(f'Then, the checkpoint is removed from memory\nthrough garbage collection.'	,      font_size=2_4		)
           step_a.move_to([2, 2, 0]		)

           self.play(Write(A	,      run_time=3		)		)

           self.play(
               FadeOut(A	,      A	,      *A	,      *A		)	,      )

           self.wait()
