

'''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					:			Tuple							=    50  # max width of layer names
_a					:			str							=    70  # max width of quantizer names


def    lowerCamelCase__		(    SCREAMING_SNAKE_CASE   : Optional[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__		(    SCREAMING_SNAKE_CASE   : 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=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__		(    SCREAMING_SNAKE_CASE   : Optional[int]							, SCREAMING_SNAKE_CASE   : List[Any]							, SCREAMING_SNAKE_CASE   : Any=False							, SCREAMING_SNAKE_CASE   : Optional[int]=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(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__		(    SCREAMING_SNAKE_CASE   : Tuple    ):
  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__		(    SCREAMING_SNAKE_CASE   : List[str]							, SCREAMING_SNAKE_CASE   : int    ):
  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__		(    SCREAMING_SNAKE_CASE   : Any							, SCREAMING_SNAKE_CASE   : str    ):

  def fusea(SCREAMING_SNAKE_CASE   : List[Any]							, SCREAMING_SNAKE_CASE   : Optional[Any]							, SCREAMING_SNAKE_CASE   : List[str]    ):
    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__		(    SCREAMING_SNAKE_CASE   : Any							, SCREAMING_SNAKE_CASE   : 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=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__		(    SCREAMING_SNAKE_CASE   : Optional[Any]    ):

  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__		(    SCREAMING_SNAKE_CASE   : Dict    ):

  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__		(    SCREAMING_SNAKE_CASE   : str							, SCREAMING_SNAKE_CASE   : Tuple=25							, SCREAMING_SNAKE_CASE   : Tuple=180							, SCREAMING_SNAKE_CASE   : int=None    ):

  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__		(    SCREAMING_SNAKE_CASE   : Optional[Any]    ):
  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__		(    SCREAMING_SNAKE_CASE   : Dict							, SCREAMING_SNAKE_CASE   : List[Any]							, SCREAMING_SNAKE_CASE   : Optional[Any]							, SCREAMING_SNAKE_CASE   : List[Any]							, SCREAMING_SNAKE_CASE   : Optional[Any]    ):
  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__		(    SCREAMING_SNAKE_CASE   : Tuple							, SCREAMING_SNAKE_CASE   : Optional[Any]							, SCREAMING_SNAKE_CASE   : List[str]="both"							, **SCREAMING_SNAKE_CASE   : Optional[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__		(    SCREAMING_SNAKE_CASE   : Optional[Any]							, SCREAMING_SNAKE_CASE   : str							, **SCREAMING_SNAKE_CASE   : List[str]    ):

  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__    )






import importlib
import shutil
import threading
import warnings
from typing import List

import fsspec
import fsspec.asyn

from . import compression
from .hffilesystem import HfFileSystem


SCREAMING_SNAKE_CASE							    =							importlib.util.find_spec('s3fs') is not None

if _has_safs:
  from .safilesystem import SaFileSystem  # noqa: F401

SCREAMING_SNAKE_CASE							    =							[
    compression.BzaFileSystem,
    compression.GzipFileSystem,
    compression.LzaFileSystem,
    compression.XzFileSystem,
    compression.ZstdFileSystem,
]

# Register custom filesystems
for fs_class in COMPRESSION_FILESYSTEMS + [HfFileSystem]:
  if fs_class.protocol in fsspec.registry and fsspec.registry[fs_class.protocol] is not fs_class:
    warnings.warn(f'''A filesystem protocol was already set for {fs_class.protocol} and will be overwritten.''')
  fsspec.register_implementation(fs_class.protocol, fs_class, clobber=True)




def  a   (lowerCAmelCase__   ):
   if "://" in dataset_path:
      __a					=  dataset_path.split("""://"""   )[1]
   return dataset_path




def  a   (lowerCAmelCase__   ):
   if fs is not None and fs.protocol != "file":
      return True
   else:
      return False




def  a   (lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__   ):
   __a					=  not is_remote_filesystem(lowerCAmelCase__   )
   if is_local:
      # LocalFileSystem.mv does copy + rm, it is more efficient to simply move a local directory
      shutil.move(fs._strip_protocol(lowerCAmelCase__   )		,				fs._strip_protocol(lowerCAmelCase__   )   )
   else:
      fs.mv(lowerCAmelCase__		,				lowerCAmelCase__		,				recursive=lowerCAmelCase__   )




def  a   ():
   if hasattr(fsspec.asyn		,				"""reset_lock"""   ):
      # for future fsspec>2022.05.0
      fsspec.asyn.reset_lock()
   else:
      __a					=  None
      __a					=  None
      __a					=  threading.Lock()



