
"""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


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


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





def    _a (						lowerCamelCase							):

 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    _a (						lowerCamelCase,     lowerCamelCase							):
 # 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		:     Any          =				pipeline_test_mapping[task]["""test"""]
 lowerCamelCase		:     Union[str, Any]          =				None

 if framework == "pt":
  lowerCamelCase		:     List[str]          =				getattr(SCREAMING_SNAKE_CASE_,     """model_mapping""",     SCREAMING_SNAKE_CASE_							)
 elif framework == "tf":
  lowerCamelCase		:     Dict          =				getattr(SCREAMING_SNAKE_CASE_,     """tf_model_mapping""",     SCREAMING_SNAKE_CASE_							)

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

 # cache the results
 lowerCamelCase		:     int          =				mapping
 return mapping





def    _a (						lowerCamelCase,     lowerCamelCase							):
 lowerCamelCase		:     int          =				get_framework(SCREAMING_SNAKE_CASE_							)
 if framework is None:
  return None
 lowerCamelCase		:     int          =				get_mapping_for_task(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_							)
 if mapping is None:
  return None

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

 return model_class





def    _a (						lowerCamelCase							):
 lowerCamelCase		:     Dict          =				[(task, get_model_for_pipeline_test(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_							)) for task in pipeline_test_mapping]
 lowerCamelCase		:     str          =				sorted([(task, model) for task, model in mapping if model is not None],     key=lambda lowerCamelCase							: x[0]							)

 return dict(SCREAMING_SNAKE_CASE_							)





def    _a (						lowerCamelCase							):
 lowerCamelCase		:     Optional[int]          =				get_framework(SCREAMING_SNAKE_CASE_							)
 if framework == "pt":
  lowerCamelCase		:     List[str]          =				"""torch"""
 lowerCamelCase		:     Optional[Any]          =				"""{}"""

 lowerCamelCase		:     Optional[int]          =				get_pipeline_model_mapping(SCREAMING_SNAKE_CASE_							)
 if len(SCREAMING_SNAKE_CASE_							) == 0:
  return ""

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

 return text





def    _a (						lowerCamelCase							):
 lowerCamelCase		:     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    _a (						lowerCamelCase							):
 lowerCamelCase		:     Optional[Any]          =				[x for x in get_test_classes(SCREAMING_SNAKE_CASE_							) if is_valid_test_class(SCREAMING_SNAKE_CASE_							)]

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

 return target_test_class





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

 return end_idx





def    _a (						lowerCamelCase,     lowerCamelCase=False							):
 if getattr(SCREAMING_SNAKE_CASE_,     """pipeline_model_mapping""",     SCREAMING_SNAKE_CASE_							) is not None:
  if not overwrite:
   return "", -1

 lowerCamelCase		:     List[str]          =				get_pipeline_model_mapping_string(SCREAMING_SNAKE_CASE_							)
 if len(SCREAMING_SNAKE_CASE_							) == 0:
  return "", -1
 lowerCamelCase		:     Union[str, Any]          =				line_to_add + """\n"""

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

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

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

 # Add indentation
 lowerCamelCase		:     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!)
 lowerCamelCase		:     Union[str, Any]          =				class_lines[: target_idx + 1] + [line_to_add] + class_lines[target_idx + 1 :]
 # Remove the lines that are marked to be removed
 lowerCamelCase		:     str          =				[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		:     List[str]          =				inspect.getsourcelines(inspect.getmodule(SCREAMING_SNAKE_CASE_							)							)[0]
 # Be careful with the 1-off between line numbers and array indices
 lowerCamelCase		:     str          =				module_lines[: class_start_line_no - 1] + class_lines + module_lines[class_end_line_no:]
 lowerCamelCase		:     Optional[int]          =				"""""".join(SCREAMING_SNAKE_CASE_							)

 lowerCamelCase		:     Union[str, 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    _a (						lowerCamelCase,     lowerCamelCase=False							):
 lowerCamelCase		:     Optional[Any]          =				find_test_class(SCREAMING_SNAKE_CASE_							)
 if test_class:
  add_pipeline_model_mapping(SCREAMING_SNAKE_CASE_,     overwrite=SCREAMING_SNAKE_CASE_							)


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

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






import argparse

import torch
from flax.training.checkpoints import restore_checkpoint

from transformers import FNetConfig, FNetForPreTraining
from transformers.utils import logging


logging.set_verbosity_info()





def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase							):
 # Initialise PyTorch model
 lowerCamelCase		:     List[str]          =				FNetConfig.from_json_file(lowerCamelCase							)
 print(F'''Building PyTorch model from configuration: {config}'''							)
 lowerCamelCase		:     Any          =				FNetForPreTraining(lowerCamelCase							)

 lowerCamelCase		:     List[Any]          =				restore_checkpoint(lowerCamelCase,     lowerCamelCase							)
 lowerCamelCase		:     List[str]          =				checkpoint_dict["""target"""]

 # Embeddings
 # Position IDs
 lowerCamelCase		:     Optional[Any]          =				fnet_pretraining_model.state_dict()

 lowerCamelCase		:     str          =				state_dict["""fnet.embeddings.position_ids"""]
 lowerCamelCase		:     List[str]          =				{"""fnet.embeddings.position_ids""": position_ids}
 # Embedding Layers
 lowerCamelCase		:     Tuple          =				torch.tensor(
     pretrained_model_params["""encoder"""]["""embedder"""]["""word"""]["""embedding"""]							)
 lowerCamelCase		:     List[Any]          =				torch.tensor(
     pretrained_model_params["""encoder"""]["""embedder"""]["""position"""]["""embedding"""][0]							)
 lowerCamelCase		:     Optional[Any]          =				torch.tensor(
     pretrained_model_params["""encoder"""]["""embedder"""]["""type"""]["""embedding"""]							)
 lowerCamelCase		:     Optional[int]          =				torch.tensor(
     pretrained_model_params["""encoder"""]["""embedder"""]["""hidden_mapping_in"""]["""kernel"""]							).T
 lowerCamelCase		:     int          =				torch.tensor(
     pretrained_model_params["""encoder"""]["""embedder"""]["""hidden_mapping_in"""]["""bias"""]							)
 lowerCamelCase		:     str          =				torch.tensor(
     pretrained_model_params["""encoder"""]["""embedder"""]["""layer_norm"""]["""scale"""]							)
 lowerCamelCase		:     Optional[int]          =				torch.tensor(
     pretrained_model_params["""encoder"""]["""embedder"""]["""layer_norm"""]["""bias"""]							)

 # Encoder Layers
 for layer in range(config.num_hidden_layers							):
  lowerCamelCase		:     List[Any]          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''encoder_{layer}''']["""mixing_layer_norm"""]["""scale"""]							)
  lowerCamelCase		:     Any          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''encoder_{layer}''']["""mixing_layer_norm"""]["""bias"""]							)

  lowerCamelCase		:     int          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''feed_forward_{layer}''']["""intermediate"""]["""kernel"""]							).T
  lowerCamelCase		:     List[str]          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''feed_forward_{layer}''']["""intermediate"""]["""bias"""]							)

  lowerCamelCase		:     List[Any]          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''feed_forward_{layer}''']["""output"""]["""kernel"""]							).T
  lowerCamelCase		:     str          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''feed_forward_{layer}''']["""output"""]["""bias"""]							)

  lowerCamelCase		:     Tuple          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''encoder_{layer}''']["""output_layer_norm"""]["""scale"""]							)
  lowerCamelCase		:     Optional[int]          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''encoder_{layer}''']["""output_layer_norm"""]["""bias"""]							)

 # Pooler Layers
 lowerCamelCase		:     Tuple          =				torch.tensor(pretrained_model_params["""encoder"""]["""pooler"""]["""kernel"""]							).T
 lowerCamelCase		:     List[Any]          =				torch.tensor(pretrained_model_params["""encoder"""]["""pooler"""]["""bias"""]							)

 # Masked LM Layers
 lowerCamelCase		:     Optional[Any]          =				torch.tensor(
     pretrained_model_params["""predictions_dense"""]["""kernel"""]							).T
 lowerCamelCase		:     Tuple          =				torch.tensor(
     pretrained_model_params["""predictions_dense"""]["""bias"""]							)
 lowerCamelCase		:     Tuple          =				torch.tensor(
     pretrained_model_params["""predictions_layer_norm"""]["""scale"""]							)
 lowerCamelCase		:     int          =				torch.tensor(
     pretrained_model_params["""predictions_layer_norm"""]["""bias"""]							)
 lowerCamelCase		:     Dict          =				torch.tensor(
     pretrained_model_params["""encoder"""]["""embedder"""]["""word"""]["""embedding"""]							)
 lowerCamelCase		:     Dict          =				torch.tensor(
     pretrained_model_params["""predictions_output"""]["""output_bias"""]							)
 lowerCamelCase		:     Any          =				torch.tensor(pretrained_model_params["""predictions_output"""]["""output_bias"""]							)

 # Seq Relationship Layers
 lowerCamelCase		:     int          =				torch.tensor(
     pretrained_model_params["""classification"""]["""output_kernel"""]							)
 lowerCamelCase		:     Optional[Any]          =				torch.tensor(
     pretrained_model_params["""classification"""]["""output_bias"""]							)

 # Load State Dict
 fnet_pretraining_model.load_state_dict(lowerCamelCase							)

 # Save PreTrained
 print(F'''Saving pretrained model to {save_path}'''							)
 fnet_pretraining_model.save_pretrained(lowerCamelCase							)


if __name__ == "__main__":
      _lowerCamelCase        =argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          """--flax_checkpoint_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
      )
      parser.add_argument(
          """--fnet_config_file""",
          default=None,
          type=str,
          required=True,
          help=(
              """The config json file corresponding to the pre-trained FNet model. \n"""
              """This specifies the model architecture."""
          ),
      )
      parser.add_argument("""--save_path""", default=None, type=str, required=True, help="""Path to the output model.""")
      _lowerCamelCase        =parser.parse_args()
      convert_flax_checkpoint_to_pytorch(args.flax_checkpoint_path, args.fnet_config_file, args.save_path)




