




'''simple docstring'''
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


A				: Optional[int]					      =			{}
for task, _ in pipeline_test_mapping.items():
  A				: int					      =			{"""pt""": None, """tf""": None}


# DO **NOT** add item to this set (unless the reason is approved)
A				: Union[str, Any]					      =			{
    """tests/models/esm/test_modeling_esmfold.py""",  # The pipeline test mapping is added to `test_modeling_esm.py`
}
def 							snake_case_    (  a__						:      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 							snake_case_    (  a__						:      List[str]       ,a__						:      int     ):



    """simple docstring"""




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

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

    if framework == "pt":
        __lowercase     							=		getattr(UpperCamelCase__       ,"""model_mapping"""       ,UpperCamelCase__     )
    elif framework == "tf":
        __lowercase     							=		getattr(UpperCamelCase__       ,"""tf_model_mapping"""       ,UpperCamelCase__     )

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

    # cache the results
    __lowercase     							=		mapping
    return mapping
def 							snake_case_    (  a__						:      Any       ,a__						:      str     ):



    """simple docstring"""




    __lowercase     							=		get_framework(UpperCamelCase__     )
    if framework is None:
        return None
    __lowercase     							=		get_mapping_for_task(UpperCamelCase__       ,UpperCamelCase__     )
    if mapping is None:
        return None

    __lowercase     							=		list({model_class.config_class for model_class in test_class.all_model_classes}     )
    if len(UpperCamelCase__     ) != 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]       ,UpperCamelCase__     )
    if isinstance(UpperCamelCase__       ,(tuple, list)     ):
        __lowercase     							=		sorted(UpperCamelCase__       ,key=lambda a__     : x.__name__     )

    return model_class
def 							snake_case_    (  a__						:      Optional[Any]     ):



    """simple docstring"""




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

    return dict(UpperCamelCase__     )
def 							snake_case_    (  a__						:      List[str]     ):



    """simple docstring"""




    __lowercase     							=		get_framework(UpperCamelCase__     )
    if framework == "pt":
        __lowercase     							=		"""torch"""
    __lowercase     							=		"""{}"""

    __lowercase     							=		get_pipeline_model_mapping(UpperCamelCase__     )
    if len(UpperCamelCase__     ) == 0:
        return ""

    __lowercase     							=		[]
    for task, model_classes in mapping.items():
        if isinstance(UpperCamelCase__       ,(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(UpperCamelCase__     ) + """}"""
    __lowercase     							=		f'pipeline_model_mapping = {text} if is_{framework}_available() else {default_value}'

    return text
def 							snake_case_    (  a__						:      Dict     ):



    """simple docstring"""




    __lowercase     							=		{"""ModelTesterMixin""", """TFModelTesterMixin""", """FlaxModelTesterMixin"""}
    if not issubclass(UpperCamelCase__       ,unittest.TestCase     ):
        return False
    return len(base_class_names.intersection([x.__name__ for x in test_class.__bases__]     )     ) > 0
def 							snake_case_    (  a__						:      Optional[Any]     ):



    """simple docstring"""




    __lowercase     							=		[x for x in get_test_classes(UpperCamelCase__     ) if is_valid_test_class(UpperCamelCase__     )]

    __lowercase     							=		None
    for test_class in test_classes:
        # If a test class has defined `pipeline_model_mapping`, let's take it
        if getattr(UpperCamelCase__       ,"""pipeline_model_mapping"""       ,UpperCamelCase__     ) 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(UpperCamelCase__     ) > 0:
        __lowercase     							=		sorted(UpperCamelCase__       ,key=lambda a__     : (len(x.__name__     ), x.__name__)     )[0]

    return target_test_class
def 							snake_case_    (  a__						:      Dict       ,a__						:      Dict       ,a__						:      Optional[int]     ):



    """simple docstring"""




    __lowercase     							=		start_idx
    for idx, line in enumerate(lines[start_idx:]     ):
        __lowercase     							=		len(UpperCamelCase__     ) - 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 							snake_case_    (  a__						:      Union[str, Any]       ,a__						:      Union[str, Any]=False     ):



    """simple docstring"""




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

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

    # The code defined the class `test_class`
    __lowercase							,__lowercase     							=		inspect.getsourcelines(UpperCamelCase__     )
    # `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(UpperCamelCase__     ):
        if line.lstrip().startswith("""class """     ):
            __lowercase     							=		class_lines[idx:]
            class_start_line_no += idx
            break
    __lowercase     							=		class_start_line_no + len(UpperCamelCase__     ) - 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(UpperCamelCase__     ):
        if line.strip().startswith("""all_model_classes = """     ):
            __lowercase     							=		len(UpperCamelCase__     ) - len(line.lstrip()     )
            __lowercase     							=		idx
        elif line.strip().startswith("""all_generative_model_classes = """     ):
            __lowercase     							=		len(UpperCamelCase__     ) - len(line.lstrip()     )
            __lowercase     							=		idx
        elif line.strip().startswith("""pipeline_model_mapping = """     ):
            __lowercase     							=		len(UpperCamelCase__     ) - 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(UpperCamelCase__       ,UpperCamelCase__       ,UpperCamelCase__     )

    # 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(UpperCamelCase__     )
        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(UpperCamelCase__       ,UpperCamelCase__     )
            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(UpperCamelCase__       ,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(UpperCamelCase__     ):
            # 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(UpperCamelCase__     )     )[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(UpperCamelCase__     )

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

    return line_to_add
def 							snake_case_    (  a__						:      List[str]       ,a__						:      Tuple=False     ):



    """simple docstring"""




    __lowercase     							=		find_test_class(UpperCamelCase__     )
    if test_class:
        add_pipeline_model_mapping(UpperCamelCase__       ,overwrite=UpperCamelCase__     )


if __name__ == "__main__":
  A				: Union[str, Any]					      =			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`.""",
  )
  A				: List[Any]					      =			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.""")

  A				: Tuple					      =			[]
  if args.test_file:
    A				: Optional[int]					      =			[args.test_file]
  else:
    A				: Optional[int]					      =			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)





"""simple docstring"""




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__lowerCamelCase									=	{
    "configuration_layoutlm": ["LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP", "LayoutLMConfig", "LayoutLMOnnxConfig"],
    "tokenization_layoutlm": ["LayoutLMTokenizer"],
}

try:
				if not is_tokenizers_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__lowerCamelCase									=	["LayoutLMTokenizerFast"]

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__lowerCamelCase									=	[
				    "LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST",
				    "LayoutLMForMaskedLM",
				    "LayoutLMForSequenceClassification",
				    "LayoutLMForTokenClassification",
				    "LayoutLMForQuestionAnswering",
				    "LayoutLMModel",
				    "LayoutLMPreTrainedModel",
				]

try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__lowerCamelCase									=	[
				    "TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST",
				    "TFLayoutLMForMaskedLM",
				    "TFLayoutLMForSequenceClassification",
				    "TFLayoutLMForTokenClassification",
				    "TFLayoutLMForQuestionAnswering",
				    "TFLayoutLMMainLayer",
				    "TFLayoutLMModel",
				    "TFLayoutLMPreTrainedModel",
				]


if TYPE_CHECKING:
				from .configuration_layoutlm import LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP, LayoutLMConfig, LayoutLMOnnxConfig
				from .tokenization_layoutlm import LayoutLMTokenizer

				try:
								if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_layoutlm_fast import LayoutLMTokenizerFast

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_layoutlm import (
								    LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
								    LayoutLMForMaskedLM,
								    LayoutLMForQuestionAnswering,
								    LayoutLMForSequenceClassification,
								    LayoutLMForTokenClassification,
								    LayoutLMModel,
								    LayoutLMPreTrainedModel,
								)
				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_layoutlm import (
								    TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
								    TFLayoutLMForMaskedLM,
								    TFLayoutLMForQuestionAnswering,
								    TFLayoutLMForSequenceClassification,
								    TFLayoutLMForTokenClassification,
								    TFLayoutLMMainLayer,
								    TFLayoutLMModel,
								    TFLayoutLMPreTrainedModel,
								)

else:
				import sys

				__lowerCamelCase									=	_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





