


'''simple docstring'''





import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


UpperCAmelCase__  			=				logging.get_logger(__name__)






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Union[str, Any] ):






			"""simple docstring"""




			__A=	r'\w+[.]\d+'
			__A=	re.findall(__lowerCAmelCase,__lowerCAmelCase )
			for pat in pats:
						__A=	key.replace(__lowerCAmelCase,'_'.join(pat.split('.' ) ) )
			return key






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: int,_SCREAMING_SNAKE_CASE					: Optional[Any],_SCREAMING_SNAKE_CASE					: int ):






			"""simple docstring"""




			__A=	pt_tuple_key[:-1] + ('scale',)
			if (
			    any('norm' in str_ for str_ in pt_tuple_key )
			    and (pt_tuple_key[-1] == "bias")
			    and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
			    and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
			):
						__A=	pt_tuple_key[:-1] + ('scale',)
						return renamed_pt_tuple_key, pt_tensor
			elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
						__A=	pt_tuple_key[:-1] + ('scale',)
						return renamed_pt_tuple_key, pt_tensor

			# embedding
			if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
						__A=	pt_tuple_key[:-1] + ('embedding',)
						return renamed_pt_tuple_key, pt_tensor

			# conv layer
			__A=	pt_tuple_key[:-1] + ('kernel',)
			if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
						__A=	pt_tensor.transpose(2,3,1,0 )
						return renamed_pt_tuple_key, pt_tensor

			# linear layer
			__A=	pt_tuple_key[:-1] + ('kernel',)
			if pt_tuple_key[-1] == "weight":
						__A=	pt_tensor.T
						return renamed_pt_tuple_key, pt_tensor

			# old PyTorch layer norm weight
			__A=	pt_tuple_key[:-1] + ('weight',)
			if pt_tuple_key[-1] == "gamma":
						return renamed_pt_tuple_key, pt_tensor

			# old PyTorch layer norm bias
			__A=	pt_tuple_key[:-1] + ('bias',)
			if pt_tuple_key[-1] == "beta":
						return renamed_pt_tuple_key, pt_tensor

			return pt_tuple_key, pt_tensor






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Optional[int],_SCREAMING_SNAKE_CASE					: str,_SCREAMING_SNAKE_CASE					: Union[str, Any]=42 ):






			"""simple docstring"""




			__A=	{k: v.numpy() for k, v in pt_state_dict.items()}

			# Step 2: Since the model is stateless, get random Flax params
			__A=	flax_model.init_weights(PRNGKey(__lowerCAmelCase ) )

			__A=	flatten_dict(__lowerCAmelCase )
			__A=	{}

			# Need to change some parameters name to match Flax names
			for pt_key, pt_tensor in pt_state_dict.items():
						__A=	rename_key(__lowerCAmelCase )
						__A=	tuple(renamed_pt_key.split('.' ) )

						# Correctly rename weight parameters
						__A,	__A=	rename_key_and_reshape_tensor(__lowerCAmelCase,__lowerCAmelCase,__lowerCAmelCase )

						if flax_key in random_flax_state_dict:
									if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
												raise ValueError(
												    f"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
												    f"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.""" )

        # also add unexpected weight so that warning is thrown
						__A=	jnp.asarray(__lowerCAmelCase )

			return unflatten_dict(__lowerCAmelCase )





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


_A  =	logging.get_logger(__name__)

_A  =	{
    "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   (      __a			):
				_lowercase     						='''bloom'''
				_lowercase     						=['''past_key_values''']
				_lowercase     						={
				    '''num_hidden_layers''': '''n_layer''',
				    '''num_attention_heads''': '''n_head''',
				}




				def __init__(   self							,							_UpperCamelCase=250_880							,							_UpperCamelCase=64							,							_UpperCamelCase=2							,							_UpperCamelCase=8							,							_UpperCamelCase=1e-5							,							_UpperCamelCase=0.02							,							_UpperCamelCase=True							,							_UpperCamelCase=1							,							_UpperCamelCase=2							,							_UpperCamelCase=False							,							_UpperCamelCase=0.0							,							_UpperCamelCase=0.0							,							_UpperCamelCase=1							,							_UpperCamelCase=False							,							**_UpperCamelCase							,							)							-> str:
							lowerCAmelCase_         =      vocab_size
							# Backward compatibility with n_embed kwarg
							lowerCAmelCase_         =      kwargs.pop("n_embed"							,							_UpperCamelCase	)
							lowerCAmelCase_         =      hidden_size if n_embed is None else n_embed
							lowerCAmelCase_         =      n_layer
							lowerCAmelCase_         =      n_head
							lowerCAmelCase_         =      layer_norm_epsilon
							lowerCAmelCase_         =      initializer_range
							lowerCAmelCase_         =      use_cache
							lowerCAmelCase_         =      pretraining_tp
							lowerCAmelCase_         =      apply_residual_connection_post_layernorm
							lowerCAmelCase_         =      hidden_dropout
							lowerCAmelCase_         =      attention_dropout

							lowerCAmelCase_         =      bos_token_id
							lowerCAmelCase_         =      eos_token_id
							lowerCAmelCase_         =      slow_but_exact

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





class       _lowerCAmelCase   (      __a			):
				_lowercase     						=version.parse('''1.12'''			)




				def __init__(   self							,							_UpperCamelCase							,							_UpperCamelCase = "default"							,							_UpperCamelCase = None							,							_UpperCamelCase = False							,							)							-> int:
							super().__init__(_UpperCamelCase							,							task=_UpperCamelCase							,							patching_specs=_UpperCamelCase							,							use_past=_UpperCamelCase	)
							if not getattr(self._config							,							"pad_token_id"							,							_UpperCamelCase	):
										# TODO: how to do that better?
										lowerCAmelCase_         =      0




				@property
				def 	__a		(   self	)							-> Mapping[str, Mapping[int, str]]:
							lowerCAmelCase_         =      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_(_UpperCamelCase							,							direction="inputs"							,							inverted_values_shape=_UpperCamelCase	)
										lowerCAmelCase_         =      {0: "batch", 1: "past_sequence + sequence"}
							else:
										lowerCAmelCase_         =      {0: "batch", 1: "sequence"}

							return common_inputs




				@property
				def 	__a		(   self	)							-> int:
							return self._config.n_layer




				@property
				def 	__a		(   self	)							-> int:
							return self._config.n_head




				@property
				def 	__a		(   self	)							-> float:
							return 1e-3




				def 	__a		(   self							,							_UpperCamelCase							,							_UpperCamelCase = -1							,							_UpperCamelCase = -1							,							_UpperCamelCase = False							,							_UpperCamelCase = None							,							)							-> Mapping[str, Any]:
							lowerCAmelCase_         =      super(_UpperCamelCase							,							self	).generate_dummy_inputs(
							    _UpperCamelCase							,							batch_size=_UpperCamelCase							,							seq_length=_UpperCamelCase							,							is_pair=_UpperCamelCase							,							framework=_UpperCamelCase	)

							# We need to order the input in the way they appears in the forward()
							lowerCAmelCase_         =      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

													lowerCAmelCase_      ,    lowerCAmelCase_         =      common_inputs["input_ids"].shape
													# Not using the same length for past_key_values
													lowerCAmelCase_         =      seqlen + 2
													lowerCAmelCase_         =      self._config.hidden_size // self.num_attention_heads
													lowerCAmelCase_         =      (
													    batch * self.num_attention_heads,
													    head_dim,
													    past_key_values_length,
													)
													lowerCAmelCase_         =      (
													    batch * self.num_attention_heads,
													    past_key_values_length,
													    head_dim,
													)
													lowerCAmelCase_         =      [
													    (torch.zeros(_UpperCamelCase	), torch.zeros(_UpperCamelCase	)) for _ in range(self.num_layers	)
													]

							lowerCAmelCase_         =      common_inputs["attention_mask"]
							if self.use_past:
										lowerCAmelCase_         =      ordered_inputs["attention_mask"].dtype
										lowerCAmelCase_         =      torch.cat(
										    [ordered_inputs["attention_mask"], torch.ones(_UpperCamelCase							,							_UpperCamelCase							,							dtype=_UpperCamelCase	)]							,							dim=1	)

							return ordered_inputs




				@property
				def 	__a		(   self	)							-> int:
							return 13



