



from queue import Queue
from typing import TYPE_CHECKING, Optional


if TYPE_CHECKING:
							from ..models.auto import AutoTokenizer
class 			_snake_case     :




				def 							A__ (  self   :  List[str],   __lowercase   :  Optional[int]  ):
							raise NotImplementedError()




				def 							A__ (  self   :  List[Any]  ):
							raise NotImplementedError()
class 			_snake_case     (   lowercase__):




				def __init__(  self   :  List[str],   __lowercase   :  "AutoTokenizer",   __lowercase   :  bool = False,   **__lowercase   :  List[Any]  ):
							lowercase__   					=			tokenizer
							lowercase__   					=			skip_prompt
							lowercase__   					=			decode_kwargs

							# variables used in the streaming process
							lowercase__   					=			[]
							lowercase__   					=			0
							lowercase__   					=			True




				def 							A__ (  self   :  Union[str, Any],   __lowercase   :  Any  ):
							if len(value.shape  ) > 1 and value.shape[0] > 1:
										raise ValueError("TextStreamer only supports batch size 1"  )
							elif len(value.shape  ) > 1:
										lowercase__   					=			value[0]

							if self.skip_prompt and self.next_tokens_are_prompt:
										lowercase__   					=			False
										return

							# Add the new token to the cache and decodes the entire thing.
							self.token_cache.extend(value.tolist()  )
							lowercase__   					=			self.tokenizer.decode(self.token_cache,   **self.decode_kwargs  )

							# After the symbol for a new line, we flush the cache.
							if text.endswith("\n"  ):
										lowercase__   					=			text[self.print_len :]
										lowercase__   					=			[]
										lowercase__   					=			0
							# If the last token is a CJK character, we print the characters.
							elif len(__lowercase  ) > 0 and self._is_chinese_char(ord(text[-1]  )  ):
										lowercase__   					=			text[self.print_len :]
										self.print_len += len(__lowercase  )
							# Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words,
							# which may change with the subsequent token -- there are probably smarter ways to do this!)
							else:
										lowercase__   					=			text[self.print_len : text.rfind(" "  ) + 1]
										self.print_len += len(__lowercase  )

							self.on_finalized_text(__lowercase  )




				def 							A__ (  self   :  Optional[Any]  ):
							# Flush the cache, if it exists
							if len(self.token_cache  ) > 0:
										lowercase__   					=			self.tokenizer.decode(self.token_cache,   **self.decode_kwargs  )
										lowercase__   					=			text[self.print_len :]
										lowercase__   					=			[]
										lowercase__   					=			0
							else:
										lowercase__   					=			""

							lowercase__   					=			True
							self.on_finalized_text(__lowercase,   stream_end=__lowercase  )




				def 							A__ (  self   :  Optional[Any],   __lowercase   :  str,   __lowercase   :  bool = False  ):
							print(__lowercase,   flush=__lowercase,   end="" if not stream_end else None  )




				def 							A__ (  self   :  Optional[int],   __lowercase   :  Union[str, Any]  ):
							# This defines a "chinese character" as anything in the CJK Unicode block:
							#   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
							#
							# Note that the CJK Unicode block is NOT all Japanese and Korean characters,
							# despite its name. The modern Korean Hangul alphabet is a different block,
							# as is Japanese Hiragana and Katakana. Those alphabets are used to write
							# space-separated words, so they are not treated specially and handled
							# like the all of the other languages.
							if (
							    (cp >= 0x4_e_0_0 and cp <= 0x9_f_f_f)
							    or (cp >= 0x3_4_0_0 and cp <= 0x4_d_b_f)  #
							    or (cp >= 0x2_0_0_0_0 and cp <= 0x2_a_6_d_f)  #
							    or (cp >= 0x2_a_7_0_0 and cp <= 0x2_b_7_3_f)  #
							    or (cp >= 0x2_b_7_4_0 and cp <= 0x2_b_8_1_f)  #
							    or (cp >= 0x2_b_8_2_0 and cp <= 0x2_c_e_a_f)  #
							    or (cp >= 0xf_9_0_0 and cp <= 0xf_a_f_f)
							    or (cp >= 0x2_f_8_0_0 and cp <= 0x2_f_a_1_f)  #
							):  #
										return True

							return False
class 			_snake_case     (   lowercase__):




				def __init__(  self   :  Dict,   __lowercase   :  "AutoTokenizer",   __lowercase   :  bool = False,   __lowercase   :  Optional[float] = None,   **__lowercase   :  str  ):
							super().__init__(__lowercase,   __lowercase,   **__lowercase  )
							lowercase__   					=			Queue()
							lowercase__   					=			None
							lowercase__   					=			timeout




				def 							A__ (  self   :  List[str],   __lowercase   :  str,   __lowercase   :  bool = False  ):
							self.text_queue.put(__lowercase,   timeout=self.timeout  )
							if stream_end:
										self.text_queue.put(self.stop_signal,   timeout=self.timeout  )




				def __iter__(  self   :  str  ):
							return self




				def 							A__ (  self   :  Dict  ):
							lowercase__   					=			self.text_queue.get(timeout=self.timeout  )
							if value == self.stop_signal:
										raise StopIteration()
							else:
										return value








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


lowercase_						=       logging.get_logger(__name__)



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



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					):
			lowercase__   					=			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)
			):
						lowercase__   					=			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:
						lowercase__   					=			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:
						lowercase__   					=			pt_tuple_key[:-1] + ("embedding",)
						return renamed_pt_tuple_key, pt_tensor

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

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

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

			# old PyTorch layer norm bias
			lowercase__   					=			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 __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_=42					):
			# Step 1: Convert pytorch tensor to numpy
			lowercase__   					=			{k: v.numpy() for k, v in pt_state_dict.items()}

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

			lowercase__   					=			flatten_dict(SCREAMING_SNAKE_CASE_					)
			lowercase__   					=			{}

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

						# Correctly rename weight parameters
						lowercase__					,    lowercase__   					=			rename_key_and_reshape_tensor(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					)

						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
						lowercase__   					=			jnp.asarray(SCREAMING_SNAKE_CASE_					)

			return unflatten_dict(SCREAMING_SNAKE_CASE_					)




