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


lowercase_         =       logging.getLogger(__name__)

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




def       A_      (       lowercase   )  ->				List[Any]:


     """simple docstring"""







     UpperCAmelCase_    :     Tuple            =	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   )  ->				Tuple:


     """simple docstring"""







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

     UpperCAmelCase_    :     List[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    ,			lowercase    ,			lowercase=False    ,			lowercase=False   )  ->				Optional[int]:


     """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   )  ->				List[Any]:


     """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    ,			lowercase   )  ->				int:


     """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    ,			lowercase   )  ->				Optional[Any]:


     """simple docstring"""







     def fusea(lowercase    ,			lowercase    ,			lowercase   ):
          for mod in [qq, qk, qv]:
               if not hasattr(lowercase    ,			"""_amax"""   ):
                    print("""          WARNING: NO AMAX BUFFER"""   )
                    return
          UpperCAmelCase_    :     Union[str, Any]            =	qq._amax.detach().item()
          UpperCAmelCase_    :     int            =	qk._amax.detach().item()
          UpperCAmelCase_    :     Tuple            =	qv._amax.detach().item()

          UpperCAmelCase_    :     str            =	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    ,			lowercase   )  ->				Any:


     """simple docstring"""







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


     """simple docstring"""







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




def       A_      (       lowercase   )  ->				Union[str, Any]:


     """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_    :     Optional[int]            =	set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis   )
               UpperCAmelCase_    :     Optional[Any]            =	set(range(len(mod.weight.size()   )   )   ) - axis_set
               UpperCAmelCase_    :     Tuple            =	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_    :     int            =	amax




def       A_      (       lowercase    ,			lowercase=25    ,			lowercase=180    ,			lowercase=None   )  ->				Dict:


     """simple docstring"""







     if ignore is None:
          UpperCAmelCase_    :     List[Any]            =	[]
     elif not isinstance(lowercase    ,			lowercase   ):
          UpperCAmelCase_    :     int            =	[ignore]

     UpperCAmelCase_    :     Union[str, Any]            =	0
     for name, mod in model.named_modules():
          if not hasattr(lowercase    ,			"""weight"""   ):
               continue
          UpperCAmelCase_    :     Union[str, Any]            =	max(lowercase    ,			len(lowercase   )   )

     for name, mod in model.named_modules():
          UpperCAmelCase_    :     List[Any]            =	getattr(lowercase    ,			"""_input_quantizer"""    ,			lowercase   )
          UpperCAmelCase_    :     List[str]            =	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_    :     Dict            =	f'''Act:{input_q.extra_repr()}'''
          UpperCAmelCase_    :     str            =	f'''Wgt:{weight_q.extra_repr()}'''
          UpperCAmelCase_    :     List[str]            =	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   )  ->				Tuple:


     """simple docstring"""







     UpperCAmelCase_    :     int            =	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    ,			lowercase    ,			lowercase    ,			lowercase    ,			lowercase   )  ->				str:


     """simple docstring"""







     UpperCAmelCase_    :     str            =	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    ,			lowercase    ,			lowercase="both"    ,			**lowercase   )  ->				Tuple:


     """simple docstring"""







     UpperCAmelCase_    :     Union[str, 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(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    ,			lowercase    ,			**lowercase   )  ->				Optional[int]:


     """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_    :     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"""







import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
       import torch

       from transformers.generation import DisjunctiveConstraint



@require_torch
class 		UpperCAmelCase_						(unittest.TestCase    ):
 """simple docstring"""




 def 	a   (			self							:			Any			)->       int:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[Any]            =	[[1, 2, 4], [1, 2, 3, 4]]
      UpperCAmelCase_    :     Union[str, Any]            =	DisjunctiveConstraint(a_			)
      self.assertTrue(isinstance(dc.token_ids  , a_			)			)

      with self.assertRaises(a_			):
           DisjunctiveConstraint(torch.LongTensor([[1, 2, 4], [1, 2, 3]]			)			)

      with self.assertRaises(a_			):
           DisjunctiveConstraint([torch.LongTensor([1, 2, 4]			), torch.LongTensor([1, 2, 3, 4, 5]			)]			)




 def 	a   (			self							:			List[str]			)->       List[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Tuple            =	[[1, 2], [1, 2, 3, 4]]

      with self.assertRaises(a_			):
           DisjunctiveConstraint(a_			)  # fails here




 def 	a   (			self							:			Dict			)->       List[str]:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[int]            =	[[1, 2, 3], [1, 2, 4]]

      UpperCAmelCase_    :     Dict            =	DisjunctiveConstraint(a_			)

      UpperCAmelCase_       ,UpperCAmelCase_       ,UpperCAmelCase_    :     int            =	dc.update(1			)
      UpperCAmelCase_    :     str            =	stepped is True and completed is False and reset is False
      self.assertTrue(a_			)
      self.assertTrue(not dc.completed			)
      self.assertTrue(dc.current_seq == [1]			)

      UpperCAmelCase_       ,UpperCAmelCase_       ,UpperCAmelCase_    :     Optional[int]            =	dc.update(2			)
      UpperCAmelCase_    :     Tuple            =	stepped is True and completed is False and reset is False
      self.assertTrue(a_			)
      self.assertTrue(not dc.completed			)
      self.assertTrue(dc.current_seq == [1, 2]			)

      UpperCAmelCase_       ,UpperCAmelCase_       ,UpperCAmelCase_    :     Union[str, Any]            =	dc.update(3			)
      UpperCAmelCase_    :     str            =	stepped is True and completed is True and reset is False
      self.assertTrue(a_			)
      self.assertTrue(dc.completed			)  # Completed!
      self.assertTrue(dc.current_seq == [1, 2, 3]			)






 def 	a   (			self							:			List[Any]			)->       List[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[int]            =	[[1, 2, 3], [1, 2, 4, 5], [1, 2, 5]]

      UpperCAmelCase_    :     Tuple            =	DisjunctiveConstraint(a_			)

      UpperCAmelCase_       ,UpperCAmelCase_       ,UpperCAmelCase_    :     List[Any]            =	dc.update(1			)
      self.assertTrue(not dc.completed			)
      self.assertTrue(dc.current_seq == [1]			)

      UpperCAmelCase_       ,UpperCAmelCase_       ,UpperCAmelCase_    :     Tuple            =	dc.update(2			)
      self.assertTrue(not dc.completed			)
      self.assertTrue(dc.current_seq == [1, 2]			)

      UpperCAmelCase_       ,UpperCAmelCase_       ,UpperCAmelCase_    :     List[str]            =	dc.update(4			)
      self.assertTrue(not dc.completed			)
      self.assertTrue(dc.current_seq == [1, 2, 4]			)

      UpperCAmelCase_       ,UpperCAmelCase_       ,UpperCAmelCase_    :     str            =	dc.update(5			)
      self.assertTrue(dc.completed			)  # Completed!
      self.assertTrue(dc.current_seq == [1, 2, 4, 5]			)

      dc.reset()

      UpperCAmelCase_       ,UpperCAmelCase_       ,UpperCAmelCase_    :     Any            =	dc.update(1			)
      self.assertTrue(not dc.completed			)
      self.assertTrue(dc.remaining() == 3			)
      self.assertTrue(dc.current_seq == [1]			)

      UpperCAmelCase_       ,UpperCAmelCase_       ,UpperCAmelCase_    :     Any            =	dc.update(2			)
      self.assertTrue(not dc.completed			)
      self.assertTrue(dc.remaining() == 2			)
      self.assertTrue(dc.current_seq == [1, 2]			)

      UpperCAmelCase_       ,UpperCAmelCase_       ,UpperCAmelCase_    :     List[str]            =	dc.update(5			)
      self.assertTrue(dc.completed			)  # Completed!
      self.assertTrue(dc.remaining() == 0			)
      self.assertTrue(dc.current_seq == [1, 2, 5]			)

