






'''simple docstring'''







from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


__A   :							Dict     	=						logging.get_logger(__name__)

__A   :							Optional[int]     	=						{
    'facebook/blenderbot-3B': 'https://huggingface.co/facebook/blenderbot-3B/resolve/main/config.json',
    # See all Blenderbot models at https://huggingface.co/models?filter=blenderbot
}






class        a__     (      lowerCamelCase    ):
     '''simple docstring'''







     _snake_case							   =	'blenderbot'
     _snake_case							   =	['past_key_values']
     _snake_case							   =	{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}






     def __init__(self     :						Tuple					,					lowercase__     :						Union[str, Any]=8008					,					lowercase__     :						str=128					,					lowercase__     :						Dict=2					,					lowercase__     :						Dict=10240					,					lowercase__     :						Dict=32					,					lowercase__     :						List[Any]=24					,					lowercase__     :						List[str]=10240					,					lowercase__     :						Optional[int]=32					,					lowercase__     :						str=0.0					,					lowercase__     :						str=0.0					,					lowercase__     :						List[Any]=True					,					lowercase__     :						Union[str, Any]=True					,					lowercase__     :						List[str]="gelu"					,					lowercase__     :						Any=2560					,					lowercase__     :						Optional[Any]=0.1					,					lowercase__     :						str=0.0					,					lowercase__     :						List[Any]=0.0					,					lowercase__     :						Union[str, Any]=0.02					,					lowercase__     :						Optional[int]=1					,					lowercase__     :						Any=False					,					lowercase__     :						List[Any]=0					,					lowercase__     :						Dict=1					,					lowercase__     :						List[str]=2					,					lowercase__     :						Dict=3					,					lowercase__     :						str=2					,					**lowercase__     :						List[str]					,					)	->							List[str]:
           __A		 =   vocab_size
           __A		 =   max_position_embeddings
           __A		 =   d_model
           __A		 =   encoder_ffn_dim
           __A		 =   encoder_layers
           __A		 =   encoder_attention_heads
           __A		 =   decoder_ffn_dim
           __A		 =   decoder_layers
           __A		 =   decoder_attention_heads
           __A		 =   dropout
           __A		 =   attention_dropout
           __A		 =   activation_dropout
           __A		 =   activation_function
           __A		 =   init_std
           __A		 =   encoder_layerdrop
           __A		 =   decoder_layerdrop
           __A		 =   use_cache
           __A		 =   encoder_layers
           __A		 =   scale_embedding  # scale factor will be sqrt(d_model) if True

           super().__init__(
               pad_token_id=lowercase__					,					bos_token_id=lowercase__					,					eos_token_id=lowercase__					,					is_encoder_decoder=lowercase__					,					decoder_start_token_id=lowercase__					,					encoder_no_repeat_ngram_size=lowercase__					,					forced_eos_token_id=lowercase__					,					**lowercase__					,					)







class        a__     (      lowerCamelCase    ):
     '''simple docstring'''







     @property
     def      _UpperCamelCase (self     :						Tuple       )	->							Mapping[str, Mapping[int, str]]:
           if self.task in ["default", "seq2seq-lm"]:
                 __A		 =   OrderedDict(
                     [
                         ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                         ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                     ]       )
                 if self.use_past:
                       __A		 =   {0: "batch"}
                       __A		 =   {0: "batch", 1: "past_decoder_sequence + sequence"}
                 else:
                       __A		 =   {0: "batch", 1: "decoder_sequence"}
                       __A		 =   {0: "batch", 1: "decoder_sequence"}
                 if self.use_past:
                       self.fill_with_past_key_values_(lowercase__					,					direction="inputs"       )
           elif self.task == "causal-lm":
                 __A		 =   OrderedDict(
                     [
                         ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                         ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                     ]       )
                 if self.use_past:
                       __A ,					__A		 =   self.num_layers
                       for i in range(lowercase__       ):
                             __A		 =   {0: "batch", 2: "past_sequence + sequence"}
                             __A		 =   {0: "batch", 2: "past_sequence + sequence"}
           else:
                 __A		 =   OrderedDict(
                     [
                         ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                         ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                         ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
                         ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
                     ]       )

           return common_inputs






     @property
     # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
     def      _UpperCamelCase (self     :						Dict       )	->							Mapping[str, Mapping[int, str]]:
           if self.task in ["default", "seq2seq-lm"]:
                 __A		 =   super().outputs
           else:
                 __A		 =   super(lowercase__					,					self       ).outputs
                 if self.use_past:
                       __A ,					__A		 =   self.num_layers
                       for i in range(lowercase__       ):
                             __A		 =   {0: "batch", 2: "past_sequence + sequence"}
                             __A		 =   {0: "batch", 2: "past_sequence + sequence"}
           return common_outputs






     def      _UpperCamelCase (self     :						str					,					lowercase__     :						PreTrainedTokenizer					,					lowercase__     :						int = -1					,					lowercase__     :						int = -1					,					lowercase__     :						bool = False					,					lowercase__     :						Optional[TensorType] = None					,					)	->							Mapping[str, Any]:
           __A		 =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
               lowercase__					,					lowercase__					,					lowercase__					,					lowercase__					,					lowercase__       )
           # Generate decoder inputs
           __A		 =   seq_length if not self.use_past else 1
           __A		 =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
               lowercase__					,					lowercase__					,					lowercase__					,					lowercase__					,					lowercase__       )
           __A		 =   {f"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}
           __A		 =   dict(**lowercase__					,					**lowercase__       )

           if self.use_past:
                 if not is_torch_available():
                       raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."       )
                 else:
                       import torch
                 __A ,					__A		 =   common_inputs["input_ids"].shape
                 __A		 =   common_inputs["decoder_input_ids"].shape[1]
                 __A ,					__A		 =   self.num_attention_heads
                 __A		 =   (
                     batch,
                     num_encoder_attention_heads,
                     encoder_seq_length,
                     self._config.hidden_size // num_encoder_attention_heads,
                 )
                 __A		 =   decoder_seq_length
                 __A		 =   (
                     batch,
                     num_decoder_attention_heads,
                     decoder_past_length,
                     self._config.hidden_size // num_decoder_attention_heads,
                 )
                 __A		 =   torch.cat(
                     [common_inputs["decoder_attention_mask"], torch.ones(lowercase__					,					lowercase__       )]					,					dim=1       )
                 __A		 =   []
                 __A ,					__A		 =   self.num_layers

                 for _ in range(lowercase__       ):
                       common_inputs["past_key_values"].append(
                           (
                               torch.zeros(lowercase__       ),
                               torch.zeros(lowercase__       ),
                               torch.zeros(lowercase__       ),
                               torch.zeros(lowercase__       ),
                           )       )
           return common_inputs






     def      _UpperCamelCase (self     :						Union[str, Any]					,					lowercase__     :						PreTrainedTokenizer					,					lowercase__     :						int = -1					,					lowercase__     :						int = -1					,					lowercase__     :						bool = False					,					lowercase__     :						Optional[TensorType] = None					,					)	->							Mapping[str, Any]:
           __A		 =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
               lowercase__					,					lowercase__					,					lowercase__					,					lowercase__					,					lowercase__       )

           if self.use_past:
                 if not is_torch_available():
                       raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."       )
                 else:
                       import torch
                 __A ,					__A		 =   common_inputs["input_ids"].shape
                 __A		 =   seqlen
                 __A ,					__A		 =   self.num_layers
                 __A ,					__A		 =   self.num_attention_heads
                 __A		 =   (
                     batch,
                     num_encoder_attention_heads,
                     past_key_values_length,
                     self._config.hidden_size // num_encoder_attention_heads,
                 )
                 __A		 =   common_inputs["attention_mask"].dtype
                 __A		 =   torch.cat(
                     [common_inputs["attention_mask"], torch.ones(lowercase__					,					lowercase__					,					dtype=lowercase__       )]					,					dim=1       )
                 __A		 =   [
                     (torch.zeros(lowercase__       ), torch.zeros(lowercase__       )) for _ in range(lowercase__       )
                 ]
           return common_inputs






     def      _UpperCamelCase (self     :						str					,					lowercase__     :						PreTrainedTokenizer					,					lowercase__     :						int = -1					,					lowercase__     :						int = -1					,					lowercase__     :						bool = False					,					lowercase__     :						Optional[TensorType] = None					,					)	->							Mapping[str, Any]:
           # Copied from OnnxConfig.generate_dummy_inputs
           # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
           # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
           __A		 =   compute_effective_axis_dimension(
               lowercase__					,					fixed_dimension=OnnxConfig.default_fixed_batch					,					num_token_to_add=0       )

           # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
           __A		 =   tokenizer.num_special_tokens_to_add(lowercase__       )
           __A		 =   compute_effective_axis_dimension(
               lowercase__					,					fixed_dimension=OnnxConfig.default_fixed_sequence					,					num_token_to_add=lowercase__       )

           # Generate dummy inputs according to compute batch and sequence
           __A		 =   [" ".join([tokenizer.unk_token]       ) * seq_length] * batch_size
           __A		 =   dict(tokenizer(lowercase__					,					return_tensors=lowercase__       )       )
           return common_inputs






     def      _UpperCamelCase (self     :						int					,					lowercase__     :						PreTrainedTokenizer					,					lowercase__     :						int = -1					,					lowercase__     :						int = -1					,					lowercase__     :						bool = False					,					lowercase__     :						Optional[TensorType] = None					,					)	->							Mapping[str, Any]:
           if self.task in ["default", "seq2seq-lm"]:
                 __A		 =   self._generate_dummy_inputs_for_default_and_seqaseq_lm(
                     lowercase__					,					batch_size=lowercase__					,					seq_length=lowercase__					,					is_pair=lowercase__					,					framework=lowercase__       )

           elif self.task == "causal-lm":
                 __A		 =   self._generate_dummy_inputs_for_causal_lm(
                     lowercase__					,					batch_size=lowercase__					,					seq_length=lowercase__					,					is_pair=lowercase__					,					framework=lowercase__       )
           else:
                 __A		 =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
                     lowercase__					,					batch_size=lowercase__					,					seq_length=lowercase__					,					is_pair=lowercase__					,					framework=lowercase__       )

           return common_inputs






     def      _UpperCamelCase (self     :						List[Any]					,					lowercase__     :						List[str]					,					lowercase__     :						Dict					,					lowercase__     :						int					,					lowercase__     :						List[str]       )	->							Any:
           if self.task in ["default", "seq2seq-lm"]:
                 __A		 =   super()._flatten_past_key_values_(lowercase__					,					lowercase__					,					lowercase__					,					lowercase__       )
           else:
                 __A		 =   super(lowercase__					,					self       )._flatten_past_key_values_(
                     lowercase__					,					lowercase__					,					lowercase__					,					lowercase__       )






     def      _UpperCamelCase (self     :						List[Any]					,					lowercase__     :						Mapping[str, Mapping[int, str]]					,					lowercase__     :						str       )	->							List[Any]:
           if direction not in ["inputs", "outputs"]:
                 raise ValueError(f"direction must either be \"inputs\" or \"outputs\", but {direction} was given"       )

           __A		 =   "past_key_values" if direction == "inputs" else "present"
           __A ,					__A		 =   self.num_layers

           __A		 =   "past_encoder_sequence"
           __A		 =   "past_decoder_sequence" if direction == "inputs" else "past_decoder_sequence + sequence"

           for i in range(lowercase__       ):
                 __A		 =   {0: "batch", 2: decoder_sequence}
                 __A		 =   {0: "batch", 2: decoder_sequence}
                 __A		 =   {0: "batch", 2: encoder_sequence}
                 __A		 =   {0: "batch", 2: encoder_sequence}










'''simple docstring'''







import torch

from diffusers import KDPMaAncestralDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest






class        a__     (      lowerCamelCase    ):
     '''simple docstring'''







     _snake_case							   =	(KDPMaAncestralDiscreteScheduler,)
     _snake_case							   =	10






     def      _UpperCamelCase (self     :						int					,					**lowercase__     :						Optional[int]       )	->							Dict:
           __A		 =   {
               "num_train_timesteps": 1100,
               "beta_start": 0.00_01,
               "beta_end": 0.02,
               "beta_schedule": "linear",
           }

           config.update(**lowercase__       )
           return config






     def      _UpperCamelCase (self     :						str       )	->							Optional[Any]:
           for timesteps in [10, 50, 100, 1000]:
                 self.check_over_configs(num_train_timesteps=lowercase__       )






     def      _UpperCamelCase (self     :						List[Any]       )	->							int:
           for beta_start, beta_end in zip([0.0_00_01, 0.00_01, 0.0_01]					,					[0.00_02, 0.0_02, 0.02]       ):
                 self.check_over_configs(beta_start=lowercase__					,					beta_end=lowercase__       )






     def      _UpperCamelCase (self     :						int       )	->							Tuple:
           for schedule in ["linear", "scaled_linear"]:
                 self.check_over_configs(beta_schedule=lowercase__       )






     def      _UpperCamelCase (self     :						int       )	->							Any:
           if torch_device == "mps":
                 return
           __A		 =   self.scheduler_classes[0]
           __A		 =   self.get_scheduler_config()
           __A		 =   scheduler_class(**lowercase__       )

           scheduler.set_timesteps(self.num_inference_steps       )

           __A		 =   torch.manual_seed(0       )

           __A		 =   self.dummy_model()
           __A		 =   self.dummy_sample_deter * scheduler.init_noise_sigma
           __A		 =   sample.to(lowercase__       )

           for i, t in enumerate(scheduler.timesteps       ):
                 __A		 =   scheduler.scale_model_input(lowercase__					,					lowercase__       )

                 __A		 =   model(lowercase__					,					lowercase__       )

                 __A		 =   scheduler.step(lowercase__					,					lowercase__					,					lowercase__					,					generator=lowercase__       )
                 __A		 =   output.prev_sample

           __A		 =   torch.sum(torch.abs(lowercase__       )       )
           __A		 =   torch.mean(torch.abs(lowercase__       )       )

           assert abs(result_sum.item() - 1_38_49.38_77       ) < 1e-2
           assert abs(result_mean.item() - 18.03_31       ) < 5e-3






     def      _UpperCamelCase (self     :						Optional[Any]       )	->							str:
           for prediction_type in ["epsilon", "v_prediction"]:
                 self.check_over_configs(prediction_type=lowercase__       )






     def      _UpperCamelCase (self     :						List[Any]       )	->							List[Any]:
           if torch_device == "mps":
                 return
           __A		 =   self.scheduler_classes[0]
           __A		 =   self.get_scheduler_config(prediction_type="v_prediction"       )
           __A		 =   scheduler_class(**lowercase__       )

           scheduler.set_timesteps(self.num_inference_steps       )

           __A		 =   self.dummy_model()
           __A		 =   self.dummy_sample_deter * scheduler.init_noise_sigma
           __A		 =   sample.to(lowercase__       )

           __A		 =   torch.manual_seed(0       )

           for i, t in enumerate(scheduler.timesteps       ):
                 __A		 =   scheduler.scale_model_input(lowercase__					,					lowercase__       )

                 __A		 =   model(lowercase__					,					lowercase__       )

                 __A		 =   scheduler.step(lowercase__					,					lowercase__					,					lowercase__					,					generator=lowercase__       )
                 __A		 =   output.prev_sample

           __A		 =   torch.sum(torch.abs(lowercase__       )       )
           __A		 =   torch.mean(torch.abs(lowercase__       )       )

           assert abs(result_sum.item() - 3_28.99_70       ) < 1e-2
           assert abs(result_mean.item() - 0.42_84       ) < 1e-3






     def      _UpperCamelCase (self     :						Optional[Any]       )	->							int:
           if torch_device == "mps":
                 return
           __A		 =   self.scheduler_classes[0]
           __A		 =   self.get_scheduler_config()
           __A		 =   scheduler_class(**lowercase__       )

           scheduler.set_timesteps(self.num_inference_steps					,					device=lowercase__       )
           __A		 =   torch.manual_seed(0       )

           __A		 =   self.dummy_model()
           __A		 =   self.dummy_sample_deter.to(lowercase__       ) * scheduler.init_noise_sigma

           for t in scheduler.timesteps:
                 __A		 =   scheduler.scale_model_input(lowercase__					,					lowercase__       )

                 __A		 =   model(lowercase__					,					lowercase__       )

                 __A		 =   scheduler.step(lowercase__					,					lowercase__					,					lowercase__					,					generator=lowercase__       )
                 __A		 =   output.prev_sample

           __A		 =   torch.sum(torch.abs(lowercase__       )       )
           __A		 =   torch.mean(torch.abs(lowercase__       )       )

           assert abs(result_sum.item() - 1_38_49.38_18       ) < 1e-1
           assert abs(result_mean.item() - 18.03_31       ) < 1e-3


