






'''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


__A							:							int  							=		logging.getLogger(__name__)

__A							:							List[Any]  							=		50  # max width of layer names
__A							:							Union[str, Any]  							=		70  # max width of quantizer names






def    lowerCamelCase_     (    lowercase__):
 lowerCamelCase__  =     parser.add_argument_group("quant_trainer arguments")
 group.add_argument("--wprec"				,     type=lowercase__				,     default=8				,     help="weight precision")
 group.add_argument("--aprec"				,     type=lowercase__				,     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=lowercase__				,     nargs="+"				,     help="disable quantizers by keyword")
 group.add_argument("--quant-disable-layer-module"				,     type=lowercase__				,     help="disable quantizers by keyword under layer.")
 group.add_argument("--quant-enable-layer-module"				,     type=lowercase__				,     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=lowercase__				,     type=lowercase__				,     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=lowercase__				,     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    lowerCamelCase_     (    lowercase__):

 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=lowercase__)
 lowerCamelCase__  =     QuantDescriptor(num_bits=args.wprec				,     axis=(None if args.quant_per_tensor else (0,)))
 quant_nn.QuantLinear.set_default_quant_desc_input(lowercase__)
 quant_nn.QuantLinear.set_default_quant_desc_weight(lowercase__)






def    lowerCamelCase_     (    lowercase__				,     lowercase__				,     lowercase__=False				,     lowercase__=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(lowercase__				,     ["embeddings"]				,     which="weight"				,     _disabled=lowercase__)

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

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

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

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

  if args.recalibrate_weights:
   recalibrate_weights(lowercase__)

  if args.fuse_qkv:
   fuse_qkv(lowercase__				,     lowercase__)

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

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






def    lowerCamelCase_     (    lowercase__):
 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    lowerCamelCase_     (    lowercase__				,     lowercase__):
 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(lowercase__)






def    lowerCamelCase_     (    lowercase__				,     lowercase__):

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

  lowerCamelCase__  =     max(lowercase__				,     lowercase__				,     lowercase__)
  qq._amax.fill_(lowercase__)
  qk._amax.fill_(lowercase__)
  qv._amax.fill_(lowercase__)
  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    lowerCamelCase_     (    lowercase__				,     lowercase__):

 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=lowercase__)
   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    lowerCamelCase_     (    lowercase__):

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






def    lowerCamelCase_     (    lowercase__):

 for name, mod in model.named_modules():
  if hasattr(lowercase__				,     "_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=lowercase__				,     keepdims=lowercase__).detach()
   logger.info(F'''RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}''')
   lowerCamelCase__  =     amax






def    lowerCamelCase_     (    lowercase__				,     lowercase__=25				,     lowercase__=180				,     lowercase__=None):

 if ignore is None:
  lowerCamelCase__  =     []
 elif not isinstance(lowercase__				,     lowercase__):
  lowerCamelCase__  =     [ignore]

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

 for name, mod in model.named_modules():
  lowerCamelCase__  =     getattr(lowercase__				,     "_input_quantizer"				,     lowercase__)
  lowerCamelCase__  =     getattr(lowercase__				,     "_weight_quantizer"				,     lowercase__)
  if not hasattr(lowercase__				,     "weight"):
   continue
  if type(lowercase__) in ignore:
   continue
  if [True for s in ignore if type(lowercase__) 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(lowercase__) <= line_width:
   logger.info(lowercase__)
  else:
   logger.info(F'''{name:{name_width}} {act_str}''')
   logger.info(F'''{"  ":{name_width}} {wgt_str}''')






def    lowerCamelCase_     (    lowercase__):
 lowerCamelCase__  =     0
 for name, mod in model.named_modules():
  if isinstance(lowercase__				,     pytorch_quantization.nn.TensorQuantizer):
   print(F'''{name:80} {mod}''')
   count += 1
 print(F'''{count} TensorQuantizers found in model''')






def    lowerCamelCase_     (    lowercase__				,     lowercase__				,     lowercase__				,     lowercase__				,     lowercase__):
 lowerCamelCase__  =     getattr(lowercase__				,     lowercase__				,     lowercase__)
 if quantizer_mod is not None:
  assert hasattr(lowercase__				,     lowercase__)
  setattr(lowercase__				,     lowercase__				,     lowercase__)
 else:
  logger.warning(F'''{name} has no {quantizer}''')






def    lowerCamelCase_     (    lowercase__				,     lowercase__				,     lowercase__="both"				,     **lowercase__):
 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(lowercase__				,     lowercase__				,     "_input_quantizer"				,     lowercase__				,     lowercase__)
  if which in ["weight", "both"]:
   set_quantizer(lowercase__				,     lowercase__				,     "_weight_quantizer"				,     lowercase__				,     lowercase__)
 logger.info(lowercase__)






def    lowerCamelCase_     (    lowercase__				,     lowercase__				,     **lowercase__):

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









'''simple docstring'''





from typing import TYPE_CHECKING

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


__A							:							Tuple  							=		{
    """configuration_deberta""": ["""DEBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP""", """DebertaConfig""", """DebertaOnnxConfig"""],
    """tokenization_deberta""": ["""DebertaTokenizer"""],
}

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __A							:							int  							=		["""DebertaTokenizerFast"""]

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __A							:							str  							=		[
     """DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """DebertaForMaskedLM""",
     """DebertaForQuestionAnswering""",
     """DebertaForSequenceClassification""",
     """DebertaForTokenClassification""",
     """DebertaModel""",
     """DebertaPreTrainedModel""",
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __A							:							Union[str, Any]  							=		[
     """TF_DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """TFDebertaForMaskedLM""",
     """TFDebertaForQuestionAnswering""",
     """TFDebertaForSequenceClassification""",
     """TFDebertaForTokenClassification""",
     """TFDebertaModel""",
     """TFDebertaPreTrainedModel""",
 ]


if TYPE_CHECKING:
 from .configuration_deberta import DEBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, DebertaConfig, DebertaOnnxConfig
 from .tokenization_deberta import DebertaTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_deberta_fast import DebertaTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_deberta import (
      DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
      DebertaForMaskedLM,
      DebertaForQuestionAnswering,
      DebertaForSequenceClassification,
      DebertaForTokenClassification,
      DebertaModel,
      DebertaPreTrainedModel,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_deberta import (
      TF_DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFDebertaForMaskedLM,
      TFDebertaForQuestionAnswering,
      TFDebertaForSequenceClassification,
      TFDebertaForTokenClassification,
      TFDebertaModel,
      TFDebertaPreTrainedModel,
  )


else:
 import sys

 __A							:							List[Any]  							=		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

