


'''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      :List[str]			    =       logging.getLogger(__name__)

__UpperCAmelCase      :Optional[int]			    =       5_0  # max width of layer names
__UpperCAmelCase      :Dict			    =       7_0  # max width of quantizer names




def      _a     (						_lowercase :     List[str]  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Optional[Any]  							=      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      _a     (						_lowercase :     Optional[int]  ):





     '''simple docstring'''






     if args.calibrator == "max":
          __UpperCAmelCase     :			Union[str, Any]  							=      '''max'''
     elif args.calibrator == "percentile":
          if args.percentile is None:
               raise ValueError('''Specify --percentile when using percentile calibrator'''  )
          __UpperCAmelCase     :			Union[str, Any]  							=      '''histogram'''
     elif args.calibrator == "mse":
          __UpperCAmelCase     :			Dict  							=      '''histogram'''
     else:
          raise ValueError(F'Invalid calibrator {args.calibrator}'  )

     __UpperCAmelCase     :			str  							=      QuantDescriptor(num_bits=args.aprec      , calib_method=_lowercase  )
     __UpperCAmelCase     :			Union[str, Any]  							=      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      _a     (						_lowercase :     List[str]      , _lowercase :     Optional[int]      , _lowercase :     List[str]=False      , _lowercase :     Union[str, Any]=False  ):





     '''simple docstring'''






     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      _a     (						_lowercase :     Optional[int]  ):





     '''simple docstring'''






     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     (						_lowercase :     str      , _lowercase :     Any  ):





     '''simple docstring'''






     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      _a     (						_lowercase :     Dict      , _lowercase :     Optional[Any]  ):





     '''simple docstring'''






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

          __UpperCAmelCase     :			Tuple  							=      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      _a     (						_lowercase :     List[str]      , _lowercase :     List[Any]  ):





     '''simple docstring'''






     for name, mod in model.named_modules():
          if name.endswith('''.output.dense'''  ) and not name.endswith('''attention.output.dense'''  ):
               __UpperCAmelCase     :			Dict  							=      mod._input_quantizer._amax.data.detach().item()
               mod._input_quantizer._amax.data.detach().clamp_(max=_lowercase  )
               __UpperCAmelCase     :			int  							=      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     (						_lowercase :     Dict  ):





     '''simple docstring'''






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




def      _a     (						_lowercase :     List[str]  ):





     '''simple docstring'''






     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)
               __UpperCAmelCase     :			Any  							=      set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis  )
               __UpperCAmelCase     :			Optional[int]  							=      set(range(len(mod.weight.size()  )  )  ) - axis_set
               __UpperCAmelCase     :			Union[str, Any]  							=      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()}'  )
               __UpperCAmelCase     :			List[Any]  							=      amax




def      _a     (						_lowercase :     Tuple      , _lowercase :     int=25      , _lowercase :     Any=180      , _lowercase :     Tuple=None  ):





     '''simple docstring'''






     if ignore is None:
          __UpperCAmelCase     :			Optional[int]  							=      []
     elif not isinstance(_lowercase      , _lowercase  ):
          __UpperCAmelCase     :			Union[str, Any]  							=      [ignore]

     __UpperCAmelCase     :			int  							=      0
     for name, mod in model.named_modules():
          if not hasattr(_lowercase      , '''weight'''  ):
               continue
          __UpperCAmelCase     :			Tuple  							=      max(_lowercase      , len(_lowercase  )  )

     for name, mod in model.named_modules():
          __UpperCAmelCase     :			str  							=      getattr(_lowercase      , '''_input_quantizer'''      , _lowercase  )
          __UpperCAmelCase     :			Optional[Any]  							=      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
          __UpperCAmelCase     :			Tuple  							=      F'Act:{input_q.extra_repr()}'
          __UpperCAmelCase     :			Optional[int]  							=      F'Wgt:{weight_q.extra_repr()}'
          __UpperCAmelCase     :			Union[str, Any]  							=      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      _a     (						_lowercase :     str  ):





     '''simple docstring'''






     __UpperCAmelCase     :			List[Any]  							=      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      _a     (						_lowercase :     List[Any]      , _lowercase :     Optional[int]      , _lowercase :     Union[str, Any]      , _lowercase :     Dict      , _lowercase :     int  ):





     '''simple docstring'''






     __UpperCAmelCase     :			List[Any]  							=      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      _a     (						_lowercase :     Optional[int]      , _lowercase :     Optional[Any]      , _lowercase :     List[str]="both"      , **_lowercase :     Dict  ):





     '''simple docstring'''






     __UpperCAmelCase     :			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(_lowercase      , _lowercase      , '''_input_quantizer'''      , _lowercase      , _lowercase  )
          if which in ["weight", "both"]:
               set_quantizer(_lowercase      , _lowercase      , '''_weight_quantizer'''      , _lowercase      , _lowercase  )
     logger.info(_lowercase  )




def      _a     (						_lowercase :     Union[str, Any]      , _lowercase :     List[str]      , **_lowercase :     Union[str, Any]  ):





     '''simple docstring'''






     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  ):
                         __UpperCAmelCase     :			Optional[int]  							=      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_torch_available


__UpperCAmelCase      :Union[str, Any]			    =       {
    "configuration_rag": ["RagConfig"],
    "retrieval_rag": ["RagRetriever"],
    "tokenization_rag": ["RagTokenizer"],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase      :Dict			    =       [
          "RagModel",
          "RagPreTrainedModel",
          "RagSequenceForGeneration",
          "RagTokenForGeneration",
      ]

try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase      :Dict			    =       [
          "TFRagModel",
          "TFRagPreTrainedModel",
          "TFRagSequenceForGeneration",
          "TFRagTokenForGeneration",
      ]


if TYPE_CHECKING:
      from .configuration_rag import RagConfig
      from .retrieval_rag import RagRetriever
      from .tokenization_rag import RagTokenizer

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_rag import RagModel, RagPreTrainedModel, RagSequenceForGeneration, RagTokenForGeneration

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_rag import (
                TFRagModel,
                TFRagPreTrainedModel,
                TFRagSequenceForGeneration,
                TFRagTokenForGeneration,
            )

else:
      import sys

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