




'''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							:				Any			       =	{}
for task, _ in pipeline_test_mapping.items():
  A							:				Optional[Any]			       =	{"""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__		:     Optional[Any]      ):



    """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__		:     Union[str, Any]   ,a__		:     str      ):



    """simple docstring"""




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

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

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

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

    # cache the results
    __lowercase          =						mapping
    return mapping
def     snake_case_  (       a__		:     Optional[int]   ,a__		:     Optional[Any]      ):



    """simple docstring"""




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

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

    return model_class
def     snake_case_  (       a__		:     Any      ):



    """simple docstring"""




    __lowercase          =						[(task, get_model_for_pipeline_test(a__   ,a__      )) 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(a__      )
def     snake_case_  (       a__		:     Optional[int]      ):



    """simple docstring"""




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

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

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

    return text
def     snake_case_  (       a__		:     List[Any]      ):



    """simple docstring"""




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



    """simple docstring"""




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

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

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



    """simple docstring"""




    __lowercase          =						start_idx
    for idx, line in enumerate(lines[start_idx:]      ):
        __lowercase          =						len(a__      ) - 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__		:     int   ,a__		:     List[Any]=False      ):



    """simple docstring"""




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

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

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

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

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

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



    """simple docstring"""




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


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


A							:				Tuple			     =	{
    """configuration_dpr""": ["""DPR_PRETRAINED_CONFIG_ARCHIVE_MAP""", """DPRConfig"""],
    """tokenization_dpr""": [
        """DPRContextEncoderTokenizer""",
        """DPRQuestionEncoderTokenizer""",
        """DPRReaderOutput""",
        """DPRReaderTokenizer""",
    ],
}


try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  A							:				str			     =	[
      """DPRContextEncoderTokenizerFast""",
      """DPRQuestionEncoderTokenizerFast""",
      """DPRReaderTokenizerFast""",
  ]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  A							:				Tuple			     =	[
      """DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """DPRContextEncoder""",
      """DPRPretrainedContextEncoder""",
      """DPRPreTrainedModel""",
      """DPRPretrainedQuestionEncoder""",
      """DPRPretrainedReader""",
      """DPRQuestionEncoder""",
      """DPRReader""",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  A							:				Tuple			     =	[
      """TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """TFDPRContextEncoder""",
      """TFDPRPretrainedContextEncoder""",
      """TFDPRPretrainedQuestionEncoder""",
      """TFDPRPretrainedReader""",
      """TFDPRQuestionEncoder""",
      """TFDPRReader""",
  ]


if TYPE_CHECKING:
  from .configuration_dpr import DPR_PRETRAINED_CONFIG_ARCHIVE_MAP, DPRConfig
  from .tokenization_dpr import (
      DPRContextEncoderTokenizer,
      DPRQuestionEncoderTokenizer,
      DPRReaderOutput,
      DPRReaderTokenizer,
  )

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_dpr_fast import (
        DPRContextEncoderTokenizerFast,
        DPRQuestionEncoderTokenizerFast,
        DPRReaderTokenizerFast,
    )

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_dpr import (
        DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
        DPRContextEncoder,
        DPRPretrainedContextEncoder,
        DPRPreTrainedModel,
        DPRPretrainedQuestionEncoder,
        DPRPretrainedReader,
        DPRQuestionEncoder,
        DPRReader,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_dpr import (
        TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFDPRContextEncoder,
        TFDPRPretrainedContextEncoder,
        TFDPRPretrainedQuestionEncoder,
        TFDPRPretrainedReader,
        TFDPRQuestionEncoder,
        TFDPRReader,
    )

else:
  import sys

  A							:				int			     =	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




