
"""simple docstring"""




import unittest

from transformers import ErnieMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    ErnieMForInformationExtraction,
	    ErnieMForMultipleChoice,
	    ErnieMForQuestionAnswering,
	    ErnieMForSequenceClassification,
	    ErnieMForTokenClassification,
	    ErnieMModel,
	)
	from transformers.models.ernie_m.modeling_ernie_m import ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST




class 					UpperCamelCase     :

				"""simple docstring"""



				def __init__(    self		,     lowercase__		,     lowercase__=13		,     lowercase__=7		,     lowercase__=True		,     lowercase__=True		,     lowercase__=True		,     lowercase__=99		,     lowercase__=32		,     lowercase__=5		,     lowercase__=4		,     lowercase__=37		,     lowercase__="gelu"		,     lowercase__=0.1		,     lowercase__=0.1		,     lowercase__=512		,     lowercase__=16		,     lowercase__=2		,     lowercase__=0.02		,     lowercase__=3		,     lowercase__=4		,     lowercase__=None		,     )   ->   Tuple:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     parent
									SCREAMING_SNAKE_CASE           =     batch_size
									SCREAMING_SNAKE_CASE           =     seq_length
									SCREAMING_SNAKE_CASE           =     is_training
									SCREAMING_SNAKE_CASE           =     use_input_mask
									SCREAMING_SNAKE_CASE           =     use_labels
									SCREAMING_SNAKE_CASE           =     vocab_size
									SCREAMING_SNAKE_CASE           =     hidden_size
									SCREAMING_SNAKE_CASE           =     num_hidden_layers
									SCREAMING_SNAKE_CASE           =     num_attention_heads
									SCREAMING_SNAKE_CASE           =     intermediate_size
									SCREAMING_SNAKE_CASE           =     hidden_act
									SCREAMING_SNAKE_CASE           =     hidden_dropout_prob
									SCREAMING_SNAKE_CASE           =     attention_probs_dropout_prob
									SCREAMING_SNAKE_CASE           =     max_position_embeddings
									SCREAMING_SNAKE_CASE           =     type_vocab_size
									SCREAMING_SNAKE_CASE           =     type_sequence_label_size
									SCREAMING_SNAKE_CASE           =     initializer_range
									SCREAMING_SNAKE_CASE           =     num_labels
									SCREAMING_SNAKE_CASE           =     num_choices
									SCREAMING_SNAKE_CASE           =     scope


				def 							A (    self      )   ->   List[str]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     ids_tensor([self.batch_size, self.seq_length]		,     self.vocab_size      )

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

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

									SCREAMING_SNAKE_CASE           =     self.get_config()

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


				def 							A (    self      )   ->   Tuple:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     ids_tensor([self.batch_size, self.seq_length]		,     self.vocab_size      )

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

									return config, input_ids, input_mask


				def 							A (    self      )   ->   Optional[Any]:
									"""simple docstring"""







									return ErnieMConfig(
									    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		,     )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__      )   ->   Tuple:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     ErnieMModel(config=lowercase__      )
									model.to(lowercase__      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     model(lowercase__		,     return_dict=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__      )   ->   Tuple:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     ErnieMForQuestionAnswering(config=lowercase__      )
									model.to(lowercase__      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     model(
									    lowercase__		,     attention_mask=lowercase__		,     start_positions=lowercase__		,     end_positions=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__      )   ->   Optional[int]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     ErnieMForInformationExtraction(config=lowercase__      )
									model.to(lowercase__      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     torch.ones_like(lowercase__		,     dtype=torch.floataa      )
									SCREAMING_SNAKE_CASE           =     model(
									    lowercase__		,     attention_mask=lowercase__		,     start_positions=lowercase__		,     end_positions=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__      )   ->   List[str]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.num_labels
									SCREAMING_SNAKE_CASE           =     ErnieMForSequenceClassification(lowercase__      )
									model.to(lowercase__      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     model(lowercase__		,     attention_mask=lowercase__		,     labels=lowercase__      )
									self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.num_labels)      )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__      )   ->   List[str]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.num_labels
									SCREAMING_SNAKE_CASE           =     ErnieMForTokenClassification(config=lowercase__      )
									model.to(lowercase__      )
									model.eval()
									input_ids.to(lowercase__      )
									input_mask.to(lowercase__      )
									token_labels.to(lowercase__      )

									SCREAMING_SNAKE_CASE           =     model(lowercase__		,     attention_mask=lowercase__		,     labels=lowercase__      )

									self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.seq_length, self.num_labels)      )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__      )   ->   List[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.num_choices
									SCREAMING_SNAKE_CASE           =     ErnieMForMultipleChoice(config=lowercase__      )
									model.to(lowercase__      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     input_ids.unsqueeze(1      ).expand(-1		,     self.num_choices		,     -1      ).contiguous()
									SCREAMING_SNAKE_CASE           =     input_mask.unsqueeze(1      ).expand(-1		,     self.num_choices		,     -1      ).contiguous()
									SCREAMING_SNAKE_CASE           =     model(
									    lowercase__		,     attention_mask=lowercase__		,     labels=lowercase__		,     )
									self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.num_choices)      )




				def 							A (    self      )   ->   Optional[int]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.prepare_config_and_inputs()
									(
									    (
									    SCREAMING_SNAKE_CASE
									)     ,   (
									    SCREAMING_SNAKE_CASE
									)     ,   (
									    SCREAMING_SNAKE_CASE
									)     ,   (
									    SCREAMING_SNAKE_CASE
									)     ,   (
									    SCREAMING_SNAKE_CASE
									)     ,   (
									    SCREAMING_SNAKE_CASE
									)     ,   
									)           =     config_and_inputs
									SCREAMING_SNAKE_CASE           =     {'input_ids': input_ids, 'attention_mask': input_mask}
									return config, inputs_dict






@require_torch
class 					UpperCamelCase     (	__magic_name__  ,				__magic_name__  ,				unittest.TestCase				):

				"""simple docstring"""



				UpperCAmelCase_    :				Any						  =    (
				    (
				        ErnieMModel,
				        ErnieMForMultipleChoice,
				        ErnieMForQuestionAnswering,
				        ErnieMForSequenceClassification,
				        ErnieMForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				UpperCAmelCase_    :				Optional[Any]						  =    ()
				UpperCAmelCase_    :				List[str]						  =    (
				    {
				        "feature-extraction": ErnieMModel,
				        "question-answering": ErnieMForQuestionAnswering,
				        "text-classification": ErnieMForSequenceClassification,
				        "token-classification": ErnieMForTokenClassification,
				        "zero-shot": ErnieMForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				UpperCAmelCase_    :				Dict						  =    False


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__      )   ->   List[str]:
									"""simple docstring"""







									if pipeline_test_casse_name == "QAPipelineTests":
														return True

									return False


				def 							A (    self      )   ->   List[str]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     ErnieMModelTester(self      )
									SCREAMING_SNAKE_CASE           =     ConfigTester(self		,     config_class=lowercase__		,     hidden_size=37      )


				def 							A (    self      )   ->   str:
									"""simple docstring"""







									self.config_tester.run_common_tests()


				def 							A (    self      )   ->   Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*lowercase__      )


				def 							A (    self      )   ->   Any:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
														SCREAMING_SNAKE_CASE           =     type
														self.model_tester.create_and_check_model(*lowercase__      )


				def 							A (    self      )   ->   Tuple:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_multiple_choice(*lowercase__      )


				def 							A (    self      )   ->   List[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*lowercase__      )


				def 							A (    self      )   ->   Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*lowercase__      )


				def 							A (    self      )   ->   int:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_information_extraction(*lowercase__      )


				def 							A (    self      )   ->   Any:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*lowercase__      )




				@slow
				def 							A (    self      )   ->   Tuple:
									"""simple docstring"""







									for model_name in ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														SCREAMING_SNAKE_CASE           =     ErnieMModel.from_pretrained(lowercase__      )
														self.assertIsNotNone(lowercase__      )






@require_torch
class 					UpperCamelCase     (	unittest.TestCase				):

				"""simple docstring"""



				@slow
				def 							A (    self      )   ->   int:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     ErnieMModel.from_pretrained('susnato/ernie-m-base_pytorch'      )
									model.eval()
									SCREAMING_SNAKE_CASE           =     torch.tensor([[0, 1, 2, 3, 4, 5]]      )
									SCREAMING_SNAKE_CASE           =     model(lowercase__      )[0]

									# TODO Replace vocab size
									SCREAMING_SNAKE_CASE           =     768

									SCREAMING_SNAKE_CASE           =     torch.Size((1, 6, hidden_size)      )
									self.assertEqual(output.shape		,     lowercase__      )

									SCREAMING_SNAKE_CASE           =     torch.tensor(
									    [[[-0.0_012, 0.1_245, -0.0_214], [-0.0_742, 0.0_244, -0.0_771], [-0.0_333, 0.1_164, -0.1_554]]]      )

									self.assertTrue(torch.allclose(output[:, :3, :3]		,     lowercase__		,     atol=1E-3      )      )


"""simple docstring"""




import os
from typing import Optional, Union

import numpy as np
from huggingface_hub import hf_hub_download

from ... import AutoTokenizer
from ...utils import logging


snake_case			       =		'block_records.npy'


snake_case			       =		logging.get_logger(__name__)

def   UpperCamelCase_  (		SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_      ):
					import tensorflow.compat.va as tf

					SCREAMING_SNAKE_CASE           =     tf.data.TFRecordDataset(SCREAMING_SNAKE_CASE_,					buffer_size=5_1_2 * 1_0_2_4 * 1_0_2_4      )
					SCREAMING_SNAKE_CASE           =     blocks_dataset.batch(SCREAMING_SNAKE_CASE_,					drop_remainder=SCREAMING_SNAKE_CASE_      )
					SCREAMING_SNAKE_CASE           =     next(blocks_dataset.take(1      ).as_numpy_iterator()      )

					return np_record




class 					UpperCamelCase     :

				"""simple docstring"""



				def __init__(    self		,     lowercase__		,     lowercase__		,     lowercase__=2		,     lowercase__=1000		,     lowercase__=100		,     lowercase__=100000		,     )   ->   Dict:
									"""simple docstring"""







									from scann.scann_ops.py.scann_ops_pybind import builder as Builder

									SCREAMING_SNAKE_CASE           =     Builder(db=lowercase__		,     num_neighbors=lowercase__		,     distance_measure='dot_product'      )
									SCREAMING_SNAKE_CASE           =     builder.tree(
									    num_leaves=lowercase__		,     num_leaves_to_search=lowercase__		,     training_sample_size=lowercase__      )
									SCREAMING_SNAKE_CASE           =     builder.score_ah(dimensions_per_block=lowercase__      )

									SCREAMING_SNAKE_CASE           =     builder.build()




				def 							A (    self		,     lowercase__      )   ->   str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE           =     self.searcher.search_batched(question_projection.detach().cpu()      )
									return retrieved_block_ids.astype('int64'      )






class 					UpperCamelCase     :

				"""simple docstring"""



				def __init__(    self		,     lowercase__		,     lowercase__      )   ->   Dict:
									"""simple docstring"""







									super().__init__()
									SCREAMING_SNAKE_CASE           =     block_records
									SCREAMING_SNAKE_CASE           =     tokenizer


				def __call__(    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__=None		,     lowercase__="pt"      )   ->   Optional[int]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     np.take(self.block_records		,     indices=lowercase__		,     axis=0      )

									SCREAMING_SNAKE_CASE           =     self.tokenizer.decode(question_input_ids[0]		,     skip_special_tokens=lowercase__      )

									SCREAMING_SNAKE_CASE           =     []
									SCREAMING_SNAKE_CASE           =     []
									for retrieved_block in retrieved_blocks:
														text.append(lowercase__      )
														text_pair.append(retrieved_block.decode()      )

									SCREAMING_SNAKE_CASE           =     self.tokenizer(
									    lowercase__		,     lowercase__		,     padding=lowercase__		,     truncation=lowercase__		,     return_special_tokens_mask=lowercase__		,     max_length=lowercase__      )
									SCREAMING_SNAKE_CASE           =     concat_inputs.convert_to_tensors(lowercase__      )

									if answer_ids is not None:
														return self.block_has_answer(lowercase__		,     lowercase__      ) + (concat_inputs_tensors,)
									else:
														return (None, None, None, concat_inputs_tensors)


				@classmethod
				def 							A (    cls		,     lowercase__		,     *lowercase__		,     **lowercase__      )   ->   Optional[Any]:
									"""simple docstring"""







									if os.path.isdir(lowercase__      ):
														SCREAMING_SNAKE_CASE           =     os.path.join(lowercase__		,     _REALM_BLOCK_RECORDS_FILENAME      )
									else:
														SCREAMING_SNAKE_CASE           =     hf_hub_download(
														    repo_id=lowercase__		,     filename=_REALM_BLOCK_RECORDS_FILENAME		,     **lowercase__      )
									SCREAMING_SNAKE_CASE           =     np.load(lowercase__		,     allow_pickle=lowercase__      )

									SCREAMING_SNAKE_CASE           =     AutoTokenizer.from_pretrained(lowercase__		,     *lowercase__		,     **lowercase__      )

									return cls(lowercase__		,     lowercase__      )


				def 							A (    self		,     lowercase__      )   ->   List[str]:
									"""simple docstring"""







									np.save(os.path.join(lowercase__		,     _REALM_BLOCK_RECORDS_FILENAME      )		,     self.block_records      )
									# save tokenizer
									self.tokenizer.save_pretrained(lowercase__      )




				def 							A (    self		,     lowercase__		,     lowercase__      )   ->   Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     []
									SCREAMING_SNAKE_CASE           =     []
									SCREAMING_SNAKE_CASE           =     []
									SCREAMING_SNAKE_CASE           =     0

									for input_id in concat_inputs.input_ids:
														SCREAMING_SNAKE_CASE           =     input_id.tolist()
														# Check answers between two [SEP] tokens
														SCREAMING_SNAKE_CASE           =     input_id_list.index(self.tokenizer.sep_token_id      )
														SCREAMING_SNAKE_CASE           =     first_sep_idx + 1 + input_id_list[first_sep_idx + 1 :].index(self.tokenizer.sep_token_id      )

														start_pos.append([]      )
														end_pos.append([]      )
														for answer in answer_ids:
																			for idx in range(first_sep_idx + 1		,     lowercase__      ):
																								if answer[0] == input_id_list[idx]:
																													if input_id_list[idx : idx + len(lowercase__      )] == answer:
																																		start_pos[-1].append(lowercase__      )
																																		end_pos[-1].append(idx + len(lowercase__      ) - 1      )

														if len(start_pos[-1]      ) == 0:
																			has_answers.append(lowercase__      )
														else:
																			has_answers.append(lowercase__      )
																			if len(start_pos[-1]      ) > max_answers:
																								SCREAMING_SNAKE_CASE           =     len(start_pos[-1]      )

        # Pad -1 to max_answers
									for start_pos_, end_pos_ in zip(lowercase__		,     lowercase__      ):
														if len(start_pos_      ) < max_answers:
																			SCREAMING_SNAKE_CASE           =     [-1] * (max_answers - len(start_pos_      ))
																			start_pos_ += padded
																			end_pos_ += padded
									return has_answers, start_pos, end_pos
