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 ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
__UpperCamelCase	: Optional[int]    		=  logging.get_logger(__name__)
def    A			(					_lowercase	,						_lowercase=False		):
						SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							[]
						for i in range(config.num_hidden_layers		):
												# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
												rename_keys.append((f"""blocks.{i}.norm1.weight""", f"""vit.encoder.layer.{i}.layernorm_before.weight""")		)
												rename_keys.append((f"""blocks.{i}.norm1.bias""", f"""vit.encoder.layer.{i}.layernorm_before.bias""")		)
												rename_keys.append((f"""blocks.{i}.attn.proj.weight""", f"""vit.encoder.layer.{i}.attention.output.dense.weight""")		)
												rename_keys.append((f"""blocks.{i}.attn.proj.bias""", f"""vit.encoder.layer.{i}.attention.output.dense.bias""")		)
												rename_keys.append((f"""blocks.{i}.norm2.weight""", f"""vit.encoder.layer.{i}.layernorm_after.weight""")		)
												rename_keys.append((f"""blocks.{i}.norm2.bias""", f"""vit.encoder.layer.{i}.layernorm_after.bias""")		)
												rename_keys.append((f"""blocks.{i}.mlp.fc1.weight""", f"""vit.encoder.layer.{i}.intermediate.dense.weight""")		)
												rename_keys.append((f"""blocks.{i}.mlp.fc1.bias""", f"""vit.encoder.layer.{i}.intermediate.dense.bias""")		)
												rename_keys.append((f"""blocks.{i}.mlp.fc2.weight""", f"""vit.encoder.layer.{i}.output.dense.weight""")		)
												rename_keys.append((f"""blocks.{i}.mlp.fc2.bias""", f"""vit.encoder.layer.{i}.output.dense.bias""")		)

						# projection layer + position embeddings
						rename_keys.extend(
						    [
						        ('''cls_token''', '''vit.embeddings.cls_token'''),
						        ('''patch_embed.proj.weight''', '''vit.embeddings.patch_embeddings.projection.weight'''),
						        ('''patch_embed.proj.bias''', '''vit.embeddings.patch_embeddings.projection.bias'''),
						        ('''pos_embed''', '''vit.embeddings.position_embeddings'''),
						    ]		)

						if base_model:
												# layernorm + pooler
												rename_keys.extend(
												    [
												        ('''norm.weight''', '''layernorm.weight'''),
												        ('''norm.bias''', '''layernorm.bias'''),
												    ]		)

												# if just the base model, we should remove "vit" from all keys that start with "vit"
												SCREAMING_SNAKE_CASE							:					List[str]           =							[(pair[0], pair[1][4:]) if pair[1].startswith('''vit'''		) else pair for pair in rename_keys]
						else:
												# layernorm + classification head
												rename_keys.extend(
												    [
												        ('''norm.weight''', '''vit.layernorm.weight'''),
												        ('''norm.bias''', '''vit.layernorm.bias'''),
												        ('''head.weight''', '''classifier.weight'''),
												        ('''head.bias''', '''classifier.bias'''),
												    ]		)

						return rename_keys
def    A			(					_lowercase	,						_lowercase	,						_lowercase=False		):

						for i in range(config.num_hidden_layers		):
												if base_model:
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							""
												else:
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							"vit."
												# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
												SCREAMING_SNAKE_CASE							:					Any           =							state_dict.pop(f"""blocks.{i}.attn.qkv.weight"""		)
												SCREAMING_SNAKE_CASE							:					List[str]           =							state_dict.pop(f"""blocks.{i}.attn.qkv.bias"""		)
												# next, add query, keys and values (in that order) to the state dict
												SCREAMING_SNAKE_CASE							:					List[Any]           =							in_proj_weight[
												    : config.hidden_size, :
												]
												SCREAMING_SNAKE_CASE							:					Tuple           =							in_proj_bias[: config.hidden_size]
												SCREAMING_SNAKE_CASE							:					Any           =							in_proj_weight[
												    config.hidden_size : config.hidden_size * 2, :
												]
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							in_proj_bias[
												    config.hidden_size : config.hidden_size * 2
												]
												SCREAMING_SNAKE_CASE							:					Dict           =							in_proj_weight[
												    -config.hidden_size :, :
												]
												SCREAMING_SNAKE_CASE							:					Any           =							in_proj_bias[-config.hidden_size :]
def    A			(					_lowercase		):
						SCREAMING_SNAKE_CASE							:					Optional[int]           =							["head.weight", "head.bias"]
						for k in ignore_keys:
												state_dict.pop(lowercase_	,						lowercase_		)
def    A			(					_lowercase	,						_lowercase	,						_lowercase		):
						SCREAMING_SNAKE_CASE							:					Dict           =							dct.pop(lowercase_		)
						SCREAMING_SNAKE_CASE							:					List[str]           =							val
def    A			(					):
						SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							"http://images.cocodataset.org/val2017/000000039769.jpg"
						SCREAMING_SNAKE_CASE							:					Any           =							Image.open(requests.get(lowercase_	,						stream=lowercase_		).raw		)
						return im
@torch.no_grad()
def    A			(					_lowercase	,						_lowercase	,						_lowercase=True		):
						SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							ViTConfig()
						# patch_size
						if model_name[-1] == "8":
												SCREAMING_SNAKE_CASE							:					List[Any]           =							8
						# set labels if required
						if not base_model:
												SCREAMING_SNAKE_CASE							:					Tuple           =							1_000
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							"huggingface/label-files"
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							"imagenet-1k-id2label.json"
												SCREAMING_SNAKE_CASE							:					int           =							json.load(open(hf_hub_download(lowercase_	,						lowercase_	,						repo_type='''dataset'''		)	,						'''r'''		)		)
												SCREAMING_SNAKE_CASE							:					int           =							{int(lowercase_		): v for k, v in idalabel.items()}
												SCREAMING_SNAKE_CASE							:					int           =							idalabel
												SCREAMING_SNAKE_CASE							:					List[str]           =							{v: k for k, v in idalabel.items()}
						# size of the architecture
						if model_name in ["dino_vits8", "dino_vits16"]:
												SCREAMING_SNAKE_CASE							:					Dict           =							384
												SCREAMING_SNAKE_CASE							:					Any           =							1_536
												SCREAMING_SNAKE_CASE							:					List[Any]           =							12
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							6

						# load original model from torch hub
						SCREAMING_SNAKE_CASE							:					List[str]           =							torch.hub.load('''facebookresearch/dino:main'''	,						lowercase_		)
						original_model.eval()

						# load state_dict of original model, remove and rename some keys
						SCREAMING_SNAKE_CASE							:					int           =							original_model.state_dict()
						if base_model:
												remove_classification_head_(lowercase_		)
						SCREAMING_SNAKE_CASE							:					Optional[Any]           =							create_rename_keys(lowercase_	,						base_model=lowercase_		)
						for src, dest in rename_keys:
												rename_key(lowercase_	,						lowercase_	,						lowercase_		)
						read_in_q_k_v(lowercase_	,						lowercase_	,						lowercase_		)

						# load HuggingFace model
						if base_model:
												SCREAMING_SNAKE_CASE							:					Dict           =							ViTModel(lowercase_	,						add_pooling_layer=lowercase_		).eval()
						else:
												SCREAMING_SNAKE_CASE							:					List[str]           =							ViTForImageClassification(lowercase_		).eval()
						model.load_state_dict(lowercase_		)

						# Check outputs on an image, prepared by ViTImageProcessor
						SCREAMING_SNAKE_CASE							:					Tuple           =							ViTImageProcessor()
						SCREAMING_SNAKE_CASE							:					str           =							image_processor(images=prepare_img()	,						return_tensors='''pt'''		)
						SCREAMING_SNAKE_CASE							:					str           =							encoding["pixel_values"]
						SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							model(lowercase_		)

						if base_model:
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							original_model(lowercase_		)
												assert torch.allclose(lowercase_	,						outputs.last_hidden_state[:, 0, :]	,						atol=1e-1		)
						else:
												SCREAMING_SNAKE_CASE							:					str           =							original_model(lowercase_		)
												assert logits.shape == outputs.logits.shape
												assert torch.allclose(lowercase_	,						outputs.logits	,						atol=1e-3		)

						Path(lowercase_		).mkdir(exist_ok=lowercase_		)
						print(f"""Saving model {model_name} to {pytorch_dump_folder_path}"""		)
						model.save_pretrained(lowercase_		)
						print(f"""Saving image processor to {pytorch_dump_folder_path}"""		)
						image_processor.save_pretrained(lowercase_		)


if __name__ == "__main__":
					__UpperCamelCase	: Tuple    		=  argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '--model_name',
					    default='dino_vitb16',
					    type=str,
					    help='Name of the model trained with DINO you\'d like to convert.',
					)
					parser.add_argument(
					    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
					)
					parser.add_argument(
					    '--base_model',
					    action='store_true',
					    help='Whether to only convert the base model (no projection head weights).',
					)

					parser.set_defaults(base_model=True)
					__UpperCamelCase	: Optional[Any]    		=  parser.parse_args()
					convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)






import argparse
import json
import logging
import os
import sys
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, get_gpu_count, slow


__UpperCamelCase	: str    		=  [
    os.path.join(os.path.dirname(__file__), dirname)
    for dirname in [
        'text-classification',
        'language-modeling',
        'summarization',
        'token-classification',
        'question-answering',
    ]
]
sys.path.extend(SRC_DIRS)


if SRC_DIRS is not None:
					import run_clm_flax
					import run_flax_glue
					import run_flax_ner
					import run_mlm_flax
					import run_qa
					import run_summarization_flax
					import run_ta_mlm_flax


logging.basicConfig(level=logging.DEBUG)

__UpperCamelCase	: int    		=  logging.getLogger()
def    A			(					):
						SCREAMING_SNAKE_CASE							:					str           =							argparse.ArgumentParser()
						parser.add_argument('''-f'''		)
						SCREAMING_SNAKE_CASE							:					List[str]           =							parser.parse_args()
						return args.f
def    A			(					_lowercase	,						_lowercase="eval"		):
						SCREAMING_SNAKE_CASE							:					Dict           =							os.path.join(_lowercase	,						f"""{split}_results.json"""		)
						if os.path.exists(_lowercase		):
												with open(_lowercase	,						'''r'''		) as f:
																		return json.load(_lowercase		)
						raise ValueError(f"""can't find {path}"""		)


__UpperCamelCase	: Optional[Any]    		=  logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)





class       lowercase__		(     UpperCamelCase_):


						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					int           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Tuple           =							f"""
            run_glue.py
            --model_name_or_path distilbert-base-uncased
            --output_dir {tmp_dir}
            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv
            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --learning_rate=1e-4
            --eval_steps=2
            --warmup_steps=2
            --seed=42
            --max_seq_length=128
            """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_flax_glue.main()
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_accuracy''']					,	0.75		)


						@slow
						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Dict           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					str           =							f"""
            run_clm_flax.py
            --model_name_or_path distilgpt2
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --do_train
            --do_eval
            --block_size 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --num_train_epochs 2
            --logging_steps 2 --eval_steps 2
            --output_dir {tmp_dir}
            --overwrite_output_dir
            """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_clm_flax.main()
																		SCREAMING_SNAKE_CASE							:					Dict           =							get_results(UpperCamelCase__		)
																		self.assertLess(result['''eval_perplexity''']					,	100		)


						@slow
						def  __A ( self			:       Any		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[str]           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							f"""
            run_summarization.py
            --model_name_or_path t5-small
            --train_file tests/fixtures/tests_samples/xsum/sample.json
            --validation_file tests/fixtures/tests_samples/xsum/sample.json
            --test_file tests/fixtures/tests_samples/xsum/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --num_train_epochs=3
            --warmup_steps=8
            --do_train
            --do_eval
            --do_predict
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --predict_with_generate
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_summarization_flax.main()
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							get_results(UpperCamelCase__					,	split='''test'''		)
																		self.assertGreaterEqual(result['''test_rouge1''']					,	10		)
																		self.assertGreaterEqual(result['''test_rouge2''']					,	2		)
																		self.assertGreaterEqual(result['''test_rougeL''']					,	7		)
																		self.assertGreaterEqual(result['''test_rougeLsum''']					,	7		)


						@slow
						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[Any]           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Dict           =							f"""
            run_mlm.py
            --model_name_or_path distilroberta-base
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --max_seq_length 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --logging_steps 2 --eval_steps 2
            --do_train
            --do_eval
            --num_train_epochs=1
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_mlm_flax.main()
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							get_results(UpperCamelCase__		)
																		self.assertLess(result['''eval_perplexity''']					,	42		)


						@slow
						def  __A ( self			:       List[Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Dict           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							f"""
            run_t5_mlm_flax.py
            --model_name_or_path t5-small
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --do_train
            --do_eval
            --max_seq_length 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --num_train_epochs 2
            --logging_steps 2 --eval_steps 2
            --output_dir {tmp_dir}
            --overwrite_output_dir
            """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_ta_mlm_flax.main()
																		SCREAMING_SNAKE_CASE							:					Optional[int]           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_accuracy''']					,	0.42		)


						@slow
						def  __A ( self			:       Union[str, Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Tuple           =							7 if get_gpu_count() > 1 else 2

												SCREAMING_SNAKE_CASE							:					Dict           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Any           =							f"""
            run_flax_ner.py
            --model_name_or_path bert-base-uncased
            --train_file tests/fixtures/tests_samples/conll/sample.json
            --validation_file tests/fixtures/tests_samples/conll/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --do_train
            --do_eval
            --warmup_steps=2
            --learning_rate=2e-4
            --logging_steps 2 --eval_steps 2
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=2
            --num_train_epochs={epochs}
            --seed 7
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_flax_ner.main()
																		SCREAMING_SNAKE_CASE							:					List[str]           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_accuracy''']					,	0.75		)
																		self.assertGreaterEqual(result['''eval_f1''']					,	0.3		)



						@slow
						def  __A ( self			:       List[str]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[str]           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							f"""
            run_qa.py
            --model_name_or_path bert-base-uncased
            --version_2_with_negative
            --train_file tests/fixtures/tests_samples/SQUAD/sample.json
            --validation_file tests/fixtures/tests_samples/SQUAD/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --num_train_epochs=3
            --warmup_steps=2
            --do_train
            --do_eval
            --logging_steps 2 --eval_steps 2
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_qa.main()
																		SCREAMING_SNAKE_CASE							:					str           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_f1''']					,	30		)
																		self.assertGreaterEqual(result['''eval_exact''']					,	30		)





