


import numpy as np
import torch
from torch.utils.data import Dataset

from utils import logger






class     UpperCAmelCase__       (					A_		):



    """simple docstring"""






    def __init__(			self					,   A_					,   A_    )  ->       Union[str, Any]:
         __UpperCamelCase    							=params

         __UpperCamelCase    							=np.array(A_    )
         __UpperCamelCase    							=np.array([len(A_    ) for t in data]    )

         self.check()
         self.remove_long_sequences()
         self.remove_empty_sequences()
         self.remove_unknown_sequences()
         self.check()
         self.print_statistics()






    def __getitem__(			self					,   A_    )  ->       List[str]:
         return (self.token_ids[index], self.lengths[index])






    def __len__(			self    )  ->       Optional[Any]:
         return len(self.lengths    )






    def _a					(			self    )  ->       Any:
         assert len(self.token_ids    ) == len(self.lengths    )
         assert all(self.lengths[i] == len(self.token_ids[i]    ) for i in range(len(self.lengths    )    )    )






    def _a					(			self    )  ->       Union[str, Any]:
         __UpperCamelCase    							=self.params.max_model_input_size
         __UpperCamelCase    							=self.lengths > max_len
         logger.info(f'Splitting {sum(A_    )} too long sequences.'    )

         def divide_chunks(A_					,   A_    ):
              return [l[i : i + n] for i in range(0					,   len(A_    )					,   A_    )]

         __UpperCamelCase    							=[]
         __UpperCamelCase    							=[]
         if self.params.mlm:
              __UpperCamelCase							,      __UpperCamelCase    							=self.params.special_tok_ids['cls_token'], self.params.special_tok_ids['sep_token']
         else:
              __UpperCamelCase							,      __UpperCamelCase    							=self.params.special_tok_ids['bos_token'], self.params.special_tok_ids['eos_token']

         for seq_, len_ in zip(self.token_ids					,   self.lengths    ):
              assert (seq_[0] == cls_id) and (seq_[-1] == sep_id), seq_
              if len_ <= max_len:
                   new_tok_ids.append(seq_    )
                   new_lengths.append(len_    )
              else:
                   __UpperCamelCase    							=[]
                   for sub_s in divide_chunks(seq_					,   max_len - 2    ):
                        if sub_s[0] != cls_id:
                             __UpperCamelCase    							=np.insert(A_					,   0					,   A_    )
                        if sub_s[-1] != sep_id:
                             __UpperCamelCase    							=np.insert(A_					,   len(A_    )					,   A_    )
                        assert len(A_    ) <= max_len
                        assert (sub_s[0] == cls_id) and (sub_s[-1] == sep_id), sub_s
                        sub_seqs.append(A_    )

                   new_tok_ids.extend(A_    )
                   new_lengths.extend([len(A_    ) for l in sub_seqs]    )

         __UpperCamelCase    							=np.array(A_    )
         __UpperCamelCase    							=np.array(A_    )






    def _a					(			self    )  ->       Optional[int]:
         __UpperCamelCase    							=len(self    )
         __UpperCamelCase    							=self.lengths > 11
         __UpperCamelCase    							=self.token_ids[indices]
         __UpperCamelCase    							=self.lengths[indices]
         __UpperCamelCase    							=len(self    )
         logger.info(f'Remove {init_size - new_size} too short (<=11 tokens) sequences.'    )






    def _a					(			self    )  ->       int:
         if "unk_token" not in self.params.special_tok_ids:
              return
         else:
              __UpperCamelCase    							=self.params.special_tok_ids['unk_token']
         __UpperCamelCase    							=len(self    )
         __UpperCamelCase    							=np.array([np.count_nonzero(a == unk_token_id    ) for a in self.token_ids]    )
         __UpperCamelCase    							=(unk_occs / self.lengths) < 0.5
         __UpperCamelCase    							=self.token_ids[indices]
         __UpperCamelCase    							=self.lengths[indices]
         __UpperCamelCase    							=len(self    )
         logger.info(f'Remove {init_size - new_size} sequences with a high level of unknown tokens (50%).'    )






    def _a					(			self    )  ->       int:
         if not self.params.is_master:
              return
         logger.info(f'{len(self    )} sequences'    )
         # data_len = sum(self.lengths)
         # nb_unique_tokens = len(Counter(list(chain(*self.token_ids))))
         # logger.info(f'{data_len} tokens ({nb_unique_tokens} unique)')

         # unk_idx = self.params.special_tok_ids['unk_token']
         # nb_unknown = sum([(t==unk_idx).sum() for t in self.token_ids])
         # logger.info(f'{nb_unknown} unknown tokens (covering {100*nb_unknown/data_len:.2f}% of the data)')






    def _a					(			self					,   A_    )  ->       int:
         __UpperCamelCase    							=[t[0] for t in batch]
         __UpperCamelCase    							=[t[1] for t in batch]
         assert len(A_    ) == len(A_    )

         # Max for paddings
         __UpperCamelCase    							=max(A_    )

         # Pad token ids
         if self.params.mlm:
              __UpperCamelCase    							=self.params.special_tok_ids['pad_token']
         else:
              __UpperCamelCase    							=self.params.special_tok_ids['unk_token']
         __UpperCamelCase    							=[list(t.astype(A_    )    ) + [pad_idx] * (max_seq_len_ - len(A_    )) for t in token_ids]
         assert len(tk_    ) == len(A_    )
         assert all(len(A_    ) == max_seq_len_ for t in tk_    )

         __UpperCamelCase    							=torch.tensor(tk_    )  # (bs, max_seq_len_)
         __UpperCamelCase    							=torch.tensor(A_    )  # (bs)
         return tk_t, lg_t






"""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 			UpperCamelCase__:





						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						,)    -> Dict:
													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   )    -> Optional[Any]:
													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   )    -> Tuple:
													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   )    -> Union[str, Any]:
													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]:
													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:
													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 			UpperCamelCase__(	__A						,  __A						,  unittest.TestCase       ):
						lowerCAmelCase__     :  Optional[int]  							=						(
						    (
						        TFDPRContextEncoder,
						        TFDPRQuestionEncoder,
						        TFDPRReader,
						    )
						    if is_tf_available()
						    else ()
						)
						lowerCAmelCase__     :  List[str]  							=						{'feature-extraction': TFDPRQuestionEncoder} if is_tf_available() else {}

						lowerCAmelCase__     :  Tuple  							=						False
						lowerCAmelCase__     :  Optional[int]  							=						False
						lowerCAmelCase__     :  List[str]  							=						False
						lowerCAmelCase__     :  int  							=						False
						lowerCAmelCase__     :  str  							=						False





						def  snake_case__		(				self   )    -> str:
													A__		   =							TFDPRModelTester(self   )
													A__		   =							ConfigTester(self						,config_class=__UpperCAmelCase						,hidden_size=37   )





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





						def  snake_case__		(				self   )    -> int:
													A__		   =							self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_dpr_context_encoder(*__UpperCAmelCase   )





						def  snake_case__		(				self   )    -> Optional[Any]:
													A__		   =							self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_dpr_question_encoder(*__UpperCAmelCase   )





						def  snake_case__		(				self   )    -> List[str]:
													A__		   =							self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_dpr_reader(*__UpperCAmelCase   )





						@slow
						def  snake_case__		(				self   )    -> int:
													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 			UpperCamelCase__(	unittest.TestCase       ):





						@slow
						def  snake_case__		(				self   )    -> Optional[Any]:
													A__		   =							TFDPRQuestionEncoder.from_pretrained('facebook/dpr-question_encoder-single-nq-base'   )

													A__		   =							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__		   =							model(__UpperCAmelCase   )[0]  # embedding shape = (1, 768)
													# compare the actual values for a slice.
													A__		   =							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   )   )





