




import argparse
import os

import align
import numpy as np
import requests
import tensorflow as tf
import torch
from PIL import Image
from tokenizer import Tokenizer

from transformers import (
    AlignConfig,
    AlignModel,
    AlignProcessor,
    BertConfig,
    BertTokenizer,
    EfficientNetConfig,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
__A						:		Dict          =						logging.get_logger(__name__)





def      UpperCAmelCase  (						UpperCamelCase__   )     ->     List[str]:







				'''simple docstring'''


				__lowerCAmelCase					=  tf.image.resize(UpperCamelCase__						,			(3_46, 3_46)   )
				__lowerCAmelCase					=  tf.image.crop_to_bounding_box(UpperCamelCase__						,			(3_46 - 2_89) // 2						,			(3_46 - 2_89) // 2						,			2_89						,			2_89   )
				return image





def      UpperCAmelCase  (						)     ->     Optional[Any]:







				'''simple docstring'''


				__lowerCAmelCase					=  EfficientNetConfig.from_pretrained("""google/efficientnet-b7"""   )
				__lowerCAmelCase					=  2_89
				__lowerCAmelCase					=  6_40
				__lowerCAmelCase					=  {"""0""": """LABEL_0""", """1""": """LABEL_1"""}
				__lowerCAmelCase					=  {"""LABEL_0""": 0, """LABEL_1""": 1}
				__lowerCAmelCase					=  []

				__lowerCAmelCase					=  BertConfig()
				__lowerCAmelCase					=  AlignConfig.from_text_vision_configs(
				    text_config=UpperCamelCase__						,			vision_config=UpperCamelCase__						,			projection_dim=6_40   )
				return config





def      UpperCAmelCase  (						)     ->     Dict:







				'''simple docstring'''


				__lowerCAmelCase					=  """http://images.cocodataset.org/val2017/000000039769.jpg"""
				__lowerCAmelCase					=  Image.open(requests.get(UpperCamelCase__						,			stream=UpperCamelCase__   ).raw   )
				return im





def      UpperCAmelCase  (						)     ->     Optional[Any]:







				'''simple docstring'''


				__lowerCAmelCase					=  EfficientNetImageProcessor(
				    do_center_crop=UpperCamelCase__						,			rescale_factor=1 / 127.5						,			rescale_offset=UpperCamelCase__						,			do_normalize=UpperCamelCase__						,			include_top=UpperCamelCase__						,			resample=Image.BILINEAR						,			)
				__lowerCAmelCase					=  BertTokenizer.from_pretrained("""bert-base-uncased"""   )
				__lowerCAmelCase					=  64
				__lowerCAmelCase					=  AlignProcessor(image_processor=UpperCamelCase__						,			tokenizer=UpperCamelCase__   )
				return processor





def      UpperCAmelCase  (						UpperCamelCase__   )     ->     Tuple:







				'''simple docstring'''


				__lowerCAmelCase					=  [v.split("""_"""   )[0].split("""block"""   )[1] for v in original_param_names if v.startswith("""block"""   )]
				__lowerCAmelCase					=  list(set(UpperCamelCase__   )   )
				__lowerCAmelCase					=  sorted(UpperCamelCase__   )
				__lowerCAmelCase					=  len(UpperCamelCase__   )
				__lowerCAmelCase					=  {b: str(UpperCamelCase__   ) for b, i in zip(UpperCamelCase__						,			range(UpperCamelCase__   )   )}

				__lowerCAmelCase					=  []
				rename_keys.append(("""stem_conv/kernel:0""", """embeddings.convolution.weight""")   )
				rename_keys.append(("""stem_bn/gamma:0""", """embeddings.batchnorm.weight""")   )
				rename_keys.append(("""stem_bn/beta:0""", """embeddings.batchnorm.bias""")   )
				rename_keys.append(("""stem_bn/moving_mean:0""", """embeddings.batchnorm.running_mean""")   )
				rename_keys.append(("""stem_bn/moving_variance:0""", """embeddings.batchnorm.running_var""")   )

				for b in block_names:
								__lowerCAmelCase					=  block_name_mapping[b]
								rename_keys.append((F'''block{b}_expand_conv/kernel:0''', F'''encoder.blocks.{hf_b}.expansion.expand_conv.weight''')   )
								rename_keys.append((F'''block{b}_expand_bn/gamma:0''', F'''encoder.blocks.{hf_b}.expansion.expand_bn.weight''')   )
								rename_keys.append((F'''block{b}_expand_bn/beta:0''', F'''encoder.blocks.{hf_b}.expansion.expand_bn.bias''')   )
								rename_keys.append(
								    (F'''block{b}_expand_bn/moving_mean:0''', F'''encoder.blocks.{hf_b}.expansion.expand_bn.running_mean''')   )
								rename_keys.append(
								    (F'''block{b}_expand_bn/moving_variance:0''', F'''encoder.blocks.{hf_b}.expansion.expand_bn.running_var''')   )
								rename_keys.append(
								    (F'''block{b}_dwconv/depthwise_kernel:0''', F'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight''')   )
								rename_keys.append((F'''block{b}_bn/gamma:0''', F'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight''')   )
								rename_keys.append((F'''block{b}_bn/beta:0''', F'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias''')   )
								rename_keys.append(
								    (F'''block{b}_bn/moving_mean:0''', F'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean''')   )
								rename_keys.append(
								    (F'''block{b}_bn/moving_variance:0''', F'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var''')   )

								rename_keys.append((F'''block{b}_se_reduce/kernel:0''', F'''encoder.blocks.{hf_b}.squeeze_excite.reduce.weight''')   )
								rename_keys.append((F'''block{b}_se_reduce/bias:0''', F'''encoder.blocks.{hf_b}.squeeze_excite.reduce.bias''')   )
								rename_keys.append((F'''block{b}_se_expand/kernel:0''', F'''encoder.blocks.{hf_b}.squeeze_excite.expand.weight''')   )
								rename_keys.append((F'''block{b}_se_expand/bias:0''', F'''encoder.blocks.{hf_b}.squeeze_excite.expand.bias''')   )
								rename_keys.append(
								    (F'''block{b}_project_conv/kernel:0''', F'''encoder.blocks.{hf_b}.projection.project_conv.weight''')   )
								rename_keys.append((F'''block{b}_project_bn/gamma:0''', F'''encoder.blocks.{hf_b}.projection.project_bn.weight''')   )
								rename_keys.append((F'''block{b}_project_bn/beta:0''', F'''encoder.blocks.{hf_b}.projection.project_bn.bias''')   )
								rename_keys.append(
								    (F'''block{b}_project_bn/moving_mean:0''', F'''encoder.blocks.{hf_b}.projection.project_bn.running_mean''')   )
								rename_keys.append(
								    (F'''block{b}_project_bn/moving_variance:0''', F'''encoder.blocks.{hf_b}.projection.project_bn.running_var''')   )

				__lowerCAmelCase					=  {}
				for item in rename_keys:
								if item[0] in original_param_names:
												__lowerCAmelCase					=  """vision_model.""" + item[1]

    # BERT text encoder
				__lowerCAmelCase					=  []
				__lowerCAmelCase					=  """tf_bert_model/bert"""
				__lowerCAmelCase					=  """text_model"""
				for i in range(12   ):
								rename_keys.append(
								    (
								        F'''{old}/encoder/layer_._{i}/attention/self/query/kernel:0''',
								        F'''{new}.encoder.layer.{i}.attention.self.query.weight''',
								    )   )
								rename_keys.append(
								    (
								        F'''{old}/encoder/layer_._{i}/attention/self/query/bias:0''',
								        F'''{new}.encoder.layer.{i}.attention.self.query.bias''',
								    )   )
								rename_keys.append(
								    (
								        F'''{old}/encoder/layer_._{i}/attention/self/key/kernel:0''',
								        F'''{new}.encoder.layer.{i}.attention.self.key.weight''',
								    )   )
								rename_keys.append(
								    (
								        F'''{old}/encoder/layer_._{i}/attention/self/key/bias:0''',
								        F'''{new}.encoder.layer.{i}.attention.self.key.bias''',
								    )   )
								rename_keys.append(
								    (
								        F'''{old}/encoder/layer_._{i}/attention/self/value/kernel:0''',
								        F'''{new}.encoder.layer.{i}.attention.self.value.weight''',
								    )   )
								rename_keys.append(
								    (
								        F'''{old}/encoder/layer_._{i}/attention/self/value/bias:0''',
								        F'''{new}.encoder.layer.{i}.attention.self.value.bias''',
								    )   )
								rename_keys.append(
								    (
								        F'''{old}/encoder/layer_._{i}/attention/output/dense/kernel:0''',
								        F'''{new}.encoder.layer.{i}.attention.output.dense.weight''',
								    )   )
								rename_keys.append(
								    (
								        F'''{old}/encoder/layer_._{i}/attention/output/dense/bias:0''',
								        F'''{new}.encoder.layer.{i}.attention.output.dense.bias''',
								    )   )
								rename_keys.append(
								    (
								        F'''{old}/encoder/layer_._{i}/attention/output/LayerNorm/gamma:0''',
								        F'''{new}.encoder.layer.{i}.attention.output.LayerNorm.weight''',
								    )   )
								rename_keys.append(
								    (
								        F'''{old}/encoder/layer_._{i}/attention/output/LayerNorm/beta:0''',
								        F'''{new}.encoder.layer.{i}.attention.output.LayerNorm.bias''',
								    )   )
								rename_keys.append(
								    (
								        F'''{old}/encoder/layer_._{i}/intermediate/dense/kernel:0''',
								        F'''{new}.encoder.layer.{i}.intermediate.dense.weight''',
								    )   )
								rename_keys.append(
								    (
								        F'''{old}/encoder/layer_._{i}/intermediate/dense/bias:0''',
								        F'''{new}.encoder.layer.{i}.intermediate.dense.bias''',
								    )   )
								rename_keys.append(
								    (F'''{old}/encoder/layer_._{i}/output/dense/kernel:0''', F'''{new}.encoder.layer.{i}.output.dense.weight''')   )
								rename_keys.append(
								    (F'''{old}/encoder/layer_._{i}/output/dense/bias:0''', F'''{new}.encoder.layer.{i}.output.dense.bias''')   )
								rename_keys.append(
								    (F'''{old}/encoder/layer_._{i}/output/LayerNorm/gamma:0''', F'''{new}.encoder.layer.{i}.output.LayerNorm.weight''')   )
								rename_keys.append(
								    (F'''{old}/encoder/layer_._{i}/output/LayerNorm/beta:0''', F'''{new}.encoder.layer.{i}.output.LayerNorm.bias''')   )

				rename_keys.append((F'''{old}/embeddings/word_embeddings/weight:0''', F'''{new}.embeddings.word_embeddings.weight''')   )
				rename_keys.append(
				    (F'''{old}/embeddings/position_embeddings/embeddings:0''', F'''{new}.embeddings.position_embeddings.weight''')   )
				rename_keys.append(
				    (F'''{old}/embeddings/token_type_embeddings/embeddings:0''', F'''{new}.embeddings.token_type_embeddings.weight''')   )
				rename_keys.append((F'''{old}/embeddings/LayerNorm/gamma:0''', F'''{new}.embeddings.LayerNorm.weight''')   )
				rename_keys.append((F'''{old}/embeddings/LayerNorm/beta:0''', F'''{new}.embeddings.LayerNorm.bias''')   )

				rename_keys.append((F'''{old}/pooler/dense/kernel:0''', F'''{new}.pooler.dense.weight''')   )
				rename_keys.append((F'''{old}/pooler/dense/bias:0''', F'''{new}.pooler.dense.bias''')   )
				rename_keys.append(("""dense/kernel:0""", """text_projection.weight""")   )
				rename_keys.append(("""dense/bias:0""", """text_projection.bias""")   )
				rename_keys.append(("""dense/bias:0""", """text_projection.bias""")   )
				rename_keys.append(("""temperature:0""", """temperature""")   )

				for item in rename_keys:
								if item[0] in original_param_names:
												__lowerCAmelCase					=  item[1]
				return key_mapping





def      UpperCAmelCase  (						UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__   )     ->     Optional[Any]:







				'''simple docstring'''


				list(hf_params.keys()   )

				for key, value in tf_params.items():
								if key not in key_mapping:
												continue

								__lowerCAmelCase					=  key_mapping[key]
								if "_conv" in key and "kernel" in key:
												__lowerCAmelCase					=  torch.from_numpy(UpperCamelCase__   ).permute(3						,			2						,			0						,			1   )
								elif "embeddings" in key:
												__lowerCAmelCase					=  torch.from_numpy(UpperCamelCase__   )
								elif "depthwise_kernel" in key:
												__lowerCAmelCase					=  torch.from_numpy(UpperCamelCase__   ).permute(2						,			3						,			0						,			1   )
								elif "kernel" in key:
												__lowerCAmelCase					=  torch.from_numpy(np.transpose(UpperCamelCase__   )   )
								elif "temperature" in key:
												__lowerCAmelCase					=  value
								elif "bn/gamma" or "bn/beta" in key:
												__lowerCAmelCase					=  torch.from_numpy(np.transpose(UpperCamelCase__   )   ).squeeze()
								else:
												__lowerCAmelCase					=  torch.from_numpy(UpperCamelCase__   )

								# Replace HF parameters with original TF model parameters
								hf_params[hf_key].copy_(UpperCamelCase__   )





@torch.no_grad()
def      UpperCAmelCase  (						UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__   )     ->     List[str]:







				'''simple docstring'''


				__lowerCAmelCase					=  64
				__lowerCAmelCase					=  Tokenizer(UpperCamelCase__   )
				__lowerCAmelCase					=  align.Align("""efficientnet-b7"""						,			"""bert-base"""						,			6_40						,			UpperCamelCase__						,			tok.get_vocab_size()   )
				original_model.compile()
				original_model.load_weights(UpperCamelCase__   )

				__lowerCAmelCase					=  original_model.trainable_variables
				__lowerCAmelCase					=  original_model.non_trainable_variables
				__lowerCAmelCase					=  {param.name: param.numpy() for param in tf_params}
				for param in tf_non_train_params:
								__lowerCAmelCase					=  param.numpy()
				__lowerCAmelCase					=  list(tf_params.keys()   )

				# Load HuggingFace model
				__lowerCAmelCase					=  get_align_config()
				__lowerCAmelCase					=  AlignModel(UpperCamelCase__   ).eval()
				__lowerCAmelCase					=  hf_model.state_dict()

				# Create src-to-dst parameter name mapping dictionary
				print("""Converting parameters..."""   )
				__lowerCAmelCase					=  rename_keys(UpperCamelCase__   )
				replace_params(UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__   )

				# Initialize processor
				__lowerCAmelCase					=  get_processor()
				__lowerCAmelCase					=  processor(
				    images=prepare_img()						,			text="""A picture of a cat"""						,			padding="""max_length"""						,			max_length=64						,			return_tensors="""pt"""   )

				# HF model inference
				hf_model.eval()
				with torch.no_grad():
								__lowerCAmelCase					=  hf_model(**UpperCamelCase__   )

				__lowerCAmelCase					=  outputs.image_embeds.detach().numpy()
				__lowerCAmelCase					=  outputs.text_embeds.detach().numpy()

				# Original model inference
				__lowerCAmelCase					=  False
				__lowerCAmelCase					=  EfficientNetImageProcessor(
				    do_center_crop=UpperCamelCase__						,			do_rescale=UpperCamelCase__						,			do_normalize=UpperCamelCase__						,			include_top=UpperCamelCase__						,			resample=Image.BILINEAR						,			)
				__lowerCAmelCase					=  tf_image_processor(images=prepare_img()						,			return_tensors="""tf"""						,			data_format="""channels_last"""   )["""pixel_values"""]
				__lowerCAmelCase					=  tok(tf.constant(["""A picture of a cat"""]   )   )

				__lowerCAmelCase					=  original_model.image_encoder(UpperCamelCase__						,			training=UpperCamelCase__   )
				__lowerCAmelCase					=  original_model.text_encoder(UpperCamelCase__						,			training=UpperCamelCase__   )

				__lowerCAmelCase					=  tf.nn.la_normalize(UpperCamelCase__						,			axis=-1   )
				__lowerCAmelCase					=  tf.nn.la_normalize(UpperCamelCase__						,			axis=-1   )

				# Check whether original and HF model outputs match  -> np.allclose
				if not np.allclose(UpperCamelCase__						,			UpperCamelCase__						,			atol=1e-3   ):
								raise ValueError("""The predicted image features are not the same."""   )
				if not np.allclose(UpperCamelCase__						,			UpperCamelCase__						,			atol=1e-3   ):
								raise ValueError("""The predicted text features are not the same."""   )
				print("""Model outputs match!"""   )

				if save_model:
								# Create folder to save model
								if not os.path.isdir(UpperCamelCase__   ):
												os.mkdir(UpperCamelCase__   )
								# Save converted model and image processor
								hf_model.save_pretrained(UpperCamelCase__   )
								processor.save_pretrained(UpperCamelCase__   )

				if push_to_hub:
								# Push model and image processor to hub
								print("""Pushing converted ALIGN to the hub..."""   )
								processor.push_to_hub("""align-base"""   )
								hf_model.push_to_hub("""align-base"""   )


if __name__ == "__main__":
	__A						:		str          =						argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    "--checkpoint_path",
	    default="./weights/model-weights",
	    type=str,
	    help="Path to the pretrained TF ALIGN checkpoint.",
	)
	parser.add_argument(
	    "--pytorch_dump_folder_path",
	    default="hf_model",
	    type=str,
	    help="Path to the output PyTorch model directory.",
	)
	parser.add_argument("--save_model", action="store_true", help="Save model to local")
	parser.add_argument("--push_to_hub", action="store_true", help="Push model and image processor to the hub")

	__A						:		Union[str, Any]          =						parser.parse_args()
	convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)










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


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


# DO **NOT** add item to this set (unless the reason is approved)
UpperCAmelCase_          =				{
    "tests/models/esm/test_modeling_esmfold.py",  # The pipeline test mapping is added to `test_modeling_esm.py`
}
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Tuple					)->Union[str, Any]:

      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 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Any  ,					_SCREAMING_SNAKE_CASE :						str					)->Dict:
      # Use the cached results
      if PIPELINE_TEST_MAPPING[task].get(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					) is not None:
            return PIPELINE_TEST_MAPPING[task][framework]

      _lowerCAmelCase			      =	pipeline_test_mapping[task]['''test''']
      _lowerCAmelCase			      =	None

      if framework == "pt":
            _lowerCAmelCase			      =	getattr(_SCREAMING_SNAKE_CASE  ,					'''model_mapping'''  ,					_SCREAMING_SNAKE_CASE					)
      elif framework == "tf":
            _lowerCAmelCase			      =	getattr(_SCREAMING_SNAKE_CASE  ,					'''tf_model_mapping'''  ,					_SCREAMING_SNAKE_CASE					)

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

      # cache the results
      _lowerCAmelCase			      =	mapping
      return mapping
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Union[str, Any]  ,					_SCREAMING_SNAKE_CASE :						Union[str, Any]					)->List[Any]:
      _lowerCAmelCase			      =	get_framework(_SCREAMING_SNAKE_CASE					)
      if framework is None:
            return None
      _lowerCAmelCase			      =	get_mapping_for_task(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      if mapping is None:
            return None

      _lowerCAmelCase			      =	list({model_class.config_class for model_class in test_class.all_model_classes}					)
      if len(_SCREAMING_SNAKE_CASE					) != 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.
      _lowerCAmelCase			      =	mapping.get(config_classes[0]  ,					_SCREAMING_SNAKE_CASE					)
      if isinstance(_SCREAMING_SNAKE_CASE  ,					(tuple, list)					):
            _lowerCAmelCase			      =	sorted(_SCREAMING_SNAKE_CASE  ,					key=lambda _SCREAMING_SNAKE_CASE					: x.__name__					)

      return model_class
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						int					)->Dict:
      _lowerCAmelCase			      =	[(task, get_model_for_pipeline_test(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)) for task in pipeline_test_mapping]
      _lowerCAmelCase			      =	sorted([(task, model) for task, model in mapping if model is not None]  ,					key=lambda _SCREAMING_SNAKE_CASE					: x[0]					)

      return dict(_SCREAMING_SNAKE_CASE					)
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Union[str, Any]					)->Dict:
      _lowerCAmelCase			      =	get_framework(_SCREAMING_SNAKE_CASE					)
      if framework == "pt":
            _lowerCAmelCase			      =	'''torch'''
      _lowerCAmelCase			      =	'''{}'''

      _lowerCAmelCase			      =	get_pipeline_model_mapping(_SCREAMING_SNAKE_CASE					)
      if len(_SCREAMING_SNAKE_CASE					) == 0:
            return ""

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

      return text
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Optional[Any]					)->int:
      _lowerCAmelCase			      =	{'''ModelTesterMixin''', '''TFModelTesterMixin''', '''FlaxModelTesterMixin'''}
      if not issubclass(_SCREAMING_SNAKE_CASE  ,					unittest.TestCase					):
            return False
      return len(base_class_names.intersection([x.__name__ for x in test_class.__bases__]					)					) > 0
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Any					)->int:
      _lowerCAmelCase			      =	[x for x in get_test_classes(_SCREAMING_SNAKE_CASE					) if is_valid_test_class(_SCREAMING_SNAKE_CASE					)]

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

      return target_test_class
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						List[str]  ,					_SCREAMING_SNAKE_CASE :						Tuple  ,					_SCREAMING_SNAKE_CASE :						Optional[int]					)->Union[str, Any]:
      _lowerCAmelCase			      =	start_idx
      for idx, line in enumerate(lines[start_idx:]					):
            _lowerCAmelCase			      =	len(_SCREAMING_SNAKE_CASE					) - len(line.lstrip()					)
            if idx == 0 or indent > indent_level or (indent == indent_level and line.strip() == ")"):
                  _lowerCAmelCase			      =	start_idx + idx
            elif idx > 0 and indent <= indent_level:
                  # Outside the definition block of `pipeline_model_mapping`
                  break

      return end_idx
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Dict  ,					_SCREAMING_SNAKE_CASE :						Tuple=False					)->Tuple:
      if getattr(_SCREAMING_SNAKE_CASE  ,					'''pipeline_model_mapping'''  ,					_SCREAMING_SNAKE_CASE					) is not None:
            if not overwrite:
                  return "", -1

      _lowerCAmelCase			      =	get_pipeline_model_mapping_string(_SCREAMING_SNAKE_CASE					)
      if len(_SCREAMING_SNAKE_CASE					) == 0:
            return "", -1
      _lowerCAmelCase			      =	line_to_add + '''\n'''

      # The code defined the class `test_class`
      _lowerCAmelCase      , _lowerCAmelCase			      =	inspect.getsourcelines(_SCREAMING_SNAKE_CASE					)
      # `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(_SCREAMING_SNAKE_CASE					):
            if line.lstrip().startswith('''class '''					):
                  _lowerCAmelCase			      =	class_lines[idx:]
                  class_start_line_no += idx
                  break
      _lowerCAmelCase			      =	class_start_line_no + len(_SCREAMING_SNAKE_CASE					) - 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`.
      _lowerCAmelCase			      =	None
      # The indent level of the line at `class_lines[start_idx]` (if defined)
      _lowerCAmelCase			      =	0
      # To record if `pipeline_model_mapping` is found in `test_class`.
      _lowerCAmelCase			      =	None
      for idx, line in enumerate(_SCREAMING_SNAKE_CASE					):
            if line.strip().startswith('''all_model_classes = '''					):
                  _lowerCAmelCase			      =	len(_SCREAMING_SNAKE_CASE					) - len(line.lstrip()					)
                  _lowerCAmelCase			      =	idx
            elif line.strip().startswith('''all_generative_model_classes = '''					):
                  _lowerCAmelCase			      =	len(_SCREAMING_SNAKE_CASE					) - len(line.lstrip()					)
                  _lowerCAmelCase			      =	idx
            elif line.strip().startswith('''pipeline_model_mapping = '''					):
                  _lowerCAmelCase			      =	len(_SCREAMING_SNAKE_CASE					) - len(line.lstrip()					)
                  _lowerCAmelCase			      =	idx
                  _lowerCAmelCase			      =	line
                  break

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

      # 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
      _lowerCAmelCase			      =	re.compile(r'''\s(is_\S+?_available\(\))\s'''					)
      for line in class_lines[start_idx : end_idx + 1]:
            _lowerCAmelCase			      =	r.search(_SCREAMING_SNAKE_CASE					)
            if backend_condition is not None:
                  # replace the leading and trailing whitespaces to the space character " ".
                  _lowerCAmelCase			      =	''' ''' + backend_condition[0][1:-1] + ''' '''
                  _lowerCAmelCase			      =	r.sub(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
                  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`.
            _lowerCAmelCase			      =	end_idx
      else:
            # `pipeline_model_mapping` is defined. The target index is set to be one **BEFORE** its start index.
            _lowerCAmelCase			      =	start_idx - 1
            # mark the lines of the currently existing `pipeline_model_mapping` to be removed.
            for idx in range(_SCREAMING_SNAKE_CASE  ,					end_idx + 1					):
                  # These lines are going to be removed before writing to the test file.
                  _lowerCAmelCase			      =	None  # noqa

    # Make sure the test class is a subclass of `PipelineTesterMixin`.
      _lowerCAmelCase			      =	[x.__name__ for x in test_class.__bases__]
      if "PipelineTesterMixin" not in parent_classes:
            # Put `PipelineTesterMixin` just before `unittest.TestCase`
            _lowerCAmelCase			      =	[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'''					)
            _lowerCAmelCase			      =	''', '''.join(_parent_classes					)
            for idx, line in enumerate(_SCREAMING_SNAKE_CASE					):
                  # 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					):
                              _lowerCAmelCase			      =	None  # noqa
                        break
        # Add the new, one-line, class declaration for `test_class`
            _lowerCAmelCase			      =	f'''class {test_class.__name__}({parent_classes}):\n'''

      # Add indentation
      _lowerCAmelCase			      =	''' ''' * indent_level + line_to_add
      # Insert `pipeline_model_mapping` to `class_lines`.
      # (The line at `target_idx` should be kept by definition!)
      _lowerCAmelCase			      =	class_lines[: target_idx + 1] + [line_to_add] + class_lines[target_idx + 1 :]
      # Remove the lines that are marked to be removed
      _lowerCAmelCase			      =	[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)
      _lowerCAmelCase			      =	inspect.getsourcelines(inspect.getmodule(_SCREAMING_SNAKE_CASE					)					)[0]
      # Be careful with the 1-off between line numbers and array indices
      _lowerCAmelCase			      =	module_lines[: class_start_line_no - 1] + class_lines + module_lines[class_end_line_no:]
      _lowerCAmelCase			      =	''''''.join(_SCREAMING_SNAKE_CASE					)

      _lowerCAmelCase			      =	inspect.getsourcefile(_SCREAMING_SNAKE_CASE					)
      with open(_SCREAMING_SNAKE_CASE  ,					'''w'''  ,					encoding='''UTF-8'''  ,					newline='''\n'''					) as fp:
            fp.write(_SCREAMING_SNAKE_CASE					)

      return line_to_add
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Optional[int]  ,					_SCREAMING_SNAKE_CASE :						Optional[Any]=False					)->str:
      _lowerCAmelCase			      =	find_test_class(_SCREAMING_SNAKE_CASE					)
      if test_class:
            add_pipeline_model_mapping(_SCREAMING_SNAKE_CASE  ,					overwrite=_SCREAMING_SNAKE_CASE					)


if __name__ == "__main__":
    UpperCAmelCase_          =				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`.",
    )
    UpperCAmelCase_          =				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.")

    UpperCAmelCase_          =				[]
    if args.test_file:
        UpperCAmelCase_          =				[args.test_file]
    else:
        UpperCAmelCase_          =				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)