


import contextlib
import csv
import json
import os
import sqlitea
import tarfile
import textwrap
import zipfile

import pyarrow as pa
import pyarrow.parquet as pq
import pytest

import datasets
import datasets.config


@pytest.fixture(scope="session" )
def     lowercase	(    ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						10
		SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						datasets.Features(
		    {
		        "tokens": datasets.Sequence(datasets.Value("string" ) ),
		        "labels": datasets.Sequence(datasets.ClassLabel(names=["negative", "positive"] ) ),
		        "answers": datasets.Sequence(
		            {
		                "text": datasets.Value("string" ),
		                "answer_start": datasets.Value("int32" ),
		            } ),
		        "id": datasets.Value("int64" ),
		    } )
		SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						datasets.Dataset.from_dict(
		    {
		        "tokens": [["foo"] * 5] * n,
		        "labels": [[1] * 5] * n,
		        "answers": [{"answer_start": [97], "text": ["1976"]}] * 10,
		        "id": list(range(a ) ),
		    }     , features=a     , )
		return dataset


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:int				     =						str(tmp_path_factory.mktemp("data" ) / "file.arrow" )
		dataset.map(cache_file_name=a )
		return filename


# FILE_CONTENT + files


SCREAMING_SNAKE_CASE__     					=			"\\n    Text data.\n    Second line of data."


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:List[str]				     =						tmp_path_factory.mktemp("data" ) / "file.txt"
		SCREAMING_SNAKE_CASE_				:List[Any]				     =						FILE_CONTENT
		with open(a     , "w" ) as f:
				f.write(a )
		return filename


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		import bza

		SCREAMING_SNAKE_CASE_				:List[Any]				     =						tmp_path_factory.mktemp("data" ) / "file.txt.bz2"
		SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						bytes(a     , "utf-8" )
		with bza.open(a     , "wb" ) as f:
				f.write(a )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		import gzip

		SCREAMING_SNAKE_CASE_				:int				     =						str(tmp_path_factory.mktemp("data" ) / "file.txt.gz" )
		SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						bytes(a     , "utf-8" )
		with gzip.open(a     , "wb" ) as f:
				f.write(a )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		if datasets.config.LZ4_AVAILABLE:
				import lza.frame

				SCREAMING_SNAKE_CASE_				:Dict				     =						tmp_path_factory.mktemp("data" ) / "file.txt.lz4"
				SCREAMING_SNAKE_CASE_				:str				     =						bytes(a     , "utf-8" )
				with lza.frame.open(a     , "wb" ) as f:
						f.write(a )
				return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a ):



		'''simple docstring'''




		if datasets.config.PY7ZR_AVAILABLE:
				import pyazr

				SCREAMING_SNAKE_CASE_				:List[Any]				     =						tmp_path_factory.mktemp("data" ) / "file.txt.7z"
				with pyazr.SevenZipFile(a     , "w" ) as archive:
						archive.write(a     , arcname=os.path.basename(a ) )
				return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a ):



		'''simple docstring'''




		import tarfile

		SCREAMING_SNAKE_CASE_				:List[Any]				     =						tmp_path_factory.mktemp("data" ) / "file.txt.tar"
		with tarfile.TarFile(a     , "w" ) as f:
				f.add(a     , arcname=os.path.basename(a ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		import lzma

		SCREAMING_SNAKE_CASE_				:List[Any]				     =						tmp_path_factory.mktemp("data" ) / "file.txt.xz"
		SCREAMING_SNAKE_CASE_				:List[str]				     =						bytes(a     , "utf-8" )
		with lzma.open(a     , "wb" ) as f:
				f.write(a )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a ):



		'''simple docstring'''




		import zipfile

		SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						tmp_path_factory.mktemp("data" ) / "file.txt.zip"
		with zipfile.ZipFile(a     , "w" ) as f:
				f.write(a     , arcname=os.path.basename(a ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		if datasets.config.ZSTANDARD_AVAILABLE:
				import zstandard as zstd

				SCREAMING_SNAKE_CASE_				:List[Any]				     =						tmp_path_factory.mktemp("data" ) / "file.txt.zst"
				SCREAMING_SNAKE_CASE_				:List[str]				     =						bytes(a     , "utf-8" )
				with zstd.open(a     , "wb" ) as f:
						f.write(a )
				return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:List[str]				     =						tmp_path_factory.mktemp("data" ) / "file.xml"
		SCREAMING_SNAKE_CASE_				:Tuple				     =						textwrap.dedent(
		    "\\n    <?xml version=\"1.0\" encoding=\"UTF-8\" ?>\n    <tmx version=\"1.4\">\n      <header segtype=\"sentence\" srclang=\"ca\" />\n      <body>\n        <tu>\n          <tuv xml:lang=\"ca\"><seg>Contingut 1</seg></tuv>\n          <tuv xml:lang=\"en\"><seg>Content 1</seg></tuv>\n        </tu>\n        <tu>\n          <tuv xml:lang=\"ca\"><seg>Contingut 2</seg></tuv>\n          <tuv xml:lang=\"en\"><seg>Content 2</seg></tuv>\n        </tu>\n        <tu>\n          <tuv xml:lang=\"ca\"><seg>Contingut 3</seg></tuv>\n          <tuv xml:lang=\"en\"><seg>Content 3</seg></tuv>\n        </tu>\n        <tu>\n          <tuv xml:lang=\"ca\"><seg>Contingut 4</seg></tuv>\n          <tuv xml:lang=\"en\"><seg>Content 4</seg></tuv>\n        </tu>\n        <tu>\n          <tuv xml:lang=\"ca\"><seg>Contingut 5</seg></tuv>\n          <tuv xml:lang=\"en\"><seg>Content 5</seg></tuv>\n        </tu>\n      </body>\n    </tmx>" )
		with open(a     , "w" ) as f:
				f.write(a )
		return filename


SCREAMING_SNAKE_CASE__     					=			[
    {"col_1": "0", "col_2": 0, "col_3": 0.0},
    {"col_1": "1", "col_2": 1, "col_3": 1.0},
    {"col_1": "2", "col_2": 2, "col_3": 2.0},
    {"col_1": "3", "col_2": 3, "col_3": 3.0},
]
SCREAMING_SNAKE_CASE__     					=			[
    {"col_1": "4", "col_2": 4, "col_3": 4.0},
    {"col_1": "5", "col_2": 5, "col_3": 5.0},
]
SCREAMING_SNAKE_CASE__     					=			{
    "col_1": ["0", "1", "2", "3"],
    "col_2": [0, 1, 2, 3],
    "col_3": [0.0, 1.0, 2.0, 3.0],
}

SCREAMING_SNAKE_CASE__     					=			[
    {"col_3": 0.0, "col_1": "0", "col_2": 0},
    {"col_3": 1.0, "col_1": "1", "col_2": 1},
]

SCREAMING_SNAKE_CASE__     					=			[
    {"col_1": "s0", "col_2": 0, "col_3": 0.0},
    {"col_1": "s1", "col_2": 1, "col_3": 1.0},
    {"col_1": "s2", "col_2": 2, "col_3": 2.0},
    {"col_1": "s3", "col_2": 3, "col_3": 3.0},
]


@pytest.fixture(scope="session" )
def     lowercase	(    ):



		'''simple docstring'''




		return DATA_DICT_OF_LISTS


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:List[str]				     =						datasets.Dataset.from_dict(a )
		SCREAMING_SNAKE_CASE_				:Any				     =						str(tmp_path_factory.mktemp("data" ) / "dataset.arrow" )
		dataset.map(cache_file_name=a )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:List[str]				     =						str(tmp_path_factory.mktemp("data" ) / "dataset.sqlite" )
		with contextlib.closing(sqlitea.connect(a ) ) as con:
				SCREAMING_SNAKE_CASE_				:Dict				     =						con.cursor()
				cur.execute("CREATE TABLE dataset(col_1 text, col_2 int, col_3 real)" )
				for item in DATA:
						cur.execute("INSERT INTO dataset(col_1, col_2, col_3) VALUES (?, ?, ?)"     , tuple(item.values() ) )
				con.commit()
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:int				     =						str(tmp_path_factory.mktemp("data" ) / "dataset.csv" )
		with open(a     , "w"     , newline="" ) as f:
				SCREAMING_SNAKE_CASE_				:Dict				     =						csv.DictWriter(a     , fieldnames=["col_1", "col_2", "col_3"] )
				writer.writeheader()
				for item in DATA:
						writer.writerow(a )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Optional[int]				     =						str(tmp_path_factory.mktemp("data" ) / "dataset2.csv" )
		with open(a     , "w"     , newline="" ) as f:
				SCREAMING_SNAKE_CASE_				:str				     =						csv.DictWriter(a     , fieldnames=["col_1", "col_2", "col_3"] )
				writer.writeheader()
				for item in DATA:
						writer.writerow(a )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a ):



		'''simple docstring'''




		import bza

		SCREAMING_SNAKE_CASE_				:Tuple				     =						tmp_path_factory.mktemp("data" ) / "dataset.csv.bz2"
		with open(a     , "rb" ) as f:
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						f.read()
		# data = bytes(FILE_CONTENT, "utf-8")
		with bza.open(a     , "wb" ) as f:
				f.write(a )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:str				     =						tmp_path_factory.mktemp("data" ) / "dataset.csv.zip"
		with zipfile.ZipFile(a     , "w" ) as f:
				f.write(a     , arcname=os.path.basename(a ) )
				f.write(a     , arcname=os.path.basename(a ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:List[Any]				     =						tmp_path_factory.mktemp("data" ) / "dataset.csv.zip"
		with zipfile.ZipFile(a     , "w" ) as f:
				f.write(a     , arcname=os.path.basename(csv_path.replace(".csv"     , ".CSV" ) ) )
				f.write(a     , arcname=os.path.basename(csva_path.replace(".csv"     , ".CSV" ) ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:List[Any]				     =						tmp_path_factory.mktemp("data" ) / "dataset_with_dir.csv.zip"
		with zipfile.ZipFile(a     , "w" ) as f:
				f.write(a     , arcname=os.path.join("main_dir"     , os.path.basename(a ) ) )
				f.write(a     , arcname=os.path.join("main_dir"     , os.path.basename(a ) ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						str(tmp_path_factory.mktemp("data" ) / "dataset.parquet" )
		SCREAMING_SNAKE_CASE_				:str				     =						pa.schema(
		    {
		        "col_1": pa.string(),
		        "col_2": pa.intaa(),
		        "col_3": pa.floataa(),
		    } )
		with open(a     , "wb" ) as f:
				SCREAMING_SNAKE_CASE_				:int				     =						pq.ParquetWriter(a     , schema=a )
				SCREAMING_SNAKE_CASE_				:Any				     =						pa.Table.from_pydict({k: [DATA[i][k] for i in range(len(a ) )] for k in DATA[0]}     , schema=a )
				writer.write_table(a )
				writer.close()
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Tuple				     =						str(tmp_path_factory.mktemp("data" ) / "dataset.json" )
		SCREAMING_SNAKE_CASE_				:Tuple				     =						{"data": DATA}
		with open(a     , "w" ) as f:
				json.dump(a     , a )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Tuple				     =						str(tmp_path_factory.mktemp("data" ) / "dataset.json" )
		SCREAMING_SNAKE_CASE_				:int				     =						{"data": DATA_DICT_OF_LISTS}
		with open(a     , "w" ) as f:
				json.dump(a     , a )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Dict				     =						str(tmp_path_factory.mktemp("data" ) / "dataset.jsonl" )
		with open(a     , "w" ) as f:
				for item in DATA:
						f.write(json.dumps(a ) + "\n" )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:List[str]				     =						str(tmp_path_factory.mktemp("data" ) / "dataset2.jsonl" )
		with open(a     , "w" ) as f:
				for item in DATA:
						f.write(json.dumps(a ) + "\n" )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:int				     =						str(tmp_path_factory.mktemp("data" ) / "dataset_312.jsonl" )
		with open(a     , "w" ) as f:
				for item in DATA_312:
						f.write(json.dumps(a ) + "\n" )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:List[str]				     =						str(tmp_path_factory.mktemp("data" ) / "dataset-str.jsonl" )
		with open(a     , "w" ) as f:
				for item in DATA_STR:
						f.write(json.dumps(a ) + "\n" )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a ):



		'''simple docstring'''




		import gzip

		SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						str(tmp_path_factory.mktemp("data" ) / "dataset.txt.gz" )
		with open(a     , "rb" ) as orig_file:
				with gzip.open(a     , "wb" ) as zipped_file:
						zipped_file.writelines(a )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a ):



		'''simple docstring'''




		import gzip

		SCREAMING_SNAKE_CASE_				:List[Any]				     =						str(tmp_path_factory.mktemp("data" ) / "dataset.jsonl.gz" )
		with open(a     , "rb" ) as orig_file:
				with gzip.open(a     , "wb" ) as zipped_file:
						zipped_file.writelines(a )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Optional[int]				     =						tmp_path_factory.mktemp("data" ) / "dataset.jsonl.zip"
		with zipfile.ZipFile(a     , "w" ) as f:
				f.write(a     , arcname=os.path.basename(a ) )
				f.write(a     , arcname=os.path.basename(a ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a     , a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Dict				     =						tmp_path_factory.mktemp("data" ) / "dataset_nested.jsonl.zip"
		with zipfile.ZipFile(a     , "w" ) as f:
				f.write(a     , arcname=os.path.join("nested"     , os.path.basename(a ) ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Tuple				     =						tmp_path_factory.mktemp("data" ) / "dataset_with_dir.jsonl.zip"
		with zipfile.ZipFile(a     , "w" ) as f:
				f.write(a     , arcname=os.path.join("main_dir"     , os.path.basename(a ) ) )
				f.write(a     , arcname=os.path.join("main_dir"     , os.path.basename(a ) ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:List[str]				     =						tmp_path_factory.mktemp("data" ) / "dataset.jsonl.tar"
		with tarfile.TarFile(a     , "w" ) as f:
				f.add(a     , arcname=os.path.basename(a ) )
				f.add(a     , arcname=os.path.basename(a ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a     , a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Dict				     =						tmp_path_factory.mktemp("data" ) / "dataset_nested.jsonl.tar"
		with tarfile.TarFile(a     , "w" ) as f:
				f.add(a     , arcname=os.path.join("nested"     , os.path.basename(a ) ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Any				     =						["0", "1", "2", "3"]
		SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						str(tmp_path_factory.mktemp("data" ) / "dataset.txt" )
		with open(a     , "w" ) as f:
				for item in data:
						f.write(item + "\n" )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Any				     =						["0", "1", "2", "3"]
		SCREAMING_SNAKE_CASE_				:str				     =						str(tmp_path_factory.mktemp("data" ) / "dataset2.txt" )
		with open(a     , "w" ) as f:
				for item in data:
						f.write(item + "\n" )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Any				     =						["0", "1", "2", "3"]
		SCREAMING_SNAKE_CASE_				:int				     =						tmp_path_factory.mktemp("data" ) / "dataset.abc"
		with open(a     , "w" ) as f:
				for item in data:
						f.write(item + "\n" )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:List[Any]				     =						tmp_path_factory.mktemp("data" ) / "dataset.text.zip"
		with zipfile.ZipFile(a     , "w" ) as f:
				f.write(a     , arcname=os.path.basename(a ) )
				f.write(a     , arcname=os.path.basename(a ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:List[str]				     =						tmp_path_factory.mktemp("data" ) / "dataset_with_dir.text.zip"
		with zipfile.ZipFile(a     , "w" ) as f:
				f.write(a     , arcname=os.path.join("main_dir"     , os.path.basename(a ) ) )
				f.write(a     , arcname=os.path.join("main_dir"     , os.path.basename(a ) ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Dict				     =						tmp_path_factory.mktemp("data" ) / "dataset.ext.zip"
		with zipfile.ZipFile(a     , "w" ) as f:
				f.write(a     , arcname=os.path.basename("unsupported.ext" ) )
				f.write(a     , arcname=os.path.basename("unsupported_2.ext" ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Tuple				     =						"\n".join(["First", "Second\u2029with Unicode new line", "Third"] )
		SCREAMING_SNAKE_CASE_				:List[str]				     =						str(tmp_path_factory.mktemp("data" ) / "dataset_with_unicode_new_lines.txt" )
		with open(a     , "w"     , encoding="utf-8" ) as f:
				f.write(a )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    ):



		'''simple docstring'''




		return os.path.join("tests"     , "features"     , "data"     , "test_image_rgb.jpg" )


@pytest.fixture(scope="session" )
def     lowercase	(    ):



		'''simple docstring'''




		return os.path.join("tests"     , "features"     , "data"     , "test_audio_44100.wav" )


@pytest.fixture(scope="session" )
def     lowercase	(    a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Tuple				     =						tmp_path_factory.mktemp("data" ) / "dataset.img.zip"
		with zipfile.ZipFile(a     , "w" ) as f:
				f.write(a     , arcname=os.path.basename(a ) )
				f.write(a     , arcname=os.path.basename(a ).replace(".jpg"     , "2.jpg" ) )
		return path


@pytest.fixture(scope="session" )
def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:int				     =						tmp_path_factory.mktemp("data_dir" )

		(data_dir / "subdir").mkdir()
		with open(data_dir / "subdir" / "train.txt"     , "w" ) as f:
				f.write("foo\n" * 10 )
		with open(data_dir / "subdir" / "test.txt"     , "w" ) as f:
				f.write("bar\n" * 10 )
		# hidden file
		with open(data_dir / "subdir" / ".test.txt"     , "w" ) as f:
				f.write("bar\n" * 10 )

		# hidden directory
		(data_dir / ".subdir").mkdir()
		with open(data_dir / ".subdir" / "train.txt"     , "w" ) as f:
				f.write("foo\n" * 10 )
		with open(data_dir / ".subdir" / "test.txt"     , "w" ) as f:
				f.write("bar\n" * 10 )

		return data_dir










import argparse
import json

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

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)


def     lowercase	(    a     , a=False ):



		'''simple docstring'''




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

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

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

				# if just the base model, we should remove "vit" from all keys that start with "vit"
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						[(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     lowercase	(    a     , a     , a=False ):



		'''simple docstring'''




		for i in range(config.num_hidden_layers ):
				if base_model:
						SCREAMING_SNAKE_CASE_				:List[Any]				     =						""
				else:
						SCREAMING_SNAKE_CASE_				:List[Any]				     =						"vit."
				# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
				SCREAMING_SNAKE_CASE_				:str				     =						state_dict.pop(F"module.blocks.{i}.attn.qkv.weight" )
				SCREAMING_SNAKE_CASE_				:Dict				     =						state_dict.pop(F"module.blocks.{i}.attn.qkv.bias" )
				# next, add query, keys and values (in that order) to the state dict
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						in_proj_weight[
				    : config.hidden_size, :
				]
				SCREAMING_SNAKE_CASE_				:Tuple				     =						in_proj_bias[: config.hidden_size]
				SCREAMING_SNAKE_CASE_				:Dict				     =						in_proj_weight[
				    config.hidden_size : config.hidden_size * 2, :
				]
				SCREAMING_SNAKE_CASE_				:Any				     =						in_proj_bias[
				    config.hidden_size : config.hidden_size * 2
				]
				SCREAMING_SNAKE_CASE_				:str				     =						in_proj_weight[
				    -config.hidden_size :, :
				]
				SCREAMING_SNAKE_CASE_				:int				     =						in_proj_bias[-config.hidden_size :]


def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Optional[int]				     =						["head.weight", "head.bias"]
		for k in ignore_keys:
				state_dict.pop(a     , a )


def     lowercase	(    a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:int				     =						[
		    "module.fc.fc1.weight",
		    "module.fc.fc1.bias",
		    "module.fc.bn1.weight",
		    "module.fc.bn1.bias",
		    "module.fc.bn1.running_mean",
		    "module.fc.bn1.running_var",
		    "module.fc.bn1.num_batches_tracked",
		    "module.fc.fc2.weight",
		    "module.fc.fc2.bias",
		    "module.fc.bn2.weight",
		    "module.fc.bn2.bias",
		    "module.fc.bn2.running_mean",
		    "module.fc.bn2.running_var",
		    "module.fc.bn2.num_batches_tracked",
		    "module.fc.fc3.weight",
		    "module.fc.fc3.bias",
		]
		for k in ignore_keys:
				state_dict.pop(a     , a )


def     lowercase	(    a     , a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						dct.pop(a )
		SCREAMING_SNAKE_CASE_				:Dict				     =						val


def     lowercase	(    a     , a ):



		'''simple docstring'''




		SCREAMING_SNAKE_CASE_				:List[str]				     =						ViTMSNConfig()
		SCREAMING_SNAKE_CASE_				:Tuple				     =						1000

		SCREAMING_SNAKE_CASE_				:List[str]				     =						"datasets/huggingface/label-files"
		SCREAMING_SNAKE_CASE_				:int				     =						"imagenet-1k-id2label.json"
		SCREAMING_SNAKE_CASE_				:str				     =						json.load(open(hf_hub_download(a     , a )     , "r" ) )
		SCREAMING_SNAKE_CASE_				:List[str]				     =						{int(a ): v for k, v in idalabel.items()}
		SCREAMING_SNAKE_CASE_				:List[Any]				     =						idalabel
		SCREAMING_SNAKE_CASE_				:Tuple				     =						{v: k for k, v in idalabel.items()}

		if "s16" in checkpoint_url:
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						384
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						1536
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						6
		elif "l16" in checkpoint_url:
				SCREAMING_SNAKE_CASE_				:str				     =						1024
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						4096
				SCREAMING_SNAKE_CASE_				:Dict				     =						24
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						16
				SCREAMING_SNAKE_CASE_				:Dict				     =						0.1
		elif "b4" in checkpoint_url:
				SCREAMING_SNAKE_CASE_				:Dict				     =						4
		elif "l7" in checkpoint_url:
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						7
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						1024
				SCREAMING_SNAKE_CASE_				:str				     =						4096
				SCREAMING_SNAKE_CASE_				:Tuple				     =						24
				SCREAMING_SNAKE_CASE_				:Dict				     =						16
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						0.1

		SCREAMING_SNAKE_CASE_				:List[Any]				     =						ViTMSNModel(a )

		SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						torch.hub.load_state_dict_from_url(a     , map_location="cpu" )["target_encoder"]

		SCREAMING_SNAKE_CASE_				:Tuple				     =						ViTImageProcessor(size=config.image_size )

		remove_projection_head(a )
		SCREAMING_SNAKE_CASE_				:int				     =						create_rename_keys(a     , base_model=a )

		for src, dest in rename_keys:
				rename_key(a     , a     , a )
		read_in_q_k_v(a     , a     , base_model=a )

		model.load_state_dict(a )
		model.eval()

		SCREAMING_SNAKE_CASE_				:List[Any]				     =						"http://images.cocodataset.org/val2017/000000039769.jpg"

		SCREAMING_SNAKE_CASE_				:Dict				     =						Image.open(requests.get(a     , stream=a ).raw )
		SCREAMING_SNAKE_CASE_				:Optional[int]				     =						ViTImageProcessor(
		    size=config.image_size     , image_mean=a     , image_std=a )
		SCREAMING_SNAKE_CASE_				:Any				     =						image_processor(images=a     , return_tensors="pt" )

		# forward pass
		torch.manual_seed(2 )
		SCREAMING_SNAKE_CASE_				:Tuple				     =						model(**a )
		SCREAMING_SNAKE_CASE_				:Optional[int]				     =						outputs.last_hidden_state

		# The following Colab Notebook was used to generate these outputs:
		# https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
		if "s16" in checkpoint_url:
				SCREAMING_SNAKE_CASE_				:Dict				     =						torch.tensor([[-1.0_915, -1.4_876, -1.1_809]] )
		elif "b16" in checkpoint_url:
				SCREAMING_SNAKE_CASE_				:Dict				     =						torch.tensor([[14.2_889, -18.9_045, 11.7_281]] )
		elif "l16" in checkpoint_url:
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						torch.tensor([[41.5_028, -22.8_681, 45.6_475]] )
		elif "b4" in checkpoint_url:
				SCREAMING_SNAKE_CASE_				:List[str]				     =						torch.tensor([[-4.3_868, 5.2_932, -0.4_137]] )
		else:
				SCREAMING_SNAKE_CASE_				:List[str]				     =						torch.tensor([[-0.1_792, -0.6_465, 2.4_263]] )

		# verify logits
		assert torch.allclose(last_hidden_state[:, 0, :3]     , a     , atol=1e-4 )

		print(F"Saving model to {pytorch_dump_folder_path}" )
		model.save_pretrained(a )

		print(F"Saving image processor to {pytorch_dump_folder_path}" )
		image_processor.save_pretrained(a )


if __name__ == "__main__":
		SCREAMING_SNAKE_CASE__     					=			argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    "--checkpoint_url",
		    default="https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar",
		    type=str,
		    help="URL of the checkpoint you'd like to convert.",
		)
		parser.add_argument(
		    "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
		)

		SCREAMING_SNAKE_CASE__     					=			parser.parse_args()
		convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)






