





import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import BeitConfig, BeitForImageClassification, BeitForMaskedImageModeling, BeitImageProcessor
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
_a						=   logging.get_logger(__name__)





def lowerCAmelCase__(__snake_case ,__snake_case=False ,__snake_case=False							)       ->   List[Any]:







	'''simple docstring'''







	lowerCamelCase__              =	'''backbone.''' if is_semantic else ''''''

	lowerCamelCase__              =	[]
	for i in range(config.num_hidden_layers							):
		# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
		rename_keys.append((F'{prefix}blocks.{i}.norm1.weight', F'beit.encoder.layer.{i}.layernorm_before.weight')							)
		rename_keys.append((F'{prefix}blocks.{i}.norm1.bias', F'beit.encoder.layer.{i}.layernorm_before.bias')							)
		rename_keys.append(
		    (F'{prefix}blocks.{i}.attn.proj.weight', F'beit.encoder.layer.{i}.attention.output.dense.weight')							)
		rename_keys.append(
		    (F'{prefix}blocks.{i}.attn.proj.bias', F'beit.encoder.layer.{i}.attention.output.dense.bias')							)
		rename_keys.append((F'{prefix}blocks.{i}.norm2.weight', F'beit.encoder.layer.{i}.layernorm_after.weight')							)
		rename_keys.append((F'{prefix}blocks.{i}.norm2.bias', F'beit.encoder.layer.{i}.layernorm_after.bias')							)
		rename_keys.append((F'{prefix}blocks.{i}.mlp.fc1.weight', F'beit.encoder.layer.{i}.intermediate.dense.weight')							)
		rename_keys.append((F'{prefix}blocks.{i}.mlp.fc1.bias', F'beit.encoder.layer.{i}.intermediate.dense.bias')							)
		rename_keys.append((F'{prefix}blocks.{i}.mlp.fc2.weight', F'beit.encoder.layer.{i}.output.dense.weight')							)
		rename_keys.append((F'{prefix}blocks.{i}.mlp.fc2.bias', F'beit.encoder.layer.{i}.output.dense.bias')							)

	# projection layer + position embeddings
	rename_keys.extend(
	    [
	        (F'{prefix}cls_token', '''beit.embeddings.cls_token'''),
	        (F'{prefix}patch_embed.proj.weight', '''beit.embeddings.patch_embeddings.projection.weight'''),
	        (F'{prefix}patch_embed.proj.bias', '''beit.embeddings.patch_embeddings.projection.bias'''),
	        (F'{prefix}pos_embed', '''beit.embeddings.position_embeddings'''),
	    ]							)

	if has_lm_head:
		# mask token + layernorm
		rename_keys.extend(
		    [
		        ('''mask_token''', '''beit.embeddings.mask_token'''),
		        ('''norm.weight''', '''layernorm.weight'''),
		        ('''norm.bias''', '''layernorm.bias'''),
		    ]							)
	else:
		# layernorm + classification head
		rename_keys.extend(
		    [
		        ('''fc_norm.weight''', '''beit.pooler.layernorm.weight'''),
		        ('''fc_norm.bias''', '''beit.pooler.layernorm.bias'''),
		        ('''head.weight''', '''classifier.weight'''),
		        ('''head.bias''', '''classifier.bias'''),
		    ]							)

	return rename_keys





def lowerCAmelCase__(__snake_case ,__snake_case ,__snake_case=False ,__snake_case=False							)       ->   Optional[int]:







	'''simple docstring'''







	for i in range(config.num_hidden_layers							):
		lowerCamelCase__              =	'''backbone.''' if is_semantic else ''''''
		# queries, keys and values
		lowerCamelCase__              =	state_dict.pop(F'{prefix}blocks.{i}.attn.qkv.weight'							)
		lowerCamelCase__              =	state_dict.pop(F'{prefix}blocks.{i}.attn.q_bias'							)
		lowerCamelCase__              =	state_dict.pop(F'{prefix}blocks.{i}.attn.v_bias'							)

		lowerCamelCase__              =	in_proj_weight[
		    : config.hidden_size, :
		]
		lowerCamelCase__              =	q_bias
		lowerCamelCase__              =	in_proj_weight[
		    config.hidden_size : config.hidden_size * 2, :
		]
		lowerCamelCase__              =	in_proj_weight[
		    -config.hidden_size :, :
		]
		lowerCamelCase__              =	v_bias

		# gamma_1 and gamma_2
		# we call them lambda because otherwise they are renamed when using .from_pretrained
		lowerCamelCase__              =	state_dict.pop(F'{prefix}blocks.{i}.gamma_1'							)
		lowerCamelCase__              =	state_dict.pop(F'{prefix}blocks.{i}.gamma_2'							)

		lowerCamelCase__              =	gamma_a
		lowerCamelCase__              =	gamma_a





def lowerCAmelCase__(__snake_case ,__snake_case ,__snake_case							)       ->   int:







	'''simple docstring'''







	lowerCamelCase__              =	dct.pop(__snake_case							)
	lowerCamelCase__              =	val





def lowerCAmelCase__()       ->   List[Any]:







	'''simple docstring'''







	lowerCamelCase__              =	'''http://images.cocodataset.org/val2017/000000039769.jpg'''
	lowerCamelCase__              =	Image.open(requests.get(__snake_case ,stream=__snake_case							).raw							)
	return im





@torch.no_grad()
def lowerCAmelCase__(__snake_case ,__snake_case ,__snake_case=False							)       ->   Union[str, Any]:







	'''simple docstring'''







	lowerCamelCase__              =	False if '''rvlcdip''' in checkpoint_url else True
	lowerCamelCase__              =	BeitConfig(use_absolute_position_embeddings=__snake_case ,use_mask_token=__snake_case							)

	# size of the architecture
	if "large" in checkpoint_url or "dit-l" in checkpoint_url:
		lowerCamelCase__              =	1024
		lowerCamelCase__              =	4096
		lowerCamelCase__              =	24
		lowerCamelCase__              =	16

	# labels
	if "rvlcdip" in checkpoint_url:
		lowerCamelCase__              =	16
		lowerCamelCase__              =	'''huggingface/label-files'''
		lowerCamelCase__              =	'''rvlcdip-id2label.json'''
		lowerCamelCase__              =	json.load(open(hf_hub_download(__snake_case ,__snake_case ,repo_type='''dataset'''							) ,'''r'''							)							)
		lowerCamelCase__              =	{int(__snake_case							): v for k, v in idalabel.items()}
		lowerCamelCase__              =	idalabel
		lowerCamelCase__              =	{v: k for k, v in idalabel.items()}

	# load state_dict of original model, remove and rename some keys
	lowerCamelCase__              =	torch.hub.load_state_dict_from_url(__snake_case ,map_location='''cpu'''							)['''model''']

	lowerCamelCase__              =	create_rename_keys(__snake_case ,has_lm_head=__snake_case							)
	for src, dest in rename_keys:
		rename_key(__snake_case ,__snake_case ,__snake_case							)
	read_in_q_k_v(__snake_case ,__snake_case ,has_lm_head=__snake_case							)

	# load HuggingFace model
	lowerCamelCase__              =	BeitForMaskedImageModeling(__snake_case							) if has_lm_head else BeitForImageClassification(__snake_case							)
	model.eval()
	model.load_state_dict(__snake_case							)

	# Check outputs on an image
	lowerCamelCase__              =	BeitImageProcessor(
	    size=config.image_size ,resample=PILImageResampling.BILINEAR ,do_center_crop=__snake_case							)
	lowerCamelCase__              =	prepare_img()

	lowerCamelCase__              =	image_processor(images=__snake_case ,return_tensors='''pt'''							)
	lowerCamelCase__              =	encoding['''pixel_values''']

	lowerCamelCase__              =	model(__snake_case							)
	lowerCamelCase__              =	outputs.logits

	# verify logits
	lowerCamelCase__              =	[1, 16] if '''rvlcdip''' in checkpoint_url else [1, 196, 8192]
	assert logits.shape == torch.Size(__snake_case							), "Shape of logits not as expected"

	Path(__snake_case							).mkdir(exist_ok=__snake_case							)
	print(F'Saving model to {pytorch_dump_folder_path}'							)
	model.save_pretrained(__snake_case							)
	print(F'Saving image processor to {pytorch_dump_folder_path}'							)
	image_processor.save_pretrained(__snake_case							)

	if push_to_hub:
		if has_lm_head:
			lowerCamelCase__              =	'''dit-base''' if '''base''' in checkpoint_url else '''dit-large'''
		else:
			lowerCamelCase__              =	'''dit-base-finetuned-rvlcdip''' if '''dit-b''' in checkpoint_url else '''dit-large-finetuned-rvlcdip'''
		image_processor.push_to_hub(
		    repo_path_or_name=Path(__snake_case ,__snake_case							) ,organization='''nielsr''' ,commit_message='''Add image processor''' ,use_temp_dir=__snake_case ,)
		model.push_to_hub(
		    repo_path_or_name=Path(__snake_case ,__snake_case							) ,organization='''nielsr''' ,commit_message='''Add model''' ,use_temp_dir=__snake_case ,)


if __name__ == "__main__":
			_a						=   argparse.ArgumentParser()

			parser.add_argument(
			    "--checkpoint_url",
			    default="https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-base-224-p16-500k-62d53a.pth",
			    type=str,
			    help="URL to the original PyTorch checkpoint (.pth file).",
			)
			parser.add_argument(
			    "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
			)
			parser.add_argument(
			    "--push_to_hub",
			    action="store_true",
			)
			_a						=   parser.parse_args()
			convert_dit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)










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

import os
from pathlib import Path





def lowerCAmelCase__(__snake_case ,__snake_case ,__snake_case ,__snake_case							)       ->   Any:







	'''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''': [2_8.3, 2_7.5_2],
	    '''wmt16-en-de-dist-6-1''': [2_7.4, 2_7.1_1],
	    '''wmt16-en-de-12-1''': [2_6.9, 2_5.7_5],
	}
	lowerCamelCase__              =	F'{src_lang}-{tgt_lang}'

	lowerCamelCase__              =	F'\n---\nlanguage:\n- {src_lang}\n- {tgt_lang}\nthumbnail:\ntags:\n- translation\n- wmt16\n- allenai\nlicense: apache-2.0\ndatasets:\n- wmt16\nmetrics:\n- bleu\n---\n\n# FSMT\n\n## Model description\n\nThis is a ported version of fairseq-based [wmt16 transformer](https://github.com/jungokasai/deep-shallow/) for {src_lang}-{tgt_lang}.\n\nFor more details, please, see [Deep Encoder, Shallow Decoder: Reevaluating the Speed-Quality Tradeoff in Machine Translation](https://arxiv.org/abs/2006.10369).\n\nAll 3 models are available:\n\n* [wmt16-en-de-dist-12-1](https://huggingface.co/allenai/wmt16-en-de-dist-12-1)\n* [wmt16-en-de-dist-6-1](https://huggingface.co/allenai/wmt16-en-de-dist-6-1)\n* [wmt16-en-de-12-1](https://huggingface.co/allenai/wmt16-en-de-12-1)\n\n\n## Intended uses & limitations\n\n#### How to use\n\n```python\nfrom transformers import FSMTForConditionalGeneration, FSMTTokenizer\nmname = "allenai/{model_name}"\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\n## Training data\n\nPretrained weights were left identical to the original model released by allenai. For more details, please, see the [paper](https://arxiv.org/abs/2006.10369).\n\n## Eval results\n\nHere are the BLEU scores:\n\nmodel   | fairseq | transformers\n-------|---------|----------\n{model_name}  | {scores[model_name][0]} | {scores[model_name][1]}\n\nThe 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.\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=5\nmkdir -p $DATA_DIR\nsacrebleu -t wmt16 -l $PAIR --echo src > $DATA_DIR/val.source\nsacrebleu -t wmt16 -l $PAIR --echo ref > $DATA_DIR/val.target\necho $PAIR\nPYTHONPATH="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\n```\n\n## Data Sources\n\n- [training, etc.](http://www.statmt.org/wmt16/)\n- [test set](http://matrix.statmt.org/test_sets/newstest2016.tgz?1504722372)\n\n\n### BibTeX entry and citation info\n\n```\n@misc{{kasai2020deep,\n    title={{Deep Encoder, Shallow Decoder: Reevaluating the Speed-Quality Tradeoff in Machine Translation}},\n    author={{Jungo Kasai and Nikolaos Pappas and Hao Peng and James Cross and Noah A. Smith}},\n    year={{2020}},\n    eprint={{2006.10369}},\n    archivePrefix={{arXiv}},\n    primaryClass={{cs.CL}}\n}}\n```\n\n'
	model_card_dir.mkdir(parents=__snake_case ,exist_ok=__snake_case							)
	lowerCamelCase__              =	os.path.join(__snake_case ,'''README.md'''							)
	print(F'Generating {path}'							)
	with open(__snake_case ,'''w''' ,encoding='''utf-8'''							) as f:
		f.write(__snake_case							)


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



