
'''simple docstring'''





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    					=   logging.getLogger(__name__)

UpperCAmelCase    					=   50  # max width of layer names
UpperCAmelCase    					=   70  # max width of quantizer names

def      A (       A_							:					List[Any]       ):
						snake_case       :     Dict =    parser.add_argument_group('''quant_trainer arguments'''       )
						group.add_argument('''--wprec'''    ,							type=A_    ,							default=8    ,							help='''weight precision'''       )
						group.add_argument('''--aprec'''    ,							type=A_    ,							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=A_    ,							nargs='''+'''    ,							help='''disable quantizers by keyword'''       )
						group.add_argument('''--quant-disable-layer-module'''    ,							type=A_    ,							help='''disable quantizers by keyword under layer.'''       )
						group.add_argument('''--quant-enable-layer-module'''    ,							type=A_    ,							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=A_    ,							type=A_    ,							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=A_    ,							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      A (       A_							:					Union[str, Any]       ):

						if args.calibrator == "max":
												snake_case       :     Tuple =    '''max'''
						elif args.calibrator == "percentile":
												if args.percentile is None:
																		raise ValueError('''Specify --percentile when using percentile calibrator'''       )
												snake_case       :     List[Any] =    '''histogram'''
						elif args.calibrator == "mse":
												snake_case       :     int =    '''histogram'''
						else:
												raise ValueError(F"""Invalid calibrator {args.calibrator}"""       )

						snake_case       :     List[Any] =    QuantDescriptor(num_bits=args.aprec    ,							calib_method=A_       )
						snake_case       :     List[Any] =    QuantDescriptor(num_bits=args.wprec    ,							axis=(None if args.quant_per_tensor else (0,))       )
						quant_nn.QuantLinear.set_default_quant_desc_input(A_       )
						quant_nn.QuantLinear.set_default_quant_desc_weight(A_       )

def      A (       A_							:					Any    ,							A_							:					Union[str, Any]    ,							A_							:					Dict=False    ,							A_							:					Optional[Any]=False       ):
						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(A_    ,							['''embeddings''']    ,							which='''weight'''    ,							_disabled=A_       )

												if args.quant_disable:
																		set_quantizer_by_name(A_    ,							['''''']    ,							_disabled=A_       )

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

												if args.quant_disable_layer_module:
																		set_quantizer_by_name(A_    ,							[R'''layer.\d+.''' + args.quant_disable_layer_module]    ,							_disabled=A_       )

												if args.quant_enable_layer_module:
																		set_quantizer_by_name(A_    ,							[R'''layer.\d+.''' + args.quant_enable_layer_module]    ,							_disabled=A_       )

												if args.recalibrate_weights:
																		recalibrate_weights(A_       )

												if args.fuse_qkv:
																		fuse_qkv(A_    ,							A_       )

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

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

def      A (       A_							:					Union[str, 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      A (       A_							:					List[Any]    ,							A_							:					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(A_       )

def      A (       A_							:					int    ,							A_							:					Optional[Any]       ):

						def fusea(A_							:					List[Any]    ,							A_							:					Dict    ,							A_							:					Optional[Any]       ):
												for mod in [qq, qk, qv]:
																		if not hasattr(A_    ,							'''_amax'''       ):
																								print('''          WARNING: NO AMAX BUFFER'''       )
																								return
												snake_case       :     List[str] =    qq._amax.detach().item()
												snake_case       :     Any =    qk._amax.detach().item()
												snake_case       :     Union[str, Any] =    qv._amax.detach().item()

												snake_case       :     Union[str, Any] =    max(A_    ,							A_    ,							A_       )
												qq._amax.fill_(A_       )
												qk._amax.fill_(A_       )
												qv._amax.fill_(A_       )
												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      A (       A_							:					Union[str, Any]    ,							A_							:					Tuple       ):

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

def      A (       A_							:					Optional[Any]       ):

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

def      A (       A_							:					Optional[Any]       ):

						for name, mod in model.named_modules():
												if hasattr(A_    ,							'''_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)
																		snake_case       :     Optional[Any] =    set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis       )
																		snake_case       :     Optional[Any] =    set(range(len(mod.weight.size()       )       )       ) - axis_set
																		snake_case       :     str =    pytorch_quantization.utils.reduce_amax(mod.weight    ,							axis=A_    ,							keepdims=A_       ).detach()
																		logger.info(F"""RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}"""       )
																		snake_case       :     Any =    amax

def      A (       A_							:					List[str]    ,							A_							:					Any=25    ,							A_							:					List[str]=180    ,							A_							:					List[Any]=None       ):

						if ignore is None:
												snake_case       :     Tuple =    []
						elif not isinstance(A_    ,							A_       ):
												snake_case       :     str =    [ignore]

						snake_case       :     str =    0
						for name, mod in model.named_modules():
												if not hasattr(A_    ,							'''weight'''       ):
																		continue
												snake_case       :     Optional[Any] =    max(A_    ,							len(A_       )       )

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

def      A (       A_							:					Union[str, Any]       ):
						snake_case       :     Dict =    0
						for name, mod in model.named_modules():
												if isinstance(A_    ,							pytorch_quantization.nn.TensorQuantizer       ):
																		print(F"""{name:80} {mod}"""       )
																		count += 1
						print(F"""{count} TensorQuantizers found in model"""       )

def      A (       A_							:					List[Any]    ,							A_							:					Optional[int]    ,							A_							:					int    ,							A_							:					str    ,							A_							:					int       ):
						snake_case       :     Tuple =    getattr(A_    ,							A_    ,							A_       )
						if quantizer_mod is not None:
												assert hasattr(A_    ,							A_       )
												setattr(A_    ,							A_    ,							A_       )
						else:
												logger.warning(F"""{name} has no {quantizer}"""       )

def      A (       A_							:					List[Any]    ,							A_							:					List[Any]    ,							A_							:					Optional[int]="both"    ,							**A_							:					Dict       ):
						snake_case       :     int =    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(A_    ,							A_    ,							'''_input_quantizer'''    ,							A_    ,							A_       )
												if which in ["weight", "both"]:
																		set_quantizer(A_    ,							A_    ,							'''_weight_quantizer'''    ,							A_    ,							A_       )
						logger.info(A_       )





def      A (       A_							:					List[str]    ,							A_							:					int    ,							**A_							:					Optional[int]       ):

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



'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase    					=   {
    "configuration_convbert": ["CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "ConvBertConfig", "ConvBertOnnxConfig"],
    "tokenization_convbert": ["ConvBertTokenizer"],
}

try:
						if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						UpperCAmelCase    					=   ["ConvBertTokenizerFast"]

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						UpperCAmelCase    					=   [
						    "CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
						    "ConvBertForMaskedLM",
						    "ConvBertForMultipleChoice",
						    "ConvBertForQuestionAnswering",
						    "ConvBertForSequenceClassification",
						    "ConvBertForTokenClassification",
						    "ConvBertLayer",
						    "ConvBertModel",
						    "ConvBertPreTrainedModel",
						    "load_tf_weights_in_convbert",
						]


try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						UpperCAmelCase    					=   [
						    "TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
						    "TFConvBertForMaskedLM",
						    "TFConvBertForMultipleChoice",
						    "TFConvBertForQuestionAnswering",
						    "TFConvBertForSequenceClassification",
						    "TFConvBertForTokenClassification",
						    "TFConvBertLayer",
						    "TFConvBertModel",
						    "TFConvBertPreTrainedModel",
						]


if TYPE_CHECKING:
						from .configuration_convbert import CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, ConvBertConfig, ConvBertOnnxConfig
						from .tokenization_convbert import ConvBertTokenizer

						try:
												if not is_tokenizers_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_convbert_fast import ConvBertTokenizerFast

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_convbert import (
												    CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
												    ConvBertForMaskedLM,
												    ConvBertForMultipleChoice,
												    ConvBertForQuestionAnswering,
												    ConvBertForSequenceClassification,
												    ConvBertForTokenClassification,
												    ConvBertLayer,
												    ConvBertModel,
												    ConvBertPreTrainedModel,
												    load_tf_weights_in_convbert,
												)

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_convbert import (
												    TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
												    TFConvBertForMaskedLM,
												    TFConvBertForMultipleChoice,
												    TFConvBertForQuestionAnswering,
												    TFConvBertForSequenceClassification,
												    TFConvBertForTokenClassification,
												    TFConvBertLayer,
												    TFConvBertModel,
												    TFConvBertPreTrainedModel,
												)


else:
						import sys

						UpperCAmelCase    					=   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

