
'''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_           =    {}
for task, _ in pipeline_test_mapping.items():
  a_           =    {"pt": None, "tf": None}


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

def     UpperCamelCase_  (						__SCREAMING_SNAKE_CASE							):
    """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     UpperCamelCase_  (						__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							):
    """simple docstring"""






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

    snake_case_			:     List[str]        =    pipeline_test_mapping[task]["test"]
    snake_case_			:     Dict        =    None

    if framework == "pt":
        snake_case_			:     Union[str, Any]        =    getattr(__SCREAMING_SNAKE_CASE,		"model_mapping",		__SCREAMING_SNAKE_CASE							)
    elif framework == "tf":
        snake_case_			:     str        =    getattr(__SCREAMING_SNAKE_CASE,		"tf_model_mapping",		__SCREAMING_SNAKE_CASE							)

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

    # cache the results
    snake_case_			:     Union[str, Any]        =    mapping
    return mapping

def     UpperCamelCase_  (						__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							):
    """simple docstring"""






    snake_case_			:     int        =    get_framework(__SCREAMING_SNAKE_CASE							)
    if framework is None:
        return None
    snake_case_			:     List[Any]        =    get_mapping_for_task(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							)
    if mapping is None:
        return None

    snake_case_			:     List[Any]        =    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.
    snake_case_			:     List[Any]        =    mapping.get(config_classes[0],		__SCREAMING_SNAKE_CASE							)
    if isinstance(__SCREAMING_SNAKE_CASE,		(tuple, list)							):
        snake_case_			:     Dict        =    sorted(__SCREAMING_SNAKE_CASE,		key=lambda __SCREAMING_SNAKE_CASE							: x.__name__							)

    return model_class

def     UpperCamelCase_  (						__SCREAMING_SNAKE_CASE							):
    """simple docstring"""






    snake_case_			:     Union[str, Any]        =    [(task, get_model_for_pipeline_test(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							)) for task in pipeline_test_mapping]
    snake_case_			:     str        =    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							):
    """simple docstring"""






    snake_case_			:     str        =    get_framework(__SCREAMING_SNAKE_CASE							)
    if framework == "pt":
        snake_case_			:     int        =    "torch"
    snake_case_			:     Optional[int]        =    "{}"

    snake_case_			:     str        =    get_pipeline_model_mapping(__SCREAMING_SNAKE_CASE							)
    if len(__SCREAMING_SNAKE_CASE							) == 0:
        return ""

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

    return text

def     UpperCamelCase_  (						__SCREAMING_SNAKE_CASE							):
    """simple docstring"""






    snake_case_			:     Optional[int]        =    {"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							):
    """simple docstring"""






    snake_case_			:     Union[str, Any]        =    [x for x in get_test_classes(__SCREAMING_SNAKE_CASE							) if is_valid_test_class(__SCREAMING_SNAKE_CASE							)]

    snake_case_			:     Tuple        =    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:
            snake_case_			:     Tuple        =    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:
        snake_case_			:     List[Any]        =    sorted(__SCREAMING_SNAKE_CASE,		key=lambda __SCREAMING_SNAKE_CASE							: (len(x.__name__							), x.__name__)							)[0]

    return target_test_class

def     UpperCamelCase_  (						__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							):
    """simple docstring"""






    snake_case_			:     Optional[Any]        =    start_idx
    for idx, line in enumerate(lines[start_idx:]							):
        snake_case_			:     List[Any]        =    len(__SCREAMING_SNAKE_CASE							) - len(line.lstrip()							)
        if idx == 0 or indent > indent_level or (indent == indent_level and line.strip() == ")"):
            snake_case_			:     List[Any]        =    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,		__SCREAMING_SNAKE_CASE=False							):
    """simple docstring"""






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

    snake_case_			:     Optional[Any]        =    get_pipeline_model_mapping_string(__SCREAMING_SNAKE_CASE							)
    if len(__SCREAMING_SNAKE_CASE							) == 0:
        return "", -1
    snake_case_			:     Tuple        =    line_to_add + "\n"

    # The code defined the class `test_class`
    snake_case_      ,			snake_case_			:     List[str]        =    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 "							):
            snake_case_			:     Tuple        =    class_lines[idx:]
            class_start_line_no += idx
            break
    snake_case_			:     Optional[Any]        =    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`.
    snake_case_			:     Optional[Any]        =    None
    # The indent level of the line at `class_lines[start_idx]` (if defined)
    snake_case_			:     Optional[Any]        =    0
    # To record if `pipeline_model_mapping` is found in `test_class`.
    snake_case_			:     Any        =    None
    for idx, line in enumerate(__SCREAMING_SNAKE_CASE							):
        if line.strip().startswith("all_model_classes = "							):
            snake_case_			:     Optional[int]        =    len(__SCREAMING_SNAKE_CASE							) - len(line.lstrip()							)
            snake_case_			:     int        =    idx
        elif line.strip().startswith("all_generative_model_classes = "							):
            snake_case_			:     Any        =    len(__SCREAMING_SNAKE_CASE							) - len(line.lstrip()							)
            snake_case_			:     Tuple        =    idx
        elif line.strip().startswith("pipeline_model_mapping = "							):
            snake_case_			:     Dict        =    len(__SCREAMING_SNAKE_CASE							) - len(line.lstrip()							)
            snake_case_			:     List[str]        =    idx
            snake_case_			:     int        =    line
            break

    if start_idx is None:
        return "", -1
    # Find the ending index (inclusive) of the above found block.
    snake_case_			:     Optional[int]        =    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
    snake_case_			:     List[Any]        =    re.compile(r"\s(is_\S+?_available\(\))\s"							)
    for line in class_lines[start_idx : end_idx + 1]:
        snake_case_			:     str        =    r.search(__SCREAMING_SNAKE_CASE							)
        if backend_condition is not None:
            # replace the leading and trailing whitespaces to the space character " ".
            snake_case_			:     List[Any]        =    " " + backend_condition[0][1:-1] + " "
            snake_case_			:     Union[str, Any]        =    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`.
        snake_case_			:     List[Any]        =    end_idx
    else:
        # `pipeline_model_mapping` is defined. The target index is set to be one **BEFORE** its start index.
        snake_case_			:     str        =    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.
            snake_case_			:     int        =    None  # noqa

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

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

    snake_case_			:     Optional[Any]        =    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,		__SCREAMING_SNAKE_CASE=False							):
    """simple docstring"""






    snake_case_			:     Dict        =    find_test_class(__SCREAMING_SNAKE_CASE							)
    if test_class:
        add_pipeline_model_mapping(__SCREAMING_SNAKE_CASE,		overwrite=__SCREAMING_SNAKE_CASE							)


if __name__ == "__main__":
  a_           =    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_           =    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_           =    []
  if args.test_file:
    a_           =    [args.test_file]
  else:
    a_           =    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'''
import argparse
import json

from transformers import GPTNeoConfig, GPTNeoForCausalLM, load_tf_weights_in_gpt_neo
from transformers.utils import logging


logging.set_verbosity_info()

def     UpperCamelCase_  (						__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							):
    """simple docstring"""






    snake_case_			:     Optional[Any]        =    json.load(open(__SCREAMING_SNAKE_CASE,		"r"							)							)
    snake_case_			:     Optional[int]        =    GPTNeoConfig(
        hidden_size=config_json["n_embd"],		num_layers=config_json["n_layer"],		num_heads=config_json["n_head"],		attention_types=config_json["attention_types"],		max_position_embeddings=config_json["n_positions"],		resid_dropout=config_json["res_dropout"],		embed_dropout=config_json["embed_dropout"],		attention_dropout=config_json["attn_dropout"],		)
    print(f'Building PyTorch model from configuration: {config}'							)
    snake_case_			:     Optional[Any]        =    GPTNeoForCausalLM(__SCREAMING_SNAKE_CASE							)

    # Load weights from tf checkpoint
    load_tf_weights_in_gpt_neo(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							)

    # Save pytorch-model
    print(f'Save PyTorch model to {pytorch_dump_path}'							)
    model.save_pretrained(__SCREAMING_SNAKE_CASE							)


if __name__ == "__main__":
  a_           =    argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
  )
  parser.add_argument(
      "--config_file",
      default=None,
      type=str,
      required=True,
      help=(
          "The config json file corresponding to the pre-trained mesh-tf model. \n"
          "This specifies the model architecture."
      ),
  )
  parser.add_argument(
      "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
  )
  a_           =    parser.parse_args()
  convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path)




