





'''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:					Tuple							=				logging.getLogger(__name__)

_A:					Optional[int]							=				50  # max width of layer names
_A:					List[Any]							=				70  # max width of quantizer names







def _lowerCAmelCase     (      _lowerCAmelCase    )->     List[Any]:
 __UpperCAmelCase					=   parser.add_argument_group('quant_trainer arguments'    )
 group.add_argument('--wprec'       ,    type=_lowerCamelCase       ,    default=8       ,    help='weight precision'    )
 group.add_argument('--aprec'       ,    type=_lowerCamelCase       ,    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=_lowerCamelCase       ,    nargs='+'       ,    help='disable quantizers by keyword'    )
 group.add_argument('--quant-disable-layer-module'       ,    type=_lowerCamelCase       ,    help='disable quantizers by keyword under layer.'    )
 group.add_argument('--quant-enable-layer-module'       ,    type=_lowerCamelCase       ,    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=_lowerCamelCase       ,    type=_lowerCamelCase       ,    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=_lowerCamelCase       ,    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     (      _lowerCAmelCase    )->     Dict:





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







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase=False       ,    _lowerCAmelCase=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(_lowerCamelCase       ,    ['embeddings']       ,    which='weight'       ,    _disabled=_lowerCamelCase    )

  if args.quant_disable:
   set_quantizer_by_name(_lowerCamelCase       ,    ['']       ,    _disabled=_lowerCamelCase    )

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

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

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

  if args.recalibrate_weights:
   recalibrate_weights(_lowerCamelCase    )

  if args.fuse_qkv:
   fuse_qkv(_lowerCamelCase       ,    _lowerCamelCase    )

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

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







def _lowerCAmelCase     (      _lowerCAmelCase    )->     Optional[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 _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase    )->     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(_lowerCamelCase    )







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase    )->     str:





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

  __UpperCAmelCase					=   max(_lowerCamelCase       ,    _lowerCamelCase       ,    _lowerCamelCase    )
  qq._amax.fill_(_lowerCamelCase    )
  qk._amax.fill_(_lowerCamelCase    )
  qv._amax.fill_(_lowerCamelCase    )
  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     (      _lowerCAmelCase       ,    _lowerCAmelCase    )->     Tuple:





 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=_lowerCamelCase    )
   __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 _lowerCAmelCase     (      _lowerCAmelCase    )->     Dict:





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







def _lowerCAmelCase     (      _lowerCAmelCase    )->     int:





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







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase=25       ,    _lowerCAmelCase=1_80       ,    _lowerCAmelCase=None    )->     Any:





 if ignore is None:
  __UpperCAmelCase					=   []
 elif not isinstance(_lowerCamelCase       ,    _lowerCamelCase    ):
  __UpperCAmelCase					=   [ignore]

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

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







def _lowerCAmelCase     (      _lowerCAmelCase    )->     Tuple:
 __UpperCAmelCase					=   0
 for name, mod in model.named_modules():
  if isinstance(_lowerCamelCase       ,    pytorch_quantization.nn.TensorQuantizer    ):
   print(F'{name:80} {mod}'    )
   count += 1
 print(F'{count} TensorQuantizers found in model'    )







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase    )->     Dict:
 __UpperCAmelCase					=   getattr(_lowerCamelCase       ,    _lowerCamelCase       ,    _lowerCamelCase    )
 if quantizer_mod is not None:
  assert hasattr(_lowerCamelCase       ,    _lowerCamelCase    )
  setattr(_lowerCamelCase       ,    _lowerCamelCase       ,    _lowerCamelCase    )
 else:
  logger.warning(F'{name} has no {quantizer}'    )







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







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase       ,    **_lowerCAmelCase    )->     Optional[Any]:





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












'''simple docstring'''




from ...configuration_utils import PretrainedConfig
from ...utils import logging


_A:					Union[str, Any]							=				logging.get_logger(__name__)

_A:					List[str]							=				{
    """weiweishi/roc-bert-base-zh""": """https://huggingface.co/weiweishi/roc-bert-base-zh/resolve/main/config.json""",
}






class UpperCAmelCase							(		UpperCAmelCase_							):
    _A :			int						   =       """roc_bert"""



    def __init__(  self     ,			__A=30_522     ,			__A=768     ,			__A=12     ,			__A=12     ,			__A=3_072     ,			__A="gelu"     ,			__A=0.1     ,			__A=0.1     ,			__A=512     ,			__A=2     ,			__A=0.0_2     ,			__A=1E-12     ,			__A=True     ,			__A=0     ,			__A="absolute"     ,			__A=None     ,			__A=True     ,			__A=True     ,			__A=768     ,			__A=910     ,			__A=512     ,			__A=24_858     ,			__A=True     ,			**__A     ,			):
     __UpperCAmelCase					=   vocab_size
     __UpperCAmelCase					=   max_position_embeddings
     __UpperCAmelCase					=   hidden_size
     __UpperCAmelCase					=   num_hidden_layers
     __UpperCAmelCase					=   num_attention_heads
     __UpperCAmelCase					=   intermediate_size
     __UpperCAmelCase					=   hidden_act
     __UpperCAmelCase					=   hidden_dropout_prob
     __UpperCAmelCase					=   attention_probs_dropout_prob
     __UpperCAmelCase					=   initializer_range
     __UpperCAmelCase					=   type_vocab_size
     __UpperCAmelCase					=   layer_norm_eps
     __UpperCAmelCase					=   use_cache
     __UpperCAmelCase					=   enable_pronunciation
     __UpperCAmelCase					=   enable_shape
     __UpperCAmelCase					=   pronunciation_embed_dim
     __UpperCAmelCase					=   pronunciation_vocab_size
     __UpperCAmelCase					=   shape_embed_dim
     __UpperCAmelCase					=   shape_vocab_size
     __UpperCAmelCase					=   concat_input
     __UpperCAmelCase					=   position_embedding_type
     __UpperCAmelCase					=   classifier_dropout
     super().__init__(pad_token_id=__A     ,			**__A					)





