






import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, 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 (
       MODEL_FOR_PRETRAINING_MAPPING,
       MegatronBertForCausalLM,
       MegatronBertForMaskedLM,
       MegatronBertForMultipleChoice,
       MegatronBertForNextSentencePrediction,
       MegatronBertForPreTraining,
       MegatronBertForQuestionAnswering,
       MegatronBertForSequenceClassification,
       MegatronBertForTokenClassification,
       MegatronBertModel,
   )


class 			SCREAMING_SNAKE_CASE      :
     '''simple docstring'''


     def __init__(					self							:						str       ,			UpperCAmelCase_							:						Optional[int]       ,			UpperCAmelCase_							:						Optional[Any]=13       ,			UpperCAmelCase_							:						Tuple=7       ,			UpperCAmelCase_							:						Tuple=True       ,			UpperCAmelCase_							:						int=True       ,			UpperCAmelCase_							:						Tuple=True       ,			UpperCAmelCase_							:						Any=True       ,			UpperCAmelCase_							:						int=99       ,			UpperCAmelCase_							:						int=64       ,			UpperCAmelCase_							:						Optional[int]=32       ,			UpperCAmelCase_							:						Optional[int]=5       ,			UpperCAmelCase_							:						str=4       ,			UpperCAmelCase_							:						Optional[int]=37       ,			UpperCAmelCase_							:						Union[str, Any]="gelu"       ,			UpperCAmelCase_							:						Any=0.1       ,			UpperCAmelCase_							:						Tuple=0.1       ,			UpperCAmelCase_							:						Optional[Any]=512       ,			UpperCAmelCase_							:						Optional[int]=16       ,			UpperCAmelCase_							:						Union[str, Any]=2       ,			UpperCAmelCase_							:						Optional[Any]=0.02       ,			UpperCAmelCase_							:						Union[str, Any]=3       ,			UpperCAmelCase_							:						Dict=4       ,			UpperCAmelCase_							:						Any=None       ,			):
       SCREAMING_SNAKE_CASE							:   Any						    =     parent
       SCREAMING_SNAKE_CASE							:   Any						    =     batch_size
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     seq_length
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     is_training
       SCREAMING_SNAKE_CASE							:   Tuple						    =     use_input_mask
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     use_token_type_ids
       SCREAMING_SNAKE_CASE							:   Any						    =     use_labels
       SCREAMING_SNAKE_CASE							:   Tuple						    =     vocab_size
       SCREAMING_SNAKE_CASE							:   str						    =     hidden_size
       SCREAMING_SNAKE_CASE							:   Dict						    =     embedding_size
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     num_hidden_layers
       SCREAMING_SNAKE_CASE							:   Tuple						    =     num_attention_heads
       SCREAMING_SNAKE_CASE							:   Tuple						    =     intermediate_size
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     hidden_act
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     hidden_dropout_prob
       SCREAMING_SNAKE_CASE							:   str						    =     attention_probs_dropout_prob
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     max_position_embeddings
       SCREAMING_SNAKE_CASE							:   str						    =     type_vocab_size
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     type_sequence_label_size
       SCREAMING_SNAKE_CASE							:   Dict						    =     initializer_range
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     num_labels
       SCREAMING_SNAKE_CASE							:   List[str]						    =     num_choices
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     scope


     def 					_A							(					self							:						List[Any]   ):
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     ids_tensor([self.batch_size, self.seq_length]       ,			self.vocab_size   )

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

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

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

       SCREAMING_SNAKE_CASE							:   int						    =     self.get_config()

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


     def 					_A							(					self							:						Union[str, Any]   ):
       return MegatronBertConfig(
           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       ,			embedding_size=self.embedding_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       ,			)


     def 					_A							(					self							:						str       ,			UpperCAmelCase_							:						Dict       ,			UpperCAmelCase_							:						Any       ,			UpperCAmelCase_							:						Tuple       ,			UpperCAmelCase_							:						Tuple       ,			UpperCAmelCase_							:						Optional[Any]       ,			UpperCAmelCase_							:						Tuple       ,			UpperCAmelCase_							:						int   ):
       SCREAMING_SNAKE_CASE							:   Any						    =     MegatronBertModel(config=UpperCAmelCase_   )
       model.to(UpperCAmelCase_   )
       model.eval()
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     model(UpperCAmelCase_       ,			attention_mask=UpperCAmelCase_       ,			token_type_ids=UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   str						    =     model(UpperCAmelCase_       ,			token_type_ids=UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   Tuple						    =     model(UpperCAmelCase_   )

       self.parent.assertEqual(result.last_hidden_state.shape       ,			(self.batch_size, self.seq_length, self.hidden_size)   )
       self.parent.assertEqual(result.pooler_output.shape       ,			(self.batch_size, self.hidden_size)   )


     def 					_A							(					self							:						Optional[int]       ,			UpperCAmelCase_							:						List[Any]       ,			UpperCAmelCase_							:						Optional[int]       ,			UpperCAmelCase_							:						List[str]       ,			UpperCAmelCase_							:						Any       ,			UpperCAmelCase_							:						List[str]       ,			UpperCAmelCase_							:						int       ,			UpperCAmelCase_							:						str   ):
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     MegatronBertForMaskedLM(config=UpperCAmelCase_   )
       model.to(UpperCAmelCase_   )
       model.eval()
       SCREAMING_SNAKE_CASE							:   Tuple						    =     model(UpperCAmelCase_       ,			attention_mask=UpperCAmelCase_       ,			token_type_ids=UpperCAmelCase_       ,			labels=UpperCAmelCase_   )
       self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.seq_length, self.vocab_size)   )


     def 					_A							(					self							:						Optional[int]       ,			UpperCAmelCase_							:						List[Any]       ,			UpperCAmelCase_							:						Dict       ,			UpperCAmelCase_							:						Optional[Any]       ,			UpperCAmelCase_							:						List[str]       ,			UpperCAmelCase_							:						int       ,			UpperCAmelCase_							:						str       ,			UpperCAmelCase_							:						Dict   ):
       SCREAMING_SNAKE_CASE							:   Tuple						    =     MegatronBertForCausalLM(config=UpperCAmelCase_   )
       model.to(UpperCAmelCase_   )
       model.eval()
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     model(UpperCAmelCase_       ,			attention_mask=UpperCAmelCase_       ,			token_type_ids=UpperCAmelCase_       ,			labels=UpperCAmelCase_   )
       self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.seq_length, self.vocab_size)   )


     def 					_A							(					self							:						Optional[Any]       ,			UpperCAmelCase_							:						str       ,			UpperCAmelCase_							:						List[str]       ,			UpperCAmelCase_							:						str       ,			UpperCAmelCase_							:						Optional[Any]       ,			UpperCAmelCase_							:						Optional[Any]       ,			UpperCAmelCase_							:						int       ,			UpperCAmelCase_							:						Tuple   ):
       SCREAMING_SNAKE_CASE							:   Dict						    =     MegatronBertForNextSentencePrediction(config=UpperCAmelCase_   )
       model.to(UpperCAmelCase_   )
       model.eval()
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     model(
           UpperCAmelCase_       ,			attention_mask=UpperCAmelCase_       ,			token_type_ids=UpperCAmelCase_       ,			labels=UpperCAmelCase_       ,			)
       self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, 2)   )


     def 					_A							(					self							:						List[str]       ,			UpperCAmelCase_							:						Tuple       ,			UpperCAmelCase_							:						str       ,			UpperCAmelCase_							:						int       ,			UpperCAmelCase_							:						Dict       ,			UpperCAmelCase_							:						List[Any]       ,			UpperCAmelCase_							:						int       ,			UpperCAmelCase_							:						List[str]   ):
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     MegatronBertForPreTraining(config=UpperCAmelCase_   )
       model.to(UpperCAmelCase_   )
       model.eval()
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     model(
           UpperCAmelCase_       ,			attention_mask=UpperCAmelCase_       ,			token_type_ids=UpperCAmelCase_       ,			labels=UpperCAmelCase_       ,			next_sentence_label=UpperCAmelCase_       ,			)
       self.parent.assertEqual(result.prediction_logits.shape       ,			(self.batch_size, self.seq_length, self.vocab_size)   )
       self.parent.assertEqual(result.seq_relationship_logits.shape       ,			(self.batch_size, 2)   )


     def 					_A							(					self							:						Tuple       ,			UpperCAmelCase_							:						List[str]       ,			UpperCAmelCase_							:						Any       ,			UpperCAmelCase_							:						int       ,			UpperCAmelCase_							:						Any       ,			UpperCAmelCase_							:						str       ,			UpperCAmelCase_							:						str       ,			UpperCAmelCase_							:						List[Any]   ):
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     MegatronBertForQuestionAnswering(config=UpperCAmelCase_   )
       model.to(UpperCAmelCase_   )
       model.eval()
       SCREAMING_SNAKE_CASE							:   Dict						    =     model(
           UpperCAmelCase_       ,			attention_mask=UpperCAmelCase_       ,			token_type_ids=UpperCAmelCase_       ,			start_positions=UpperCAmelCase_       ,			end_positions=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)   )


     def 					_A							(					self							:						Any       ,			UpperCAmelCase_							:						Tuple       ,			UpperCAmelCase_							:						Optional[Any]       ,			UpperCAmelCase_							:						Any       ,			UpperCAmelCase_							:						List[str]       ,			UpperCAmelCase_							:						Optional[int]       ,			UpperCAmelCase_							:						int       ,			UpperCAmelCase_							:						Tuple   ):
       SCREAMING_SNAKE_CASE							:   str						    =     self.num_labels
       SCREAMING_SNAKE_CASE							:   int						    =     MegatronBertForSequenceClassification(UpperCAmelCase_   )
       model.to(UpperCAmelCase_   )
       model.eval()
       SCREAMING_SNAKE_CASE							:   int						    =     model(UpperCAmelCase_       ,			attention_mask=UpperCAmelCase_       ,			token_type_ids=UpperCAmelCase_       ,			labels=UpperCAmelCase_   )
       self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.num_labels)   )


     def 					_A							(					self							:						List[Any]       ,			UpperCAmelCase_							:						Tuple       ,			UpperCAmelCase_							:						str       ,			UpperCAmelCase_							:						Tuple       ,			UpperCAmelCase_							:						Any       ,			UpperCAmelCase_							:						str       ,			UpperCAmelCase_							:						str       ,			UpperCAmelCase_							:						List[Any]   ):
       SCREAMING_SNAKE_CASE							:   int						    =     self.num_labels
       SCREAMING_SNAKE_CASE							:   Any						    =     MegatronBertForTokenClassification(config=UpperCAmelCase_   )
       model.to(UpperCAmelCase_   )
       model.eval()
       SCREAMING_SNAKE_CASE							:   List[str]						    =     model(UpperCAmelCase_       ,			attention_mask=UpperCAmelCase_       ,			token_type_ids=UpperCAmelCase_       ,			labels=UpperCAmelCase_   )
       self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.seq_length, self.num_labels)   )


     def 					_A							(					self							:						Optional[Any]       ,			UpperCAmelCase_							:						int       ,			UpperCAmelCase_							:						List[str]       ,			UpperCAmelCase_							:						List[Any]       ,			UpperCAmelCase_							:						str       ,			UpperCAmelCase_							:						List[str]       ,			UpperCAmelCase_							:						List[Any]       ,			UpperCAmelCase_							:						str   ):
       SCREAMING_SNAKE_CASE							:   int						    =     self.num_choices
       SCREAMING_SNAKE_CASE							:   Any						    =     MegatronBertForMultipleChoice(config=UpperCAmelCase_   )
       model.to(UpperCAmelCase_   )
       model.eval()
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     input_ids.unsqueeze(1   ).expand(-1       ,			self.num_choices       ,			-1   ).contiguous()
       SCREAMING_SNAKE_CASE							:   List[str]						    =     token_type_ids.unsqueeze(1   ).expand(-1       ,			self.num_choices       ,			-1   ).contiguous()
       SCREAMING_SNAKE_CASE							:   Any						    =     input_mask.unsqueeze(1   ).expand(-1       ,			self.num_choices       ,			-1   ).contiguous()
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     model(
           UpperCAmelCase_       ,			attention_mask=UpperCAmelCase_       ,			token_type_ids=UpperCAmelCase_       ,			labels=UpperCAmelCase_       ,			)
       self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.num_choices)   )






     def 					_A							(					self							:						Optional[int]   ):
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     self.prepare_config_and_inputs()
       (
           (
           SCREAMING_SNAKE_CASE
       )    ,	(
           SCREAMING_SNAKE_CASE
       )    ,	(
           SCREAMING_SNAKE_CASE
       )    ,	(
           SCREAMING_SNAKE_CASE
       )    ,	(
           SCREAMING_SNAKE_CASE
       )    ,	(
           SCREAMING_SNAKE_CASE
       )    ,	(
           SCREAMING_SNAKE_CASE
       )    ,	
       )							:   Optional[int]						    =     config_and_inputs
       SCREAMING_SNAKE_CASE							:   str						    =     {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
       return config, inputs_dict





@require_torch
class 			SCREAMING_SNAKE_CASE      (      lowerCAmelCase   ,		lowerCAmelCase   ,		unittest.TestCase       ):
     '''simple docstring'''


     UpperCamelCase_				:		List[str]								=						(
         (
             MegatronBertModel,
             MegatronBertForMaskedLM,
             MegatronBertForCausalLM,
             MegatronBertForMultipleChoice,
             MegatronBertForNextSentencePrediction,
             MegatronBertForPreTraining,
             MegatronBertForQuestionAnswering,
             MegatronBertForSequenceClassification,
             MegatronBertForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     UpperCamelCase_				:		int								=						(
         {
             '''feature-extraction''': MegatronBertModel,
             '''fill-mask''': MegatronBertForMaskedLM,
             '''question-answering''': MegatronBertForQuestionAnswering,
             '''text-classification''': MegatronBertForSequenceClassification,
             '''text-generation''': MegatronBertForCausalLM,
             '''token-classification''': MegatronBertForTokenClassification,
             '''zero-shot''': MegatronBertForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     UpperCamelCase_				:		List[Any]								=						True
     # test_resize_embeddings = False
     UpperCamelCase_				:		str								=						False


     def 					_A							(					self							:						List[Any]       ,			UpperCAmelCase_							:						str       ,			UpperCAmelCase_							:						Any       ,			UpperCAmelCase_							:						str=False   ):
       SCREAMING_SNAKE_CASE							:   Any						    =     super()._prepare_for_class(UpperCAmelCase_       ,			UpperCAmelCase_       ,			return_labels=UpperCAmelCase_   )

       if return_labels:
         if model_class in get_values(UpperCAmelCase_   ):
           SCREAMING_SNAKE_CASE							:   List[Any]						    =     torch.zeros(
               (self.model_tester.batch_size, self.model_tester.seq_length)       ,			dtype=torch.long       ,			device=UpperCAmelCase_   )
           SCREAMING_SNAKE_CASE							:   Tuple						    =     torch.zeros(
               self.model_tester.batch_size       ,			dtype=torch.long       ,			device=UpperCAmelCase_   )
       return inputs_dict


     def 					_A							(					self							:						Optional[Any]   ):
       SCREAMING_SNAKE_CASE							:   int						    =     MegatronBertModelTester(self   )
       SCREAMING_SNAKE_CASE							:   str						    =     ConfigTester(self       ,			config_class=UpperCAmelCase_       ,			hidden_size=37   )


     def 					_A							(					self							:						Any   ):
       self.config_tester.run_common_tests()


     def 					_A							(					self							:						List[Any]   ):
       SCREAMING_SNAKE_CASE							:   str						    =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_model(*UpperCAmelCase_   )


     def 					_A							(					self							:						List[str]   ):
       SCREAMING_SNAKE_CASE							:   str						    =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_masked_lm(*UpperCAmelCase_   )


     def 					_A							(					self							:						Optional[int]   ):
       SCREAMING_SNAKE_CASE							:   Tuple						    =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*UpperCAmelCase_   )


     def 					_A							(					self							:						List[Any]   ):
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*UpperCAmelCase_   )


     def 					_A							(					self							:						Any   ):
       SCREAMING_SNAKE_CASE							:   int						    =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_pretraining(*UpperCAmelCase_   )


     def 					_A							(					self							:						Dict   ):
       SCREAMING_SNAKE_CASE							:   Any						    =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_question_answering(*UpperCAmelCase_   )


     def 					_A							(					self							:						int   ):
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*UpperCAmelCase_   )






     def 					_A							(					self							:						Optional[Any]   ):
       SCREAMING_SNAKE_CASE							:   str						    =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_token_classification(*UpperCAmelCase_   )





def 	lowerCamelCase__      (  lowercase					):

  """simple docstring"""






  return torch.tensor(
      lowercase    ,  dtype=torch.long    ,  device=lowercase    ,  )


snake_case     					=		1e-4


@require_torch
@require_sentencepiece
@require_tokenizers
class 			SCREAMING_SNAKE_CASE      (      unittest.TestCase       ):
     '''simple docstring'''


     @slow
     @unittest.skip("Model is not available."   )
     def 					_A							(					self							:						Dict   ):
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     "nvidia/megatron-bert-uncased-345m"
       if "MYDIR" in os.environ:
         SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     os.path.join(os.environ["MYDIR"]       ,			UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     MegatronBertModel.from_pretrained(UpperCAmelCase_   )
       model.to(UpperCAmelCase_   )
       model.half()
       SCREAMING_SNAKE_CASE							:   Dict						    =     _long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]   )
       with torch.no_grad():
         SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     model(UpperCAmelCase_   )[0]
       SCREAMING_SNAKE_CASE							:   str						    =     torch.Size((1, 9, 1024)   )
       self.assertEqual(output.shape       ,			UpperCAmelCase_   )

       SCREAMING_SNAKE_CASE							:   Tuple						    =     [-0.6_040, -0.2_517, -0.1_025, 0.3_420, -0.6_758, -0.0_017, -0.1_089, -0.1_990, 0.5_728]
       for ii in range(3   ):
         for jj in range(3   ):
           SCREAMING_SNAKE_CASE							:   Optional[int]						    =     output[0, ii, jj]
           SCREAMING_SNAKE_CASE							:   Tuple						    =     expected[3 * ii + jj]
           SCREAMING_SNAKE_CASE							:   str						    =     "ii={} jj={} a={} b={}".format(UpperCAmelCase_       ,			UpperCAmelCase_       ,			UpperCAmelCase_       ,			UpperCAmelCase_   )
           self.assertTrue(math.isclose(UpperCAmelCase_       ,			UpperCAmelCase_       ,			rel_tol=UpperCAmelCase_       ,			abs_tol=UpperCAmelCase_   )       ,			msg=UpperCAmelCase_   )









import json
import os
import pickle
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
					import faiss







@require_faiss
class 					lowercase							(			lowercase_     ):




				def        a				(   self			):
											snake_case_     							=       tempfile.mkdtemp()
											snake_case_     							=       8

											# DPR tok
											snake_case_     							=       [
											    '[UNK]',
											    '[CLS]',
											    '[SEP]',
											    '[PAD]',
											    '[MASK]',
											    'want',
											    '##want',
											    '##ed',
											    'wa',
											    'un',
											    'runn',
											    '##ing',
											    ',',
											    'low',
											    'lowest',
											]
											snake_case_     							=       os.path.join(self.tmpdirname ,  'dpr_tokenizer'			)
											os.makedirs(snake_case ,  exist_ok=snake_case			)
											snake_case_     							=       os.path.join(snake_case ,  DPR_VOCAB_FILES_NAMES['vocab_file']			)
											with open(self.vocab_file ,  'w' ,  encoding='utf-8'			) as vocab_writer:
																		vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]			)			)

											# BART tok
											snake_case_     							=       [
											    'l',
											    'o',
											    'w',
											    'e',
											    'r',
											    's',
											    't',
											    'i',
											    'd',
											    'n',
											    '\u0120',
											    '\u0120l',
											    '\u0120n',
											    '\u0120lo',
											    '\u0120low',
											    'er',
											    '\u0120lowest',
											    '\u0120newer',
											    '\u0120wider',
											    '<unk>',
											]
											snake_case_     							=       dict(zip(snake_case ,  range(len(snake_case			)			)			)			)
											snake_case_     							=       ['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
											snake_case_     							=       {'unk_token': '<unk>'}

											snake_case_     							=       os.path.join(self.tmpdirname ,  'bart_tokenizer'			)
											os.makedirs(snake_case ,  exist_ok=snake_case			)
											snake_case_     							=       os.path.join(snake_case ,  BART_VOCAB_FILES_NAMES['vocab_file']			)
											snake_case_     							=       os.path.join(snake_case ,  BART_VOCAB_FILES_NAMES['merges_file']			)
											with open(self.vocab_file ,  'w' ,  encoding='utf-8'			) as fp:
																		fp.write(json.dumps(snake_case			) + '\n'			)
											with open(self.merges_file ,  'w' ,  encoding='utf-8'			) as fp:
																		fp.write('\n'.join(snake_case			)			)




				def        a				(   self			):
											return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname ,  'dpr_tokenizer'			)			)




				def        a				(   self			):
											return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname ,  'dpr_tokenizer'			)			)




				def        a				(   self			):
											return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname ,  'bart_tokenizer'			)			)




				def        a				(   self			):
											shutil.rmtree(self.tmpdirname			)




				def        a				(   self			):
											snake_case_     							=       Dataset.from_dict(
											    {
											        'id': ['0', '1'],
											        'text': ['foo', 'bar'],
											        'title': ['Foo', 'Bar'],
											        'embeddings': [np.ones(self.retrieval_vector_size			), 2 * np.ones(self.retrieval_vector_size			)],
											    }			)
											dataset.add_faiss_index('embeddings' ,  string_factory='Flat' ,  metric_type=faiss.METRIC_INNER_PRODUCT			)
											return dataset




				def        a				(   self			):
											snake_case_     							=       self.get_dummy_dataset()
											snake_case_     							=       RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size ,  question_encoder=DPRConfig().to_dict() ,  generator=BartConfig().to_dict() ,  )
											with patch('transformers.models.rag.retrieval_rag.load_dataset'			) as mock_load_dataset:
																		snake_case_     							=       dataset
																		snake_case_     							=       RagRetriever(
																		    snake_case ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer() ,  )
											return retriever




				def        a				(   self ,  snake_case			):
											snake_case_     							=       self.get_dummy_dataset()
											snake_case_     							=       RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size ,  question_encoder=DPRConfig().to_dict() ,  generator=BartConfig().to_dict() ,  index_name='custom' ,  )
											if from_disk:
																		snake_case_     							=       os.path.join(self.tmpdirname ,  'dataset'			)
																		snake_case_     							=       os.path.join(self.tmpdirname ,  'index.faiss'			)
																		dataset.get_index('embeddings'			).save(os.path.join(self.tmpdirname ,  'index.faiss'			)			)
																		dataset.drop_index('embeddings'			)
																		dataset.save_to_disk(os.path.join(self.tmpdirname ,  'dataset'			)			)
																		del dataset
																		snake_case_     							=       RagRetriever(
																		    snake_case ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer() ,  )
											else:
																		snake_case_     							=       RagRetriever(
																		    snake_case ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer() ,  index=CustomHFIndex(config.retrieval_vector_size ,  snake_case			) ,  )
											return retriever




				def        a				(   self			):
											snake_case_     							=       Dataset.from_dict(
											    {
											        'id': ['0', '1'],
											        'text': ['foo', 'bar'],
											        'title': ['Foo', 'Bar'],
											        'embeddings': [np.ones(self.retrieval_vector_size + 1			), 2 * np.ones(self.retrieval_vector_size + 1			)],
											    }			)
											dataset.add_faiss_index('embeddings' ,  string_factory='Flat' ,  metric_type=faiss.METRIC_INNER_PRODUCT			)

											snake_case_     							=       os.path.join(self.tmpdirname ,  'hf_bert_base.hnswSQ8_correct_phi_128.c_index'			)
											dataset.save_faiss_index('embeddings' ,  index_file_name + '.index.dpr'			)
											pickle.dump(dataset['id'] ,  open(index_file_name + '.index_meta.dpr' ,  'wb'			)			)

											snake_case_     							=       os.path.join(self.tmpdirname ,  'psgs_w100.tsv.pkl'			)
											snake_case_     							=       {sample['id']: [sample['text'], sample['title']] for sample in dataset}
											pickle.dump(snake_case ,  open(snake_case ,  'wb'			)			)

											snake_case_     							=       RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size ,  question_encoder=DPRConfig().to_dict() ,  generator=BartConfig().to_dict() ,  index_name='legacy' ,  index_path=self.tmpdirname ,  )
											snake_case_     							=       RagRetriever(
											    snake_case ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer()			)
											return retriever




				def        a				(   self			):
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_canonical_hf_index_retriever()
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_							,       snake_case_							,       snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=snake_case			)
											self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)			)
											self.assertEqual(len(snake_case			) ,  2			)
											self.assertEqual(sorted(doc_dicts[0]			) ,  ['embeddings', 'id', 'text', 'title']			)
											self.assertEqual(len(doc_dicts[0]['id']			) ,  snake_case			)
											self.assertEqual(doc_dicts[0]['id'][0] ,  '1'			)  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]['id'][0] ,  '0'			)  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist() ,  [[1], [0]]			)




				def        a				(   self			):
											snake_case_     							=       self.get_dummy_canonical_hf_index_retriever()
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		with patch('transformers.models.rag.retrieval_rag.load_dataset'			) as mock_load_dataset:
																									snake_case_     							=       self.get_dummy_dataset()
																									retriever.save_pretrained(snake_case			)
																									snake_case_     							=       RagRetriever.from_pretrained(snake_case			)
																									self.assertIsInstance(snake_case ,  snake_case			)
																									snake_case_     							=       np.array(
																									    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
																									snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=1			)
																									self.assertTrue(out is not None			)




				def        a				(   self			):
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_							,       snake_case_							,       snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=snake_case			)
											self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)			)
											self.assertEqual(len(snake_case			) ,  2			)
											self.assertEqual(sorted(doc_dicts[0]			) ,  ['embeddings', 'id', 'text', 'title']			)
											self.assertEqual(len(doc_dicts[0]['id']			) ,  snake_case			)
											self.assertEqual(doc_dicts[0]['id'][0] ,  '1'			)  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]['id'][0] ,  '0'			)  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist() ,  [[1], [0]]			)




				def        a				(   self			):
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		retriever.save_pretrained(snake_case			)
																		snake_case_     							=       RagRetriever.from_pretrained(snake_case			)
																		self.assertIsInstance(snake_case ,  snake_case			)
																		snake_case_     							=       np.array(
																		    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
																		snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=1			)
																		self.assertTrue(out is not None			)




				def        a				(   self			):
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_							,       snake_case_							,       snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=snake_case			)
											self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)			)
											self.assertEqual(len(snake_case			) ,  2			)
											self.assertEqual(sorted(doc_dicts[0]			) ,  ['embeddings', 'id', 'text', 'title']			)
											self.assertEqual(len(doc_dicts[0]['id']			) ,  snake_case			)
											self.assertEqual(doc_dicts[0]['id'][0] ,  '1'			)  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]['id'][0] ,  '0'			)  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist() ,  [[1], [0]]			)




				def        a				(   self			):
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		retriever.save_pretrained(snake_case			)
																		snake_case_     							=       RagRetriever.from_pretrained(snake_case			)
																		self.assertIsInstance(snake_case ,  snake_case			)
																		snake_case_     							=       np.array(
																		    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
																		snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=1			)
																		self.assertTrue(out is not None			)




				def        a				(   self			):
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_legacy_index_retriever()
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_							,       snake_case_							,       snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=snake_case			)
											self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)			)
											self.assertEqual(len(snake_case			) ,  2			)
											self.assertEqual(sorted(doc_dicts[0]			) ,  ['text', 'title']			)
											self.assertEqual(len(doc_dicts[0]['text']			) ,  snake_case			)
											self.assertEqual(doc_dicts[0]['text'][0] ,  'bar'			)  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]['text'][0] ,  'foo'			)  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist() ,  [[1], [0]]			)




				def        a				(   self			):
											snake_case_     							=       self.get_dummy_legacy_index_retriever()
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		retriever.save_pretrained(snake_case			)
																		snake_case_     							=       RagRetriever.from_pretrained(snake_case			)
																		self.assertIsInstance(snake_case ,  snake_case			)
																		snake_case_     							=       np.array(
																		    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
																		snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=1			)
																		self.assertTrue(out is not None			)




				@require_torch
				@require_tokenizers
				@require_sentencepiece
				def        a				(   self			):
											import torch

											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_canonical_hf_index_retriever()
											snake_case_     							=       [[5, 7], [10, 11]]
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_     							=       retriever(snake_case ,  snake_case ,  prefix=retriever.config.generator.prefix ,  n_docs=snake_case			)
											snake_case_							,       snake_case_							,       snake_case_     							=       (
											    out['context_input_ids'],
											    out['context_attention_mask'],
											    out['retrieved_doc_embeds'],
											)
											self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)			)
											self.assertIsInstance(snake_case ,  snake_case			)
											self.assertIsInstance(snake_case ,  snake_case			)
											self.assertIsInstance(snake_case ,  np.ndarray			)

											snake_case_     							=       retriever(
											    snake_case ,  snake_case ,  prefix=retriever.config.generator.prefix ,  n_docs=snake_case ,  return_tensors='pt' ,  )
											snake_case_							,       snake_case_							,       snake_case_							,       snake_case_     							=       (  # noqa: F841
											    out['context_input_ids'],
											    out['context_attention_mask'],
											    out['retrieved_doc_embeds'],
											    out['doc_ids'],
											)
											self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)			)
											self.assertIsInstance(snake_case ,  torch.Tensor			)
											self.assertIsInstance(snake_case ,  torch.Tensor			)
											self.assertIsInstance(snake_case ,  torch.Tensor			)




				@require_torch
				@require_tokenizers
				@require_sentencepiece
				def        a				(   self			):
											snake_case_     							=       self.get_dpr_ctx_encoder_tokenizer()
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											retriever.set_ctx_encoder_tokenizer(snake_case			)

											snake_case_     							=       [[5, 7], [10, 11]]
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_     							=       retriever(snake_case ,  snake_case ,  prefix=retriever.config.generator.prefix ,  n_docs=snake_case			)

											self.assertEqual(
											    len(snake_case			) ,  6			)  # check whether the retriever output consist of 6 attributes including tokenized docs
											self.assertEqual(
											    all(k in out for k in ('tokenized_doc_ids', 'tokenized_doc_attention_mask')			) ,  snake_case			)  # check for doc token related keys in dictionary.




