



# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
import difflib
from collections.abc import Mapping
from typing import Any, Callable, Dict





class 		UpperCAmelCase    (    snake_case_					):
   pass






def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						str  ,					_SCREAMING_SNAKE_CASE :						Dict[str, Callable]  ,					_SCREAMING_SNAKE_CASE :						Any=None  ,					_SCREAMING_SNAKE_CASE :						Optional[int]=False					)->Optional[Any]:
      try:
            _lowerCAmelCase			      =	ast.parse(_SCREAMING_SNAKE_CASE					)
      except SyntaxError as e:
            print('''The code generated by the agent is not valid.\n'''  ,					_SCREAMING_SNAKE_CASE					)
            return
      if state is None:
            _lowerCAmelCase			      =	{}
      _lowerCAmelCase			      =	None
      for idx, node in enumerate(expression.body					):
            try:
                  _lowerCAmelCase			      =	evaluate_ast(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
            except InterpretorError as e:
                  _lowerCAmelCase			      =	f'''Evaluation of the code stopped at line {idx} before the end because of the following error'''
                  if chat_mode:
                        msg += (
                            f'''. Copy paste the following error message and send it back to the agent:\nI get an error: \'{e}\''''
                        )
                  else:
                        msg += f''':\n{e}'''
                  print(_SCREAMING_SNAKE_CASE					)
                  break
            if line_result is not None:
                  _lowerCAmelCase			      =	line_result

      return result
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						ast.AST  ,					_SCREAMING_SNAKE_CASE :						Dict[str, Any]  ,					_SCREAMING_SNAKE_CASE :						Dict[str, Callable]					)->Tuple:
      if isinstance(_SCREAMING_SNAKE_CASE  ,					ast.Assign					):
            # Assignement -> we evaluate the assignement which should update the state
            # We return the variable assigned as it may be used to determine the final result.
            return evaluate_assign(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.Call					):
            # Function call -> we return the value of the function call
            return evaluate_call(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.Constant					):
            # Constant -> just return the value
            return expression.value
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.Dict					):
            # Dict -> evaluate all keys and values
            _lowerCAmelCase			      =	[evaluate_ast(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					) for k in expression.keys]
            _lowerCAmelCase			      =	[evaluate_ast(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					) for v in expression.values]
            return dict(zip(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)					)
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.Expr					):
            # Expression -> evaluate the content
            return evaluate_ast(expression.value  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.For					):
            # For loop -> execute the loop
            return evaluate_for(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.FormattedValue					):
            # Formatted value (part of f-string) -> evaluate the content and return
            return evaluate_ast(expression.value  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.If					):
            # If -> execute the right branch
            return evaluate_if(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      elif hasattr(_SCREAMING_SNAKE_CASE  ,					'''Index'''					) and isinstance(_SCREAMING_SNAKE_CASE  ,					ast.Index					):
            return evaluate_ast(expression.value  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.JoinedStr					):
            return "".join([str(evaluate_ast(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)					) for v in expression.values]					)
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.List					):
            # List -> evaluate all elements
            return [evaluate_ast(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					) for elt in expression.elts]
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.Name					):
            # Name -> pick up the value in the state
            return evaluate_name(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.Subscript					):
            # Subscript -> return the value of the indexing
            return evaluate_subscript(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      else:
            # For now we refuse anything else. Let's add things as we need them.
            raise InterpretorError(f'''{expression.__class__.__name__} is not supported.'''					)
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Optional[int]  ,					_SCREAMING_SNAKE_CASE :						Optional[int]  ,					_SCREAMING_SNAKE_CASE :						int					)->Union[str, Any]:
      _lowerCAmelCase			      =	assign.targets
      _lowerCAmelCase			      =	evaluate_ast(assign.value  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)

      if len(_SCREAMING_SNAKE_CASE					) == 1:
            _lowerCAmelCase			      =	result
      else:
            if len(_SCREAMING_SNAKE_CASE					) != len(_SCREAMING_SNAKE_CASE					):
                  raise InterpretorError(f'''Expected {len(_SCREAMING_SNAKE_CASE					)} values but got {len(_SCREAMING_SNAKE_CASE					)}.'''					)
            for var_name, r in zip(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					):
                  _lowerCAmelCase			      =	r
      return result
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						List[Any]  ,					_SCREAMING_SNAKE_CASE :						Dict  ,					_SCREAMING_SNAKE_CASE :						Union[str, Any]					)->str:
      if not isinstance(call.func  ,					ast.Name					):
            raise InterpretorError(
                f'''It is not permitted to evaluate other functions than the provided tools (tried to execute {call.func} of '''
                f'''type {type(call.func					)}.'''					)
      _lowerCAmelCase			      =	call.func.id
      if func_name not in tools:
            raise InterpretorError(
                f'''It is not permitted to evaluate other functions than the provided tools (tried to execute {call.func.id}).'''					)

      _lowerCAmelCase			      =	tools[func_name]
      # Todo deal with args
      _lowerCAmelCase			      =	[evaluate_ast(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					) for arg in call.args]
      _lowerCAmelCase			      =	{keyword.arg: evaluate_ast(keyword.value  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					) for keyword in call.keywords}
      return func(*_SCREAMING_SNAKE_CASE  ,					**_SCREAMING_SNAKE_CASE					)
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Optional[int]  ,					_SCREAMING_SNAKE_CASE :						Union[str, Any]  ,					_SCREAMING_SNAKE_CASE :						str					)->Any:
      _lowerCAmelCase			      =	evaluate_ast(subscript.slice  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      _lowerCAmelCase			      =	evaluate_ast(subscript.value  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      if isinstance(_SCREAMING_SNAKE_CASE  ,					(list, tuple)					):
            return value[int(_SCREAMING_SNAKE_CASE					)]
      if index in value:
            return value[index]
      if isinstance(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					) and isinstance(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					):
            _lowerCAmelCase			      =	difflib.get_close_matches(_SCREAMING_SNAKE_CASE  ,					list(value.keys()					)					)
            if len(_SCREAMING_SNAKE_CASE					) > 0:
                  return value[close_matches[0]]

      raise InterpretorError(f'''Could not index {value} with \'{index}\'.'''					)
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						str  ,					_SCREAMING_SNAKE_CASE :						int  ,					_SCREAMING_SNAKE_CASE :						str					)->Optional[int]:
      if name.id in state:
            return state[name.id]
      _lowerCAmelCase			      =	difflib.get_close_matches(name.id  ,					list(state.keys()					)					)
      if len(_SCREAMING_SNAKE_CASE					) > 0:
            return state[close_matches[0]]
      raise InterpretorError(f'''The variable `{name.id}` is not defined.'''					)
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Dict  ,					_SCREAMING_SNAKE_CASE :						List[str]  ,					_SCREAMING_SNAKE_CASE :						List[Any]					)->Optional[int]:
      if len(condition.ops					) > 1:
            raise InterpretorError('''Cannot evaluate conditions with multiple operators'''					)

      _lowerCAmelCase			      =	evaluate_ast(condition.left  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      _lowerCAmelCase			      =	condition.ops[0]
      _lowerCAmelCase			      =	evaluate_ast(condition.comparators[0]  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)

      if isinstance(_SCREAMING_SNAKE_CASE  ,					ast.Eq					):
            return left == right
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.NotEq					):
            return left != right
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.Lt					):
            return left < right
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.LtE					):
            return left <= right
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.Gt					):
            return left > right
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.GtE					):
            return left >= right
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.Is					):
            return left is right
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.IsNot					):
            return left is not right
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.In					):
            return left in right
      elif isinstance(_SCREAMING_SNAKE_CASE  ,					ast.NotIn					):
            return left not in right
      else:
            raise InterpretorError(f'''Operator not supported: {comparator}'''					)
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Tuple  ,					_SCREAMING_SNAKE_CASE :						Optional[int]  ,					_SCREAMING_SNAKE_CASE :						Any					)->int:
      _lowerCAmelCase			      =	None
      if evaluate_condition(if_statement.test  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					):
            for line in if_statement.body:
                  _lowerCAmelCase			      =	evaluate_ast(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
                  if line_result is not None:
                        _lowerCAmelCase			      =	line_result
      else:
            for line in if_statement.orelse:
                  _lowerCAmelCase			      =	evaluate_ast(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
                  if line_result is not None:
                        _lowerCAmelCase			      =	line_result
      return result


def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						List[str]  ,					_SCREAMING_SNAKE_CASE :						Tuple  ,					_SCREAMING_SNAKE_CASE :						List[str]					)->Optional[int]:
      _lowerCAmelCase			      =	None
      _lowerCAmelCase			      =	evaluate_ast(for_loop.iter  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
      for counter in iterator:
            _lowerCAmelCase			      =	counter
            for expression in for_loop.body:
                  _lowerCAmelCase			      =	evaluate_ast(_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE  ,					_SCREAMING_SNAKE_CASE					)
                  if line_result is not None:
                        _lowerCAmelCase			      =	line_result
      return result




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)