






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


SCREAMING_SNAKE_CASE__     =       logging.getLogger(__name__)

SCREAMING_SNAKE_CASE__     =       5_0  # max width of layer names
SCREAMING_SNAKE_CASE__     =       7_0  # max width of quantizer names







def    lowercase__			(				__UpperCamelCase   )->			int:
      UpperCamelCase      = 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   )->			Union[str, Any]:

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

      UpperCamelCase      = QuantDescriptor(num_bits=args.aprec		,						calib_method=__UpperCamelCase   )
      UpperCamelCase      = 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   )->			List[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   )->			Any:
      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   )->			str:

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

            UpperCamelCase      = 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   )->			Optional[Any]:

      for name, mod in model.named_modules():
            if name.endswith(""".output.dense"""   ) and not name.endswith("""attention.output.dense"""   ):
                  UpperCamelCase      = mod._input_quantizer._amax.data.detach().item()
                  mod._input_quantizer._amax.data.detach().clamp_(max=__UpperCamelCase   )
                  UpperCamelCase      = 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   )->			List[str]:

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







def    lowercase__			(				__UpperCamelCase   )->			int:

      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)
                  UpperCamelCase      = set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis   )
                  UpperCamelCase      = set(range(len(mod.weight.size()   )   )   ) - axis_set
                  UpperCamelCase      = 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()}"   )
                  UpperCamelCase      = amax







def    lowercase__			(				__UpperCamelCase		,						__UpperCamelCase=25		,						__UpperCamelCase=180		,						__UpperCamelCase=None   )->			str:

      if ignore is None:
            UpperCamelCase      = []
      elif not isinstance(__UpperCamelCase		,						__UpperCamelCase   ):
            UpperCamelCase      = [ignore]

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

      for name, mod in model.named_modules():
            UpperCamelCase      = getattr(__UpperCamelCase		,						"""_input_quantizer"""		,						__UpperCamelCase   )
            UpperCamelCase      = 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
            UpperCamelCase      = F"Act:{input_q.extra_repr()}"
            UpperCamelCase      = F"Wgt:{weight_q.extra_repr()}"
            UpperCamelCase      = 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   )->			List[Any]:
      UpperCamelCase      = 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   )->			Any:
      UpperCamelCase      = 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   )->			Optional[int]:
      UpperCamelCase      = 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   )->			Any:

      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   ):
                              UpperCamelCase      = F"Warning: changing {name:{name_width}}"
                              for k, v in kwargs.items():
                                    s += F" {k}={v}"
                                    setattr(__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase   )
                              logger.info(__UpperCamelCase   )








'''simple docstring'''

from typing import TYPE_CHECKING

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


SCREAMING_SNAKE_CASE__     =       {
    'configuration_mobilebert': [
        'MOBILEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'MobileBertConfig',
        'MobileBertOnnxConfig',
    ],
    'tokenization_mobilebert': ['MobileBertTokenizer'],
}

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

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     SCREAMING_SNAKE_CASE__     =       [
         'MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
         'MobileBertForMaskedLM',
         'MobileBertForMultipleChoice',
         'MobileBertForNextSentencePrediction',
         'MobileBertForPreTraining',
         'MobileBertForQuestionAnswering',
         'MobileBertForSequenceClassification',
         'MobileBertForTokenClassification',
         'MobileBertLayer',
         'MobileBertModel',
         'MobileBertPreTrainedModel',
         'load_tf_weights_in_mobilebert',
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     SCREAMING_SNAKE_CASE__     =       [
         'TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
         'TFMobileBertForMaskedLM',
         'TFMobileBertForMultipleChoice',
         'TFMobileBertForNextSentencePrediction',
         'TFMobileBertForPreTraining',
         'TFMobileBertForQuestionAnswering',
         'TFMobileBertForSequenceClassification',
         'TFMobileBertForTokenClassification',
         'TFMobileBertMainLayer',
         'TFMobileBertModel',
         'TFMobileBertPreTrainedModel',
     ]


if TYPE_CHECKING:
     from .configuration_mobilebert import (
         MOBILEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
         MobileBertConfig,
         MobileBertOnnxConfig,
     )
     from .tokenization_mobilebert import MobileBertTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_mobilebert_fast import MobileBertTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_mobilebert import (
              MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
              MobileBertForMaskedLM,
              MobileBertForMultipleChoice,
              MobileBertForNextSentencePrediction,
              MobileBertForPreTraining,
              MobileBertForQuestionAnswering,
              MobileBertForSequenceClassification,
              MobileBertForTokenClassification,
              MobileBertLayer,
              MobileBertModel,
              MobileBertPreTrainedModel,
              load_tf_weights_in_mobilebert,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_mobilebert import (
              TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFMobileBertForMaskedLM,
              TFMobileBertForMultipleChoice,
              TFMobileBertForNextSentencePrediction,
              TFMobileBertForPreTraining,
              TFMobileBertForQuestionAnswering,
              TFMobileBertForSequenceClassification,
              TFMobileBertForTokenClassification,
              TFMobileBertMainLayer,
              TFMobileBertModel,
              TFMobileBertPreTrainedModel,
          )

else:
     import sys

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