



'''simple docstring'''

import argparse
import json
import os
from collections import OrderedDict

import torch

from transformers import LukeConfig, LukeForMaskedLM, MLukeTokenizer, XLMRobertaTokenizer
from transformers.tokenization_utils_base import AddedToken






@torch.no_grad()
def    __lowercase	(_SCREAMING_SNAKE_CASE   :int     ,						_SCREAMING_SNAKE_CASE   :Tuple     ,						_SCREAMING_SNAKE_CASE   :Optional[Any]     ,						_SCREAMING_SNAKE_CASE   :List[Any]     ,						_SCREAMING_SNAKE_CASE   :List[Any]							):


				with open(UpperCamelCase__							) as metadata_file:
								SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			json.load(UpperCamelCase__							)
				SCREAMING_SNAKE_CASE	:    Optional[Any]          =			LukeConfig(use_entity_aware_attention=UpperCamelCase__     ,						**metadata['''model_config''']							)

				# Load in the weights from the checkpoint_path
				SCREAMING_SNAKE_CASE	:    Tuple          =			torch.load(UpperCamelCase__     ,						map_location='''cpu'''							)['''module''']

				# Load the entity vocab file
				SCREAMING_SNAKE_CASE	:    List[str]          =			load_original_entity_vocab(UpperCamelCase__							)
				# add an entry for [MASK2]
				SCREAMING_SNAKE_CASE	:    List[str]          =			max(entity_vocab.values()							) + 1
				config.entity_vocab_size += 1

				SCREAMING_SNAKE_CASE	:    Tuple          =			XLMRobertaTokenizer.from_pretrained(metadata['''model_config''']['''bert_model_name''']							)

				# Add special tokens to the token vocabulary for downstream tasks
				SCREAMING_SNAKE_CASE	:    Optional[Any]          =			AddedToken('''<ent>'''     ,						lstrip=UpperCamelCase__     ,						rstrip=UpperCamelCase__							)
				SCREAMING_SNAKE_CASE	:    Optional[int]          =			AddedToken('''<ent2>'''     ,						lstrip=UpperCamelCase__     ,						rstrip=UpperCamelCase__							)
				tokenizer.add_special_tokens({'''additional_special_tokens''': [entity_token_a, entity_token_a]}							)
				config.vocab_size += 2

				print(F'''Saving tokenizer to {pytorch_dump_folder_path}'''							)
				tokenizer.save_pretrained(UpperCamelCase__							)
				with open(os.path.join(UpperCamelCase__     ,						'''tokenizer_config.json'''							)     ,						'''r'''							) as f:
								SCREAMING_SNAKE_CASE	:    Dict          =			json.load(UpperCamelCase__							)
				SCREAMING_SNAKE_CASE	:    str          =			'''MLukeTokenizer'''
				with open(os.path.join(UpperCamelCase__     ,						'''tokenizer_config.json'''							)     ,						'''w'''							) as f:
								json.dump(UpperCamelCase__     ,						UpperCamelCase__							)

				with open(os.path.join(UpperCamelCase__     ,						MLukeTokenizer.vocab_files_names['''entity_vocab_file''']							)     ,						'''w'''							) as f:
								json.dump(UpperCamelCase__     ,						UpperCamelCase__							)

				SCREAMING_SNAKE_CASE	:    int          =			MLukeTokenizer.from_pretrained(UpperCamelCase__							)

				# Initialize the embeddings of the special tokens
				SCREAMING_SNAKE_CASE	:    Optional[int]          =			tokenizer.convert_tokens_to_ids(['''@''']							)[0]
				SCREAMING_SNAKE_CASE	:    Tuple          =			tokenizer.convert_tokens_to_ids(['''#''']							)[0]

				SCREAMING_SNAKE_CASE	:    Any          =			state_dict['''embeddings.word_embeddings.weight''']
				SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			word_emb[ent_init_index].unsqueeze(0							)
				SCREAMING_SNAKE_CASE	:    Any          =			word_emb[enta_init_index].unsqueeze(0							)
				SCREAMING_SNAKE_CASE	:    str          =			torch.cat([word_emb, ent_emb, enta_emb]							)
				# add special tokens for 'entity_predictions.bias'
				for bias_name in ["lm_head.decoder.bias", "lm_head.bias"]:
								SCREAMING_SNAKE_CASE	:    Dict          =			state_dict[bias_name]
								SCREAMING_SNAKE_CASE	:    int          =			decoder_bias[ent_init_index].unsqueeze(0							)
								SCREAMING_SNAKE_CASE	:    Optional[Any]          =			decoder_bias[enta_init_index].unsqueeze(0							)
								SCREAMING_SNAKE_CASE	:    Dict          =			torch.cat([decoder_bias, ent_decoder_bias, enta_decoder_bias]							)

				# Initialize the query layers of the entity-aware self-attention mechanism
				for layer_index in range(config.num_hidden_layers							):
								for matrix_name in ["query.weight", "query.bias"]:
												SCREAMING_SNAKE_CASE	:    Dict          =			F'''encoder.layer.{layer_index}.attention.self.'''
												SCREAMING_SNAKE_CASE	:    Tuple          =			state_dict[prefix + matrix_name]
												SCREAMING_SNAKE_CASE	:    int          =			state_dict[prefix + matrix_name]
												SCREAMING_SNAKE_CASE	:    Optional[int]          =			state_dict[prefix + matrix_name]

    # Initialize the embedding of the [MASK2] entity using that of the [MASK] entity for downstream tasks
				SCREAMING_SNAKE_CASE	:    str          =			state_dict['''entity_embeddings.entity_embeddings.weight''']
				SCREAMING_SNAKE_CASE	:    List[str]          =			entity_emb[entity_vocab['''[MASK]''']].unsqueeze(0							)
				SCREAMING_SNAKE_CASE	:    str          =			torch.cat([entity_emb, entity_mask_emb]							)
				# add [MASK2] for 'entity_predictions.bias'
				SCREAMING_SNAKE_CASE	:    Dict          =			state_dict['''entity_predictions.bias''']
				SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			entity_prediction_bias[entity_vocab['''[MASK]''']].unsqueeze(0							)
				SCREAMING_SNAKE_CASE	:    List[str]          =			torch.cat([entity_prediction_bias, entity_mask_bias]							)

				SCREAMING_SNAKE_CASE	:    Optional[int]          =			LukeForMaskedLM(config=UpperCamelCase__							).eval()

				state_dict.pop('''entity_predictions.decoder.weight'''							)
				state_dict.pop('''lm_head.decoder.weight'''							)
				state_dict.pop('''lm_head.decoder.bias'''							)
				SCREAMING_SNAKE_CASE	:    List[str]          =			OrderedDict()
				for key, value in state_dict.items():
								if not (key.startswith('''lm_head'''							) or key.startswith('''entity_predictions'''							)):
												SCREAMING_SNAKE_CASE	:    Optional[Any]          =			state_dict[key]
								else:
												SCREAMING_SNAKE_CASE	:    List[Any]          =			state_dict[key]

				SCREAMING_SNAKE_CASE       ,       SCREAMING_SNAKE_CASE	:    str          =			model.load_state_dict(UpperCamelCase__     ,						strict=UpperCamelCase__							)

				if set(UpperCamelCase__							) != {"luke.embeddings.position_ids"}:
								raise ValueError(F'''Unexpected unexpected_keys: {unexpected_keys}'''							)
				if set(UpperCamelCase__							) != {
				    "lm_head.decoder.weight",
				    "lm_head.decoder.bias",
				    "entity_predictions.decoder.weight",
				}:
								raise ValueError(F'''Unexpected missing_keys: {missing_keys}'''							)

				model.tie_weights()
				assert (model.luke.embeddings.word_embeddings.weight == model.lm_head.decoder.weight).all()
				assert (model.luke.entity_embeddings.entity_embeddings.weight == model.entity_predictions.decoder.weight).all()

				# Check outputs
				SCREAMING_SNAKE_CASE	:    List[str]          =			MLukeTokenizer.from_pretrained(UpperCamelCase__     ,						task='''entity_classification'''							)

				SCREAMING_SNAKE_CASE	:    Optional[Any]          =			'''ISO 639-3 uses the code fas for the dialects spoken across Iran and アフガニスタン (Afghanistan).'''
				SCREAMING_SNAKE_CASE	:    Optional[int]          =			(0, 9)
				SCREAMING_SNAKE_CASE	:    Any          =			tokenizer(UpperCamelCase__     ,						entity_spans=[span]     ,						return_tensors='''pt'''							)

				SCREAMING_SNAKE_CASE	:    str          =			model(**UpperCamelCase__							)

				# Verify word hidden states
				if model_size == "large":
								raise NotImplementedError
				else:  # base
								SCREAMING_SNAKE_CASE	:    Optional[Any]          =			torch.Size((1, 33, 7_68)							)
								SCREAMING_SNAKE_CASE	:    List[Any]          =			torch.tensor([[0.0892, 0.0596, -0.2819], [0.0134, 0.1199, 0.0573], [-0.0169, 0.0927, 0.0644]]							)

				if not (outputs.last_hidden_state.shape == expected_shape):
								raise ValueError(
								    F'''Outputs.last_hidden_state.shape is {outputs.last_hidden_state.shape}, Expected shape is {expected_shape}'''							)
				if not torch.allclose(outputs.last_hidden_state[0, :3, :3]     ,						UpperCamelCase__     ,						atol=1E-4							):
								raise ValueError

				# Verify entity hidden states
				if model_size == "large":
								raise NotImplementedError
				else:  # base
								SCREAMING_SNAKE_CASE	:    Tuple          =			torch.Size((1, 1, 7_68)							)
								SCREAMING_SNAKE_CASE	:    Any          =			torch.tensor([[-0.1482, 0.0609, 0.0322]]							)

				if not (outputs.entity_last_hidden_state.shape == expected_shape):
								raise ValueError(
								    F'''Outputs.entity_last_hidden_state.shape is {outputs.entity_last_hidden_state.shape}, Expected shape is'''
								    F''' {expected_shape}'''							)
				if not torch.allclose(outputs.entity_last_hidden_state[0, :3, :3]     ,						UpperCamelCase__     ,						atol=1E-4							):
								raise ValueError

				# Verify masked word/entity prediction
				SCREAMING_SNAKE_CASE	:    str          =			MLukeTokenizer.from_pretrained(UpperCamelCase__							)
				SCREAMING_SNAKE_CASE	:    Tuple          =			'''Tokyo is the capital of <mask>.'''
				SCREAMING_SNAKE_CASE	:    int          =			(24, 30)
				SCREAMING_SNAKE_CASE	:    int          =			tokenizer(UpperCamelCase__     ,						entity_spans=[span]     ,						return_tensors='''pt'''							)

				SCREAMING_SNAKE_CASE	:    Tuple          =			model(**UpperCamelCase__							)

				SCREAMING_SNAKE_CASE	:    Tuple          =			encoding['''input_ids'''][0].tolist()
				SCREAMING_SNAKE_CASE	:    List[Any]          =			input_ids.index(tokenizer.convert_tokens_to_ids('''<mask>'''							)							)
				SCREAMING_SNAKE_CASE	:    Any          =			outputs.logits[0][mask_position_id].argmax(dim=-1							)
				assert "Japan" == tokenizer.decode(UpperCamelCase__							)

				SCREAMING_SNAKE_CASE	:    Dict          =			outputs.entity_logits[0][0].argmax().item()
				SCREAMING_SNAKE_CASE	:    Union[str, Any]          =			[
				    entity for entity, entity_id in tokenizer.entity_vocab.items() if entity_id == predicted_entity_id
				]
				assert [e for e in multilingual_predicted_entities if e.startswith('''en:'''							)][0] == "en:Japan"

				# Finally, save our PyTorch model and tokenizer
				print('''Saving PyTorch model to {}'''.format(UpperCamelCase__							)							)
				model.save_pretrained(UpperCamelCase__							)






def    __lowercase	(_SCREAMING_SNAKE_CASE   :str							):
				SCREAMING_SNAKE_CASE	:    int          =			['''[MASK]''', '''[PAD]''', '''[UNK]''']

				SCREAMING_SNAKE_CASE	:    Dict          =			[json.loads(UpperCamelCase__							) for line in open(UpperCamelCase__							)]

				SCREAMING_SNAKE_CASE	:    int          =			{}
				for entry in data:
								SCREAMING_SNAKE_CASE	:    Dict          =			entry['''id''']
								for entity_name, language in entry["entities"]:
												if entity_name in SPECIAL_TOKENS:
																SCREAMING_SNAKE_CASE	:    Tuple          =			entity_id
																break
												SCREAMING_SNAKE_CASE	:    List[str]          =			F'''{language}:{entity_name}'''
												SCREAMING_SNAKE_CASE	:    str          =			entity_id
				return new_mapping


if __name__ == "__main__":
	snake_case_      =  argparse.ArgumentParser()
	# Required parameters
	parser.add_argument("""--checkpoint_path""", type=str, help="""Path to a pytorch_model.bin file.""")
	parser.add_argument(
	    """--metadata_path""", default=None, type=str, help="""Path to a metadata.json file, defining the configuration."""
	)
	parser.add_argument(
	    """--entity_vocab_path""",
	    default=None,
	    type=str,
	    help="""Path to an entity_vocab.tsv file, containing the entity vocabulary.""",
	)
	parser.add_argument(
	    """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to where to dump the output PyTorch model."""
	)
	parser.add_argument(
	    """--model_size""", default="""base""", type=str, choices=["""base""", """large"""], help="""Size of the model to be converted."""
	)
	snake_case_      =  parser.parse_args()
	convert_luke_checkpoint(
	    args.checkpoint_path,
	    args.metadata_path,
	    args.entity_vocab_path,
	    args.pytorch_dump_folder_path,
	    args.model_size,
	)










# 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)


