



'''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 (  _a		):






    """simple docstring"""





    def __init__(  self      :					List[str]  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					Optional[Any]=1_3  ,						UpperCamelCase__      :					Dict=7  ,						UpperCamelCase__      :					int=True  ,						UpperCamelCase__      :					Optional[Any]=True  ,						UpperCamelCase__      :					Dict=False  ,						UpperCamelCase__      :					int=True  ,						UpperCamelCase__      :					List[str]=9_9  ,						UpperCamelCase__      :					Optional[int]=3_2  ,						UpperCamelCase__      :					int=5  ,						UpperCamelCase__      :					str=4  ,						UpperCamelCase__      :					int=3_7  ,						UpperCamelCase__      :					Optional[int]="gelu"  ,						UpperCamelCase__      :					str=0.1  ,						UpperCamelCase__      :					Tuple=0.1  ,						UpperCamelCase__      :					int=5_1_2  ,						UpperCamelCase__      :					List[Any]=1_6  ,						UpperCamelCase__      :					List[Any]=2  ,						UpperCamelCase__      :					Optional[Any]=0.0_2  ,						UpperCamelCase__      :					str=3  ,						UpperCamelCase__      :					int=4  ,						UpperCamelCase__      :					Optional[int]=None  ,						):


       """simple docstring"""







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






    def  A       (  self      :					str ):


       """simple docstring"""







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

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

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

       UpperCamelCase					=  self.get_config()

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






    def  A       (  self      :					Tuple ):


       """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  A       (  self      :					Dict  ,						UpperCamelCase__      :					Union[str, Any]  ,						UpperCamelCase__      :					Optional[int]  ,						UpperCamelCase__      :					Tuple  ,						UpperCamelCase__      :					Tuple  ,						UpperCamelCase__      :					Any  ,						UpperCamelCase__      :					List[str] ):


       """simple docstring"""







       UpperCamelCase					=  DistilBertModel(config=UpperCamelCase__ )
       model.to(UpperCamelCase__ )
       model.eval()
       UpperCamelCase					=  model(UpperCamelCase__  ,						UpperCamelCase__ )
       UpperCamelCase					=  model(UpperCamelCase__ )
       self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, self.seq_length, self.hidden_size) )






    def  A       (  self      :					List[Any]  ,						UpperCamelCase__      :					List[str]  ,						UpperCamelCase__      :					str  ,						UpperCamelCase__      :					Tuple  ,						UpperCamelCase__      :					Optional[Any]  ,						UpperCamelCase__      :					Union[str, Any]  ,						UpperCamelCase__      :					Tuple ):


       """simple docstring"""







       UpperCamelCase					=  DistilBertForMaskedLM(config=UpperCamelCase__ )
       model.to(UpperCamelCase__ )
       model.eval()
       UpperCamelCase					=  model(UpperCamelCase__  ,						attention_mask=UpperCamelCase__  ,						labels=UpperCamelCase__ )
       self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length, self.vocab_size) )






    def  A       (  self      :					Tuple  ,						UpperCamelCase__      :					str  ,						UpperCamelCase__      :					Any  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					Optional[Any]  ,						UpperCamelCase__      :					List[Any]  ,						UpperCamelCase__      :					List[Any] ):


       """simple docstring"""







       UpperCamelCase					=  DistilBertForQuestionAnswering(config=UpperCamelCase__ )
       model.to(UpperCamelCase__ )
       model.eval()
       UpperCamelCase					=  model(
           UpperCamelCase__  ,						attention_mask=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      :					Union[str, Any]  ,						UpperCamelCase__      :					Tuple  ,						UpperCamelCase__      :					str  ,						UpperCamelCase__      :					int  ,						UpperCamelCase__      :					Union[str, Any]  ,						UpperCamelCase__      :					Optional[int]  ,						UpperCamelCase__      :					int ):


       """simple docstring"""







       UpperCamelCase					=  self.num_labels
       UpperCamelCase					=  DistilBertForSequenceClassification(UpperCamelCase__ )
       model.to(UpperCamelCase__ )
       model.eval()
       UpperCamelCase					=  model(UpperCamelCase__  ,						attention_mask=UpperCamelCase__  ,						labels=UpperCamelCase__ )
       self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.num_labels) )






    def  A       (  self      :					List[Any]  ,						UpperCamelCase__      :					List[str]  ,						UpperCamelCase__      :					Any  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					Any  ,						UpperCamelCase__      :					List[Any]  ,						UpperCamelCase__      :					Optional[int] ):


       """simple docstring"""







       UpperCamelCase					=  self.num_labels
       UpperCamelCase					=  DistilBertForTokenClassification(config=UpperCamelCase__ )
       model.to(UpperCamelCase__ )
       model.eval()

       UpperCamelCase					=  model(UpperCamelCase__  ,						attention_mask=UpperCamelCase__  ,						labels=UpperCamelCase__ )
       self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length, self.num_labels) )






    def  A       (  self      :					List[Any]  ,						UpperCamelCase__      :					Optional[int]  ,						UpperCamelCase__      :					List[Any]  ,						UpperCamelCase__      :					List[Any]  ,						UpperCamelCase__      :					Tuple  ,						UpperCamelCase__      :					Optional[Any]  ,						UpperCamelCase__      :					Union[str, Any] ):


       """simple docstring"""







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






    def  A       (  self      :					Optional[int] ):


       """simple docstring"""







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

@require_torch
class  SCREAMING_SNAKE_CASE (  _a   ,						_a   ,						unittest.TestCase		):






    """simple docstring"""





    _SCREAMING_SNAKE_CASE	  =  (
        (
            DistilBertModel,
            DistilBertForMaskedLM,
            DistilBertForMultipleChoice,
            DistilBertForQuestionAnswering,
            DistilBertForSequenceClassification,
            DistilBertForTokenClassification,
        )
        if is_torch_available()
        else None
    )
    _SCREAMING_SNAKE_CASE	  =  (
        {
            """feature-extraction""": DistilBertModel,
            """fill-mask""": DistilBertForMaskedLM,
            """question-answering""": DistilBertForQuestionAnswering,
            """text-classification""": DistilBertForSequenceClassification,
            """token-classification""": DistilBertForTokenClassification,
            """zero-shot""": DistilBertForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    _SCREAMING_SNAKE_CASE	  =  True
    _SCREAMING_SNAKE_CASE	  =  True
    _SCREAMING_SNAKE_CASE	  =  True
    _SCREAMING_SNAKE_CASE	  =  True






    def  A       (  self      :					Tuple ):


       """simple docstring"""







       UpperCamelCase					=  DistilBertModelTester(self )
       UpperCamelCase					=  ConfigTester(self  ,						config_class=UpperCamelCase__  ,						dim=3_7 )






    def  A       (  self      :					List[str] ):


       """simple docstring"""







       self.config_tester.run_common_tests()






    def  A       (  self      :					Tuple ):


       """simple docstring"""







       UpperCamelCase					=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_model(*UpperCamelCase__ )






    def  A       (  self      :					Optional[Any] ):


       """simple docstring"""







       UpperCamelCase					=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_masked_lm(*UpperCamelCase__ )






    def  A       (  self      :					Tuple ):


       """simple docstring"""







       UpperCamelCase					=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_question_answering(*UpperCamelCase__ )






    def  A       (  self      :					str ):


       """simple docstring"""







       UpperCamelCase					=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_sequence_classification(*UpperCamelCase__ )






    def  A       (  self      :					Union[str, Any] ):


       """simple docstring"""







       UpperCamelCase					=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_token_classification(*UpperCamelCase__ )






    def  A       (  self      :					Union[str, Any] ):


       """simple docstring"""







       UpperCamelCase					=  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_multiple_choice(*UpperCamelCase__ )






    @slow
    def  A       (  self      :					Tuple ):


       """simple docstring"""







       for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          UpperCamelCase					=  DistilBertModel.from_pretrained(UpperCamelCase__ )
          self.assertIsNotNone(UpperCamelCase__ )






    @slow
    @require_torch_gpu
    def  A       (  self      :					Optional[Any] ):


       """simple docstring"""







       UpperCamelCase				,      UpperCamelCase					=  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					=  True
          UpperCamelCase					=  model_class(config=UpperCamelCase__ )

          UpperCamelCase					=  self._prepare_for_class(UpperCamelCase__  ,						UpperCamelCase__ )
          UpperCamelCase					=  torch.jit.trace(
              UpperCamelCase__  ,						(inputs_dict['input_ids'].to('cpu' ), inputs_dict['attention_mask'].to('cpu' )) )

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

@require_torch
class  SCREAMING_SNAKE_CASE (  unittest.TestCase		):






    """simple docstring"""





    @slow
    def  A       (  self      :					str ):


       """simple docstring"""







       UpperCamelCase					=  DistilBertModel.from_pretrained('distilbert-base-uncased' )
       UpperCamelCase					=  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					=  torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]] )
       with torch.no_grad():
          UpperCamelCase					=  model(UpperCamelCase__  ,						attention_mask=UpperCamelCase__ )[0]
       UpperCamelCase					=  torch.Size((1, 1_1, 7_6_8) )
       self.assertEqual(output.shape  ,						UpperCamelCase__ )
       UpperCamelCase					=  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]  ,						UpperCamelCase__  ,						atol=1E-4 ) )




"""simple docstring"""





import importlib
import math
import os
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, Optional, Tuple, Union

import flax
import jax.numpy as jnp

from ..utils import BaseOutput


__UpperCamelCase  :      Optional[Any]							      =		'''scheduler_config.json'''







class 				SCREAMING_SNAKE_CASE					(		a_       ):






					"""simple docstring"""



					lowercase__				 =       1
					lowercase__				 =       2
					lowercase__				 =       3
					lowercase__				 =       4
					lowercase__				 =       5







@dataclass
class 				SCREAMING_SNAKE_CASE					(		a_       ):






					"""simple docstring"""



					lowercase__				 =       42







class 				SCREAMING_SNAKE_CASE					:






					"""simple docstring"""



					lowercase__				 =       SCHEDULER_CONFIG_NAME
					lowercase__				 =       ["dtype"]
					lowercase__				 =       []
					lowercase__				 =       True







					@classmethod
					def        __lowerCAmelCase (       cls			:				List[Any]   ,lowercase_			:				Dict[str, Any] = None   ,lowercase_			:				Optional[str] = None   ,lowercase_			:				Optional[int]=False   ,**lowercase_			:				Any   ,):
						lowerCAmelCase__				,lowerCAmelCase__	:						Dict  				=	cls.load_config(
						    pretrained_model_name_or_path=lowercase_   ,subfolder=lowercase_   ,return_unused_kwargs=lowercase_   ,**lowercase_   ,)
						lowerCAmelCase__				,lowerCAmelCase__	:						Union[str, Any]  				=	cls.from_config(lowercase_   ,return_unused_kwargs=lowercase_   ,**lowercase_		)

						if hasattr(lowercase_   ,'''create_state'''		) and getattr(lowercase_   ,'''has_state'''   ,lowercase_		):
							lowerCAmelCase__	:						List[Any]  				=	scheduler.create_state()

						if return_unused_kwargs:
							return scheduler, state, unused_kwargs

						return scheduler, state







					def        __lowerCAmelCase (       self			:				Tuple   ,lowercase_			:				Union[str, os.PathLike]   ,lowercase_			:				bool = False   ,**lowercase_			:				str		):
						self.save_config(save_directory=lowercase_   ,push_to_hub=lowercase_   ,**lowercase_		)







					@property
					def        __lowerCAmelCase (       self			:				List[str]		):
						return self._get_compatibles()







					@classmethod
					def        __lowerCAmelCase (       cls			:				List[Any]		):
						lowerCAmelCase__	:						Tuple  				=	list(set([cls.__name__] + cls._compatibles		)		)
						lowerCAmelCase__	:						Tuple  				=	importlib.import_module(__name__.split('''.'''		)[0]		)
						lowerCAmelCase__	:						Union[str, Any]  				=	[
						    getattr(lowercase_   ,lowercase_		) for c in compatible_classes_str if hasattr(lowercase_   ,lowercase_		)
						]
						return compatible_classes

def   __SCREAMING_SNAKE_CASE  (      A_						,     A_ ):
	assert len(A_ ) >= x.ndim
	return jnp.broadcast_to(x.reshape(x.shape + (1,) * (len(A_ ) - x.ndim) )						,     A_ )

def   __SCREAMING_SNAKE_CASE  (      A_						,     A_=0.999						,     A_=jnp.floataa ):

	def alpha_bar(A_ ):
		return math.cos((time_step + 0.008) / 1.008 * math.pi / 2 ) ** 2

	lowerCAmelCase__	:						Optional[Any]  				=	[]
	for i in range(A_ ):
		lowerCAmelCase__	:						str  				=	i / num_diffusion_timesteps
		lowerCAmelCase__	:						Union[str, Any]  				=	(i + 1) / num_diffusion_timesteps
		betas.append(min(1 - alpha_bar(A_ ) / alpha_bar(A_ )						,     A_ ) )
	return jnp.array(A_						,     dtype=A_ )







@flax.struct.dataclass
class 				SCREAMING_SNAKE_CASE					:






					"""simple docstring"""



					lowercase__				 =       42
					lowercase__				 =       42
					lowercase__				 =       42







					@classmethod
					def        __lowerCAmelCase (       cls			:				Union[str, Any]   ,lowercase_			:				List[Any]		):
						lowerCAmelCase__	:						Optional[int]  				=	scheduler.config

						if config.trained_betas is not None:
							lowerCAmelCase__	:						Any  				=	jnp.asarray(config.trained_betas   ,dtype=scheduler.dtype		)
						elif config.beta_schedule == "linear":
							lowerCAmelCase__	:						Union[str, Any]  				=	jnp.linspace(config.beta_start   ,config.beta_end   ,config.num_train_timesteps   ,dtype=scheduler.dtype		)
						elif config.beta_schedule == "scaled_linear":
							# this schedule is very specific to the latent diffusion model.
							lowerCAmelCase__	:						int  				=	(
							    jnp.linspace(
							        config.beta_start**0.5   ,config.beta_end**0.5   ,config.num_train_timesteps   ,dtype=scheduler.dtype		)
							    ** 2
							)
						elif config.beta_schedule == "squaredcos_cap_v2":
							# Glide cosine schedule
							lowerCAmelCase__	:						List[Any]  				=	betas_for_alpha_bar(config.num_train_timesteps   ,dtype=scheduler.dtype		)
						else:
							raise NotImplementedError(
							    F'beta_schedule {config.beta_schedule} is not implemented for scheduler {scheduler.__class__.__name__}'		)

						lowerCAmelCase__	:						str  				=	1.0 - betas

						lowerCAmelCase__	:						Union[str, Any]  				=	jnp.cumprod(lowercase_   ,axis=0		)

						return cls(
						    alphas=lowercase_   ,betas=lowercase_   ,alphas_cumprod=lowercase_   ,)

def   __SCREAMING_SNAKE_CASE  (      A_						,     A_						,     A_						,     A_ ):
	lowerCAmelCase__	:						Any  				=	state.alphas_cumprod

	lowerCAmelCase__	:						Optional[Any]  				=	alphas_cumprod[timesteps] ** 0.5
	lowerCAmelCase__	:						Tuple  				=	sqrt_alpha_prod.flatten()
	lowerCAmelCase__	:						str  				=	broadcast_to_shape_from_left(A_						,     original_samples.shape )

	lowerCAmelCase__	:						Optional[Any]  				=	(1 - alphas_cumprod[timesteps]) ** 0.5
	lowerCAmelCase__	:						Optional[Any]  				=	sqrt_one_minus_alpha_prod.flatten()
	lowerCAmelCase__	:						Optional[int]  				=	broadcast_to_shape_from_left(A_						,     original_samples.shape )

	return sqrt_alpha_prod, sqrt_one_minus_alpha_prod

def   __SCREAMING_SNAKE_CASE  (      A_						,     A_						,     A_						,     A_ ):
	lowerCAmelCase__				,lowerCAmelCase__	:						List[Any]  				=	get_sqrt_alpha_prod(A_						,     A_						,     A_						,     A_ )
	lowerCAmelCase__	:						str  				=	sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
	return noisy_samples





def   __SCREAMING_SNAKE_CASE  (      A_						,     A_						,     A_						,     A_ ):
	lowerCAmelCase__				,lowerCAmelCase__	:						List[Any]  				=	get_sqrt_alpha_prod(A_						,     A_						,     A_						,     A_ )
	lowerCAmelCase__	:						Union[str, Any]  				=	sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
	return velocity




