




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


SCREAMING_SNAKE_CASE    =	logging.getLogger(__name__)

SCREAMING_SNAKE_CASE    =	50  # max width of layer names
SCREAMING_SNAKE_CASE    =	70  # max width of quantizer names







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		List[str]:
 A__	=    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   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		Any:



 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=lowerCAmelCase__						)
 A__	=    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   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_      ,			lowercase_=False      ,			lowercase_=False						)						->		Tuple:
 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   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		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   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_						)						->		Union[str, Any]:
 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   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_						)						->		List[str]:



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

  A__	=    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   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_						)						->		Optional[int]:



 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=lowerCAmelCase__						)
   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   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		Union[str, Any]:



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







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		Union[str, Any]:



 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)
   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=lowerCAmelCase__      ,			keepdims=lowerCAmelCase__						).detach()
   logger.info(f"""RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}"""						)
   A__	=    amax







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_=25      ,			lowercase_=1_80      ,			lowercase_=None						)						->		List[str]:



 if ignore is None:
  A__	=    []
 elif not isinstance(lowerCAmelCase__      ,			lowerCAmelCase__						):
  A__	=    [ignore]

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

 for name, mod in model.named_modules():
  A__	=    getattr(lowerCAmelCase__      ,			"_input_quantizer"      ,			lowerCAmelCase__						)
  A__	=    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
  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(lowerCAmelCase__						) <= line_width:
   logger.info(lowerCAmelCase__						)
  else:
   logger.info(f"""{name:{name_width}} {act_str}"""						)
   logger.info(f"""{"  ":{name_width}} {wgt_str}"""						)







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		List[Any]:
 A__	=    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   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_      ,			lowercase_      ,			lowercase_      ,			lowercase_						)						->		int:
 A__	=    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   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_      ,			lowercase_="both"      ,			**lowercase_						)						->		Dict:
 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(lowerCAmelCase__      ,			lowerCAmelCase__      ,			"_input_quantizer"      ,			lowerCAmelCase__      ,			lowerCAmelCase__						)
  if which in ["weight", "both"]:
   set_quantizer(lowerCAmelCase__      ,			lowerCAmelCase__      ,			"_weight_quantizer"      ,			lowerCAmelCase__      ,			lowerCAmelCase__						)
 logger.info(lowerCAmelCase__						)







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_      ,			**lowercase_						)						->		int:



 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__						):
     A__	=    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 typing import Any




class      UpperCAmelCase_  :




    def __init__(   self						:	Optional[Any]     ,			snake_case_						:	Any				)				->							List[str]:




     '''simple docstring'''

     A__	=    data
     A__	=    None




    def __repr__(   self						:	Optional[int]				)				->							str:




     '''simple docstring'''

     return F"""Node({self.data})"""






class      UpperCAmelCase_  :




    def __init__(   self						:	Dict				)				->							Any:




     '''simple docstring'''

     A__	=    None




    def __iter__(   self						:	List[Any]				)				->							Any:




     '''simple docstring'''

     A__	=    self.head
     while node:
      yield node.data
      A__	=    node.next




    def __len__(   self						:	Any				)				->							int:




     '''simple docstring'''

     return sum(1 for _ in self				)




    def __repr__(   self						:	List[str]				)				->							str:




     '''simple docstring'''

     return "->".join([str(snake_case_				) for item in self]				)




    def __getitem__(   self						:	str     ,			snake_case_						:	int				)				->							Any:




     '''simple docstring'''

     if not 0 <= index < len(self				):
      raise ValueError("list index out of range."				)
     for i, node in enumerate(self				):
      if i == index:
       return node
     return None




    def __setitem__(   self						:	Tuple     ,			snake_case_						:	int     ,			snake_case_						:	Any				)				->							None:




     '''simple docstring'''

     if not 0 <= index < len(self				):
      raise ValueError("list index out of range."				)
     A__	=    self.head
     for _ in range(snake_case_				):
      A__	=    current.next
     A__	=    data




    def 							__magic_name__							(   self						:	List[Any]     ,			snake_case_						:	Any				)				->							None:




     '''simple docstring'''

     self.insert_nth(len(self				)     ,			snake_case_				)




    def 							__magic_name__							(   self						:	Tuple     ,			snake_case_						:	Any				)				->							None:




     '''simple docstring'''

     self.insert_nth(0     ,			snake_case_				)




    def 							__magic_name__							(   self						:	Dict     ,			snake_case_						:	int     ,			snake_case_						:	Any				)				->							None:




     '''simple docstring'''

     if not 0 <= index <= len(self				):
      raise IndexError("list index out of range"				)
     A__	=    Node(snake_case_				)
     if self.head is None:
      A__	=    new_node
     elif index == 0:
      A__	=    self.head  # link new_node to head
      A__	=    new_node
     else:
      A__	=    self.head
      for _ in range(index - 1				):
       A__	=    temp.next
      A__	=    temp.next
      A__	=    new_node




    def 							__magic_name__							(   self						:	Dict				)				->							None:  # print every node data




     '''simple docstring'''

     print(self				)




    def 							__magic_name__							(   self						:	Dict				)				->							Any:




     '''simple docstring'''

     return self.delete_nth(0				)




    def 							__magic_name__							(   self						:	Optional[Any]				)				->							Any:  # delete from tail




     '''simple docstring'''

     return self.delete_nth(len(self				) - 1				)




    def 							__magic_name__							(   self						:	Any     ,			snake_case_						:	int = 0				)				->							Any:




     '''simple docstring'''

     if not 0 <= index <= len(self				) - 1:  # test if index is valid
      raise IndexError("List index out of range."				)
     A__	=    self.head  # default first node
     if index == 0:
      A__	=    self.head.next
     else:
      A__	=    self.head
      for _ in range(index - 1				):
       A__	=    temp.next
      A__	=    temp.next
      A__	=    temp.next.next
     return delete_node.data




    def 							__magic_name__							(   self						:	Dict				)				->							bool:




     '''simple docstring'''

     return self.head is None




    def 							__magic_name__							(   self						:	List[Any]				)				->							None:




     '''simple docstring'''

     A__	=    None
     A__	=    self.head

     while current:
      # Store the current node's next node.
      A__	=    current.next
      # Make the current node's next point backwards
      A__	=    prev
      # Make the previous node be the current node
      A__	=    current
      # Make the current node the next node (to progress iteration)
      A__	=    next_node
     # Return prev in order to put the head at the end
     A__	=    prev







def   _SCREAMING_SNAKE_CASE       (							)						->		None:
 A__	=    LinkedList()
 assert linked_list.is_empty() is True
 assert str(lowercase_						) == ""

 try:
  linked_list.delete_head()
  raise AssertionError  # This should not happen.
 except IndexError:
  assert True  # This should happen.

 try:
  linked_list.delete_tail()
  raise AssertionError  # This should not happen.
 except IndexError:
  assert True  # This should happen.

 for i in range(10						):
  assert len(lowercase_						) == i
  linked_list.insert_nth(lowercase_      ,			i + 1						)
 assert str(lowercase_						) == "->".join(str(lowercase_						) for i in range(1      ,			11						)						)

 linked_list.insert_head(0						)
 linked_list.insert_tail(11						)
 assert str(lowercase_						) == "->".join(str(lowercase_						) for i in range(0      ,			12						)						)

 assert linked_list.delete_head() == 0
 assert linked_list.delete_nth(9						) == 10
 assert linked_list.delete_tail() == 11
 assert len(lowercase_						) == 9
 assert str(lowercase_						) == "->".join(str(lowercase_						) for i in range(1      ,			10						)						)

 assert all(linked_list[i] == i + 1 for i in range(0      ,			9						)						) is True

 for i in range(0      ,			9						):
  A__	=    -i
 assert all(linked_list[i] == -i for i in range(0      ,			9						)						) is True

 linked_list.reverse()
 assert str(lowercase_						) == "->".join(str(lowercase_						) for i in range(-8      ,			1						)						)







def   _SCREAMING_SNAKE_CASE       (							)						->		None:
 A__	=    [
     -9,
     1_00,
     Node(77_34_51_12						),
     "dlrow olleH",
     7,
     55_55,
     0,
     -1_9_2.5_5_5_5_5,
     "Hello, world!",
     7_7.9,
     Node(10						),
     None,
     None,
     1_2.2_0,
 ]
 A__	=    LinkedList()

 for i in test_input:
  linked_list.insert_tail(lowercase_						)

 # Check if it's empty or not
 assert linked_list.is_empty() is False
 assert (
     str(lowercase_						) == "-9->100->Node(77345112)->dlrow olleH->7->5555->0->"
     "-192.55555->Hello, world!->77.9->Node(10)->None->None->12.2"
 )

 # Delete the head
 A__	=    linked_list.delete_head()
 assert result == -9
 assert (
     str(lowercase_						) == "100->Node(77345112)->dlrow olleH->7->5555->0->-192.55555->"
     "Hello, world!->77.9->Node(10)->None->None->12.2"
 )

 # Delete the tail
 A__	=    linked_list.delete_tail()
 assert result == 1_2.2
 assert (
     str(lowercase_						) == "100->Node(77345112)->dlrow olleH->7->5555->0->-192.55555->"
     "Hello, world!->77.9->Node(10)->None->None"
 )

 # Delete a node in specific location in linked list
 A__	=    linked_list.delete_nth(10						)
 assert result is None
 assert (
     str(lowercase_						) == "100->Node(77345112)->dlrow olleH->7->5555->0->-192.55555->"
     "Hello, world!->77.9->Node(10)->None"
 )

 # Add a Node instance to its head
 linked_list.insert_head(Node("Hello again, world!"						)						)
 assert (
     str(lowercase_						)
     == "Node(Hello again, world!)->100->Node(77345112)->dlrow olleH->"
     "7->5555->0->-192.55555->Hello, world!->77.9->Node(10)->None"
 )

 # Add None to its tail
 linked_list.insert_tail(lowercase_						)
 assert (
     str(lowercase_						)
     == "Node(Hello again, world!)->100->Node(77345112)->dlrow olleH->"
     "7->5555->0->-192.55555->Hello, world!->77.9->Node(10)->None->None"
 )

 # Reverse the linked list
 linked_list.reverse()
 assert (
     str(lowercase_						)
     == "None->None->Node(10)->77.9->Hello, world!->-192.55555->0->5555->"
     "7->dlrow olleH->Node(77345112)->100->Node(Hello again, world!)"
 )







def   _SCREAMING_SNAKE_CASE       (							)						->		Optional[int]:
 from doctest import testmod

 testmod()

 A__	=    LinkedList()
 linked_list.insert_head(input("Inserting 1st at head "						).strip()						)
 linked_list.insert_head(input("Inserting 2nd at head "						).strip()						)
 print("\nPrint list:"						)
 linked_list.print_list()
 linked_list.insert_tail(input("\nInserting 1st at tail "						).strip()						)
 linked_list.insert_tail(input("Inserting 2nd at tail "						).strip()						)
 print("\nPrint list:"						)
 linked_list.print_list()
 print("\nDelete head"						)
 linked_list.delete_head()
 print("Delete tail"						)
 linked_list.delete_tail()
 print("\nPrint list:"						)
 linked_list.print_list()
 print("\nReverse linked list"						)
 linked_list.reverse()
 print("\nPrint list:"						)
 linked_list.print_list()
 print("\nString representation of linked list:"						)
 print(lowercase_						)
 print("\nReading/changing Node data using indexing:"						)
 print(f"""Element at Position 1: {linked_list[1]}"""						)
 A__	=    input("Enter New Value: "						).strip()
 print("New list:"						)
 print(lowercase_						)
 print(f"""length of linked_list is : {len(lowercase_						)}"""						)


if __name__ == "__main__":
      main()





