


import os
import tempfile
import unittest

from transformers import DistilBertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, 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 (
				    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    DistilBertForMaskedLM,
				    DistilBertForMultipleChoice,
				    DistilBertForQuestionAnswering,
				    DistilBertForSequenceClassification,
				    DistilBertForTokenClassification,
				    DistilBertModel,
				)




class        lowerCamelCase      (     _UpperCamelCase   ):


	def __init__(						self							,    lowercase__							,    lowercase__=1_3							,    lowercase__=7							,    lowercase__=True							,    lowercase__=True							,    lowercase__=False							,    lowercase__=True							,    lowercase__=9_9							,    lowercase__=3_2							,    lowercase__=5							,    lowercase__=4							,    lowercase__=3_7							,    lowercase__="gelu"							,    lowercase__=0.1							,    lowercase__=0.1							,    lowercase__=5_1_2							,    lowercase__=1_6							,    lowercase__=2							,    lowercase__=0.0_2							,    lowercase__=3							,    lowercase__=4							,    lowercase__=None							,    ):
							__UpperCAmelCase :       int  						=  parent
							__UpperCAmelCase :       Union[str, Any]  						=  batch_size
							__UpperCAmelCase :       List[Any]  						=  seq_length
							__UpperCAmelCase :       List[Any]  						=  is_training
							__UpperCAmelCase :       Any  						=  use_input_mask
							__UpperCAmelCase :       List[Any]  						=  use_token_type_ids
							__UpperCAmelCase :       Union[str, Any]  						=  use_labels
							__UpperCAmelCase :       str  						=  vocab_size
							__UpperCAmelCase :       str  						=  hidden_size
							__UpperCAmelCase :       Any  						=  num_hidden_layers
							__UpperCAmelCase :       List[str]  						=  num_attention_heads
							__UpperCAmelCase :       List[str]  						=  intermediate_size
							__UpperCAmelCase :       Optional[int]  						=  hidden_act
							__UpperCAmelCase :       Optional[Any]  						=  hidden_dropout_prob
							__UpperCAmelCase :       Optional[int]  						=  attention_probs_dropout_prob
							__UpperCAmelCase :       Dict  						=  max_position_embeddings
							__UpperCAmelCase :       Optional[Any]  						=  type_vocab_size
							__UpperCAmelCase :       Any  						=  type_sequence_label_size
							__UpperCAmelCase :       List[Any]  						=  initializer_range
							__UpperCAmelCase :       Any  						=  num_labels
							__UpperCAmelCase :       List[Any]  						=  num_choices
							__UpperCAmelCase :       Tuple  						=  scope


	def    A(						self):
							__UpperCAmelCase :       List[Any]  						=  ids_tensor([self.batch_size, self.seq_length]							,    self.vocab_size)

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

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

							__UpperCAmelCase :       str  						=  self.get_config()

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


	def    A(						self):
							return DistilBertConfig(
							    vocab_size=self.vocab_size							,    dim=self.hidden_size							,    n_layers=self.num_hidden_layers							,    n_heads=self.num_attention_heads							,    hidden_dim=self.intermediate_size							,    hidden_act=self.hidden_act							,    dropout=self.hidden_dropout_prob							,    attention_dropout=self.attention_probs_dropout_prob							,    max_position_embeddings=self.max_position_embeddings							,    initializer_range=self.initializer_range							,    )


	def    A(						self							,    lowercase__							,    lowercase__							,    lowercase__							,    lowercase__							,    lowercase__							,    lowercase__):
							__UpperCAmelCase :       Tuple  						=  DistilBertModel(config=lowercase__)
							model.to(lowercase__)
							model.eval()
							__UpperCAmelCase :       List[str]  						=  model(lowercase__							,    lowercase__)
							__UpperCAmelCase :       Optional[int]  						=  model(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__):
							__UpperCAmelCase :       Dict  						=  DistilBertForMaskedLM(config=lowercase__)
							model.to(lowercase__)
							model.eval()
							__UpperCAmelCase :       Optional[Any]  						=  model(lowercase__							,    attention_mask=lowercase__							,    labels=lowercase__)
							self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.seq_length, self.vocab_size))


	def    A(						self							,    lowercase__							,    lowercase__							,    lowercase__							,    lowercase__							,    lowercase__							,    lowercase__):
							__UpperCAmelCase :       Optional[int]  						=  DistilBertForQuestionAnswering(config=lowercase__)
							model.to(lowercase__)
							model.eval()
							__UpperCAmelCase :       Optional[int]  						=  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__):
							__UpperCAmelCase :       Dict  						=  self.num_labels
							__UpperCAmelCase :       int  						=  DistilBertForSequenceClassification(lowercase__)
							model.to(lowercase__)
							model.eval()
							__UpperCAmelCase :       str  						=  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__):
							__UpperCAmelCase :       Optional[Any]  						=  self.num_labels
							__UpperCAmelCase :       Union[str, Any]  						=  DistilBertForTokenClassification(config=lowercase__)
							model.to(lowercase__)
							model.eval()

							__UpperCAmelCase :       List[Any]  						=  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__):
							__UpperCAmelCase :       Union[str, Any]  						=  self.num_choices
							__UpperCAmelCase :       Optional[Any]  						=  DistilBertForMultipleChoice(config=lowercase__)
							model.to(lowercase__)
							model.eval()
							__UpperCAmelCase :       Tuple  						=  input_ids.unsqueeze(1).expand(-1							,    self.num_choices							,    -1).contiguous()
							__UpperCAmelCase :       Optional[int]  						=  input_mask.unsqueeze(1).expand(-1							,    self.num_choices							,    -1).contiguous()
							__UpperCAmelCase :       List[str]  						=  model(
							    lowercase__							,    attention_mask=lowercase__							,    labels=lowercase__							,    )
							self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.num_choices))


	def    A(						self):
							__UpperCAmelCase :       int  						=  self.prepare_config_and_inputs()
							((__UpperCAmelCase)						,			(__UpperCAmelCase)						,			(__UpperCAmelCase)						,			(__UpperCAmelCase)						,			(__UpperCAmelCase)						,			(__UpperCAmelCase)) :       Dict  						=  config_and_inputs
							__UpperCAmelCase :       str  						=  {'''input_ids''': input_ids, '''attention_mask''': input_mask}
							return config, inputs_dict






@require_torch
class        lowerCamelCase      (     _UpperCamelCase       ,	_UpperCamelCase       ,	unittest.TestCase   ):
	_lowerCAmelCase    :		Optional[Any]        =	(
	    (
	        DistilBertModel,
	        DistilBertForMaskedLM,
	        DistilBertForMultipleChoice,
	        DistilBertForQuestionAnswering,
	        DistilBertForSequenceClassification,
	        DistilBertForTokenClassification,
	    )
	    if is_torch_available()
	    else None
	)
	_lowerCAmelCase    :		Dict        =	(
	    {
	        '''feature-extraction''': DistilBertModel,
	        '''fill-mask''': DistilBertForMaskedLM,
	        '''question-answering''': DistilBertForQuestionAnswering,
	        '''text-classification''': DistilBertForSequenceClassification,
	        '''token-classification''': DistilBertForTokenClassification,
	        '''zero-shot''': DistilBertForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	_lowerCAmelCase    :		Optional[int]        =	True
	_lowerCAmelCase    :		str        =	True
	_lowerCAmelCase    :		Any        =	True
	_lowerCAmelCase    :		int        =	True


	def    A(						self):
							__UpperCAmelCase :       int  						=  DistilBertModelTester(self)
							__UpperCAmelCase :       str  						=  ConfigTester(self							,    config_class=lowercase__							,    dim=3_7)


	def    A(						self):
							self.config_tester.run_common_tests()


	def    A(						self):
							__UpperCAmelCase :       Dict  						=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_model(*lowercase__)


	def    A(						self):
							__UpperCAmelCase :       str  						=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_masked_lm(*lowercase__)


	def    A(						self):
							__UpperCAmelCase :       int  						=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_question_answering(*lowercase__)


	def    A(						self):
							__UpperCAmelCase :       Any  						=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowercase__)


	def    A(						self):
							__UpperCAmelCase :       List[Any]  						=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_token_classification(*lowercase__)


	def    A(						self):
							__UpperCAmelCase :       Any  						=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowercase__)


	@slow
	def    A(						self):
							for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													__UpperCAmelCase :       List[str]  						=  DistilBertModel.from_pretrained(lowercase__)
													self.assertIsNotNone(lowercase__)


	@slow
	@require_torch_gpu
	def    A(						self):
							__UpperCAmelCase						,			__UpperCAmelCase :       Tuple  						=  self.model_tester.prepare_config_and_inputs_for_common()
							for model_class in self.all_model_classes:
													# BertForMultipleChoice behaves incorrectly in JIT environments.
													if model_class == DistilBertForMultipleChoice:
																			return

													__UpperCAmelCase :       str  						=  True
													__UpperCAmelCase :       str  						=  model_class(config=lowercase__)

													__UpperCAmelCase :       str  						=  self._prepare_for_class(lowercase__							,    lowercase__)
													__UpperCAmelCase :       Any  						=  torch.jit.trace(
													    lowercase__							,    (inputs_dict['''input_ids'''].to('''cpu'''), inputs_dict['''attention_mask'''].to('''cpu''')))

													with tempfile.TemporaryDirectory() as tmp:
																			torch.jit.save(lowercase__							,    os.path.join(lowercase__							,    '''traced_model.pt'''))
																			__UpperCAmelCase :       Optional[int]  						=  torch.jit.load(os.path.join(lowercase__							,    '''traced_model.pt''')							,    map_location=lowercase__)
																			loaded(inputs_dict['''input_ids'''].to(lowercase__)							,    inputs_dict['''attention_mask'''].to(lowercase__))






@require_torch
class        lowerCamelCase      (     unittest.TestCase   ):


	@slow
	def    A(						self):
							__UpperCAmelCase :       List[str]  						=  DistilBertModel.from_pretrained('''distilbert-base-uncased''')
							__UpperCAmelCase :       Tuple  						=  torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]])
							__UpperCAmelCase :       Dict  						=  torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
							with torch.no_grad():
													__UpperCAmelCase :       Dict  						=  model(lowercase__							,    attention_mask=lowercase__)[0]
							__UpperCAmelCase :       int  						=  torch.Size((1, 1_1, 7_6_8))
							self.assertEqual(output.shape							,    lowercase__)
							__UpperCAmelCase :       List[str]  						=  torch.tensor(
							    [[[-0.1_6_3_9, 0.3_2_9_9, 0.1_6_4_8], [-0.1_7_4_6, 0.3_2_8_9, 0.1_7_1_0], [-0.1_8_8_4, 0.3_3_5_7, 0.1_8_1_0]]])

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





from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase							=						logging.get_logger(__name__)

lowerCAmelCase							=						{
    """google/realm-cc-news-pretrained-embedder""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/config.json"""
    ),
    """google/realm-cc-news-pretrained-encoder""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/config.json"""
    ),
    """google/realm-cc-news-pretrained-scorer""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/config.json"""
    ),
    """google/realm-cc-news-pretrained-openqa""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/config.json"""
    ),
    """google/realm-orqa-nq-openqa""": """https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/config.json""",
    """google/realm-orqa-nq-reader""": """https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/config.json""",
    """google/realm-orqa-wq-openqa""": """https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/config.json""",
    """google/realm-orqa-wq-reader""": """https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/config.json""",
    # See all REALM models at https://huggingface.co/models?filter=realm
}




class        lowerCamelCase      (     _UpperCamelCase   ):
	_lowerCAmelCase    :		Tuple        =	'''realm'''


	def __init__(						self							,    lowercase__=3_0_5_2_2							,    lowercase__=7_6_8							,    lowercase__=1_2_8							,    lowercase__=1_2							,    lowercase__=1_2							,    lowercase__=8							,    lowercase__=3_0_7_2							,    lowercase__="gelu_new"							,    lowercase__=0.1							,    lowercase__=0.1							,    lowercase__=5_1_2							,    lowercase__=2							,    lowercase__=0.0_2							,    lowercase__=1e-12							,    lowercase__=2_5_6							,    lowercase__=1_0							,    lowercase__=1e-3							,    lowercase__=5							,    lowercase__=3_2_0							,    lowercase__=1_3_3_5_3_7_1_8							,    lowercase__=5_0_0_0							,    lowercase__=1							,    lowercase__=0							,    lowercase__=2							,    **lowercase__							,    ):
							super().__init__(pad_token_id=lowercase__							,    bos_token_id=lowercase__							,    eos_token_id=lowercase__							,    **lowercase__)

							# Common config
							__UpperCAmelCase :       Optional[Any]  						=  vocab_size
							__UpperCAmelCase :       int  						=  max_position_embeddings
							__UpperCAmelCase :       Tuple  						=  hidden_size
							__UpperCAmelCase :       Optional[Any]  						=  retriever_proj_size
							__UpperCAmelCase :       List[Any]  						=  num_hidden_layers
							__UpperCAmelCase :       Optional[Any]  						=  num_attention_heads
							__UpperCAmelCase :       int  						=  num_candidates
							__UpperCAmelCase :       Dict  						=  intermediate_size
							__UpperCAmelCase :       Optional[Any]  						=  hidden_act
							__UpperCAmelCase :       Tuple  						=  hidden_dropout_prob
							__UpperCAmelCase :       Any  						=  attention_probs_dropout_prob
							__UpperCAmelCase :       Optional[Any]  						=  initializer_range
							__UpperCAmelCase :       List[str]  						=  type_vocab_size
							__UpperCAmelCase :       Any  						=  layer_norm_eps

							# Reader config
							__UpperCAmelCase :       Optional[int]  						=  span_hidden_size
							__UpperCAmelCase :       Dict  						=  max_span_width
							__UpperCAmelCase :       int  						=  reader_layer_norm_eps
							__UpperCAmelCase :       int  						=  reader_beam_size
							__UpperCAmelCase :       Optional[int]  						=  reader_seq_len

							# Retrieval config
							__UpperCAmelCase :       Optional[int]  						=  num_block_records
							__UpperCAmelCase :       Optional[Any]  						=  searcher_beam_size

