






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


__snake_case												=  logging.getLogger(__name__)

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


def   _A   (				SCREAMING_SNAKE_CASE__  :	Optional[int]       ):
    UpperCamelCase       :Optional[int]								=   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   _A   (				SCREAMING_SNAKE_CASE__  :	Tuple       ):

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

    UpperCamelCase       :str								=   QuantDescriptor(num_bits=args.aprec     ,  calib_method=UpperCamelCase__       )
    UpperCamelCase       :Dict								=   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   _A   (				SCREAMING_SNAKE_CASE__  :	Optional[Any]     ,  SCREAMING_SNAKE_CASE__  :	Dict     ,  SCREAMING_SNAKE_CASE__  :	Optional[Any]=False     ,  SCREAMING_SNAKE_CASE__  :	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(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   _A   (				SCREAMING_SNAKE_CASE__  :	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   _A   (				SCREAMING_SNAKE_CASE__  :	str     ,  SCREAMING_SNAKE_CASE__  :	Dict       ):
    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   _A   (				SCREAMING_SNAKE_CASE__  :	List[str]     ,  SCREAMING_SNAKE_CASE__  :	List[str]       ):

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

        UpperCamelCase       :List[Any]								=   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   _A   (				SCREAMING_SNAKE_CASE__  :	int     ,  SCREAMING_SNAKE_CASE__  :	str       ):

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

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


def   _A   (				SCREAMING_SNAKE_CASE__  :	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)
            UpperCamelCase       :List[str]								=   set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis       )
            UpperCamelCase       :List[str]								=   set(range(len(mod.weight.size()       )       )       ) - axis_set
            UpperCamelCase       :Union[str, Any]								=   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       :List[Any]								=   amax


def   _A   (				SCREAMING_SNAKE_CASE__  :	Optional[int]     ,  SCREAMING_SNAKE_CASE__  :	Union[str, Any]=25     ,  SCREAMING_SNAKE_CASE__  :	Tuple=180     ,  SCREAMING_SNAKE_CASE__  :	List[Any]=None       ):

    if ignore is None:
        UpperCamelCase       :int								=   []
    elif not isinstance(UpperCamelCase__     ,  UpperCamelCase__       ):
        UpperCamelCase       :Optional[int]								=   [ignore]

    UpperCamelCase       :int								=   0
    for name, mod in model.named_modules():
        if not hasattr(UpperCamelCase__     ,  '''weight'''       ):
            continue
        UpperCamelCase       :str								=   max(UpperCamelCase__     ,  len(UpperCamelCase__       )       )

    for name, mod in model.named_modules():
        UpperCamelCase       :List[Any]								=   getattr(UpperCamelCase__     ,  '''_input_quantizer'''     ,  UpperCamelCase__       )
        UpperCamelCase       :List[Any]								=   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       :Tuple								=   F'''Act:{input_q.extra_repr()}'''
        UpperCamelCase       :Optional[Any]								=   F'''Wgt:{weight_q.extra_repr()}'''
        UpperCamelCase       :Dict								=   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   _A   (				SCREAMING_SNAKE_CASE__  :	Dict       ):
    UpperCamelCase       :Optional[int]								=   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   _A   (				SCREAMING_SNAKE_CASE__  :	List[str]     ,  SCREAMING_SNAKE_CASE__  :	Tuple     ,  SCREAMING_SNAKE_CASE__  :	List[str]     ,  SCREAMING_SNAKE_CASE__  :	Tuple     ,  SCREAMING_SNAKE_CASE__  :	Dict       ):
    UpperCamelCase       :Dict								=   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   _A   (				SCREAMING_SNAKE_CASE__  :	List[Any]     ,  SCREAMING_SNAKE_CASE__  :	Union[str, Any]     ,  SCREAMING_SNAKE_CASE__  :	List[Any]="both"     ,  **SCREAMING_SNAKE_CASE__  :	Tuple       ):
    UpperCamelCase       :Any								=   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   _A   (				SCREAMING_SNAKE_CASE__  :	Optional[int]     ,  SCREAMING_SNAKE_CASE__  :	Union[str, Any]     ,  **SCREAMING_SNAKE_CASE__  :	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       :Dict								=   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"""




import warnings

from ...utils import logging
from .image_processing_clip import CLIPImageProcessor


_snake_case          =   logging.get_logger(__name__)






class        UpperCamelCase			(	snake_case_      ):


 def __init__(     self		:						Any							,						*UpperCAmelCase__		:						Optional[Any]							,						**UpperCAmelCase__		:						Tuple							)					->       None:
   warnings.warn(
       """The class CLIPFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"""
       """ use CLIPImageProcessor instead."""							,						UpperCAmelCase__							,						)
   super().__init__(*UpperCAmelCase__							,						**UpperCAmelCase__							)

