





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

lowercase_		:					Optional[int]      						=    50  # max width of layer names
lowercase_		:					Dict      						=    70  # max width of quantizer names

def    A__(		__lowerCAmelCase			):
				_snake_case			:      Union[str, Any]	    = 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    A__(		__lowerCAmelCase			):

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

				_snake_case			:      Dict	    = QuantDescriptor(num_bits=args.aprec						,      calib_method=__lowerCAmelCase			)
				_snake_case			:      Optional[Any]	    = 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    A__(		__lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase=False						,      __lowerCAmelCase=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    A__(		__lowerCAmelCase			):
				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__(		__lowerCAmelCase						,      __lowerCAmelCase			):
				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    A__(		__lowerCAmelCase						,      __lowerCAmelCase			):

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

								_snake_case			:      int	    = 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    A__(		__lowerCAmelCase						,      __lowerCAmelCase			):

				for name, mod in model.named_modules():
								if name.endswith('.output.dense'			) and not name.endswith('attention.output.dense'			):
												_snake_case			:      Tuple	    = mod._input_quantizer._amax.data.detach().item()
												mod._input_quantizer._amax.data.detach().clamp_(max=__lowerCAmelCase			)
												_snake_case			:      List[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__(		__lowerCAmelCase			):

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

def    A__(		__lowerCAmelCase			):

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

def    A__(		__lowerCAmelCase						,      __lowerCAmelCase=25						,      __lowerCAmelCase=1_80						,      __lowerCAmelCase=None			):

				if ignore is None:
								_snake_case			:      Optional[int]	    = []
				elif not isinstance(__lowerCAmelCase						,      __lowerCAmelCase			):
								_snake_case			:      List[Any]	    = [ignore]

				_snake_case			:      List[Any]	    = 0
				for name, mod in model.named_modules():
								if not hasattr(__lowerCAmelCase						,      'weight'			):
												continue
								_snake_case			:      List[str]	    = max(__lowerCAmelCase						,      len(__lowerCAmelCase			)			)

				for name, mod in model.named_modules():
								_snake_case			:      List[str]	    = getattr(__lowerCAmelCase						,      '_input_quantizer'						,      __lowerCAmelCase			)
								_snake_case			:      Union[str, Any]	    = 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
								_snake_case			:      int	    = F'''Act:{input_q.extra_repr()}'''
								_snake_case			:      Union[str, Any]	    = F'''Wgt:{weight_q.extra_repr()}'''
								_snake_case			:      List[str]	    = 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    A__(		__lowerCAmelCase			):
				_snake_case			:      List[str]	    = 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    A__(		__lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase			):
				_snake_case			:      Tuple	    = 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    A__(		__lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase="both"						,      **__lowerCAmelCase			):
				_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(__lowerCAmelCase						,      __lowerCAmelCase						,      '_input_quantizer'						,      __lowerCAmelCase						,      __lowerCAmelCase			)
								if which in ["weight", "both"]:
												set_quantizer(__lowerCAmelCase						,      __lowerCAmelCase						,      '_weight_quantizer'						,      __lowerCAmelCase						,      __lowerCAmelCase			)
				logger.info(__lowerCAmelCase			)

def    A__(		__lowerCAmelCase						,      __lowerCAmelCase						,      **__lowerCAmelCase			):

				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			):
																				_snake_case			:      Dict	    = 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"""



import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    AutoConfig,
    AutoImageProcessor,
    CLIPConfig,
    CLIPImageProcessor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER


sys.path.append(str(Path(__file__).parent.parent.parent.parent / "utils"))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_image_processing import CustomImageProcessor  # noqa E402
class 						__UpperCAmelCase    (unittest.TestCase     ):







			'''simple docstring'''





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     List[str]								=			0





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     List[str]								=			AutoImageProcessor.from_pretrained("""openai/clip-vit-base-patch32"""							)
										self.assertIsInstance(snake_case_   ,      snake_case_							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										with tempfile.TemporaryDirectory() as tmpdirname:
																	A__       :     Optional[Any]								=			Path(snake_case_							) / """preprocessor_config.json"""
																	A__       :     Any								=			Path(snake_case_							) / """config.json"""
																	json.dump(
																	    {"""image_processor_type""": """CLIPImageProcessor""", """processor_class""": """CLIPProcessor"""}   ,      open(snake_case_   ,      """w"""							)   ,      )
																	json.dump({"""model_type""": """clip"""}   ,      open(snake_case_   ,      """w"""							)							)

																	A__       :     Any								=			AutoImageProcessor.from_pretrained(snake_case_							)
																	self.assertIsInstance(snake_case_   ,      snake_case_							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										with tempfile.TemporaryDirectory() as tmpdirname:
																	A__       :     str								=			Path(snake_case_							) / """preprocessor_config.json"""
																	A__       :     Union[str, Any]								=			Path(snake_case_							) / """config.json"""
																	json.dump(
																	    {"""feature_extractor_type""": """CLIPFeatureExtractor""", """processor_class""": """CLIPProcessor"""}   ,      open(snake_case_   ,      """w"""							)   ,      )
																	json.dump({"""model_type""": """clip"""}   ,      open(snake_case_   ,      """w"""							)							)

																	A__       :     Optional[int]								=			AutoImageProcessor.from_pretrained(snake_case_							)
																	self.assertIsInstance(snake_case_   ,      snake_case_							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										with tempfile.TemporaryDirectory() as tmpdirname:
																	A__       :     Union[str, Any]								=			CLIPConfig()

																	# Create a dummy config file with image_proceesor_type
																	A__       :     Dict								=			Path(snake_case_							) / """preprocessor_config.json"""
																	A__       :     Optional[int]								=			Path(snake_case_							) / """config.json"""
																	json.dump(
																	    {"""image_processor_type""": """CLIPImageProcessor""", """processor_class""": """CLIPProcessor"""}   ,      open(snake_case_   ,      """w"""							)   ,      )
																	json.dump({"""model_type""": """clip"""}   ,      open(snake_case_   ,      """w"""							)							)

																	# remove image_processor_type to make sure config.json alone is enough to load image processor locally
																	A__       :     List[str]								=			AutoImageProcessor.from_pretrained(snake_case_							).to_dict()

																	config_dict.pop("""image_processor_type"""							)
																	A__       :     str								=			CLIPImageProcessor(**snake_case_							)

																	# save in new folder
																	model_config.save_pretrained(snake_case_							)
																	config.save_pretrained(snake_case_							)

																	A__       :     Optional[int]								=			AutoImageProcessor.from_pretrained(snake_case_							)

																	# make sure private variable is not incorrectly saved
																	A__       :     Dict								=			json.loads(config.to_json_string()							)
																	self.assertTrue("""_processor_class""" not in dict_as_saved							)

										self.assertIsInstance(snake_case_   ,      snake_case_							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										with tempfile.TemporaryDirectory() as tmpdirname:
																	A__       :     List[Any]								=			Path(snake_case_							) / """preprocessor_config.json"""
																	json.dump(
																	    {"""image_processor_type""": """CLIPImageProcessor""", """processor_class""": """CLIPProcessor"""}   ,      open(snake_case_   ,      """w"""							)   ,      )

																	A__       :     List[str]								=			AutoImageProcessor.from_pretrained(snake_case_							)
																	self.assertIsInstance(snake_case_   ,      snake_case_							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										with self.assertRaisesRegex(
										    snake_case_   ,      """clip-base is not a local folder and is not a valid model identifier"""							):
																	A__       :     Optional[Any]								=			AutoImageProcessor.from_pretrained("""clip-base"""							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										with self.assertRaisesRegex(
										    snake_case_   ,      r"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""							):
																	A__       :     Any								=			AutoImageProcessor.from_pretrained(snake_case_   ,      revision="""aaaaaa"""							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										with self.assertRaisesRegex(
										    snake_case_   ,      """hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json."""   ,      ):
																	A__       :     Optional[Any]								=			AutoImageProcessor.from_pretrained("""hf-internal-testing/config-no-model"""							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										with self.assertRaises(snake_case_							):
																	A__       :     str								=			AutoImageProcessor.from_pretrained("""hf-internal-testing/test_dynamic_image_processor"""							)
										# If remote code is disabled, we can't load this config.
										with self.assertRaises(snake_case_							):
																	A__       :     int								=			AutoImageProcessor.from_pretrained(
																	    """hf-internal-testing/test_dynamic_image_processor"""   ,      trust_remote_code=snake_case_							)

										A__       :     Optional[int]								=			AutoImageProcessor.from_pretrained(
										    """hf-internal-testing/test_dynamic_image_processor"""   ,      trust_remote_code=snake_case_							)
										self.assertEqual(image_processor.__class__.__name__   ,      """NewImageProcessor"""							)

										# Test image processor can be reloaded.
										with tempfile.TemporaryDirectory() as tmp_dir:
																	image_processor.save_pretrained(snake_case_							)
																	A__       :     Optional[Any]								=			AutoImageProcessor.from_pretrained(snake_case_   ,      trust_remote_code=snake_case_							)
										self.assertEqual(reloaded_image_processor.__class__.__name__   ,      """NewImageProcessor"""							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										try:
																	AutoConfig.register("""custom"""   ,      snake_case_							)
																	AutoImageProcessor.register(snake_case_   ,      snake_case_							)
																	# Trying to register something existing in the Transformers library will raise an error
																	with self.assertRaises(snake_case_							):
																								AutoImageProcessor.register(snake_case_   ,      snake_case_							)

																	with tempfile.TemporaryDirectory() as tmpdirname:
																								A__       :     Tuple								=			Path(snake_case_							) / """preprocessor_config.json"""
																								A__       :     List[str]								=			Path(snake_case_							) / """config.json"""
																								json.dump(
																								    {"""feature_extractor_type""": """CLIPFeatureExtractor""", """processor_class""": """CLIPProcessor"""}   ,      open(snake_case_   ,      """w"""							)   ,      )
																								json.dump({"""model_type""": """clip"""}   ,      open(snake_case_   ,      """w"""							)							)

																								A__       :     Optional[int]								=			CustomImageProcessor.from_pretrained(snake_case_							)

																	# Now that the config is registered, it can be used as any other config with the auto-API
																	with tempfile.TemporaryDirectory() as tmp_dir:
																								image_processor.save_pretrained(snake_case_							)
																								A__       :     int								=			AutoImageProcessor.from_pretrained(snake_case_							)
																								self.assertIsInstance(snake_case_   ,      snake_case_							)

										finally:
																	if "custom" in CONFIG_MAPPING._extra_content:
																								del CONFIG_MAPPING._extra_content["custom"]
																	if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
																								del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										class 						__UpperCAmelCase    (__A     ):







													'''simple docstring'''
													_UpperCamelCase  :						Union[str, Any]					     =				True

										try:
																	AutoConfig.register("""custom"""   ,      snake_case_							)
																	AutoImageProcessor.register(snake_case_   ,      snake_case_							)
																	# If remote code is not set, the default is to use local
																	A__       :     int								=			AutoImageProcessor.from_pretrained("""hf-internal-testing/test_dynamic_image_processor"""							)
																	self.assertEqual(image_processor.__class__.__name__   ,      """NewImageProcessor"""							)
																	self.assertTrue(image_processor.is_local							)

																	# If remote code is disabled, we load the local one.
																	A__       :     Any								=			AutoImageProcessor.from_pretrained(
																	    """hf-internal-testing/test_dynamic_image_processor"""   ,      trust_remote_code=snake_case_							)
																	self.assertEqual(image_processor.__class__.__name__   ,      """NewImageProcessor"""							)
																	self.assertTrue(image_processor.is_local							)

																	# If remote is enabled, we load from the Hub
																	A__       :     Optional[Any]								=			AutoImageProcessor.from_pretrained(
																	    """hf-internal-testing/test_dynamic_image_processor"""   ,      trust_remote_code=snake_case_							)
																	self.assertEqual(image_processor.__class__.__name__   ,      """NewImageProcessor"""							)
																	self.assertTrue(not hasattr(snake_case_   ,      """is_local"""							)							)

										finally:
																	if "custom" in CONFIG_MAPPING._extra_content:
																								del CONFIG_MAPPING._extra_content["custom"]
																	if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
																								del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]





