


'''simple docstring'''


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

from packaging import version


if TYPE_CHECKING:
						from ... import PreTrainedTokenizer, TensorType

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import is_torch_available, logging


lowerCamelCase__					=   logging.get_logger(__name__)

lowerCamelCase__					=   {
    """bigscience/bloom""": """https://huggingface.co/bigscience/bloom/resolve/main/config.json""",
    """bigscience/bloom-560m""": """https://huggingface.co/bigscience/bloom-560m/blob/main/config.json""",
    """bigscience/bloom-1b1""": """https://huggingface.co/bigscience/bloom-1b1/blob/main/config.json""",
    """bigscience/bloom-1b7""": """https://huggingface.co/bigscience/bloom-1b7/blob/main/config.json""",
    """bigscience/bloom-3b""": """https://huggingface.co/bigscience/bloom-3b/blob/main/config.json""",
    """bigscience/bloom-7b1""": """https://huggingface.co/bigscience/bloom-7b1/blob/main/config.json""",
}






class 		lowerCAmelCase__	(	_UpperCAmelCase		):
			lowerCAmelCase       :		str				      =   '''bloom'''
			lowerCAmelCase       :		int				      =   ['''past_key_values''']
			lowerCAmelCase       :		Union[str, Any]				      =   {
			    '''num_hidden_layers''': '''n_layer''',
			    '''num_attention_heads''': '''n_head''',
			}




			def __init__(						self     :   str		,     lowerCamelCase__     :   Any=25_08_80		,     lowerCamelCase__     :   int=64		,     lowerCamelCase__     :   Any=2		,     lowerCamelCase__     :   int=8		,     lowerCamelCase__     :   List[str]=1E-5		,     lowerCamelCase__     :   List[Any]=0.0_2		,     lowerCamelCase__     :   Tuple=True		,     lowerCamelCase__     :   Optional[Any]=1		,     lowerCamelCase__     :   Optional[Any]=2		,     lowerCamelCase__     :   Any=False		,     lowerCamelCase__     :   int=0.0		,     lowerCamelCase__     :   Tuple=0.0		,     lowerCamelCase__     :   int=1		,     lowerCamelCase__     :   int=False		,     **lowerCamelCase__     :   int		,     )      ->Union[str, Any]:



					'''simple docstring'''







					_UpperCAmelCase						:				Tuple	    =					vocab_size
					# Backward compatibility with n_embed kwarg
					_UpperCAmelCase						:				Dict	    =					kwargs.pop("n_embed"		,     lowerCAmelCase__       )
					_UpperCAmelCase						:				str	    =					hidden_size if n_embed is None else n_embed
					_UpperCAmelCase						:				str	    =					n_layer
					_UpperCAmelCase						:				Optional[int]	    =					n_head
					_UpperCAmelCase						:				List[Any]	    =					layer_norm_epsilon
					_UpperCAmelCase						:				Dict	    =					initializer_range
					_UpperCAmelCase						:				str	    =					use_cache
					_UpperCAmelCase						:				Optional[Any]	    =					pretraining_tp
					_UpperCAmelCase						:				Dict	    =					apply_residual_connection_post_layernorm
					_UpperCAmelCase						:				int	    =					hidden_dropout
					_UpperCAmelCase						:				Union[str, Any]	    =					attention_dropout

					_UpperCAmelCase						:				Tuple	    =					bos_token_id
					_UpperCAmelCase						:				Union[str, Any]	    =					eos_token_id
					_UpperCAmelCase						:				Union[str, Any]	    =					slow_but_exact

					super().__init__(bos_token_id=lowerCAmelCase__		,     eos_token_id=lowerCAmelCase__		,     **lowerCAmelCase__       )






class 		lowerCAmelCase__	(	_UpperCAmelCase		):
			lowerCAmelCase       :		Optional[Any]				      =   version.parse("1.12"		)




			def __init__(						self     :   Optional[Any]		,     lowerCamelCase__     :   List[str]		,     lowerCamelCase__     :   List[Any] = "default"		,     lowerCamelCase__     :   List[str] = None		,     lowerCamelCase__     :   Union[str, Any] = False		,     )      ->Optional[int]:



					'''simple docstring'''







					super().__init__(lowerCAmelCase__		,     task=lowerCAmelCase__		,     patching_specs=lowerCAmelCase__		,     use_past=lowerCAmelCase__       )
					if not getattr(self._config		,     "pad_token_id"		,     lowerCAmelCase__       ):
							# TODO: how to do that better?
							_UpperCAmelCase						:				Any	    =					0




			@property
			def 							lowerCAmelCase__			(						self     :   Any       )      ->Mapping[str, Mapping[int, str]]:



					'''simple docstring'''







					_UpperCAmelCase						:				Tuple	    =					OrderedDict({"input_ids": {0: "batch", 1: "sequence"}}       )
					if self.use_past:
							# BLOOM stores values on dynamic axis 2. For more details see: https://github.com/huggingface/transformers/pull/18344
							self.fill_with_past_key_values_(lowerCAmelCase__		,     direction="inputs"		,     inverted_values_shape=lowerCAmelCase__       )
							_UpperCAmelCase						:				Dict	    =					{0: "batch", 1: "past_sequence + sequence"}
					else:
							_UpperCAmelCase						:				Any	    =					{0: "batch", 1: "sequence"}

					return common_inputs




			@property
			def 							lowerCAmelCase__			(						self     :   Optional[int]       )      ->int:



					'''simple docstring'''







					return self._config.n_layer




			@property
			def 							lowerCAmelCase__			(						self     :   Tuple       )      ->int:



					'''simple docstring'''







					return self._config.n_head




			@property
			def 							lowerCAmelCase__			(						self     :   int       )      ->float:



					'''simple docstring'''







					return 1E-3




			def 							lowerCAmelCase__			(						self     :   str		,     lowerCamelCase__     :   Any		,     lowerCamelCase__     :   Any = -1		,     lowerCamelCase__     :   List[Any] = -1		,     lowerCamelCase__     :   Optional[int] = False		,     lowerCamelCase__     :   Optional[Any] = None		,     )      ->Mapping[str, Any]:



					'''simple docstring'''







					_UpperCAmelCase						:				Optional[Any]	    =					super(lowerCAmelCase__		,     self       ).generate_dummy_inputs(
					    lowerCAmelCase__		,     batch_size=lowerCAmelCase__		,     seq_length=lowerCAmelCase__		,     is_pair=lowerCAmelCase__		,     framework=lowerCAmelCase__       )

					# We need to order the input in the way they appears in the forward()
					_UpperCAmelCase						:				List[Any]	    =					OrderedDict({"input_ids": common_inputs["input_ids"]}       )

					# Need to add the past_keys
					if self.use_past:
							if not is_torch_available():
									raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."       )
							else:
									import torch

									_UpperCAmelCase    ,			_UpperCAmelCase						:				int	    =					common_inputs["input_ids"].shape
									# Not using the same length for past_key_values
									_UpperCAmelCase						:				List[Any]	    =					seqlen + 2
									_UpperCAmelCase						:				Optional[int]	    =					self._config.hidden_size // self.num_attention_heads
									_UpperCAmelCase						:				int	    =					(
									    batch * self.num_attention_heads,
									    head_dim,
									    past_key_values_length,
									)
									_UpperCAmelCase						:				List[str]	    =					(
									    batch * self.num_attention_heads,
									    past_key_values_length,
									    head_dim,
									)
									_UpperCAmelCase						:				Optional[Any]	    =					[
									    (torch.zeros(lowerCAmelCase__       ), torch.zeros(lowerCAmelCase__       )) for _ in range(self.num_layers       )
									]

					_UpperCAmelCase						:				str	    =					common_inputs["attention_mask"]
					if self.use_past:
							_UpperCAmelCase						:				Any	    =					ordered_inputs["attention_mask"].dtype
							_UpperCAmelCase						:				str	    =					torch.cat(
							    [ordered_inputs["attention_mask"], torch.ones(lowerCAmelCase__		,     lowerCAmelCase__		,     dtype=lowerCAmelCase__       )]		,     dim=1       )

					return ordered_inputs






			@property
			def 							lowerCAmelCase__			(						self     :   Optional[int]       )      ->int:



					'''simple docstring'''







					return 13




import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest




class    _UpperCamelCase      (						_UpperCAmelCase ):


			"""simple docstring"""






			__a  : Optional[Any]			 =						(KDPMaDiscreteScheduler,)
			__a  : Dict			 =						10





			def        _SCREAMING_SNAKE_CASE			(   self      ,   **lowerCAmelCase__						)     ->					int:

							'''simple docstring'''


							__lowercase 			=							{
							    '''num_train_timesteps''': 11_00,
							    '''beta_start''': 0.0001,
							    '''beta_end''': 0.02,
							    '''beta_schedule''': '''linear''',
							}

							config.update(**lowerCAmelCase__						)
							return config





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					List[str]:

							'''simple docstring'''


							for timesteps in [10, 50, 1_00, 10_00]:
											self.check_over_configs(num_train_timesteps=lowerCAmelCase__						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Any:

							'''simple docstring'''


							for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001]      ,   [0.0002, 0.002, 0.02]						):
											self.check_over_configs(beta_start=lowerCAmelCase__      ,   beta_end=lowerCAmelCase__						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Union[str, Any]:

							'''simple docstring'''


							for schedule in ["linear", "scaled_linear"]:
											self.check_over_configs(beta_schedule=lowerCAmelCase__						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Union[str, Any]:

							'''simple docstring'''


							for prediction_type in ["epsilon", "v_prediction"]:
											self.check_over_configs(prediction_type=lowerCAmelCase__						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Tuple:

							'''simple docstring'''


							__lowercase 			=							self.scheduler_classes[0]
							__lowercase 			=							self.get_scheduler_config(prediction_type='''v_prediction'''						)
							__lowercase 			=							scheduler_class(**lowerCAmelCase__						)

							scheduler.set_timesteps(self.num_inference_steps						)

							__lowercase 			=							self.dummy_model()
							__lowercase 			=							self.dummy_sample_deter * scheduler.init_noise_sigma
							__lowercase 			=							sample.to(lowerCAmelCase__						)

							for i, t in enumerate(scheduler.timesteps						):
											__lowercase 			=							scheduler.scale_model_input(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							model(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							scheduler.step(lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)
											__lowercase 			=							output.prev_sample

							__lowercase 			=							torch.sum(torch.abs(lowerCAmelCase__						)						)
							__lowercase 			=							torch.mean(torch.abs(lowerCAmelCase__						)						)

							if torch_device in ["cpu", "mps"]:
											assert abs(result_sum.item() - 4.6934E-07						) < 1E-2
											assert abs(result_mean.item() - 6.1112E-10						) < 1E-3
							else:
											# CUDA
											assert abs(result_sum.item() - 4.693_4286_5017_0972E-07						) < 1E-2
											assert abs(result_mean.item() - 0.0002						) < 1E-3





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					str:

							'''simple docstring'''


							if torch_device == "mps":
											return
							__lowercase 			=							self.scheduler_classes[0]
							__lowercase 			=							self.get_scheduler_config()
							__lowercase 			=							scheduler_class(**lowerCAmelCase__						)

							scheduler.set_timesteps(self.num_inference_steps						)

							__lowercase 			=							self.dummy_model()
							__lowercase 			=							self.dummy_sample_deter * scheduler.init_noise_sigma
							__lowercase 			=							sample.to(lowerCAmelCase__						)

							for i, t in enumerate(scheduler.timesteps						):
											__lowercase 			=							scheduler.scale_model_input(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							model(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							scheduler.step(lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)
											__lowercase 			=							output.prev_sample

							__lowercase 			=							torch.sum(torch.abs(lowerCAmelCase__						)						)
							__lowercase 			=							torch.mean(torch.abs(lowerCAmelCase__						)						)

							if torch_device in ["cpu", "mps"]:
											assert abs(result_sum.item() - 20.4125						) < 1E-2
											assert abs(result_mean.item() - 0.0266						) < 1E-3
							else:
											# CUDA
											assert abs(result_sum.item() - 20.4125						) < 1E-2
											assert abs(result_mean.item() - 0.0266						) < 1E-3





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Dict:

							'''simple docstring'''


							if torch_device == "mps":
											return
							__lowercase 			=							self.scheduler_classes[0]
							__lowercase 			=							self.get_scheduler_config()
							__lowercase 			=							scheduler_class(**lowerCAmelCase__						)

							scheduler.set_timesteps(self.num_inference_steps      ,   device=lowerCAmelCase__						)

							__lowercase 			=							self.dummy_model()
							__lowercase 			=							self.dummy_sample_deter.to(lowerCAmelCase__						) * scheduler.init_noise_sigma

							for t in scheduler.timesteps:
											__lowercase 			=							scheduler.scale_model_input(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							model(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							scheduler.step(lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)
											__lowercase 			=							output.prev_sample

							__lowercase 			=							torch.sum(torch.abs(lowerCAmelCase__						)						)
							__lowercase 			=							torch.mean(torch.abs(lowerCAmelCase__						)						)

							if str(lowerCAmelCase__						).startswith('''cpu'''						):
											# The following sum varies between 148 and 156 on mps. Why?
											assert abs(result_sum.item() - 20.4125						) < 1E-2
											assert abs(result_mean.item() - 0.0266						) < 1E-3
							else:
											# CUDA
											assert abs(result_sum.item() - 20.4125						) < 1E-2
											assert abs(result_mean.item() - 0.0266						) < 1E-3