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


UpperCamelCase_       :    int						      =			logging.getLogger(__name__)

UpperCamelCase_       :    List[Any]						      =			50  # max width of layer names
UpperCamelCase_       :    Union[str, Any]						      =			70  # max width of quantizer names
def 				A_  (__a					):

					'''simple docstring'''





					A_								=    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 				A_  (__a					):

					'''simple docstring'''





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

					A_								=    QuantDescriptor(num_bits=args.aprec					, calib_method=__a					)
					A_								=    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 				A_  (__a					, __a					, __a=False					, __a=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 				A_  (__a					):

					'''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_  (__a					, __a					):

					'''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 				A_  (__a					, __a					):

					'''simple docstring'''





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

										A_								=    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 				A_  (__a					, __a					):

					'''simple docstring'''





					for name, mod in model.named_modules():
										if name.endswith(".output.dense"					) and not name.endswith("attention.output.dense"					):
															A_								=    mod._input_quantizer._amax.data.detach().item()
															mod._input_quantizer._amax.data.detach().clamp_(max=__a					)
															A_								=    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_  (__a					):

					'''simple docstring'''





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

					'''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)
															A_								=    set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis					)
															A_								=    set(range(len(mod.weight.size()					)					)					) - axis_set
															A_								=    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()}'					)
															A_								=    amax
def 				A_  (__a					, __a=25					, __a=180					, __a=None					):

					'''simple docstring'''





					if ignore is None:
										A_								=    []
					elif not isinstance(__a					, __a					):
										A_								=    [ignore]

					A_								=    0
					for name, mod in model.named_modules():
										if not hasattr(__a					, "weight"					):
															continue
										A_								=    max(__a					, len(__a					)					)

					for name, mod in model.named_modules():
										A_								=    getattr(__a					, "_input_quantizer"					, __a					)
										A_								=    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
										A_								=    f'Act:{input_q.extra_repr()}'
										A_								=    f'Wgt:{weight_q.extra_repr()}'
										A_								=    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 				A_  (__a					):

					'''simple docstring'''





					A_								=    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 				A_  (__a					, __a					, __a					, __a					, __a					):

					'''simple docstring'''





					A_								=    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 				A_  (__a					, __a					, __a="both"					, **__a					):

					'''simple docstring'''





					A_								=    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 				A_  (__a					, __a					, **__a					):

					'''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					):
																									A_								=    f'Warning: changing {name:{name_width}}'
																									for k, v in kwargs.items():
																														s += f' {k}={v}'
																														setattr(__a					, __a					, __a					)
																									logger.info(__a					)






"""simple docstring"""





import os
from shutil import copyfile
from typing import List, Optional, Tuple

from tokenizers import processors

from ...tokenization_utils import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
		from .tokenization_mbart import MBartTokenizer
else:
		UpperCamelCase_       :    Dict						      =			None


UpperCamelCase_       :    int						      =			logging.get_logger(__name__)


UpperCamelCase_       :    str						      =			{'''vocab_file''': '''sentencepiece.bpe.model''', '''tokenizer_file''': '''tokenizer.json'''}

UpperCamelCase_       :    List[Any]						      =			{
    '''vocab_file''': {
        '''facebook/mbart-large-en-ro''': (
            '''https://huggingface.co/facebook/mbart-large-en-ro/resolve/main/sentencepiece.bpe.model'''
        ),
        '''facebook/mbart-large-cc25''': (
            '''https://huggingface.co/facebook/mbart-large-cc25/resolve/main/sentencepiece.bpe.model'''
        ),
    },
    '''tokenizer_file''': {
        '''facebook/mbart-large-en-ro''': '''https://huggingface.co/facebook/mbart-large-en-ro/resolve/main/tokenizer.json''',
        '''facebook/mbart-large-cc25''': '''https://huggingface.co/facebook/mbart-large-cc25/resolve/main/tokenizer.json''',
    },
}

UpperCamelCase_       :    Optional[Any]						      =			{
    '''facebook/mbart-large-en-ro''': 1024,
    '''facebook/mbart-large-cc25''': 1024,
}

# fmt: off
UpperCamelCase_       :    str						      =			['''ar_AR''', '''cs_CZ''', '''de_DE''', '''en_XX''', '''es_XX''', '''et_EE''', '''fi_FI''', '''fr_XX''', '''gu_IN''', '''hi_IN''', '''it_IT''', '''ja_XX''', '''kk_KZ''', '''ko_KR''', '''lt_LT''', '''lv_LV''', '''my_MM''', '''ne_NP''', '''nl_XX''', '''ro_RO''', '''ru_RU''', '''si_LK''', '''tr_TR''', '''vi_VN''', '''zh_CN''']



class __lowerCAmelCase					(       _lowercase   ):




			"""simple docstring"""







			snake_case						   =					VOCAB_FILES_NAMES
			snake_case						   =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			snake_case						   =					PRETRAINED_VOCAB_FILES_MAP
			snake_case						   =					["input_ids", "attention_mask"]
			snake_case						   =					MBartTokenizer

			snake_case						   =					[]
			snake_case						   =					[]




			def __init__(					self    :       List[str]							, _snake_case    :       Tuple=None							, _snake_case    :       int=None							, _snake_case    :       List[Any]="<s>"							, _snake_case    :       Tuple="</s>"							, _snake_case    :       str="</s>"							, _snake_case    :       List[Any]="<s>"							, _snake_case    :       Dict="<unk>"							, _snake_case    :       str="<pad>"							, _snake_case    :       Any="<mask>"							, _snake_case    :       int=None							, _snake_case    :       Optional[int]=None							, _snake_case    :       Any=None							, **_snake_case    :       Union[str, Any]							, )   ->			List[str]:







								"""simple docstring"""
								# Mask token behave like a normal word, i.e. include the space before it
								A_								=    AddedToken(_snake_case							, lstrip=_snake_case							, rstrip=_snake_case						) if isinstance(_snake_case							, _snake_case						) else mask_token

								super().__init__(
								    vocab_file=_snake_case							, tokenizer_file=_snake_case							, bos_token=_snake_case							, eos_token=_snake_case							, sep_token=_snake_case							, cls_token=_snake_case							, unk_token=_snake_case							, pad_token=_snake_case							, mask_token=_snake_case							, src_lang=_snake_case							, tgt_lang=_snake_case							, additional_special_tokens=_snake_case							, **_snake_case							, )

								A_								=    vocab_file
								A_								=    False if not self.vocab_file else True

								A_								=    FAIRSEQ_LANGUAGE_CODES.copy()

								if additional_special_tokens is not None:
													# Only add those special tokens if they are not already there.
													_additional_special_tokens.extend(
													    [t for t in additional_special_tokens if t not in _additional_special_tokens]						)

								self.add_special_tokens({"additional_special_tokens": _additional_special_tokens}						)
								A_								=    {
								    lang_code: self.convert_tokens_to_ids(_snake_case						) for lang_code in FAIRSEQ_LANGUAGE_CODES
								}

								A_								=    src_lang if src_lang is not None else "en_XX"
								A_								=    self.convert_tokens_to_ids(self._src_lang						)
								A_								=    tgt_lang
								self.set_src_lang_special_tokens(self._src_lang						)




			@property
			def   lowerCamelCase__  (					self    :       Dict						)   ->			str:







								"""simple docstring"""
								return self._src_lang




			@src_lang.setter
			def   lowerCamelCase__  (					self    :       Tuple							, _snake_case    :       str						)   ->			None:







								"""simple docstring"""
								A_								=    new_src_lang
								self.set_src_lang_special_tokens(self._src_lang						)




			def   lowerCamelCase__  (					self    :       Optional[Any]							, _snake_case    :       List[int]							, _snake_case    :       Optional[List[int]] = None						)   ->			List[int]:







								"""simple docstring"""
								if token_ids_a is None:
													return self.prefix_tokens + token_ids_a + self.suffix_tokens
								# We don't expect to process pairs, but leave the pair logic for API consistency
								return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens




			def   lowerCamelCase__  (					self    :       Dict							, _snake_case    :       List[int]							, _snake_case    :       Optional[List[int]] = None						)   ->			List[int]:







								"""simple docstring"""
								A_								=    [self.sep_token_id]
								A_								=    [self.cls_token_id]

								if token_ids_a is None:
													return len(cls + token_ids_a + sep						) * [0]
								return len(cls + token_ids_a + sep + sep + token_ids_a + sep						) * [0]




			def   lowerCamelCase__  (					self    :       List[Any]							, _snake_case    :       str							, _snake_case    :       str							, _snake_case    :       Optional[str]							, _snake_case    :       Optional[str]							, **_snake_case    :       Optional[int]						)   ->			str:







								"""simple docstring"""
								if src_lang is None or tgt_lang is None:
													raise ValueError("Translation requires a `src_lang` and a `tgt_lang` for this model"						)
								A_								=    src_lang
								A_								=    self(_snake_case							, add_special_tokens=_snake_case							, return_tensors=_snake_case							, **_snake_case						)
								A_								=    self.convert_tokens_to_ids(_snake_case						)
								A_								=    tgt_lang_id
								return inputs




			def   lowerCamelCase__  (					self    :       Dict							, _snake_case    :       List[str]							, _snake_case    :       str = "en_XX"							, _snake_case    :       Optional[List[str]] = None							, _snake_case    :       str = "ro_RO"							, **_snake_case    :       str							, )   ->			BatchEncoding:







								"""simple docstring"""
								A_								=    src_lang
								A_								=    tgt_lang
								return super().prepare_seqaseq_batch(_snake_case							, _snake_case							, **_snake_case						)




			def   lowerCamelCase__  (					self    :       Any						)   ->			Optional[Any]:







								"""simple docstring"""
								return self.set_src_lang_special_tokens(self.src_lang						)




			def   lowerCamelCase__  (					self    :       Dict						)   ->			List[Any]:







								"""simple docstring"""
								return self.set_tgt_lang_special_tokens(self.tgt_lang						)




			def   lowerCamelCase__  (					self    :       Tuple							, _snake_case    :       List[str]						)   ->			None:







								"""simple docstring"""
								A_								=    self.convert_tokens_to_ids(_snake_case						)
								A_								=    []
								A_								=    [self.eos_token_id, self.cur_lang_code]

								A_								=    self.convert_ids_to_tokens(self.prefix_tokens						)
								A_								=    self.convert_ids_to_tokens(self.suffix_tokens						)

								A_								=    processors.TemplateProcessing(
								    single=prefix_tokens_str + ["$A"] + suffix_tokens_str							, pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str							, special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str							, self.prefix_tokens + self.suffix_tokens						)						)							, )




			def   lowerCamelCase__  (					self    :       List[str]							, _snake_case    :       str						)   ->			None:







								"""simple docstring"""
								A_								=    self.convert_tokens_to_ids(_snake_case						)
								A_								=    []
								A_								=    [self.eos_token_id, self.cur_lang_code]

								A_								=    self.convert_ids_to_tokens(self.prefix_tokens						)
								A_								=    self.convert_ids_to_tokens(self.suffix_tokens						)

								A_								=    processors.TemplateProcessing(
								    single=prefix_tokens_str + ["$A"] + suffix_tokens_str							, pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str							, special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str							, self.prefix_tokens + self.suffix_tokens						)						)							, )






			def   lowerCamelCase__  (					self    :       Union[str, Any]							, _snake_case    :       str							, _snake_case    :       Optional[str] = None						)   ->			Tuple[str]:







								"""simple docstring"""
								if not self.can_save_slow_tokenizer:
													raise ValueError(
													    "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
													    "tokenizer."						)

								if not os.path.isdir(_snake_case						):
													logger.error(F'Vocabulary path ({save_directory}) should be a directory.'						)
													return
								A_								=    os.path.join(
								    _snake_case							, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]						)

								if os.path.abspath(self.vocab_file						) != os.path.abspath(_snake_case						):
													copyfile(self.vocab_file							, _snake_case						)

								return (out_vocab_file,)




