





from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


UpperCAmelCase__  :     Dict   			=logging.get_logger(__name__)

UpperCAmelCase__  :     Any   			={
    '''google/umt5-small''': '''https://huggingface.co/google/umt5-small/resolve/main/config.json''',
    # See all umt5 models at https://huggingface.co/models?filter=umt5
}

class  __A       (  a      ):
				__A								=			"""umt5"""
				__A								=			["""past_key_values"""]


				def __init__(						self				,	UpperCAmelCase_=250112				,	UpperCAmelCase_=512				,	UpperCAmelCase_=64				,	UpperCAmelCase_=1024				,	UpperCAmelCase_=8				,	UpperCAmelCase_=None				,	UpperCAmelCase_=6				,	UpperCAmelCase_=32				,	UpperCAmelCase_=128				,	UpperCAmelCase_=0.1				,	UpperCAmelCase_=1E-6				,	UpperCAmelCase_=1.0				,	UpperCAmelCase_="gated-gelu"				,	UpperCAmelCase_=True				,	UpperCAmelCase_=True				,	UpperCAmelCase_="T5Tokenizer"				,	UpperCAmelCase_=True				,	UpperCAmelCase_=0				,	UpperCAmelCase_=1				,	UpperCAmelCase_=0				,	**UpperCAmelCase_				,	):
								super().__init__(
								    is_encoder_decoder=UpperCAmelCase_				,	tokenizer_class=UpperCAmelCase_				,	tie_word_embeddings=UpperCAmelCase_				,	pad_token_id=UpperCAmelCase_				,	eos_token_id=UpperCAmelCase_				,	decoder_start_token_id=UpperCAmelCase_				,	**UpperCAmelCase_				,	)
								lowerCamelCase			=vocab_size
								lowerCamelCase			=d_model
								lowerCamelCase			=d_kv
								lowerCamelCase			=d_ff
								lowerCamelCase			=num_layers
								lowerCamelCase			=(
								    num_decoder_layers if num_decoder_layers is not None else self.num_layers
								)  # default = symmetry
								lowerCamelCase			=num_heads
								lowerCamelCase			=relative_attention_num_buckets
								lowerCamelCase			=relative_attention_max_distance
								lowerCamelCase			=dropout_rate
								lowerCamelCase			=layer_norm_epsilon
								lowerCamelCase			=initializer_factor
								lowerCamelCase			=feed_forward_proj
								lowerCamelCase			=use_cache

								lowerCamelCase			=self.feed_forward_proj.split("""-"""							)
								lowerCamelCase			=act_info[-1]
								lowerCamelCase			=act_info[0] == """gated"""

								if len(UpperCAmelCase_							) > 1 and act_info[0] != "gated" or len(UpperCAmelCase_							) > 2:
												raise ValueError(
												    f"""`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."""
												    """Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. """
												    """'gated-gelu' or 'relu'"""							)

								if feed_forward_proj == "gated-gelu":
												lowerCamelCase			="""gelu_new"""


				@property
				def _snake_case    (						self							):
								return self.d_model


				@property
				def _snake_case    (						self							):
								return self.num_heads


				@property
				def _snake_case    (						self							):
								return self.num_layers




class  __A       (  a      ):


				@property
				# Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.inputs
				def _snake_case    (						self							):
								lowerCamelCase			={
								    """input_ids""": {0: """batch""", 1: """encoder_sequence"""},
								    """attention_mask""": {0: """batch""", 1: """encoder_sequence"""},
								}
								if self.use_past:
												lowerCamelCase			="""past_encoder_sequence + sequence"""
												lowerCamelCase			={0: """batch"""}
												lowerCamelCase			={0: """batch""", 1: """past_decoder_sequence + sequence"""}
								else:
												lowerCamelCase			={0: """batch""", 1: """decoder_sequence"""}
												lowerCamelCase			={0: """batch""", 1: """decoder_sequence"""}

								if self.use_past:
												self.fill_with_past_key_values_(UpperCAmelCase_				,	direction="""inputs"""							)

								return common_inputs


				@property
				# Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.default_onnx_opset
				def _snake_case    (						self							):
								return 13


				@property
				def _snake_case    (						self							):
								return 5E-4












import logging
import re

import pytorch_quantization
import pytorch_quantization.nn as quant_nn
import torch
from pytorch_quantization import calib
from pytorch_quantization.tensor_quant import QuantDescriptor


UpperCAmelCase__  :     Optional[int]   			=logging.getLogger(__name__)

UpperCAmelCase__  :     Tuple   			=50  # max width of layer names
UpperCAmelCase__  :     List[str]   			=70  # max width of quantizer names





def        _lowercase   (  _UpperCAmelCase			)  ->					List[str]:
				lowerCamelCase			=parser.add_argument_group("""quant_trainer arguments"""			)
				group.add_argument("""--wprec"""     , type=_UpperCAmelCase     , default=8     , help="""weight precision"""			)
				group.add_argument("""--aprec"""     , type=_UpperCAmelCase     , default=8     , help="""activation precision"""			)
				group.add_argument("""--quant-per-tensor"""     , action="""store_true"""     , help="""per tensor weight scaling"""			)
				group.add_argument("""--quant-disable"""     , action="""store_true"""     , help="""disable all quantizers"""			)
				group.add_argument("""--quant-disable-embeddings"""     , action="""store_true"""     , help="""disable all embeddings quantizers"""			)
				group.add_argument("""--quant-disable-keyword"""     , type=_UpperCAmelCase     , nargs="""+"""     , help="""disable quantizers by keyword"""			)
				group.add_argument("""--quant-disable-layer-module"""     , type=_UpperCAmelCase     , help="""disable quantizers by keyword under layer."""			)
				group.add_argument("""--quant-enable-layer-module"""     , type=_UpperCAmelCase     , help="""enable quantizers by keyword under layer"""			)
				group.add_argument("""--calibrator"""     , default="""max"""     , help="""which quantization range calibrator to use"""			)
				group.add_argument("""--percentile"""     , default=_UpperCAmelCase     , type=_UpperCAmelCase     , help="""percentile for PercentileCalibrator"""			)
				group.add_argument("""--fuse-qkv"""     , action="""store_true"""     , help="""use the same scale factor for qkv"""			)
				group.add_argument("""--clip-gelu"""     , metavar="""N"""     , type=_UpperCAmelCase     , help="""clip gelu output maximum value to N"""			)
				group.add_argument(
				    """--recalibrate-weights"""     , action="""store_true"""     , help=(
				        """recalibrate weight amaxes by taking the max of the weights."""
				        """ amaxes will be computed with the current quantization granularity (axis)."""
				    )     , )





def        _lowercase   (  _UpperCAmelCase			)  ->					Dict:

				if args.calibrator == "max":
								lowerCamelCase			="""max"""
				elif args.calibrator == "percentile":
								if args.percentile is None:
												raise ValueError("""Specify --percentile when using percentile calibrator"""			)
								lowerCamelCase			="""histogram"""
				elif args.calibrator == "mse":
								lowerCamelCase			="""histogram"""
				else:
								raise ValueError(F"""Invalid calibrator {args.calibrator}"""			)

				lowerCamelCase			=QuantDescriptor(num_bits=args.aprec     , calib_method=_UpperCAmelCase			)
				lowerCamelCase			=QuantDescriptor(num_bits=args.wprec     , axis=(None if args.quant_per_tensor else (0,))			)
				quant_nn.QuantLinear.set_default_quant_desc_input(_UpperCAmelCase			)
				quant_nn.QuantLinear.set_default_quant_desc_weight(_UpperCAmelCase			)





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase=False     , _UpperCAmelCase=False			)  ->					int:
				logger.info("""Configuring Model for Quantization"""			)
				logger.info(F"""using quantization package {pytorch_quantization.__file__}"""			)

				if not calib:
								if args.quant_disable_embeddings:
												set_quantizer_by_name(_UpperCAmelCase     , ["""embeddings"""]     , which="""weight"""     , _disabled=_UpperCAmelCase			)

								if args.quant_disable:
												set_quantizer_by_name(_UpperCAmelCase     , [""""""]     , _disabled=_UpperCAmelCase			)

								if args.quant_disable_keyword:
												set_quantizer_by_name(_UpperCAmelCase     , args.quant_disable_keyword     , _disabled=_UpperCAmelCase			)

								if args.quant_disable_layer_module:
												set_quantizer_by_name(_UpperCAmelCase     , [r"""layer.\d+.""" + args.quant_disable_layer_module]     , _disabled=_UpperCAmelCase			)

								if args.quant_enable_layer_module:
												set_quantizer_by_name(_UpperCAmelCase     , [r"""layer.\d+.""" + args.quant_enable_layer_module]     , _disabled=_UpperCAmelCase			)

								if args.recalibrate_weights:
												recalibrate_weights(_UpperCAmelCase			)

								if args.fuse_qkv:
												fuse_qkv(_UpperCAmelCase     , _UpperCAmelCase			)

				if args.clip_gelu:
								clip_gelu(_UpperCAmelCase     , args.clip_gelu			)

				# if args.local_rank in [-1, 0] and not calib:
				print_quant_summary(_UpperCAmelCase			)





def        _lowercase   (  _UpperCAmelCase			)  ->					Optional[Any]:
				logger.info("""Enabling Calibration"""			)
				for name, module in model.named_modules():
								if name.endswith("""_quantizer"""			):
												if module._calibrator is not None:
																module.disable_quant()
																module.enable_calib()
												else:
																module.disable()
												logger.info(F"""{name:80}: {module}"""			)





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase			)  ->					Optional[int]:
				logger.info("""Loading calibrated amax"""			)
				for name, module in model.named_modules():
								if name.endswith("""_quantizer"""			):
												if module._calibrator is not None:
																if isinstance(module._calibrator     , calib.MaxCalibrator			):
																				module.load_calib_amax()
																else:
																				module.load_calib_amax("""percentile"""     , percentile=args.percentile			)
																module.enable_quant()
																module.disable_calib()
												else:
																module.enable()
				model.cuda()
				print_quant_summary(_UpperCAmelCase			)





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase			)  ->					Tuple:

				def fusea(_UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase			):
								for mod in [qq, qk, qv]:
												if not hasattr(_UpperCAmelCase     , """_amax"""			):
																print("""          WARNING: NO AMAX BUFFER"""			)
																return
								lowerCamelCase			=qq._amax.detach().item()
								lowerCamelCase			=qk._amax.detach().item()
								lowerCamelCase			=qv._amax.detach().item()

								lowerCamelCase			=max(_UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase			)
								qq._amax.fill_(_UpperCAmelCase			)
								qk._amax.fill_(_UpperCAmelCase			)
								qv._amax.fill_(_UpperCAmelCase			)
								logger.info(F"""          q={q:5.2f} k={k:5.2f} v={v:5.2f} -> {amax:5.2f}"""			)

				for name, mod in model.named_modules():
								if name.endswith(""".attention.self"""			):
												logger.info(F"""FUSE_QKV: {name:{name_width}}"""			)
												fusea(mod.matmul_q_input_quantizer     , mod.matmul_k_input_quantizer     , mod.matmul_v_input_quantizer			)
												if args.quant_per_tensor:
																fusea(mod.query._weight_quantizer     , mod.key._weight_quantizer     , mod.value._weight_quantizer			)





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase			)  ->					int:

				for name, mod in model.named_modules():
								if name.endswith(""".output.dense"""			) and not name.endswith("""attention.output.dense"""			):
												lowerCamelCase			=mod._input_quantizer._amax.data.detach().item()
												mod._input_quantizer._amax.data.detach().clamp_(max=_UpperCAmelCase			)
												lowerCamelCase			=mod._input_quantizer._amax.data.detach().item()
												logger.info(F"""CLIP_GELU: {name:{name_width}} amax: {amax_init:5.2f} -> {amax:5.2f}"""			)





def        _lowercase   (  _UpperCAmelCase			)  ->					Dict:

				for name, mod in model.named_modules():
								if hasattr(_UpperCAmelCase     , """_weight_quantizer"""			) and mod._weight_quantizer.axis is not None:
												lowerCamelCase			=mod.weight.shape[0]
												lowerCamelCase			=mod._weight_quantizer._amax.detach()
												lowerCamelCase			=torch.ones(_UpperCAmelCase     , dtype=amax.dtype     , device=amax.device			) * amax
												print(F"""expanding {name} {amax} -> {mod._weight_quantizer._amax}"""			)





def        _lowercase   (  _UpperCAmelCase			)  ->					List[str]:

				for name, mod in model.named_modules():
								if hasattr(_UpperCAmelCase     , """_weight_quantizer"""			):
												if not hasattr(mod.weight_quantizer     , """_amax"""			):
																print("""RECALIB: {name:{name_width}} WARNING: NO AMAX BUFFER"""			)
																continue

												# determine which axes to reduce across
												# e.g. a 4D tensor quantized per axis 0 should reduce over (1,2,3)
												lowerCamelCase			=set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis			)
												lowerCamelCase			=set(range(len(mod.weight.size()			)			)			) - axis_set
												lowerCamelCase			=pytorch_quantization.utils.reduce_amax(mod.weight     , axis=_UpperCAmelCase     , keepdims=_UpperCAmelCase			).detach()
												logger.info(F"""RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}"""			)
												lowerCamelCase			=amax





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase=25     , _UpperCAmelCase=1_80     , _UpperCAmelCase=None			)  ->					Dict:

				if ignore is None:
								lowerCamelCase			=[]
				elif not isinstance(_UpperCAmelCase     , _UpperCAmelCase			):
								lowerCamelCase			=[ignore]

				lowerCamelCase			=0
				for name, mod in model.named_modules():
								if not hasattr(_UpperCAmelCase     , """weight"""			):
												continue
								lowerCamelCase			=max(_UpperCAmelCase     , len(_UpperCAmelCase			)			)

				for name, mod in model.named_modules():
								lowerCamelCase			=getattr(_UpperCAmelCase     , """_input_quantizer"""     , _UpperCAmelCase			)
								lowerCamelCase			=getattr(_UpperCAmelCase     , """_weight_quantizer"""     , _UpperCAmelCase			)
								if not hasattr(_UpperCAmelCase     , """weight"""			):
												continue
								if type(_UpperCAmelCase			) in ignore:
												continue
								if [True for s in ignore if type(_UpperCAmelCase			) is str and s in name]:
												continue
								lowerCamelCase			=F"""Act:{input_q.extra_repr()}"""
								lowerCamelCase			=F"""Wgt:{weight_q.extra_repr()}"""
								lowerCamelCase			=F"""{name:{name_width}} {act_str} {wgt_str}"""
								if len(_UpperCAmelCase			) <= line_width:
												logger.info(_UpperCAmelCase			)
								else:
												logger.info(F"""{name:{name_width}} {act_str}"""			)
												logger.info(F"""{"  ":{name_width}} {wgt_str}"""			)





def        _lowercase   (  _UpperCAmelCase			)  ->					Dict:
				lowerCamelCase			=0
				for name, mod in model.named_modules():
								if isinstance(_UpperCAmelCase     , pytorch_quantization.nn.TensorQuantizer			):
												print(F"""{name:80} {mod}"""			)
												count += 1
				print(F"""{count} TensorQuantizers found in model"""			)





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase			)  ->					Union[str, Any]:
				lowerCamelCase			=getattr(_UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase			)
				if quantizer_mod is not None:
								assert hasattr(_UpperCAmelCase     , _UpperCAmelCase			)
								setattr(_UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase			)
				else:
								logger.warning(F"""{name} has no {quantizer}"""			)





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase="both"     , **_UpperCAmelCase			)  ->					List[str]:
				lowerCamelCase			=F"""Warning: changing {which} quantizers of {name:{qname_width}}"""
				for k, v in kwargs.items():
								s += F""" {k}={v}"""
								if which in ["input", "both"]:
												set_quantizer(_UpperCAmelCase     , _UpperCAmelCase     , """_input_quantizer"""     , _UpperCAmelCase     , _UpperCAmelCase			)
								if which in ["weight", "both"]:
												set_quantizer(_UpperCAmelCase     , _UpperCAmelCase     , """_weight_quantizer"""     , _UpperCAmelCase     , _UpperCAmelCase			)
				logger.info(_UpperCAmelCase			)





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase     , **_UpperCAmelCase			)  ->					int:

				for name, mod in model.named_modules():
								if hasattr(_UpperCAmelCase     , """_input_quantizer"""			) or hasattr(_UpperCAmelCase     , """_weight_quantizer"""			):
												for n in names:
																if re.search(_UpperCAmelCase     , _UpperCAmelCase			):
																				set_quantizers(_UpperCAmelCase     , _UpperCAmelCase     , **_UpperCAmelCase			)
								elif name.endswith("""_quantizer"""			):
												for n in names:
																if re.search(_UpperCAmelCase     , _UpperCAmelCase			):
																				lowerCamelCase			=F"""Warning: changing {name:{name_width}}"""
																				for k, v in kwargs.items():
																								s += F""" {k}={v}"""
																								setattr(_UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase			)
																				logger.info(_UpperCAmelCase			)





