



"""simple docstring"""







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

import os
from pathlib import Path

def 				UpperCamelCase_    (    lowerCAmelCase__      :				List[Any]    ,   lowerCAmelCase__      :				Tuple    ,   lowerCAmelCase__      :				Optional[Any]						)						->   Dict:

  """simple docstring"""




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

  # BLUE scores as follows:
  # "pair": [fairseq, transformers]
  lowerCAmelCase_   :				Optional[int]								= {
      "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"],
  }
  lowerCAmelCase_   :				List[str]								= f"{src_lang}-{tgt_lang}"

  lowerCAmelCase_   :				List[Any]								= f"\n---\nlanguage: \n- {src_lang}\n- {tgt_lang}\nthumbnail:\ntags:\n- translation\n- wmt19\n- facebook\nlicense: apache-2.0\ndatasets:\n- wmt19\nmetrics:\n- bleu\n---\n\n# FSMT\n\n## Model description\n\nThis is a ported version of [fairseq wmt19 transformer](https://github.com/pytorch/fairseq/blob/master/examples/wmt19/README.md) for {src_lang}-{tgt_lang}.\n\nFor more details, please see, [Facebook FAIR\'s WMT19 News Translation Task Submission](https://arxiv.org/abs/1907.06616).\n\nThe abbreviation FSMT stands for FairSeqMachineTranslation\n\nAll four models are available:\n\n* [wmt19-en-ru](https://huggingface.co/facebook/wmt19-en-ru)\n* [wmt19-ru-en](https://huggingface.co/facebook/wmt19-ru-en)\n* [wmt19-en-de](https://huggingface.co/facebook/wmt19-en-de)\n* [wmt19-de-en](https://huggingface.co/facebook/wmt19-de-en)\n\n## Intended uses & limitations\n\n#### How to use\n\n```python\nfrom transformers import FSMTForConditionalGeneration, FSMTTokenizer\nmname = \"facebook/wmt19-{src_lang}-{tgt_lang}\"\ntokenizer = FSMTTokenizer.from_pretrained(mname)\nmodel = FSMTForConditionalGeneration.from_pretrained(mname)\n\ninput = \"{texts[src_lang]}\"\ninput_ids = tokenizer.encode(input, return_tensors=\"pt\")\noutputs = model.generate(input_ids)\ndecoded = tokenizer.decode(outputs[0], skip_special_tokens=True)\nprint(decoded) # {texts[tgt_lang]}\n\n```\n\n#### Limitations and bias\n\n- 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)\n\n## Training data\n\nPretrained weights were left identical to the original model released by fairseq. For more details, please, see the [paper](https://arxiv.org/abs/1907.06616).\n\n## Eval results\n\npair   | fairseq | transformers\n-------|---------|----------\n{pair}  | {scores[pair][0]} | {scores[pair][1]}\n\nThe score is slightly below the score reported by `fairseq`, since `transformers`` currently doesn\'t support:\n- model ensemble, therefore the best performing checkpoint was ported (``model4.pt``).\n- re-ranking\n\nThe score was calculated using this code:\n\n```bash\ngit clone https://github.com/huggingface/transformers\ncd transformers\nexport PAIR={pair}\nexport DATA_DIR=data/$PAIR\nexport SAVE_DIR=data/$PAIR\nexport BS=8\nexport NUM_BEAMS=15\nmkdir -p $DATA_DIR\nsacrebleu -t wmt19 -l $PAIR --echo src > $DATA_DIR/val.source\nsacrebleu -t wmt19 -l $PAIR --echo ref > $DATA_DIR/val.target\necho $PAIR\nPYTHONPATH=\"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\n```\nnote: fairseq reports using a beam of 50, so you should get a slightly higher score if re-run with `--num_beams 50`.\n\n## Data Sources\n\n- [training, etc.](http://www.statmt.org/wmt19/)\n- [test set](http://matrix.statmt.org/test_sets/newstest2019.tgz?1556572561)\n\n\n### BibTeX entry and citation info\n\n```bibtex\n@inproceedings{{...,\n  year={{2020}},\n  title={{Facebook FAIR\'s WMT19 News Translation Task Submission}},\n  author={{Ng, Nathan and Yee, Kyra and Baevski, Alexei and Ott, Myle and Auli, Michael and Edunov, Sergey}},\n  booktitle={{Proc. of WMT}},\n}}\n```\n\n\n## TODO\n\n- port model ensemble (fairseq uses 4 model checkpoints)\n\n"
  os.makedirs(lowerCAmelCase__    ,   exist_ok=lowerCAmelCase__						)
  lowerCAmelCase_   :				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
lowercase__		:				Any          =	Path(__file__).resolve().parent.parent.parent
lowercase__		:				Optional[Any]          =	repo_dir / '''model_cards'''

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






import argparse
import re

from flax.traverse_util import flatten_dict, unflatten_dict
from tax import checkpoints

from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration
from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model
from transformers.utils import logging


logging.set_verbosity_info()


# should not include what is already done by the `from_pt` argument
__lowerCamelCase					:	Any       =  {
    '''/attention/''': '''/0/SelfAttention/''',
    '''/self_attention/''': '''/0/SelfAttention/''',
    '''/encoder_decoder_attention/''': '''/1/EncDecAttention/''',
    '''value''': '''v''',
    '''query''': '''q''',
    '''key''': '''k''',
    '''out''': '''o''',
    '''pre_self_attention_layer_norm''': '''0/layer_norm''',
    '''pre_cross_attention_layer_norm''': '''1/layer_norm''',
    '''pre_attention_layer_norm''': '''0/layer_norm''',  # previously 1, but seems wrong
    '''token_embedder''': '''shared''',
    '''encoder_norm''': '''final_layer_norm''',
    '''decoder_norm''': '''final_layer_norm''',
    '''relpos_bias/rel_embedding''': '''block/0/layer/0/SelfAttention/relative_attention_bias/weight''',
    '''router/router_weights/w/''': '''router/classifier/''',
    '''roer/roer_weights/w/''': '''router/classifier/''',
    '''logits_dense''': '''lm_head''',
}




def       _snake_case	(   lowerCAmelCase							:   List[str]	):
		"""simple docstring"""







		SCREAMING_SNAKE_CASE_  :		int				      =    list(s_dict.keys()	)
		for key in keys:
				SCREAMING_SNAKE_CASE_  :		int				      =    R".*/layers_(\d+)"
				SCREAMING_SNAKE_CASE_  :		List[Any]				      =    key
				if re.match(lowerCAmelCase   , lowerCAmelCase	):
						SCREAMING_SNAKE_CASE_  :		Any				      =    re.sub(R"layers_(\d+)"   , R"block/\1/layer"   , lowerCAmelCase	)

				SCREAMING_SNAKE_CASE_  :		List[Any]				      =    R"(encoder|decoder)\/"

				if re.match(lowerCAmelCase   , lowerCAmelCase	):
						SCREAMING_SNAKE_CASE_  :		str				      =    re.match(lowerCAmelCase   , lowerCAmelCase	).groups()
						if groups[0] == "encoder":
								SCREAMING_SNAKE_CASE_  :		Any				      =    re.sub(R"/mlp/"   , R"/1/mlp/"   , lowerCAmelCase	)
								SCREAMING_SNAKE_CASE_  :		List[Any]				      =    re.sub(R"/pre_mlp_layer_norm/"   , R"/1/layer_norm/"   , lowerCAmelCase	)

						elif groups[0] == "decoder":
								SCREAMING_SNAKE_CASE_  :		List[str]				      =    re.sub(R"/mlp/"   , R"/2/mlp/"   , lowerCAmelCase	)
								SCREAMING_SNAKE_CASE_  :		Dict				      =    re.sub(R"/pre_mlp_layer_norm/"   , R"/2/layer_norm/"   , lowerCAmelCase	)

        # 2. Convert other classic mappings
				for old_key, temp_key in MOE_LAYER_NAME_MAPPING.items():
						if old_key in new_key:
								SCREAMING_SNAKE_CASE_  :		List[Any]				      =    new_key.replace(lowerCAmelCase   , lowerCAmelCase	)

				print(f'{key} -> {new_key}'	)
				SCREAMING_SNAKE_CASE_  :		List[Any]				      =    s_dict.pop(lowerCAmelCase	)

		if "encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
				SCREAMING_SNAKE_CASE_  :		str				      =    s_dict[
				    "encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"
				].T
		if "decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
				SCREAMING_SNAKE_CASE_  :		Optional[int]				      =    s_dict[
				    "decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"
				].T

		# 3. Take extra care of the EXPERTS layer
		for key in list(s_dict.keys()	):
				if "expert" in key:
						SCREAMING_SNAKE_CASE_  :		Union[str, Any]				      =    s_dict[key].shape[0]
						SCREAMING_SNAKE_CASE_  :		List[Any]				      =    s_dict[key]
						for idx in range(lowerCAmelCase	):
								SCREAMING_SNAKE_CASE_  :		Tuple				      =    expert_weihts[idx]
								print(f'{key} -> {key.replace("expert/"   , "nested fstring"	)}'	)

						s_dict.pop(lowerCAmelCase	)

		return s_dict


__lowerCamelCase					:	List[Any]       =  {
    '''NUM_ENCODER_LAYERS''': '''num_layers''',
    '''NUM_DECODER_LAYERS''': '''num_decoder_layers''',
    '''NUM_HEADS''': '''num_heads''',
    '''HEAD_DIM''': '''d_kv''',
    '''EMBED_DIM''': '''d_model''',
    '''MLP_DIM''': '''d_ff''',
    '''NUM_SELECTED_EXPERTS''': '''num_selected_experts''',
    '''NUM_ENCODER_SPARSE_LAYERS''': '''num_sparse_encoder_layers''',
    '''NUM_DECODER_SPARSE_LAYERS''': '''num_sparse_decoder_layers''',
    '''dense.MlpBlock.activations''': '''feed_forward_proj''',
}




def       _snake_case	(   lowerCAmelCase							:   Dict   , lowerCAmelCase							:   Optional[Any]	):
		"""simple docstring"""







		import regex as re

		with open(lowerCAmelCase   , "r"	) as f:
				SCREAMING_SNAKE_CASE_  :		Optional[Any]				      =    f.read()

		SCREAMING_SNAKE_CASE_  :		List[str]				      =    re.findall(R"(.*) = ([0-9.]*)"   , lowerCAmelCase	)
		SCREAMING_SNAKE_CASE_  :		Any				      =    {}
		for param, value in regex_match:
				if param in GIN_TO_CONFIG_MAPPING and value != "":
						SCREAMING_SNAKE_CASE_  :		int				      =    float(lowerCAmelCase	) if "." in value else int(lowerCAmelCase	)

		SCREAMING_SNAKE_CASE_  :		Union[str, Any]				      =    re.findall(R"(.*activations) = \(\'(.*)\',\)"   , lowerCAmelCase	)[0]
		SCREAMING_SNAKE_CASE_  :		List[str]				      =    str(activation[1]	)

		SCREAMING_SNAKE_CASE_  :		str				      =    num_experts
		SCREAMING_SNAKE_CASE_  :		Tuple				      =    SwitchTransformersConfig(**lowerCAmelCase	)
		return config




def       _snake_case	(   lowerCAmelCase							:   str   , lowerCAmelCase							:   str   , lowerCAmelCase							:   str=None   , lowerCAmelCase							:   Optional[Any]="./"   , lowerCAmelCase							:   Dict=8	):
		"""simple docstring"""







		print(f'Loading flax weights from : {flax_checkpoint_path}'	)
		SCREAMING_SNAKE_CASE_  :		int				      =    checkpoints.load_tax_checkpoint(lowerCAmelCase	)

		if gin_file is not None:
				SCREAMING_SNAKE_CASE_  :		int				      =    convert_gin_to_config(lowerCAmelCase   , lowerCAmelCase	)
		else:
				SCREAMING_SNAKE_CASE_  :		Dict				      =    SwitchTransformersConfig.from_pretrained(lowerCAmelCase	)

		SCREAMING_SNAKE_CASE_  :		str				      =    SwitchTransformersForConditionalGeneration(lowerCAmelCase	)

		SCREAMING_SNAKE_CASE_  :		Optional[int]				      =    flax_params["target"]
		SCREAMING_SNAKE_CASE_  :		List[str]				      =    flatten_dict(lowerCAmelCase   , sep="/"	)
		SCREAMING_SNAKE_CASE_  :		List[str]				      =    rename_keys(lowerCAmelCase	)
		SCREAMING_SNAKE_CASE_  :		Optional[int]				      =    unflatten_dict(lowerCAmelCase   , sep="/"	)

		# Load the flax params in the PT model
		load_flax_weights_in_pytorch_model(lowerCAmelCase   , lowerCAmelCase	)

		print(f'Save PyTorch model to {pytorch_dump_path}'	)
		pt_model.save_pretrained(lowerCAmelCase	)


if __name__ == "__main__":
		__lowerCamelCase					:	List[str]       =  argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '''--switch_t5x_checkpoint_path''',
		    default=None,
		    type=str,
		    required=True,
		    help=(
		        '''The config json file corresponding to the pre-trained SwitchTransformers model. \nThis specifies the'''
		        ''' model architecture. If not provided, a `gin_file` has to be provided.'''
		    ),
		)
		parser.add_argument(
		    '''--gin_file''',
		    default=None,
		    type=str,
		    required=False,
		    help='''Path to the gin config file. If not provided, a `config_file` has to be passed   ''',
		)
		parser.add_argument(
		    '''--config_name''', default=None, type=str, required=False, help='''Config name of SwitchTransformers model.'''
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', default=None, type=str, required=True, help='''Path to the output pytorch model.'''
		)
		parser.add_argument('''--num_experts''', default=8, type=int, required=False, help='''Number of experts''')
		__lowerCamelCase					:	Any       =  parser.parse_args()
		convert_flax_checkpoint_to_pytorch(
		    args.switch_tax_checkpoint_path,
		    args.config_name,
		    args.gin_file,
		    args.pytorch_dump_folder_path,
		    args.num_experts,
		)
