


import argparse
import glob
import inspect
import os
import re
import unittest

from get_test_info import get_test_classes

from tests.test_pipeline_mixin import pipeline_test_mapping


SCREAMING_SNAKE_CASE__       			=				{}
for task, _ in pipeline_test_mapping.items():
							SCREAMING_SNAKE_CASE__       			=				{"""pt""": None, """tf""": None}


# DO **NOT** add item to this set (unless the reason is approved)
SCREAMING_SNAKE_CASE__       			=				{
    """tests/models/esm/test_modeling_esmfold.py""",  # The pipeline test mapping is added to `test_modeling_esm.py`
}




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		Dict  ):







	'''simple docstring'''



	if "ModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
		return "pt"
	elif "TFModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
		return "tf"
	elif "FlaxModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
		return "flax"
	else:
		return None




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		Optional[Any]  ,			__lowerCamelCase:		List[str]  ):







	'''simple docstring'''



	if PIPELINE_TEST_MAPPING[task].get(__lowerCamelCase  ,			__lowerCamelCase  ) is not None:
		return PIPELINE_TEST_MAPPING[task][framework]

	lowercase_			 = pipeline_test_mapping[task]["test"]
	lowercase_			 = None

	if framework == "pt":
		lowercase_			 = getattr(__lowerCamelCase  ,			"model_mapping"  ,			__lowerCamelCase  )
	elif framework == "tf":
		lowercase_			 = getattr(__lowerCamelCase  ,			"tf_model_mapping"  ,			__lowerCamelCase  )

	if mapping is not None:
		lowercase_			 = dict(mapping.items()  )

	# cache the results
	lowercase_			 = mapping
	return mapping




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		str  ,			__lowerCamelCase:		Tuple  ):







	'''simple docstring'''



	lowercase_			 = get_framework(__lowerCamelCase  )
	if framework is None:
		return None
	lowercase_			 = get_mapping_for_task(__lowerCamelCase  ,			__lowerCamelCase  )
	if mapping is None:
		return None

	lowercase_			 = list({model_class.config_class for model_class in test_class.all_model_classes}  )
	if len(__lowerCamelCase  ) != 1:
		raise ValueError("There should be exactly one configuration class from `test_class.all_model_classes`."  )

	# This could be a list/tuple of model classes, but it's rare.
	lowercase_			 = mapping.get(config_classes[0]  ,			__lowerCamelCase  )
	if isinstance(__lowerCamelCase  ,			(tuple, list)  ):
		lowercase_			 = sorted(__lowerCamelCase  ,			key=lambda __lowerCamelCase  : x.__name__  )

	return model_class




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		int  ):







	'''simple docstring'''



	lowercase_			 = [(task, get_model_for_pipeline_test(__lowerCamelCase  ,			__lowerCamelCase  )) for task in pipeline_test_mapping]
	lowercase_			 = sorted([(task, model) for task, model in mapping if model is not None]  ,			key=lambda __lowerCamelCase  : x[0]  )

	return dict(__lowerCamelCase  )




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		Dict  ):







	'''simple docstring'''



	lowercase_			 = get_framework(__lowerCamelCase  )
	if framework == "pt":
		lowercase_			 = "torch"
	lowercase_			 = "{}"

	lowercase_			 = get_pipeline_model_mapping(__lowerCamelCase  )
	if len(__lowerCamelCase  ) == 0:
		return ""

	lowercase_			 = []
	for task, model_classes in mapping.items():
		if isinstance(__lowerCamelCase  ,			(tuple, list)  ):
			# A list/tuple of model classes
			lowercase_			 = "(" + ", ".join([x.__name__ for x in model_classes]  ) + ")"
		else:
			# A single model class
			lowercase_			 = model_classes.__name__
		texts.append(F'"{task}": {value}'  )
	lowercase_			 = "{" + ", ".join(__lowerCamelCase  ) + "}"
	lowercase_			 = F'pipeline_model_mapping = {text} if is_{framework}_available() else {default_value}'

	return text




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		Any  ):







	'''simple docstring'''



	lowercase_			 = {"ModelTesterMixin", "TFModelTesterMixin", "FlaxModelTesterMixin"}
	if not issubclass(__lowerCamelCase  ,			unittest.TestCase  ):
		return False
	return len(base_class_names.intersection([x.__name__ for x in test_class.__bases__]  )  ) > 0




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		Any  ):







	'''simple docstring'''



	lowercase_			 = [x for x in get_test_classes(__lowerCamelCase  ) if is_valid_test_class(__lowerCamelCase  )]

	lowercase_			 = None
	for test_class in test_classes:
		# If a test class has defined `pipeline_model_mapping`, let's take it
		if getattr(__lowerCamelCase  ,			"pipeline_model_mapping"  ,			__lowerCamelCase  ) is not None:
			lowercase_			 = test_class
			break
    # Take the test class with the shortest name (just a heuristic)
	if target_test_class is None and len(__lowerCamelCase  ) > 0:
		lowercase_			 = sorted(__lowerCamelCase  ,			key=lambda __lowerCamelCase  : (len(x.__name__  ), x.__name__)  )[0]

	return target_test_class




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		str  ,			__lowerCamelCase:		str  ,			__lowerCamelCase:		Optional[Any]  ):







	'''simple docstring'''



	lowercase_			 = start_idx
	for idx, line in enumerate(lines[start_idx:]  ):
		lowercase_			 = len(__lowerCamelCase  ) - len(line.lstrip()  )
		if idx == 0 or indent > indent_level or (indent == indent_level and line.strip() == ")"):
			lowercase_			 = start_idx + idx
		elif idx > 0 and indent <= indent_level:
			# Outside the definition block of `pipeline_model_mapping`
			break

	return end_idx




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		Any  ,			__lowerCamelCase:		Tuple=False  ):







	'''simple docstring'''



	if getattr(__lowerCamelCase  ,			"pipeline_model_mapping"  ,			__lowerCamelCase  ) is not None:
		if not overwrite:
			return "", -1

	lowercase_			 = get_pipeline_model_mapping_string(__lowerCamelCase  )
	if len(__lowerCamelCase  ) == 0:
		return "", -1
	lowercase_			 = line_to_add + "\n"

	# The code defined the class `test_class`
	lowercase_ ,	lowercase_			 = inspect.getsourcelines(__lowerCamelCase  )
	# `inspect` gives the code for an object, including decorator(s) if any.
	# We (only) need the exact line of the class definition.
	for idx, line in enumerate(__lowerCamelCase  ):
		if line.lstrip().startswith("class "  ):
			lowercase_			 = class_lines[idx:]
			class_start_line_no += idx
			break
	lowercase_			 = class_start_line_no + len(__lowerCamelCase  ) - 1

	# The index in `class_lines` that starts the definition of `all_model_classes`, `all_generative_model_classes` or
	# `pipeline_model_mapping`. This assumes they are defined in such order, and we take the start index of the last
	# block that appears in a `test_class`.
	lowercase_			 = None
	# The indent level of the line at `class_lines[start_idx]` (if defined)
	lowercase_			 = 0
	# To record if `pipeline_model_mapping` is found in `test_class`.
	lowercase_			 = None
	for idx, line in enumerate(__lowerCamelCase  ):
		if line.strip().startswith("all_model_classes = "  ):
			lowercase_			 = len(__lowerCamelCase  ) - len(line.lstrip()  )
			lowercase_			 = idx
		elif line.strip().startswith("all_generative_model_classes = "  ):
			lowercase_			 = len(__lowerCamelCase  ) - len(line.lstrip()  )
			lowercase_			 = idx
		elif line.strip().startswith("pipeline_model_mapping = "  ):
			lowercase_			 = len(__lowerCamelCase  ) - len(line.lstrip()  )
			lowercase_			 = idx
			lowercase_			 = line
			break

	if start_idx is None:
		return "", -1
	# Find the ending index (inclusive) of the above found block.
	lowercase_			 = find_block_ending(__lowerCamelCase  ,			__lowerCamelCase  ,			__lowerCamelCase  )

	# Extract `is_xxx_available()` from existing blocks: some models require specific libraries like `timm` and use
	# `is_timm_available()` instead of `is_torch_available()`.
	# Keep leading and trailing whitespaces
	lowercase_			 = re.compile(r"\s(is_\S+?_available\(\))\s"  )
	for line in class_lines[start_idx : end_idx + 1]:
		lowercase_			 = r.search(__lowerCamelCase  )
		if backend_condition is not None:
			# replace the leading and trailing whitespaces to the space character " ".
			lowercase_			 = " " + backend_condition[0][1:-1] + " "
			lowercase_			 = r.sub(__lowerCamelCase  ,			__lowerCamelCase  )
			break

	if def_line is None:
		# `pipeline_model_mapping` is not defined. The target index is set to the ending index (inclusive) of
		# `all_model_classes` or `all_generative_model_classes`.
		lowercase_			 = end_idx
	else:
		# `pipeline_model_mapping` is defined. The target index is set to be one **BEFORE** its start index.
		lowercase_			 = start_idx - 1
		# mark the lines of the currently existing `pipeline_model_mapping` to be removed.
		for idx in range(__lowerCamelCase  ,			end_idx + 1  ):
			# These lines are going to be removed before writing to the test file.
			lowercase_			 = None  # noqa

    # Make sure the test class is a subclass of `PipelineTesterMixin`.
	lowercase_			 = [x.__name__ for x in test_class.__bases__]
	if "PipelineTesterMixin" not in parent_classes:
		# Put `PipelineTesterMixin` just before `unittest.TestCase`
		lowercase_			 = [x for x in parent_classes if x != "TestCase"] + ["PipelineTesterMixin"]
		if "TestCase" in parent_classes:
			# Here we **assume** the original string is always with `unittest.TestCase`.
			_parent_classes.append("unittest.TestCase"  )
		lowercase_			 = ", ".join(_parent_classes  )
		for idx, line in enumerate(__lowerCamelCase  ):
			# Find the ending of the declaration of `test_class`
			if line.strip().endswith("):"  ):
				# mark the lines of the declaration of `test_class` to be removed
				for _idx in range(idx + 1  ):
					lowercase_			 = None  # noqa
				break
        # Add the new, one-line, class declaration for `test_class`
		lowercase_			 = F'class {test_class.__name__}({parent_classes}):\n'

	# Add indentation
	lowercase_			 = " " * indent_level + line_to_add
	# Insert `pipeline_model_mapping` to `class_lines`.
	# (The line at `target_idx` should be kept by definition!)
	lowercase_			 = class_lines[: target_idx + 1] + [line_to_add] + class_lines[target_idx + 1 :]
	# Remove the lines that are marked to be removed
	lowercase_			 = [x for x in class_lines if x is not None]

	# Move from test class to module (in order to write to the test file)
	lowercase_			 = inspect.getsourcelines(inspect.getmodule(__lowerCamelCase  )  )[0]
	# Be careful with the 1-off between line numbers and array indices
	lowercase_			 = module_lines[: class_start_line_no - 1] + class_lines + module_lines[class_end_line_no:]
	lowercase_			 = "".join(__lowerCamelCase  )

	lowercase_			 = inspect.getsourcefile(__lowerCamelCase  )
	with open(__lowerCamelCase  ,			"w"  ,			encoding="UTF-8"  ,			newline="\n"  ) as fp:
		fp.write(__lowerCamelCase  )

	return line_to_add




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		Tuple  ,			__lowerCamelCase:		List[str]=False  ):







	'''simple docstring'''



	lowercase_			 = find_test_class(__lowerCamelCase  )
	if test_class:
		add_pipeline_model_mapping(__lowerCamelCase  ,			overwrite=__lowerCamelCase  )


if __name__ == "__main__":
							SCREAMING_SNAKE_CASE__       			=				argparse.ArgumentParser()
							parser.add_argument(
							    """--test_file""", type=str, help="""A path to the test file, starting with the repository's `tests` directory."""
							)
							parser.add_argument(
							    """--all""",
							    action="""store_true""",
							    help="""If to check and modify all test files.""",
							)
							parser.add_argument(
							    """--overwrite""",
							    action="""store_true""",
							    help="""If to overwrite a test class if it has already defined `pipeline_model_mapping`.""",
							)
							SCREAMING_SNAKE_CASE__       			=				parser.parse_args()

							if not args.all and not args.test_file:
														raise ValueError("""Please specify either `test_file` or pass `--all` to check/modify all test files.""")
							elif args.all and args.test_file:
														raise ValueError("""Only one of `--test_file` and `--all` could be specified.""")

							SCREAMING_SNAKE_CASE__       			=				[]
							if args.test_file:
														SCREAMING_SNAKE_CASE__       			=				[args.test_file]
							else:
														SCREAMING_SNAKE_CASE__       			=				os.path.join("""tests""", """models""", """**""", """test_modeling_*.py""")
														for test_file in glob.glob(pattern):
																					# `Flax` is not concerned at this moment
																					if not test_file.startswith("""test_modeling_flax_"""):
																												test_files.append(test_file)

							for test_file in test_files:
														if test_file in TEST_FILE_TO_IGNORE:
																					print(f"""[SKIPPED] {test_file} is skipped as it is in `TEST_FILE_TO_IGNORE` in the file {__file__}.""")
																					continue
														add_pipeline_model_mapping_to_test_file(test_file, overwrite=args.overwrite)





import argparse

import torch
from clip import load

from transformers import CLIPConfig, CLIPModel




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		Optional[int]  ,			__lowerCamelCase:		Optional[int]  ):







	'''simple docstring'''



	lowercase_ ,	lowercase_ ,	lowercase_			 = pt_attn_layer.in_proj_weight.chunk(3  ,			dim=0  )
	lowercase_ ,	lowercase_ ,	lowercase_			 = pt_attn_layer.in_proj_bias.chunk(3  ,			dim=0  )

	lowercase_			 = pt_attn_layer.out_proj.weight
	lowercase_			 = pt_attn_layer.out_proj.bias

	lowercase_			 = q_proj
	lowercase_			 = q_proj_bias

	lowercase_			 = k_proj
	lowercase_			 = k_proj_bias

	lowercase_			 = v_proj
	lowercase_			 = v_proj_bias

	lowercase_			 = out_proj_weights
	lowercase_			 = out_proj_bias




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		int  ,			__lowerCamelCase:		List[str]  ):







	'''simple docstring'''



	copy_linear(hf_mlp.fca  ,			pt_mlp.c_fc  )
	copy_linear(hf_mlp.fca  ,			pt_mlp.c_proj  )




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		List[str]  ,			__lowerCamelCase:		Dict  ):







	'''simple docstring'''



	lowercase_			 = pt_linear.weight
	lowercase_			 = pt_linear.bias




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		int  ,			__lowerCamelCase:		Any  ):







	'''simple docstring'''



	copy_linear(hf_layer.layer_norma  ,			pt_layer.ln_a  )
	copy_linear(hf_layer.layer_norma  ,			pt_layer.ln_a  )

	# copy MLP
	copy_mlp(hf_layer.mlp  ,			pt_layer.mlp  )

	# copy attn
	copy_attn_layer(hf_layer.self_attn  ,			pt_layer.attn  )




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		str  ,			__lowerCamelCase:		Optional[int]  ):







	'''simple docstring'''



	for hf_layer, pt_layer in zip(__lowerCamelCase  ,			__lowerCamelCase  ):
		copy_layer(__lowerCamelCase  ,			__lowerCamelCase  )




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		Optional[Any]  ,			__lowerCamelCase:		str  ):







	'''simple docstring'''



	lowercase_			 = pt_model.token_embedding.weight
	lowercase_			 = pt_model.positional_embedding

	# copy layer norm
	copy_linear(hf_encoder.final_layer_norm  ,			pt_model.ln_final  )

	# copy hidden layers
	copy_layers(hf_encoder.encoder.layers  ,			pt_model.transformer.resblocks  )




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		Optional[Any]  ,			__lowerCamelCase:		Any  ):







	'''simple docstring'''



	lowercase_			 = pt_model.text_projection.data.T

	# copy text encoder
	copy_encoder(hf_model.text_model  ,			__lowerCamelCase  )




def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		List[Any]  ,			__lowerCamelCase:		List[Any]  ):







	'''simple docstring'''



	lowercase_			 = pt_model.visual.proj.data.T

	# copy layer norms
	copy_linear(hf_model.vision_model.pre_layrnorm  ,			pt_model.visual.ln_pre  )
	copy_linear(hf_model.vision_model.post_layernorm  ,			pt_model.visual.ln_post  )

	# copy embeds
	lowercase_			 = pt_model.visual.conva.weight.data
	lowercase_			 = pt_model.visual.class_embedding
	lowercase_			 = pt_model.visual.positional_embedding.data

	# copy encoder
	copy_layers(hf_model.vision_model.encoder.layers  ,			pt_model.visual.transformer.resblocks  )




@torch.no_grad()
def SCREAMING_SNAKE_CASE_  (				__lowerCamelCase:		int  ,			__lowerCamelCase:		Dict  ,			__lowerCamelCase:		Dict=None  ):







	'''simple docstring'''



	if config_path is not None:
		lowercase_			 = CLIPConfig.from_pretrained(__lowerCamelCase  )
	else:
		lowercase_			 = CLIPConfig(projection_dim=512  ,			text_config={}  ,			vision_config={}  )

	lowercase_			 = CLIPModel(__lowerCamelCase  ).eval()

	lowercase_ ,	lowercase_			 = load(__lowerCamelCase  ,			device="cpu"  ,			jit=__lowerCamelCase  )
	lowercase_			 = pt_model.eval()

	copy_text_model_and_projection(__lowerCamelCase  ,			__lowerCamelCase  )
	copy_vison_model_and_projection(__lowerCamelCase  ,			__lowerCamelCase  )
	lowercase_			 = pt_model.logit_scale

	lowercase_			 = torch.arange(0  ,			77  ).unsqueeze(0  )
	lowercase_			 = torch.randn(1  ,			3  ,			224  ,			224  )

	lowercase_			 = hf_model(input_ids=__lowerCamelCase  ,			pixel_values=__lowerCamelCase  ,			return_dict=__lowerCamelCase  )
	lowercase_			 = hf_outputs.logits_per_image
	lowercase_			 = hf_outputs.logits_per_text
	lowercase_ ,	lowercase_			 = pt_model(__lowerCamelCase  ,			__lowerCamelCase  )

	assert torch.allclose(__lowerCamelCase  ,			__lowerCamelCase  ,			atol=1E-3  )
	assert torch.allclose(__lowerCamelCase  ,			__lowerCamelCase  ,			atol=1E-3  )

	hf_model.save_pretrained(__lowerCamelCase  )


if __name__ == "__main__":
							SCREAMING_SNAKE_CASE__       			=				argparse.ArgumentParser()
							parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
							parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to fairseq checkpoint""")
							parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
							SCREAMING_SNAKE_CASE__       			=				parser.parse_args()

							convert_clip_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)

