



"""simple docstring"""

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 SCREAMING_SNAKE_CASE__					(		__snake_case      ):



      def __init__(self						,     _lowercase						,     _lowercase=13						,     _lowercase=7						,     _lowercase=True						,     _lowercase=True						,     _lowercase=False						,     _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						,     ):





             '''simple docstring'''

             __a		:		Optional[int]									=       parent
             __a		:		List[Any]									=       batch_size
             __a		:		Any									=       seq_length
             __a		:		Optional[Any]									=       is_training
             __a		:		Dict									=       use_input_mask
             __a		:		Dict									=       use_token_type_ids
             __a		:		Optional[int]									=       use_labels
             __a		:		Dict									=       vocab_size
             __a		:		Tuple									=       hidden_size
             __a		:		Optional[Any]									=       num_hidden_layers
             __a		:		str									=       num_attention_heads
             __a		:		List[Any]									=       intermediate_size
             __a		:		List[str]									=       hidden_act
             __a		:		Any									=       hidden_dropout_prob
             __a		:		Dict									=       attention_probs_dropout_prob
             __a		:		Union[str, Any]									=       max_position_embeddings
             __a		:		str									=       type_vocab_size
             __a		:		List[str]									=       type_sequence_label_size
             __a		:		Dict									=       initializer_range
             __a		:		Union[str, Any]									=       num_labels
             __a		:		Tuple									=       num_choices
             __a		:		List[str]									=       scope



      def     lowerCAmelCase__(self							):





             '''simple docstring'''

             __a		:		List[str]									=       ids_tensor([self.batch_size, self.seq_length]						,     self.vocab_size							)

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

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

             __a		:		Dict									=       self.get_config()

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



      def     lowerCAmelCase__(self							):





             '''simple docstring'''

             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     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							):





             '''simple docstring'''

             __a		:		Union[str, Any]									=       DistilBertModel(config=_lowercase							)
             model.to(_lowercase							)
             model.eval()
             __a		:		Union[str, Any]									=       model(_lowercase						,     _lowercase							)
             __a		:		int									=       model(_lowercase							)
             self.parent.assertEqual(result.last_hidden_state.shape						,     (self.batch_size, self.seq_length, self.hidden_size)							)



      def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							):





             '''simple docstring'''

             __a		:		Union[str, Any]									=       DistilBertForMaskedLM(config=_lowercase							)
             model.to(_lowercase							)
             model.eval()
             __a		:		List[Any]									=       model(_lowercase						,     attention_mask=_lowercase						,     labels=_lowercase							)
             self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.seq_length, self.vocab_size)							)



      def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							):





             '''simple docstring'''

             __a		:		Any									=       DistilBertForQuestionAnswering(config=_lowercase							)
             model.to(_lowercase							)
             model.eval()
             __a		:		Optional[Any]									=       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     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							):





             '''simple docstring'''

             __a		:		Dict									=       self.num_labels
             __a		:		str									=       DistilBertForSequenceClassification(_lowercase							)
             model.to(_lowercase							)
             model.eval()
             __a		:		str									=       model(_lowercase						,     attention_mask=_lowercase						,     labels=_lowercase							)
             self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.num_labels)							)



      def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							):





             '''simple docstring'''

             __a		:		Any									=       self.num_labels
             __a		:		Dict									=       DistilBertForTokenClassification(config=_lowercase							)
             model.to(_lowercase							)
             model.eval()

             __a		:		List[str]									=       model(_lowercase						,     attention_mask=_lowercase						,     labels=_lowercase							)
             self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.seq_length, self.num_labels)							)



      def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							):





             '''simple docstring'''

             __a		:		List[Any]									=       self.num_choices
             __a		:		Tuple									=       DistilBertForMultipleChoice(config=_lowercase							)
             model.to(_lowercase							)
             model.eval()
             __a		:		Optional[int]									=       input_ids.unsqueeze(1							).expand(-1						,     self.num_choices						,     -1							).contiguous()
             __a		:		Union[str, Any]									=       input_mask.unsqueeze(1							).expand(-1						,     self.num_choices						,     -1							).contiguous()
             __a		:		List[Any]									=       model(
                 _lowercase						,     attention_mask=_lowercase						,     labels=_lowercase						,     )
             self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.num_choices)							)





      def     lowerCAmelCase__(self							):





             '''simple docstring'''

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







@require_torch
class SCREAMING_SNAKE_CASE__					(		__snake_case     ,					__snake_case     ,					unittest.TestCase      ):
      _lowerCAmelCase							=					(
          (
              DistilBertModel,
              DistilBertForMaskedLM,
              DistilBertForMultipleChoice,
              DistilBertForQuestionAnswering,
              DistilBertForSequenceClassification,
              DistilBertForTokenClassification,
          )
          if is_torch_available()
          else None
      )
      _lowerCAmelCase							=					(
          {
              "feature-extraction": DistilBertModel,
              "fill-mask": DistilBertForMaskedLM,
              "question-answering": DistilBertForQuestionAnswering,
              "text-classification": DistilBertForSequenceClassification,
              "token-classification": DistilBertForTokenClassification,
              "zero-shot": DistilBertForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _lowerCAmelCase							=					True
      _lowerCAmelCase							=					True
      _lowerCAmelCase							=					True
      _lowerCAmelCase							=					True



      def     lowerCAmelCase__(self							):





             '''simple docstring'''

             __a		:		Any									=       DistilBertModelTester(self							)
             __a		:		Optional[Any]									=       ConfigTester(self						,     config_class=_lowercase						,     dim=37							)



      def     lowerCAmelCase__(self							):





             '''simple docstring'''

             self.config_tester.run_common_tests()



      def     lowerCAmelCase__(self							):





             '''simple docstring'''

             __a		:		Dict									=       self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_distilbert_model(*_lowercase							)



      def     lowerCAmelCase__(self							):





             '''simple docstring'''

             __a		:		List[Any]									=       self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_distilbert_for_masked_lm(*_lowercase							)



      def     lowerCAmelCase__(self							):





             '''simple docstring'''

             __a		:		Optional[int]									=       self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_distilbert_for_question_answering(*_lowercase							)



      def     lowerCAmelCase__(self							):





             '''simple docstring'''

             __a		:		Union[str, Any]									=       self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_distilbert_for_sequence_classification(*_lowercase							)



      def     lowerCAmelCase__(self							):





             '''simple docstring'''

             __a		:		Union[str, Any]									=       self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_distilbert_for_token_classification(*_lowercase							)



      def     lowerCAmelCase__(self							):





             '''simple docstring'''

             __a		:		Union[str, Any]									=       self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_distilbert_for_multiple_choice(*_lowercase							)



      @slow
      def     lowerCAmelCase__(self							):





             '''simple docstring'''

             for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    __a		:		Optional[int]									=       DistilBertModel.from_pretrained(_lowercase							)
                    self.assertIsNotNone(_lowercase							)





      @slow
      @require_torch_gpu
      def     lowerCAmelCase__(self							):





             '''simple docstring'''

             __a							,     __a		:		Any									=       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

                    __a		:		List[Any]									=       True
                    __a		:		Optional[int]									=       model_class(config=_lowercase							)

                    __a		:		Union[str, Any]									=       self._prepare_for_class(_lowercase						,     _lowercase							)
                    __a		:		List[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"""							)							)
                           __a		:		Any									=       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 SCREAMING_SNAKE_CASE__					(		unittest.TestCase      ):



      @slow
      def     lowerCAmelCase__(self							):





             '''simple docstring'''

             __a		:		int									=       DistilBertModel.from_pretrained("""distilbert-base-uncased"""							)
             __a		:		List[str]									=       torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]							)
             __a		:		Any									=       torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]							)
             with torch.no_grad():
                    __a		:		Optional[Any]									=       model(_lowercase						,     attention_mask=_lowercase							)[0]
             __a		:		Dict									=       torch.Size((1, 11, 768)							)
             self.assertEqual(output.shape						,     _lowercase							)
             __a		:		Optional[int]									=       torch.tensor(
                 [[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]							)

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






"""simple docstring"""

import importlib

import torch
import yaml
from omegaconf import OmegaConf
from taming.models.vqgan import VQModel





def       __magic_name__    (   _lowerCamelCase				: int  ,       _lowerCamelCase				: Optional[Any]=False					):
							__a		:		Dict									=       OmegaConf.load(_lowerCamelCase					)
							if display:
														print(yaml.dump(OmegaConf.to_container(_lowerCamelCase					)					)					)
							return config





def       __magic_name__    (   _lowerCamelCase				: str  ,       _lowerCamelCase				: Union[str, Any]=None  ,       _lowerCamelCase				: int=None					):
							if conf_path is None:
														__a		:		str									=       """./model_checkpoints/vqgan_only.yaml"""
							__a		:		List[Any]									=       load_config(_lowerCamelCase  ,       display=_lowerCamelCase					)
							__a		:		Dict									=       VQModel(**config.model.params					)
							if ckpt_path is None:
														__a		:		List[Any]									=       """./model_checkpoints/vqgan_only.pt"""
							__a		:		Tuple									=       torch.load(_lowerCamelCase  ,       map_location=_lowerCamelCase					)
							if ".ckpt" in ckpt_path:
														__a		:		List[str]									=       sd["""state_dict"""]
							model.load_state_dict(_lowerCamelCase  ,       strict=_lowerCamelCase					)
							model.to(_lowerCamelCase					)
							del sd
							return model





def       __magic_name__    (   _lowerCamelCase				: Optional[int]  ,       _lowerCamelCase				: List[str]					):
							__a							,     __a							,     __a		:		Tuple									=       model.encode(_lowerCamelCase					)
							print(F'''VQGAN --- {model.__class__.__name__}: latent shape: {z.shape[2:]}'''					)
							__a		:		Union[str, Any]									=       model.decode(_lowerCamelCase					)
							return xrec





def       __magic_name__    (   _lowerCamelCase				: Optional[int]  ,       _lowerCamelCase				: Union[str, Any]=False					):
							__a							,     __a		:		Optional[Any]									=       string.rsplit("""."""  ,       1					)
							if reload:
														__a		:		Optional[Any]									=       importlib.import_module(_lowerCamelCase					)
														importlib.reload(_lowerCamelCase					)
							return getattr(importlib.import_module(_lowerCamelCase  ,       package=_lowerCamelCase					)  ,       cls					)





def       __magic_name__    (   _lowerCamelCase				: Any					):
							if "target" not in config:
														raise KeyError("""Expected key `target` to instantiate."""					)
							return get_obj_from_str(config["""target"""]					)(**config.get("""params"""  ,       {}					)					)





def       __magic_name__    (   _lowerCamelCase				: Any  ,       _lowerCamelCase				: Dict  ,       _lowerCamelCase				: int=True  ,       _lowerCamelCase				: int=True					):
							__a		:		Union[str, Any]									=       instantiate_from_config(_lowerCamelCase					)
							if sd is not None:
														model.load_state_dict(_lowerCamelCase					)
							if gpu:
														model.cuda()
							if eval_mode:
														model.eval()
							return {"model": model}







def       __magic_name__    (   _lowerCamelCase				: Any  ,       _lowerCamelCase				: Union[str, Any]  ,       _lowerCamelCase				: str  ,       _lowerCamelCase				: int					):
							# load the specified checkpoint
							if ckpt:
														__a		:		List[str]									=       torch.load(_lowerCamelCase  ,       map_location="""cpu"""					)
														__a		:		Any									=       pl_sd["""global_step"""]
														print(F'''loaded model from global step {global_step}.'''					)
							else:
														__a		:		List[Any]									=       {"""state_dict""": None}
														__a		:		Any									=       None
							__a		:		Union[str, Any]									=       load_model_from_config(config.model  ,       pl_sd["""state_dict"""]  ,       gpu=_lowerCamelCase  ,       eval_mode=_lowerCamelCase					)["""model"""]
							return model, global_step

