




import collections

from .utils import ExplicitEnum, is_torch_available, logging


if is_torch_available():
							import torch


UpperCamelCase__  =		logging.get_logger(__name__)


class        __SCREAMING_SNAKE_CASE							:

					def __init__(       self				, __lowercase				, __lowercase=21				, __lowercase=[]				, __lowercase=None      ):







						'''simple docstring'''




						lowerCamelCase_							:	Union[str, Any]							     =							model
						lowerCamelCase_							:	Any							     =							trace_batch_nums
						lowerCamelCase_							:	Dict							     =							abort_after_batch_num

						# keep a LIFO buffer of frames to dump as soon as inf/nan is encountered to give context to the problem emergence
						lowerCamelCase_							:	str							     =							collections.deque([]				, __lowercase      )
						lowerCamelCase_							:	str							     =							[]
						lowerCamelCase_							:	int							     =							0
						lowerCamelCase_							:	int							     =							0
						lowerCamelCase_							:	Optional[int]							     =							False
						lowerCamelCase_							:	Optional[int]							     =							'                 '

						self.analyse_model()

						self.register_forward_hook()

					def 			lowercase_							(       self				, __lowercase=None      ):







						'''simple docstring'''




						if frame is not None:
							self.expand_frame(__lowercase      )
						self.frames.append('\n'.join(self.frame      )      )
						lowerCamelCase_							:	List[str]							     =							[]  # start a new frame

					def 			lowercase_							(       self				, __lowercase      ):







						'''simple docstring'''




						self.frame.append(__lowercase      )

					def 			lowercase_							(       self      ):







						'''simple docstring'''




						print('\n'.join(self.frames      )      )
						lowerCamelCase_							:	Dict							     =							[]

					def 			lowercase_							(       self      ):







						'''simple docstring'''




						lowerCamelCase_							:	Dict							     =							[]

					def 			lowercase_							(       self      ):







						'''simple docstring'''




						print(f"\nDetected inf/nan during batch_number={self.batch_number}"      )
						print(f"Last {len(self.frames      )} forward frames:"      )
						print(f"{'abs min':8} {'abs max':8} metadata"      )
						print('\n'.join(self.frames      )      )
						print('\n\n'      )
						lowerCamelCase_							:	List[Any]							     =							[]

					def 			lowercase_							(       self      ):







						'''simple docstring'''




						lowerCamelCase_							:	int							     =							{m: name for name, m in self.model.named_modules()}
						# self.longest_module_name = max(len(v) for v in self.module_names.values())

					def 			lowercase_							(       self				, __lowercase				, __lowercase      ):







						'''simple docstring'''




						if torch.is_tensor(__lowercase      ):
							self.expand_frame(get_abs_min_max(__lowercase				, __lowercase      )      )
							if detect_overflow(__lowercase				, __lowercase      ):
								lowerCamelCase_							:	Optional[int]							     =							True
						elif var is None:
							self.expand_frame(f"{'None':>17} {ctx}"      )
						else:
							self.expand_frame(f"{'not a tensor':>17} {ctx}"      )

					def 			lowercase_							(       self      ):







						'''simple docstring'''




						self.expand_frame(f"\n\n{self.prefix} *** Starting batch number={self.batch_number} ***"      )
						self.expand_frame(f"{'abs min':8} {'abs max':8} metadata"      )

					def 			lowercase_							(       self      ):







						'''simple docstring'''




						self.expand_frame(f"{self.prefix} *** Finished batch number={self.batch_number-1} ***\n\n"      )

					def 			lowercase_							(       self				, __lowercase				, __lowercase				, __lowercase      ):







						'''simple docstring'''




						self.expand_frame(f"{self.prefix} {self.module_names[module]} {module.__class__.__name__}"      )

						# params
						for name, p in module.named_parameters(recurse=__lowercase      ):
							self.analyse_variable(__lowercase				, __lowercase      )

						# inputs
						if isinstance(__lowercase				, __lowercase      ):
							for i, x in enumerate(__lowercase      ):
								self.analyse_variable(__lowercase				, f"input[{i}]"      )
						else:
							self.analyse_variable(__lowercase				, 'input'      )

						# outputs
						if isinstance(__lowercase				, __lowercase      ):
							for i, x in enumerate(__lowercase      ):
								# possibly a tuple of tuples
								if isinstance(__lowercase				, __lowercase      ):
									for j, y in enumerate(__lowercase      ):
										self.analyse_variable(__lowercase				, f"output[{i}][{j}]"      )
								else:
									self.analyse_variable(__lowercase				, f"output[{i}]"      )
						else:
							self.analyse_variable(__lowercase				, 'output'      )

						self.save_frame()

					def 			lowercase_							(       self      ):







						'''simple docstring'''




						self.model.apply(self._register_forward_hook      )

					def 			lowercase_							(       self				, __lowercase      ):







						'''simple docstring'''




						module.register_forward_hook(self.forward_hook      )





					def 			lowercase_							(       self				, __lowercase				, __lowercase				, __lowercase      ):







						'''simple docstring'''




						lowerCamelCase_							:	Union[str, Any]							     =							False

						lowerCamelCase_							:	int							     =							True if self.batch_number in self.trace_batch_nums else False
						if trace_mode:
							self.reset_saved_frames()

						if self.total_calls == 0:
							self.batch_start_frame()
						self.total_calls += 1

						# count batch numbers - the very first forward hook of the batch will be called when the
						# batch completes - i.e. it gets called very last - we know this batch has finished
						if module == self.model:
							self.batch_number += 1
							lowerCamelCase_							:	Tuple							     =							True

						self.create_frame(__lowercase				, __lowercase				, __lowercase      )

						# if last_frame_of_batch:
						#     self.batch_end_frame()

						if trace_mode:
							self.trace_frames()

						if last_frame_of_batch:
							self.batch_start_frame()

						if self.detected_overflow and not trace_mode:
							self.dump_saved_frames()

							# now we can abort, as it's pointless to continue running
							raise ValueError(
							    'DebugUnderflowOverflow: inf/nan detected, aborting as there is no point running further. '
							    'Please scroll up above this traceback to see the activation values prior to this event.'      )

						# abort after certain batch if requested to do so
						if self.abort_after_batch_num is not None and self.batch_number > self.abort_after_batch_num:
							raise ValueError(
							    f"DebugUnderflowOverflow: aborting after {self.batch_number} batches due to"
							    f" `abort_after_batch_num={self.abort_after_batch_num}` arg"      )






def       snake_case      (							A		,     A    )							->							Any:




	'''simple docstring'''


	lowerCamelCase_							:	Union[str, Any]							     =							var.abs()
	return f"{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}"






def       snake_case      (							A		,     A    )							->							Dict:




	'''simple docstring'''


	lowerCamelCase_							:	Optional[Any]							     =							False
	if torch.isnan(A    ).any().item():
		lowerCamelCase_							:	List[Any]							     =							True
		print(f"{ctx} has nans"    )
	if torch.isinf(A    ).any().item():
		lowerCamelCase_							:	Optional[Any]							     =							True
		print(f"{ctx} has infs"    )

	# if needed to monitor large elements can enable the following
	if 0:  # and detected:
		lowerCamelCase_							:	Optional[int]							     =							var[torch.ge(var.abs()		,     100    )]
		if naaa.numel() > 0:
			print(f"{ctx}:  n100={naaa.numel()}"    )
		lowerCamelCase_							:	Union[str, Any]							     =							var[torch.ge(var.abs()		,     1000    )]
		if naaaa.numel() > 0:
			print(f"{ctx}: n1000={naaaa.numel()}"    )
		lowerCamelCase_							:	str							     =							var[torch.ge(var.abs()		,     10000    )]
		if naaaaa.numel() > 0:
			print(f"{ctx}: n10000={naaaaa.numel()}"    )

	if 0:
		print(f"min={var.min():9.2e} max={var.max():9.2e}"    )

	if 0:
		print(f"min={var.min():9.2e} max={var.max():9.2e} var={var.var():9.2e} mean={var.mean():9.2e} ({ctx})"    )

	return detected



class        __SCREAMING_SNAKE_CASE							(	lowerCamelCase						):
					_lowerCamelCase							=				"""underflow_overflow"""
					_lowerCamelCase							=				"""tpu_metrics_debug"""






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_tokenizers_available


UpperCamelCase__  =		{}

try:
							if not is_sentencepiece_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCamelCase__  =		['CpmTokenizer']

try:
							if not is_tokenizers_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCamelCase__  =		['CpmTokenizerFast']


if TYPE_CHECKING:
							try:
														if not is_sentencepiece_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tokenization_cpm import CpmTokenizer

							try:
														if not is_tokenizers_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tokenization_cpm_fast import CpmTokenizerFast

else:
							import sys

							UpperCamelCase__  =		_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
