





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


__lowerCamelCase			:	Tuple					    =	logging.getLogger(__name__)

__lowerCamelCase			:	Optional[Any]					    =	50  # max width of layer names
__lowerCamelCase			:	List[Any]					    =	70  # max width of quantizer names



def 						A_							(	_lowerCAmelCase	)		->					Optional[int]:
 UpperCamelCase      :   Tuple			      =  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	)		->					Optional[int]:

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

 UpperCamelCase      :   Dict			      =  QuantDescriptor(num_bits=args.aprec				,  calib_method=_lowerCAmelCase	)
 UpperCamelCase      :   Optional[int]			      =  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	)		->					List[Any]:
 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	)		->					Optional[int]:
 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	)		->					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 						A_							(	_lowerCAmelCase				,  _lowerCAmelCase	)		->					Dict:

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

  UpperCamelCase      :   List[Any]			      =  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	)		->					Tuple:

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

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



def 						A_							(	_lowerCAmelCase	)		->					List[str]:

 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      :   List[str]			      =  set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis	)
   UpperCamelCase      :   List[str]			      =  set(range(len(mod.weight.size()	)	)	) - axis_set
   UpperCamelCase      :   Optional[int]			      =  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      :   List[str]			      =  amax



def 						A_							(	_lowerCAmelCase				,  _lowerCAmelCase=25				,  _lowerCAmelCase=180				,  _lowerCAmelCase=None	)		->					Dict:

 if ignore is None:
  UpperCamelCase      :   Any			      =  []
 elif not isinstance(_lowerCAmelCase				,  _lowerCAmelCase	):
  UpperCamelCase      :   Optional[int]			      =  [ignore]

 UpperCamelCase      :   List[Any]			      =  0
 for name, mod in model.named_modules():
  if not hasattr(_lowerCAmelCase				,  "weight"	):
   continue
  UpperCamelCase      :   Optional[int]			      =  max(_lowerCAmelCase				,  len(_lowerCAmelCase	)	)

 for name, mod in model.named_modules():
  UpperCamelCase      :   Tuple			      =  getattr(_lowerCAmelCase				,  "_input_quantizer"				,  _lowerCAmelCase	)
  UpperCamelCase      :   Dict			      =  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      :   int			      =  F"""Act:{input_q.extra_repr()}"""
  UpperCamelCase      :   Union[str, Any]			      =  F"""Wgt:{weight_q.extra_repr()}"""
  UpperCamelCase      :   int			      =  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	)		->					Optional[Any]:
 UpperCamelCase      :   int			      =  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	)		->					Tuple:
 UpperCamelCase      :   List[Any]			      =  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	)		->					int:
 UpperCamelCase      :   int			      =  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	)		->					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      :   Any			      =  F"""Warning: changing {name:{name_width}}"""
     for k, v in kwargs.items():
      s += F""" {k}={v}"""
      setattr(_lowerCAmelCase				,  _lowerCAmelCase				,  _lowerCAmelCase	)
     logger.info(_lowerCAmelCase	)







# Usage:
# ./gen-card-facebook-wmt19.py

import os
from pathlib import Path



def 						A_							(	_lowerCAmelCase				,  _lowerCAmelCase				,  _lowerCAmelCase	)		->					str:
 UpperCamelCase      :   Union[str, Any]			      =  {
     "en": "Machine learning is great, isn't it?",
     "ru": "Машинное обучение - это здорово, не так ли?",
     "de": "Maschinelles Lernen ist großartig, oder?",
 }

 # BLUE scores as follows:
 # "pair": [fairseq, transformers]
 UpperCamelCase      :   List[str]			      =  {
     "ru-en": ["[41.3](http://matrix.statmt.org/matrix/output/1907?run_id=6937)", "39.20"],
     "en-ru": ["[36.4](http://matrix.statmt.org/matrix/output/1914?run_id=6724)", "33.47"],
     "en-de": ["[43.1](http://matrix.statmt.org/matrix/output/1909?run_id=6862)", "42.83"],
     "de-en": ["[42.3](http://matrix.statmt.org/matrix/output/1902?run_id=6750)", "41.35"],
 }
 UpperCamelCase      :   str			      =  F"""{src_lang}-{tgt_lang}"""

 UpperCamelCase      :   Dict			      =  F"""
---
language: 
- {src_lang}
- {tgt_lang}
thumbnail:
tags:
- translation
- wmt19
- facebook
license: apache-2.0
datasets:
- wmt19
metrics:
- bleu
---

# FSMT

## Model description

This is a ported version of [fairseq wmt19 transformer](https://github.com/pytorch/fairseq/blob/master/examples/wmt19/README.md) for {src_lang}-{tgt_lang}.

For more details, please see, [Facebook FAIR's WMT19 News Translation Task Submission](https://arxiv.org/abs/1907.06616).

The abbreviation FSMT stands for FairSeqMachineTranslation

All four models are available:

* [wmt19-en-ru](https://huggingface.co/facebook/wmt19-en-ru)
* [wmt19-ru-en](https://huggingface.co/facebook/wmt19-ru-en)
* [wmt19-en-de](https://huggingface.co/facebook/wmt19-en-de)
* [wmt19-de-en](https://huggingface.co/facebook/wmt19-de-en)

## Intended uses & limitations

#### How to use

```python
from transformers import FSMTForConditionalGeneration, FSMTTokenizer
mname = \"facebook/wmt19-{src_lang}-{tgt_lang}\"
tokenizer = FSMTTokenizer.from_pretrained(mname)
model = FSMTForConditionalGeneration.from_pretrained(mname)

input = \"{texts[src_lang]}\"
input_ids = tokenizer.encode(input, return_tensors=\"pt\")
outputs = model.generate(input_ids)
decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(decoded) # {texts[tgt_lang]}

```

#### Limitations and bias

- The original (and this ported model) doesn't seem to handle well inputs with repeated sub-phrases, [content gets truncated](https://discuss.huggingface.co/t/issues-with-translating-inputs-containing-repeated-phrases/981)

## Training data

Pretrained weights were left identical to the original model released by fairseq. For more details, please, see the [paper](https://arxiv.org/abs/1907.06616).

## Eval results

pair   | fairseq | transformers
-------|---------|----------
{pair}  | {scores[pair][0]} | {scores[pair][1]}

The score is slightly below the score reported by `fairseq`, since `transformers`` currently doesn't support:
- model ensemble, therefore the best performing checkpoint was ported (``model4.pt``).
- re-ranking

The score was calculated using this code:

```bash
git clone https://github.com/huggingface/transformers
cd transformers
export PAIR={pair}
export DATA_DIR=data/$PAIR
export SAVE_DIR=data/$PAIR
export BS=8
export NUM_BEAMS=15
mkdir -p $DATA_DIR
sacrebleu -t wmt19 -l $PAIR --echo src > $DATA_DIR/val.source
sacrebleu -t wmt19 -l $PAIR --echo ref > $DATA_DIR/val.target
echo $PAIR
PYTHONPATH=\"src:examples/seq2seq\" python examples/seq2seq/run_eval.py facebook/wmt19-$PAIR $DATA_DIR/val.source $SAVE_DIR/test_translations.txt --reference_path $DATA_DIR/val.target --score_path $SAVE_DIR/test_bleu.json --bs $BS --task translation --num_beams $NUM_BEAMS
```
note: fairseq reports using a beam of 50, so you should get a slightly higher score if re-run with `--num_beams 50`.

## Data Sources

- [training, etc.](http://www.statmt.org/wmt19/)
- [test set](http://matrix.statmt.org/test_sets/newstest2019.tgz?1556572561)


### BibTeX entry and citation info

```bibtex
@inproceedings{{...,
  year={{2020}},
  title={{Facebook FAIR's WMT19 News Translation Task Submission}},
  author={{Ng, Nathan and Yee, Kyra and Baevski, Alexei and Ott, Myle and Auli, Michael and Edunov, Sergey}},
  booktitle={{Proc. of WMT}},
}}
```


## TODO

- port model ensemble (fairseq uses 4 model checkpoints)

"""
 os.makedirs(_lowerCAmelCase				,  exist_ok=_lowerCAmelCase	)
 UpperCamelCase      :   Union[str, Any]			      =  os.path.join(_lowerCAmelCase				,  "README.md"	)
 print(F"""Generating {path}"""	)
 with open(_lowerCAmelCase				,  "w"				,  encoding="utf-8"	) as f:
  f.write(_lowerCAmelCase	)


# make sure we are under the root of the project
__lowerCamelCase			:	str					    =	Path(__file__).resolve().parent.parent.parent
__lowerCamelCase			:	Optional[Any]					    =	repo_dir / """model_cards"""

for model_name in ["wmt19-ru-en", "wmt19-en-ru", "wmt19-en-de", "wmt19-de-en"]:
    __lowerCamelCase    ,      __lowerCamelCase    ,      __lowerCamelCase			:	str					    =	model_name.split("""-""")
    __lowerCamelCase			:	str					    =	model_cards_dir / """facebook""" / model_name
    write_model_card(model_card_dir, src_lang=src_lang, tgt_lang=tgt_lang)
