
'''simple docstring'''





import tempfile
import unittest

import transformers
from transformers import XGLMConfig, XGLMTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
 import jax
 import jax.numpy as jnp
 import numpy as np

 from transformers.modeling_flax_pytorch_utils import (
     convert_pytorch_state_dict_to_flax,
     load_flax_weights_in_pytorch_model,
 )
 from transformers.models.xglm.modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel


if is_torch_available():
 import torch


@require_flax
class 			UpperCamelCase__	:






 def __init__(      self	:List[str]						,					_A	:int						,					_A	:Any=14						,					_A	:Union[str, Any]=7						,					_A	:Optional[int]=True						,					_A	:List[str]=True						,					_A	:int=True						,					_A	:Any=99						,					_A	:Optional[int]=32						,					_A	:List[str]=5						,					_A	:Optional[Any]=4						,					_A	:Union[str, Any]=37						,					_A	:List[Any]="gelu"						,					_A	:int=0.1						,					_A	:Union[str, Any]=0.1						,					_A	:str=512						,					_A	:Optional[Any]=0.02						,					_A	:Tuple=None						,					)					->      Union[str, Any]:







       '''simple docstring'''



       __A													=      parent
       __A													=      batch_size
       __A													=      seq_length
       __A													=      is_training
       __A													=      use_input_mask
       __A													=      use_labels
       __A													=      vocab_size
       __A													=      d_model
       __A													=      num_hidden_layers
       __A													=      num_attention_heads
       __A													=      ffn_dim
       __A													=      activation_function
       __A													=      activation_dropout
       __A													=      attention_dropout
       __A													=      max_position_embeddings
       __A													=      initializer_range
       __A													=      None
       __A													=      0
       __A													=      2
       __A													=      1






 def 						lowercase_    (      self	:Union[str, Any]      )					->      Any:







       '''simple docstring'''



       __A													=      np.clip(ids_tensor([self.batch_size, self.seq_length]						,					self.vocab_size      )						,					3						,					self.vocab_size      )

       __A													=      None
       if self.use_input_mask:
             __A													=      random_attention_mask([self.batch_size, self.seq_length]      )

       __A													=      XGLMConfig(
           vocab_size=self.vocab_size						,					d_model=self.hidden_size						,					num_layers=self.num_hidden_layers						,					attention_heads=self.num_attention_heads						,					ffn_dim=self.ffn_dim						,					activation_function=self.activation_function						,					activation_dropout=self.activation_dropout						,					attention_dropout=self.attention_dropout						,					max_position_embeddings=self.max_position_embeddings						,					initializer_range=self.initializer_range						,					use_cache=_A						,					bos_token_id=self.bos_token_id						,					eos_token_id=self.eos_token_id						,					pad_token_id=self.pad_token_id						,					)

       return (config, input_ids, input_mask)






 def 						lowercase_    (      self	:Optional[Any]      )					->      str:







       '''simple docstring'''



       __A													=      self.prepare_config_and_inputs()
       __A						,   __A						,   __A													=      config_and_inputs
       __A													=      {'input_ids': input_ids, 'attention_mask': attention_mask}
       return config, inputs_dict






 def 						lowercase_    (      self	:str      )					->      int:







       '''simple docstring'''



       __A						,   __A						,   __A													=      self.prepare_config_and_inputs()

       __A													=      floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
       __A													=      ids_tensor([self.batch_size, self.seq_length]						,					vocab_size=2      )

       return (
           config,
           input_ids,
           attention_mask,
           encoder_hidden_states,
           encoder_attention_mask,
       )






 def 						lowercase_    (      self	:Union[str, Any]						,					_A	:Optional[int]						,					_A	:List[Any]						,					_A	:str						,					_A	:Optional[Any]      )					->      List[Any]:







       '''simple docstring'''



       __A													=      20
       __A													=      model_class_name(_A      )

       __A													=      model.init_cache(input_ids.shape[0]						,					_A      )
       __A													=      jnp.ones((input_ids.shape[0], max_decoder_length)						,					dtype='i4'      )

       __A													=      jnp.broadcast_to(
           jnp.arange(input_ids.shape[-1] - 1      )[None, :]						,					(input_ids.shape[0], input_ids.shape[-1] - 1)      )
       __A													=      model(
           input_ids[:, :-1]						,					attention_mask=_A						,					past_key_values=_A						,					position_ids=_A						,					)

       __A													=      jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,					dtype='i4'      )
       __A													=      model(
           input_ids[:, -1:]						,					attention_mask=_A						,					past_key_values=outputs_cache.past_key_values						,					position_ids=_A						,					)

       __A													=      model(_A      )

       __A													=      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
       self.parent.assertTrue(diff < 1E-3						,					msg=F'Max diff is {diff}'      )






 def 						lowercase_    (      self	:int						,					_A	:Union[str, Any]						,					_A	:int						,					_A	:Optional[int]						,					_A	:Tuple      )					->      List[str]:







       '''simple docstring'''



       __A													=      20
       __A													=      model_class_name(_A      )

       __A													=      jnp.concatenate(
           [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])      )]						,					axis=-1						,					)

       __A													=      model.init_cache(input_ids.shape[0]						,					_A      )
       __A													=      jnp.broadcast_to(
           jnp.arange(input_ids.shape[-1] - 1      )[None, :]						,					(input_ids.shape[0], input_ids.shape[-1] - 1)      )

       __A													=      model(
           input_ids[:, :-1]						,					attention_mask=_A						,					past_key_values=_A						,					position_ids=_A						,					)
       __A													=      jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,					dtype='i4'      )
       __A													=      model(
           input_ids[:, -1:]						,					past_key_values=outputs_cache.past_key_values						,					attention_mask=_A						,					position_ids=_A						,					)

       __A													=      model(_A						,					attention_mask=_A      )
       __A													=      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
       self.parent.assertTrue(diff < 1E-3						,					msg=F'Max diff is {diff}'      )


@require_sentencepiece
@require_flax
class 			UpperCamelCase__	(     SCREAMING_SNAKE_CASE   ,  SCREAMING_SNAKE_CASE   ,  unittest.TestCase):
 UpperCAmelCase__     :     Any    					=       (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
 UpperCAmelCase__     :     Optional[Any]    					=       (FlaxXGLMForCausalLM,) if is_flax_available() else ()






 def 						lowercase_    (      self	:Dict      )					->      Any:







       '''simple docstring'''



       __A													=      FlaxXGLMModelTester(self      )






 def 						lowercase_    (      self	:Dict      )					->      List[str]:







       '''simple docstring'''



       for model_class_name in self.all_model_classes:
             __A						,   __A						,   __A													=      self.model_tester.prepare_config_and_inputs()
             self.model_tester.check_use_cache_forward(_A						,					_A						,					_A						,					_A      )






 def 						lowercase_    (      self	:str      )					->      int:







       '''simple docstring'''



       for model_class_name in self.all_model_classes:
             __A						,   __A						,   __A													=      self.model_tester.prepare_config_and_inputs()
             self.model_tester.check_use_cache_forward_with_attn_mask(
                 _A						,					_A						,					_A						,					_A      )






 @slow
 def 						lowercase_    (      self	:Tuple      )					->      List[str]:







       '''simple docstring'''



       __A													=      XGLMTokenizer.from_pretrained('XGLM'						,					padding_side='left'      )
       __A													=      tokenizer(['Hello this is a long string', 'Hey']						,					return_tensors='np'						,					padding=_A						,					truncation=_A      )

       __A													=      FlaxXGLMForCausalLM.from_pretrained('facebook/xglm-564M'      )
       __A													=      1
       __A													=      False

       __A													=      jax.jit(model.generate      )

       __A													=      jit_generate(inputs['input_ids']						,					attention_mask=inputs['attention_mask']      ).sequences

       __A													=      tokenizer.batch_decode(_A						,					skip_special_tokens=_A      )

       __A													=      [
           'Hello this is a long string of questions, but I\'m not sure if I\'m',
           'Hey, I\'m a newbie to the forum and I\'',
       ]

       self.assertListEqual(_A						,					_A      )






 @is_pt_flax_cross_test
 def 						lowercase_    (      self	:int      )					->      Optional[Any]:







       '''simple docstring'''



       __A						,   __A													=      self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             with self.subTest(model_class.__name__      ):
                   # prepare inputs
                   __A													=      self._prepare_for_class(_A						,					_A      )
                   __A													=      {k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

                   # load corresponding PyTorch class
                   __A													=      model_class.__name__[4:]  # Skip the "Flax" at the beginning
                   __A													=      getattr(_A						,					_A      )

                   __A						,   __A													=      pt_inputs['input_ids'].shape
                   __A													=      np.random.randint(0						,					seq_length - 1						,					size=(batch_size,)      )
                   for batch_idx, start_index in enumerate(_A      ):
                         __A													=      0
                         __A													=      1
                         __A													=      0
                         __A													=      1
                   __A													=      pt_model_class(_A      ).eval()
                   # Flax models don't use the `use_cache` option and cache is not returned as a default.
                   # So we disable `use_cache` here for PyTorch model.
                   __A													=      False
                   __A													=      model_class(_A						,					dtype=jnp.floataa      )
                   __A													=      convert_pytorch_state_dict_to_flax(pt_model.state_dict()						,					_A      )
                   __A													=      fx_state

                   with torch.no_grad():
                         __A													=      pt_model(**_A      ).to_tuple()

                   __A													=      fx_model(**_A      ).to_tuple()
                   self.assertEqual(len(_A      )						,					len(_A      )						,					'Output lengths differ between Flax and PyTorch'      )
                   for fx_output, pt_output in zip(_A						,					_A      ):
                         self.assert_almost_equals(fx_output[:, -1]						,					pt_output[:, -1].numpy()						,					4E-2      )

                   with tempfile.TemporaryDirectory() as tmpdirname:
                         pt_model.save_pretrained(_A      )
                         __A													=      model_class.from_pretrained(_A						,					from_pt=_A      )

                   __A													=      fx_model_loaded(**_A      ).to_tuple()
                   self.assertEqual(
                       len(_A      )						,					len(_A      )						,					'Output lengths differ between Flax and PyTorch'      )
                   for fx_output_loaded, pt_output in zip(_A						,					_A      ):
                         self.assert_almost_equals(fx_output_loaded[:, -1]						,					pt_output[:, -1].numpy()						,					4E-2      )






 @is_pt_flax_cross_test
 def 						lowercase_    (      self	:List[str]      )					->      Optional[Any]:







       '''simple docstring'''



       __A						,   __A													=      self.model_tester.prepare_config_and_inputs_for_common()
       for model_class in self.all_model_classes:
             with self.subTest(model_class.__name__      ):
                   # prepare inputs
                   __A													=      self._prepare_for_class(_A						,					_A      )
                   __A													=      {k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

                   # load corresponding PyTorch class
                   __A													=      model_class.__name__[4:]  # Skip the "Flax" at the beginning
                   __A													=      getattr(_A						,					_A      )

                   __A													=      pt_model_class(_A      ).eval()
                   __A													=      False
                   __A													=      model_class(_A						,					dtype=jnp.floataa      )

                   __A													=      load_flax_weights_in_pytorch_model(_A						,					fx_model.params      )
                   __A						,   __A													=      pt_inputs['input_ids'].shape
                   __A													=      np.random.randint(0						,					seq_length - 1						,					size=(batch_size,)      )
                   for batch_idx, start_index in enumerate(_A      ):
                         __A													=      0
                         __A													=      1
                         __A													=      0
                         __A													=      1

                   # make sure weights are tied in PyTorch
                   pt_model.tie_weights()

                   with torch.no_grad():
                         __A													=      pt_model(**_A      ).to_tuple()

                   __A													=      fx_model(**_A      ).to_tuple()
                   self.assertEqual(len(_A      )						,					len(_A      )						,					'Output lengths differ between Flax and PyTorch'      )
                   for fx_output, pt_output in zip(_A						,					_A      ):
                         self.assert_almost_equals(fx_output[:, -1]						,					pt_output[:, -1].numpy()						,					4E-2      )

                   with tempfile.TemporaryDirectory() as tmpdirname:
                         fx_model.save_pretrained(_A      )
                         __A													=      pt_model_class.from_pretrained(_A						,					from_flax=_A      )

                   with torch.no_grad():
                         __A													=      pt_model_loaded(**_A      ).to_tuple()

                   self.assertEqual(
                       len(_A      )						,					len(_A      )						,					'Output lengths differ between Flax and PyTorch'      )
                   for fx_output, pt_output in zip(_A						,					_A      ):
                         self.assert_almost_equals(fx_output[:, -1]						,					pt_output[:, -1].numpy()						,					4E-2      )






 @slow
 def 						lowercase_    (      self	:Union[str, Any]      )					->      Union[str, Any]:







       '''simple docstring'''



       for model_class_name in self.all_model_classes:
             __A													=      model_class_name.from_pretrained('facebook/xglm-564M'      )
             __A													=      model(np.ones((1, 1)      )      )
             self.assertIsNotNone(_A      )












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    : int		= {}
for task, _ in pipeline_test_mapping.items():
  _lowerCAmelCase    : List[str]		= {"pt": None, "tf": None}


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



def     lowerCAmelCase		(		_lowerCAmelCase				: Union[str, 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     lowerCAmelCase		(		_lowerCAmelCase				: str    ,				_lowerCAmelCase				: Optional[Any]						):




 """simple docstring"""







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

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

 if framework == "pt":
  UpperCAmelCase__							=    getattr(_lowerCAmelCase    ,				"model_mapping"    ,				_lowerCAmelCase						)
 elif framework == "tf":
  UpperCAmelCase__							=    getattr(_lowerCAmelCase    ,				"tf_model_mapping"    ,				_lowerCAmelCase						)

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

 # cache the results
 UpperCAmelCase__							=    mapping
 return mapping



def     lowerCAmelCase		(		_lowerCAmelCase				: Tuple    ,				_lowerCAmelCase				: List[str]						):




 """simple docstring"""







 UpperCAmelCase__							=    get_framework(_lowerCAmelCase						)
 if framework is None:
  return None
 UpperCAmelCase__							=    get_mapping_for_task(_lowerCAmelCase    ,				_lowerCAmelCase						)
 if mapping is None:
  return None

 UpperCAmelCase__							=    list({model_class.config_class for model_class in test_class.all_model_classes}						)
 if len(_lowerCAmelCase						) != 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.
 UpperCAmelCase__							=    mapping.get(config_classes[0]    ,				_lowerCAmelCase						)
 if isinstance(_lowerCAmelCase    ,				(tuple, list)						):
  UpperCAmelCase__							=    sorted(_lowerCAmelCase    ,				key=lambda _lowerCAmelCase						: x.__name__						)

 return model_class



def     lowerCAmelCase		(		_lowerCAmelCase				: str						):




 """simple docstring"""







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

 return dict(_lowerCAmelCase						)



def     lowerCAmelCase		(		_lowerCAmelCase				: Optional[Any]						):




 """simple docstring"""







 UpperCAmelCase__							=    get_framework(_lowerCAmelCase						)
 if framework == "pt":
  UpperCAmelCase__							=    "torch"
 UpperCAmelCase__							=    "{}"

 UpperCAmelCase__							=    get_pipeline_model_mapping(_lowerCAmelCase						)
 if len(_lowerCAmelCase						) == 0:
  return ""

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

 return text



def     lowerCAmelCase		(		_lowerCAmelCase				: Any						):




 """simple docstring"""







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



def     lowerCAmelCase		(		_lowerCAmelCase				: Union[str, Any]						):




 """simple docstring"""







 UpperCAmelCase__							=    [x for x in get_test_classes(_lowerCAmelCase						) if is_valid_test_class(_lowerCAmelCase						)]

 UpperCAmelCase__							=    None
 for test_class in test_classes:
  # If a test class has defined `pipeline_model_mapping`, let's take it
  if getattr(_lowerCAmelCase    ,				"pipeline_model_mapping"    ,				_lowerCAmelCase						) is not None:
   UpperCAmelCase__							=    test_class
   break
    # Take the test class with the shortest name (just a heuristic)
 if target_test_class is None and len(_lowerCAmelCase						) > 0:
  UpperCAmelCase__							=    sorted(_lowerCAmelCase    ,				key=lambda _lowerCAmelCase						: (len(x.__name__						), x.__name__)						)[0]

 return target_test_class



def     lowerCAmelCase		(		_lowerCAmelCase				: int    ,				_lowerCAmelCase				: Optional[int]    ,				_lowerCAmelCase				: str						):




 """simple docstring"""







 UpperCAmelCase__							=    start_idx
 for idx, line in enumerate(lines[start_idx:]						):
  UpperCAmelCase__							=    len(_lowerCAmelCase						) - len(line.lstrip()						)
  if idx == 0 or indent > indent_level or (indent == indent_level and line.strip() == ")"):
   UpperCAmelCase__							=    start_idx + idx
  elif idx > 0 and indent <= indent_level:
   # Outside the definition block of `pipeline_model_mapping`
   break

 return end_idx



def     lowerCAmelCase		(		_lowerCAmelCase				: int    ,				_lowerCAmelCase				: str=False						):




 """simple docstring"""







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

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

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

 if start_idx is None:
  return "", -1
 # Find the ending index (inclusive) of the above found block.
 UpperCAmelCase__							=    find_block_ending(_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase						)

 # 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
 UpperCAmelCase__							=    re.compile(R"\s(is_\S+?_available\(\))\s"						)
 for line in class_lines[start_idx : end_idx + 1]:
  UpperCAmelCase__							=    r.search(_lowerCAmelCase						)
  if backend_condition is not None:
   # replace the leading and trailing whitespaces to the space character " ".
   UpperCAmelCase__							=    " " + backend_condition[0][1:-1] + " "
   UpperCAmelCase__							=    r.sub(_lowerCAmelCase    ,				_lowerCAmelCase						)
   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`.
  UpperCAmelCase__							=    end_idx
 else:
  # `pipeline_model_mapping` is defined. The target index is set to be one **BEFORE** its start index.
  UpperCAmelCase__							=    start_idx - 1
  # mark the lines of the currently existing `pipeline_model_mapping` to be removed.
  for idx in range(_lowerCAmelCase    ,				end_idx + 1						):
   # These lines are going to be removed before writing to the test file.
   UpperCAmelCase__							=    None  # noqa

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

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

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

 return line_to_add



def     lowerCAmelCase		(		_lowerCAmelCase				: Tuple    ,				_lowerCAmelCase				: List[Any]=False						):




 """simple docstring"""







 UpperCAmelCase__							=    find_test_class(_lowerCAmelCase						)
 if test_class:
  add_pipeline_model_mapping(_lowerCAmelCase    ,				overwrite=_lowerCAmelCase						)


if __name__ == "__main__":
  _lowerCAmelCase    : Tuple		= 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    : Optional[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.")

  _lowerCAmelCase    : List[str]		= []
  if args.test_file:
    _lowerCAmelCase    : Any		= [args.test_file]
  else:
    _lowerCAmelCase    : Any		= 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)





