



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

__lowercase       :List[str]       =  50  # max width of layer names
__lowercase       :str       =  70  # max width of quantizer names







def       UpperCAmelCase							(     _lowerCamelCase      :  Optional[int]     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					parser.add_argument_group("quant_trainer arguments"     )
						group.add_argument("--wprec"			,   type=__A			,   default=8			,   help="weight precision"     )
						group.add_argument("--aprec"			,   type=__A			,   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=__A			,   nargs="+"			,   help="disable quantizers by keyword"     )
						group.add_argument("--quant-disable-layer-module"			,   type=__A			,   help="disable quantizers by keyword under layer."     )
						group.add_argument("--quant-enable-layer-module"			,   type=__A			,   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=__A			,   type=__A			,   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=__A			,   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       UpperCAmelCase							(     _lowerCamelCase      :  int     ):





						'''simple docstring'''


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

						SCREAMING_SNAKE_CASE__     :				int   =					QuantDescriptor(num_bits=args.aprec			,   calib_method=__A     )
						SCREAMING_SNAKE_CASE__     :				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(__A     )
						quant_nn.QuantLinear.set_default_quant_desc_weight(__A     )







def       UpperCAmelCase							(     _lowerCamelCase      :  Dict			,   _lowerCamelCase      :  Tuple			,   _lowerCamelCase      :  str=False			,   _lowerCamelCase      :  Tuple=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(__A			,   ["embeddings"]			,   which="weight"			,   _disabled=__A     )

												if args.quant_disable:
																		set_quantizer_by_name(__A			,   [""]			,   _disabled=__A     )

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

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

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

												if args.recalibrate_weights:
																		recalibrate_weights(__A     )

												if args.fuse_qkv:
																		fuse_qkv(__A			,   __A     )

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

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







def       UpperCAmelCase							(     _lowerCamelCase      :  Optional[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       UpperCAmelCase							(     _lowerCamelCase      :  List[Any]			,   _lowerCamelCase      :  Union[str, 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(__A     )







def       UpperCAmelCase							(     _lowerCamelCase      :  List[Any]			,   _lowerCamelCase      :  Dict     ):





						'''simple docstring'''


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

												SCREAMING_SNAKE_CASE__     :				str   =					max(__A			,   __A			,   __A     )
												qq._amax.fill_(__A     )
												qk._amax.fill_(__A     )
												qv._amax.fill_(__A     )
												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       UpperCAmelCase							(     _lowerCamelCase      :  Dict			,   _lowerCamelCase      :  Optional[Any]     ):





						'''simple docstring'''


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





						'''simple docstring'''


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







def       UpperCAmelCase							(     _lowerCamelCase      :  Optional[int]     ):





						'''simple docstring'''


						for name, mod in model.named_modules():
												if hasattr(__A			,   "_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)
																		SCREAMING_SNAKE_CASE__     :				List[str]   =					set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis     )
																		SCREAMING_SNAKE_CASE__     :				List[Any]   =					set(range(len(mod.weight.size()     )     )     ) - axis_set
																		SCREAMING_SNAKE_CASE__     :				str   =					pytorch_quantization.utils.reduce_amax(mod.weight			,   axis=__A			,   keepdims=__A     ).detach()
																		logger.info(f"""RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}"""     )
																		SCREAMING_SNAKE_CASE__     :				List[str]   =					amax







def       UpperCAmelCase							(     _lowerCamelCase      :  Any			,   _lowerCamelCase      :  Any=25			,   _lowerCamelCase      :  Any=180			,   _lowerCamelCase      :  Tuple=None     ):





						'''simple docstring'''


						if ignore is None:
												SCREAMING_SNAKE_CASE__     :				Any   =					[]
						elif not isinstance(__A			,   __A     ):
												SCREAMING_SNAKE_CASE__     :				List[Any]   =					[ignore]

						SCREAMING_SNAKE_CASE__     :				int   =					0
						for name, mod in model.named_modules():
												if not hasattr(__A			,   "weight"     ):
																		continue
												SCREAMING_SNAKE_CASE__     :				List[str]   =					max(__A			,   len(__A     )     )

						for name, mod in model.named_modules():
												SCREAMING_SNAKE_CASE__     :				int   =					getattr(__A			,   "_input_quantizer"			,   __A     )
												SCREAMING_SNAKE_CASE__     :				Dict   =					getattr(__A			,   "_weight_quantizer"			,   __A     )
												if not hasattr(__A			,   "weight"     ):
																		continue
												if type(__A     ) in ignore:
																		continue
												if [True for s in ignore if type(__A     ) is str and s in name]:
																		continue
												SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					f"""Act:{input_q.extra_repr()}"""
												SCREAMING_SNAKE_CASE__     :				Dict   =					f"""Wgt:{weight_q.extra_repr()}"""
												SCREAMING_SNAKE_CASE__     :				int   =					f"""{name:{name_width}} {act_str} {wgt_str}"""
												if len(__A     ) <= line_width:
																		logger.info(__A     )
												else:
																		logger.info(f"""{name:{name_width}} {act_str}"""     )
																		logger.info(f"""{'  ':{name_width}} {wgt_str}"""     )







def       UpperCAmelCase							(     _lowerCamelCase      :  Optional[Any]     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				Dict   =					0
						for name, mod in model.named_modules():
												if isinstance(__A			,   pytorch_quantization.nn.TensorQuantizer     ):
																		print(f"""{name:80} {mod}"""     )
																		count += 1
						print(f"""{count} TensorQuantizers found in model"""     )







def       UpperCAmelCase							(     _lowerCamelCase      :  List[str]			,   _lowerCamelCase      :  Any			,   _lowerCamelCase      :  Optional[Any]			,   _lowerCamelCase      :  Optional[int]			,   _lowerCamelCase      :  Optional[Any]     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				Dict   =					getattr(__A			,   __A			,   __A     )
						if quantizer_mod is not None:
												assert hasattr(__A			,   __A     )
												setattr(__A			,   __A			,   __A     )
						else:
												logger.warning(f"""{name} has no {quantizer}"""     )







def       UpperCAmelCase							(     _lowerCamelCase      :  Any			,   _lowerCamelCase      :  List[Any]			,   _lowerCamelCase      :  List[Any]="both"			,   **_lowerCamelCase      :  int     ):





						'''simple docstring'''


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







def       UpperCAmelCase							(     _lowerCamelCase      :  str			,   _lowerCamelCase      :  List[str]			,   **_lowerCamelCase      :  Any     ):





						'''simple docstring'''


						for name, mod in model.named_modules():
												if hasattr(__A			,   "_input_quantizer"     ) or hasattr(__A			,   "_weight_quantizer"     ):
																		for n in names:
																								if re.search(__A			,   __A     ):
																														set_quantizers(__A			,   __A			,   **__A     )
												elif name.endswith("_quantizer"     ):
																		for n in names:
																								if re.search(__A			,   __A     ):
																														SCREAMING_SNAKE_CASE__     :				Optional[int]   =					f"""Warning: changing {name:{name_width}}"""
																														for k, v in kwargs.items():
																																				s += f""" {k}={v}"""
																																				setattr(__A			,   __A			,   __A     )
																														logger.info(__A     )






import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding



class    _a     (				lowercase__							):






							"""simple docstring"""

							snake_case_							 =						["image_processor", "tokenizer"]
							snake_case_							 =						"CLIPImageProcessor"
							snake_case_							 =						("CLIPTokenizer", "CLIPTokenizerFast")

							def __init__(      self				:     Any							,	a				:     List[Any]=None							,	a				:     Any=None							,	**a				:     int    )	->int:
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					None
													if "feature_extractor" in kwargs:
																			warnings.warn(
																			    "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
																			    " instead."							,	a							,	)
																			SCREAMING_SNAKE_CASE__     :				List[Any]   =					kwargs.pop("feature_extractor"    )

													SCREAMING_SNAKE_CASE__     :				int   =					image_processor if image_processor is not None else feature_extractor
													if image_processor is None:
																			raise ValueError("You need to specify an `image_processor`."    )
													if tokenizer is None:
																			raise ValueError("You need to specify a `tokenizer`."    )

													super().__init__(a							,	a    )

							def __call__(      self				:     Tuple							,	a				:     Tuple=None							,	a				:     Union[str, Any]=None							,	a				:     List[str]=None							,	**a				:     Optional[Any]    )	->Optional[Any]:

													if text is None and images is None:
																			raise ValueError("You have to specify either text or images. Both cannot be none."    )

													if text is not None:
																			SCREAMING_SNAKE_CASE__     :				str   =					self.tokenizer(a							,	return_tensors=a							,	**a    )

													if images is not None:
																			SCREAMING_SNAKE_CASE__     :				int   =					self.image_processor(a							,	return_tensors=a							,	**a    )

													if text is not None and images is not None:
																			SCREAMING_SNAKE_CASE__     :				Tuple   =					image_features.pixel_values
																			return encoding
													elif text is not None:
																			return encoding
													else:
																			return BatchEncoding(data=dict(**a    )							,	tensor_type=a    )

							def      A_  (      self				:     Optional[int]							,	*a				:     Any							,	**a				:     List[str]    )	->Any:
													return self.tokenizer.batch_decode(*a							,	**a    )

							def      A_  (      self				:     Any							,	*a				:     Optional[int]							,	**a				:     Dict    )	->Any:
													return self.tokenizer.decode(*a							,	**a    )

							@property
							def      A_  (      self				:     List[str]    )	->Union[str, Any]:
													SCREAMING_SNAKE_CASE__     :				Dict   =					self.tokenizer.model_input_names
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					self.image_processor.model_input_names
													return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names    )    )

							@property
							def      A_  (      self				:     Optional[int]    )	->List[Any]:
													warnings.warn(
													    "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."							,	a							,	)
													return self.image_processor_class






							@property
							def      A_  (      self				:     Dict    )	->str:
													warnings.warn(
													    "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."							,	a							,	)
													return self.image_processor