



'''simple docstring'''

import argparse
import json
import os

import tensorstore as ts
import torch
from flax import serialization
from flax.traverse_util import flatten_dict, unflatten_dict
from tensorflow.io import gfile

from transformers.modeling_utils import dtype_byte_size
from transformers.models.switch_transformers.convert_switch_transformers_original_flax_checkpoint_to_pytorch import (
    rename_keys,
)
from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME
from transformers.utils.hub import convert_file_size_to_int






def    __lowercase	(_SCREAMING_SNAKE_CASE   :int     ,						_SCREAMING_SNAKE_CASE   :Tuple							):


				if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 3:
								# expert layer
								SCREAMING_SNAKE_CASE	:    str          =			flax_key_tuple[:-1] + ('''weight''',)
								SCREAMING_SNAKE_CASE	:    List[Any]          =			torch.permute(UpperCamelCase__     ,						(0, 2, 1)							)
				elif flax_key_tuple[-1] == "kernel" and ".".join(UpperCamelCase__							):
								# linear layer
								SCREAMING_SNAKE_CASE	:    Dict          =			flax_key_tuple[:-1] + ('''weight''',)
								SCREAMING_SNAKE_CASE	:    Optional[Any]          =			flax_tensor.T
				elif flax_key_tuple[-1] in ["scale", "embedding"]:
								SCREAMING_SNAKE_CASE	:    str          =			flax_key_tuple[:-1] + ('''weight''',)

				return flax_key_tuple, flax_tensor






def    __lowercase	(_SCREAMING_SNAKE_CASE   :Tuple     ,						_SCREAMING_SNAKE_CASE   :Optional[Any]     ,						_SCREAMING_SNAKE_CASE   :Tuple							):


				if "metadata" in layer:
								SCREAMING_SNAKE_CASE	:    str          =			layer.split('''metadata'''							)
								SCREAMING_SNAKE_CASE	:    Optional[int]          =			''''''.join(split_layer[0]							)[:-1]
								SCREAMING_SNAKE_CASE	:    Optional[Any]          =			[tuple(('''metadata''' + split_layer[1]).split('''/'''							)							)]
				elif "kvstore" in layer:
								SCREAMING_SNAKE_CASE	:    List[str]          =			layer.split('''kvstore'''							)
								SCREAMING_SNAKE_CASE	:    Optional[Any]          =			''''''.join(split_layer[0]							)[:-1]
								SCREAMING_SNAKE_CASE	:    str          =			[tuple(('''kvstore''' + split_layer[1]).split('''/'''							)							)]

				else:
								SCREAMING_SNAKE_CASE	:    Tuple          =			layer.split('''/'''							)
								SCREAMING_SNAKE_CASE	:    Optional[Any]          =			'''/'''.join(split_layer[:-1]							)
								SCREAMING_SNAKE_CASE	:    Any          =			(split_layer[-1],)

				if "kvstore/path" in layer:
								SCREAMING_SNAKE_CASE	:    Optional[int]          =			F'''{switch_checkpoint_path}/{checkpoint_info[layer]}'''
				elif "kvstore/driver" in layer:
								SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			'''file'''
				else:
								SCREAMING_SNAKE_CASE	:    Tuple          =			checkpoint_info[layer]

				return curr_real_layer_name, split_layer, content






def    __lowercase	(_SCREAMING_SNAKE_CASE   :Optional[Any]     ,						_SCREAMING_SNAKE_CASE   :Tuple							):
				SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			rename_keys(UpperCamelCase__							)
				SCREAMING_SNAKE_CASE	:    Tuple          =			{}
				for k, v in current_block.items():
								SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			v
				SCREAMING_SNAKE_CASE	:    str          =			new_current_block
				torch.save(UpperCamelCase__     ,						UpperCamelCase__							)






def    __lowercase	(_SCREAMING_SNAKE_CASE   :Dict     ,						_SCREAMING_SNAKE_CASE   :Union[str, Any]     ,						_SCREAMING_SNAKE_CASE   :Optional[Any]     ,						_SCREAMING_SNAKE_CASE   :Optional[Any]     ,						_SCREAMING_SNAKE_CASE   :str = WEIGHTS_NAME							):
				SCREAMING_SNAKE_CASE	:    Optional[int]          =			convert_file_size_to_int(UpperCamelCase__							)
				SCREAMING_SNAKE_CASE	:    str          =			[]
				SCREAMING_SNAKE_CASE	:    str          =			{}
				SCREAMING_SNAKE_CASE	:    List[Any]          =			0
				SCREAMING_SNAKE_CASE	:    str          =			0

				os.makedirs(UpperCamelCase__     ,						exist_ok=UpperCamelCase__							)
				with gfile.GFile(switch_checkpoint_path + '''/checkpoint'''     ,						'''rb'''							) as fp:
								SCREAMING_SNAKE_CASE	:    List[str]          =			serialization.msgpack_restore(fp.read()							)['''optimizer''']['''target''']
								SCREAMING_SNAKE_CASE	:    Optional[Any]          =			flatten_dict(UpperCamelCase__     ,						sep='''/'''							)

				SCREAMING_SNAKE_CASE	:    List[Any]          =			{}
				for layer in checkpoint_info.keys():
								SCREAMING_SNAKE_CASE       ,       SCREAMING_SNAKE_CASE       ,       SCREAMING_SNAKE_CASE	:    Optional[Any]          =			get_key_and_tensorstore_dict(
								    UpperCamelCase__     ,						UpperCamelCase__     ,						UpperCamelCase__							)
								if curr_real_layer_name in all_layers:
												SCREAMING_SNAKE_CASE	:    str          =			content
								else:
												SCREAMING_SNAKE_CASE	:    Dict          =			{split_layer[-1]: content}

				for key in all_layers.keys():
								# open tensorstore file
								SCREAMING_SNAKE_CASE	:    Optional[Any]          =			ts.open(unflatten_dict(all_layers[key]							)							).result().read().result()
								SCREAMING_SNAKE_CASE	:    int          =			torch.tensor(UpperCamelCase__							)
								SCREAMING_SNAKE_CASE	:    Optional[Any]          =			raw_weights.numel() * dtype_byte_size(raw_weights.dtype							)

								# use the renaming pattern from the small conversion scripts
								SCREAMING_SNAKE_CASE       ,       SCREAMING_SNAKE_CASE	:    List[Any]          =			rename_base_flax_keys(tuple(key.split('''/'''							)							)     ,						UpperCamelCase__							)
								SCREAMING_SNAKE_CASE	:    str          =			'''/'''.join(UpperCamelCase__							)

								# If this weight is going to tip up over the maximal size, we split.
								if current_block_size + weight_size > max_shard_size:
												SCREAMING_SNAKE_CASE	:    Dict          =			os.path.join(
												    UpperCamelCase__     ,						weights_name.replace('''.bin'''     ,						F'''-{len(UpperCamelCase__							)+1:05d}-of-???.bin'''							)							)
												rename_and_save_block(UpperCamelCase__     ,						UpperCamelCase__							)
												sharded_state_dicts.append(current_block.keys()							)
												del current_block
												SCREAMING_SNAKE_CASE	:    Dict          =			{}
												SCREAMING_SNAKE_CASE	:    List[str]          =			0

								SCREAMING_SNAKE_CASE	:    Any          =			raw_weights.to(getattr(UpperCamelCase__     ,						UpperCamelCase__							)							)
								current_block_size += weight_size
								total_size += weight_size

				# Add the last block
				SCREAMING_SNAKE_CASE	:    Optional[Any]          =			os.path.join(UpperCamelCase__     ,						weights_name.replace('''.bin'''     ,						F'''-{len(UpperCamelCase__							)+1:05d}-of-???.bin'''							)							)
				rename_and_save_block(UpperCamelCase__     ,						UpperCamelCase__							)
				sharded_state_dicts.append(current_block.keys()							)

				# If we only have one shard, we return it
				if len(UpperCamelCase__							) == 1:
								return {weights_name: sharded_state_dicts[0]}, None

				# Otherwise, let's build the index
				SCREAMING_SNAKE_CASE	:    Dict          =			{}
				SCREAMING_SNAKE_CASE	:    Optional[Any]          =			{}
				for idx, shard in enumerate(UpperCamelCase__							):
								SCREAMING_SNAKE_CASE	:    str          =			weights_name.replace(
								    '''.bin'''     ,						F'''-{idx+1:05d}-of-{len(UpperCamelCase__							):05d}.bin'''							)  # len(sharded_state_dicts):05d}
								SCREAMING_SNAKE_CASE	:    Optional[int]          =			os.path.join(UpperCamelCase__     ,						weights_name.replace('''.bin'''     ,						F'''-{idx+1:05d}-of-???.bin'''							)							)
								os.rename(UpperCamelCase__     ,						os.path.join(UpperCamelCase__     ,						UpperCamelCase__							)							)
								SCREAMING_SNAKE_CASE	:    Optional[Any]          =			shard
								for key in shard:
												SCREAMING_SNAKE_CASE	:    List[str]          =			shard_file

    # Add the metadata
				SCREAMING_SNAKE_CASE	:    str          =			{'''total_size''': total_size}
				SCREAMING_SNAKE_CASE	:    Tuple          =			{'''metadata''': metadata, '''weight_map''': weight_map}

				with open(os.path.join(UpperCamelCase__     ,						UpperCamelCase__							)     ,						'''w'''     ,						encoding='''utf-8'''							) as f:
								SCREAMING_SNAKE_CASE	:    Tuple          =			json.dumps(UpperCamelCase__     ,						indent=2     ,						sort_keys=UpperCamelCase__							) + '''\n'''
								f.write(UpperCamelCase__							)

				return metadata, index


if __name__ == "__main__":
	snake_case_      =  argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    """--switch_t5x_checkpoint_path""",
	    default="""/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128/checkpoint_634600""",
	    type=str,
	    required=False,
	    help="""Path to a directory containing a folder per layer. Follows the original Google format.""",
	)
	parser.add_argument("""--max_shard_size""", default="""10GB""", required=False, help="""Max shard size""")
	parser.add_argument("""--dtype""", default="""bfloat16""", type=str, required=False, help="""dtype of the saved model""")
	parser.add_argument(
	    """--pytorch_dump_folder_path""",
	    default="""/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128-converted""",
	    type=str,
	    required=False,
	    help="""Path to the output pytorch model.""",
	)
	snake_case_      =  parser.parse_args()
	shard_on_the_fly(
	    args.switch_tax_checkpoint_path,
	    args.pytorch_dump_folder_path,
	    args.max_shard_size,
	    args.dtype,
	)






def    __lowercase	():
				from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration, TaTokenizer

				SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			SwitchTransformersConfig.from_pretrained('''google/switch-base-8'''							)
				config.save_pretrained('''/home/arthur_huggingface_co/transformers/switch_converted'''							)
				SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			SwitchTransformersForConditionalGeneration.from_pretrained(
				    '''/home/arthur_huggingface_co/transformers/switch_converted'''     ,						device_map='''auto'''							)

				SCREAMING_SNAKE_CASE	:    List[Any]          =			TaTokenizer.from_pretrained('''t5-small'''							)
				SCREAMING_SNAKE_CASE	:    Optional[int]          =			'''A <extra_id_0> walks into a bar a orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.'''

				SCREAMING_SNAKE_CASE	:    List[Any]          =			tokenizer(UpperCamelCase__     ,						return_tensors='''pt'''							).input_ids
				SCREAMING_SNAKE_CASE	:    List[Any]          =			model.generate(UpperCamelCase__     ,						decoder_start_token_id=0							)
				print(tokenizer.decode(out[0]							)							)










# Usage:
# ./gen-card-allenai-wmt16.py

import os
from pathlib import Path







def      lowerCamelCase_						(			UpperCamelCase__							: int					,	UpperCamelCase__							: Optional[int]					,	UpperCamelCase__							: int					,	UpperCamelCase__							: str )	->			List[str]:





				"""simple docstring"""


				__lowerCamelCase									=       {
				    'en': 'Machine learning is great, isn\'t it?',
				    'ru': 'Машинное обучение - это здорово, не так ли?',
				    'de': 'Maschinelles Lernen ist großartig, nicht wahr?',
				}

				# BLUE scores as follows:
				# "pair": [fairseq, transformers]
				__lowerCamelCase									=       {
				    'wmt16-en-de-dist-12-1': [28.3, 27.52],
				    'wmt16-en-de-dist-6-1': [27.4, 27.11],
				    'wmt16-en-de-12-1': [26.9, 25.75],
				}
				__lowerCamelCase									=       F"""{src_lang}-{tgt_lang}"""

				__lowerCamelCase									=       F"""
---
language:
- {src_lang}
- {tgt_lang}
thumbnail:
tags:
- translation
- wmt16
- allenai
license: apache-2.0
datasets:
- wmt16
metrics:
- bleu
---

# FSMT

## Model description

This is a ported version of fairseq-based [wmt16 transformer](https://github.com/jungokasai/deep-shallow/) for {src_lang}-{tgt_lang}.

For more details, please, see [Deep Encoder, Shallow Decoder: Reevaluating the Speed-Quality Tradeoff in Machine Translation](https://arxiv.org/abs/2006.10369).

All 3 models are available:

* [wmt16-en-de-dist-12-1](https://huggingface.co/allenai/wmt16-en-de-dist-12-1)
* [wmt16-en-de-dist-6-1](https://huggingface.co/allenai/wmt16-en-de-dist-6-1)
* [wmt16-en-de-12-1](https://huggingface.co/allenai/wmt16-en-de-12-1)


## Intended uses & limitations

#### How to use

```python
from transformers import FSMTForConditionalGeneration, FSMTTokenizer
mname = \"allenai/{model_name}\"
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


## Training data

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

## Eval results

Here are the BLEU scores:

model   | fairseq | transformers
-------|---------|----------
{model_name}  | {scores[model_name][0]} | {scores[model_name][1]}

The score is slightly below the score reported in the paper, as the researchers don't use `sacrebleu` and measure the score on tokenized outputs. `transformers` score was measured using `sacrebleu` on detokenized outputs.

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=5
mkdir -p $DATA_DIR
sacrebleu -t wmt16 -l $PAIR --echo src > $DATA_DIR/val.source
sacrebleu -t wmt16 -l $PAIR --echo ref > $DATA_DIR/val.target
echo $PAIR
PYTHONPATH=\"src:examples/seq2seq\" python examples/seq2seq/run_eval.py allenai/{model_name} $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
```

## Data Sources

- [training, etc.](http://www.statmt.org/wmt16/)
- [test set](http://matrix.statmt.org/test_sets/newstest2016.tgz?1504722372)


### BibTeX entry and citation info

```
@misc{{kasai2020deep,
    title={{Deep Encoder, Shallow Decoder: Reevaluating the Speed-Quality Tradeoff in Machine Translation}},
    author={{Jungo Kasai and Nikolaos Pappas and Hao Peng and James Cross and Noah A. Smith}},
    year={{2020}},
    eprint={{2006.10369}},
    archivePrefix={{arXiv}},
    primaryClass={{cs.CL}}
}}
```

"""
				model_card_dir.mkdir(parents=UpperCamelCase__					,	exist_ok=UpperCamelCase__ )
				__lowerCamelCase									=       os.path.join(UpperCamelCase__					,	'README.md' )
				print(F"""Generating {path}""" )
				with open(UpperCamelCase__					,	'w'					,	encoding='utf-8' ) as f:
								f.write(UpperCamelCase__ )


# make sure we are under the root of the project
__A								=       Path(__file__).resolve().parent.parent.parent
__A								=       repo_dir / "model_cards"

for model_name in ["wmt16-en-de-dist-12-1", "wmt16-en-de-dist-6-1", "wmt16-en-de-12-1"]:
					__A								=       model_cards_dir / "allenai" / model_name
					write_model_card(model_card_dir, src_lang="en", tgt_lang="de", model_name=model_name)


