




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


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


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




def     snake_case__     (			lowerCamelCase__					: List[Any]      ) -> Any:

					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__     (			lowerCamelCase__					: Dict      ,		lowerCamelCase__					: Union[str, Any]      ) -> Any:
					# Use the cached results
					if PIPELINE_TEST_MAPPING[task].get(_UpperCamelCase      ,		_UpperCamelCase      ) is not None:
										return PIPELINE_TEST_MAPPING[task][framework]

					A_			:  Union[str, Any]			=				pipeline_test_mapping[task]['''test''']
					A_			:  Union[str, Any]			=				None

					if framework == "pt":
										A_			:  Optional[int]			=				getattr(_UpperCamelCase      ,		'''model_mapping'''      ,		_UpperCamelCase      )
					elif framework == "tf":
										A_			:  Optional[int]			=				getattr(_UpperCamelCase      ,		'''tf_model_mapping'''      ,		_UpperCamelCase      )

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

					# cache the results
					A_			:  Dict			=				mapping
					return mapping




def     snake_case__     (			lowerCamelCase__					: List[str]      ,		lowerCamelCase__					: Optional[int]      ) -> Any:
					A_			:  Union[str, Any]			=				get_framework(_UpperCamelCase      )
					if framework is None:
										return None
					A_			:  Optional[int]			=				get_mapping_for_task(_UpperCamelCase      ,		_UpperCamelCase      )
					if mapping is None:
										return None

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

					return model_class




def     snake_case__     (			lowerCamelCase__					: Any      ) -> str:
					A_			:  Optional[int]			=				[(task, get_model_for_pipeline_test(_UpperCamelCase      ,		_UpperCamelCase      )) for task in pipeline_test_mapping]
					A_			:  List[Any]			=				sorted([(task, model) for task, model in mapping if model is not None]      ,		key=lambda lowerCamelCase__      : x[0]      )

					return dict(_UpperCamelCase      )




def     snake_case__     (			lowerCamelCase__					: List[Any]      ) -> Optional[Any]:
					A_			:  List[str]			=				get_framework(_UpperCamelCase      )
					if framework == "pt":
										A_			:  Dict			=				'''torch'''
					A_			:  int			=				'''{}'''

					A_			:  Union[str, Any]			=				get_pipeline_model_mapping(_UpperCamelCase      )
					if len(_UpperCamelCase      ) == 0:
										return ""

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

					return text




def     snake_case__     (			lowerCamelCase__					: Optional[int]      ) -> List[Any]:
					A_			:  Tuple			=				{'''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__     (			lowerCamelCase__					: List[Any]      ) -> Dict:
					A_			:  Optional[int]			=				[x for x in get_test_classes(_UpperCamelCase      ) if is_valid_test_class(_UpperCamelCase      )]

					A_			:  str			=				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:
															A_			:  Optional[Any]			=				test_class
															break
    # Take the test class with the shortest name (just a heuristic)
					if target_test_class is None and len(_UpperCamelCase      ) > 0:
										A_			:  Any			=				sorted(_UpperCamelCase      ,		key=lambda lowerCamelCase__      : (len(x.__name__      ), x.__name__)      )[0]

					return target_test_class




def     snake_case__     (			lowerCamelCase__					: Any      ,		lowerCamelCase__					: Any      ,		lowerCamelCase__					: Dict      ) -> Optional[int]:
					A_			:  List[str]			=				start_idx
					for idx, line in enumerate(lines[start_idx:]      ):
										A_			:  List[Any]			=				len(_UpperCamelCase      ) - len(line.lstrip()      )
										if idx == 0 or indent > indent_level or (indent == indent_level and line.strip() == ")"):
															A_			:  Union[str, Any]			=				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__     (			lowerCamelCase__					: Any      ,		lowerCamelCase__					: Tuple=False      ) -> Optional[Any]:
					if getattr(_UpperCamelCase      ,		'''pipeline_model_mapping'''      ,		_UpperCamelCase      ) is not None:
										if not overwrite:
															return "", -1

					A_			:  List[str]			=				get_pipeline_model_mapping_string(_UpperCamelCase      )
					if len(_UpperCamelCase      ) == 0:
										return "", -1
					A_			:  Tuple			=				line_to_add + '''\n'''

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

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

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

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

					A_			:  int			=				inspect.getsourcefile(_UpperCamelCase      )
					with open(_UpperCamelCase      ,		'''w'''      ,		encoding='''UTF-8'''      ,		newline='''\n'''      ) as fp:
										fp.write(_UpperCamelCase      )

					return line_to_add




def     snake_case__     (			lowerCamelCase__					: Tuple      ,		lowerCamelCase__					: Union[str, Any]=False      ) -> Optional[Any]:
					A_			:  Optional[Any]			=				find_test_class(_UpperCamelCase      )
					if test_class:
										add_pipeline_model_mapping(_UpperCamelCase      ,		overwrite=_UpperCamelCase      )


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

							snake_case__ 							=    []
							if args.test_file:
														snake_case__ 							=    [args.test_file]
							else:
														snake_case__ 							=    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 logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


__snake_case  : Optional[int]		       =    logging.getLogger(__name__)




class       __UpperCAmelCase							:





       '''simple docstring'''



       def __init__( self					)						->	List[str]:
           A_      =  False

       def    __A			( self       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)						->	str:
           if not self.initialized:
               A_      =  RagRetriever(
                   _SCREAMING_SNAKE_CASE       ,	question_encoder_tokenizer=_SCREAMING_SNAKE_CASE       ,	generator_tokenizer=_SCREAMING_SNAKE_CASE       ,	index=_SCREAMING_SNAKE_CASE       ,	init_retrieval=_SCREAMING_SNAKE_CASE       ,	)
               A_      =  True

       def    __A			( self					)						->	List[Any]:
           self.retriever.index.init_index()

       def    __A			( self					)						->	Dict:
           # delete the old self.retriever object before assigning the new index
           del self.retriever
           A_      =  False

       def    __A			( self       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)						->	Optional[int]:
           A_  ,A_      =  self.retriever._main_retrieve(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)
           A_      =  self.retriever.index.get_doc_dicts(_SCREAMING_SNAKE_CASE					)
           return doc_ids, retrieved_doc_embeds, doc_dicts




class       __UpperCAmelCase							(							_UpperCamelCase	):





       '''simple docstring'''



       def __init__( self       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE=None					)						->	List[Any]:
           if index is not None and index.is_initialized() and len(_SCREAMING_SNAKE_CASE					) > 0:
               raise ValueError(
                   '''When using Ray for distributed fine-tuning, '''
                   '''you\'ll need to provide the paths instead, '''
                   '''as the dataset and the index are loaded '''
                   '''separately. More info in examples/rag/use_own_knowledge_dataset.py '''					)

           super().__init__(
               _SCREAMING_SNAKE_CASE       ,	question_encoder_tokenizer=_SCREAMING_SNAKE_CASE       ,	generator_tokenizer=_SCREAMING_SNAKE_CASE       ,	index=_SCREAMING_SNAKE_CASE       ,	init_retrieval=_SCREAMING_SNAKE_CASE       ,	)

           A_      =  retrieval_workers
           A_      =  question_encoder_tokenizer
           A_      =  generator_tokenizer
           if len(self.retrieval_workers					) > 0:
               ray.get(
                   [
                       worker.create_rag_retriever.remote(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)
                       for worker in self.retrieval_workers
                   ]					)

       def    __A			( self					)						->	List[Any]:
           logger.info('''initializing retrieval'''					)

           if len(self.retrieval_workers					) > 0:
               ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers]					)
           else:
               # Non-distributed training. Load index into this same process.
               self.index.init_index()

       def    __A			( self       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)						->	Optional[int]:
           if len(self.retrieval_workers					) > 0:
               # Select a random retrieval actor.
               A_      =  self.retrieval_workers[random.randint(0       ,	len(self.retrieval_workers					) - 1					)]
               A_  ,A_  ,A_      =  ray.get(
                   random_worker.retrieve.remote(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)					)
           else:
               A_  ,A_      =  self._main_retrieve(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)
               A_      =  self.index.get_doc_dicts(_SCREAMING_SNAKE_CASE					)
           return retrieved_doc_embeds, doc_ids, doc_dicts

       @classmethod
       def    __A			( cls       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE=None       ,	**_SCREAMING_SNAKE_CASE					)						->	List[Any]:
           return super(_SCREAMING_SNAKE_CASE       ,	cls					).get_tokenizers(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	**_SCREAMING_SNAKE_CASE					)

       @classmethod
       def    __A			( cls       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE=None       ,	**_SCREAMING_SNAKE_CASE					)						->	int:
           A_      =  kwargs.pop('''config'''       ,	_SCREAMING_SNAKE_CASE					) or RagConfig.from_pretrained(_SCREAMING_SNAKE_CASE       ,	**_SCREAMING_SNAKE_CASE					)
           A_      =  RagTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE       ,	config=_SCREAMING_SNAKE_CASE					)
           A_      =  rag_tokenizer.question_encoder
           A_      =  rag_tokenizer.generator

           if indexed_dataset is not None:
               A_      =  '''custom'''
               A_      =  CustomHFIndex(config.retrieval_vector_size       ,	_SCREAMING_SNAKE_CASE					)
           else:
               A_      =  cls._build_index(_SCREAMING_SNAKE_CASE					)

           return cls(
               _SCREAMING_SNAKE_CASE       ,	question_encoder_tokenizer=_SCREAMING_SNAKE_CASE       ,	generator_tokenizer=_SCREAMING_SNAKE_CASE       ,	retrieval_workers=_SCREAMING_SNAKE_CASE       ,	index=_SCREAMING_SNAKE_CASE       ,	)

       def    __A			( self					)						->	Optional[int]:
           logger.info('''re-loading the new dataset with embeddings'''					)
           # access from the training loop

           ray.get([worker.clear_object.remote() for worker in self.retrieval_workers]					)

           # build the index object again
           A_      =  self._build_index(self.config					)

           ray.get(
               [
                   worker.create_rag_retriever.remote(
                       self.config       ,	self.question_encoder_tokenizer       ,	self.generator_tokenizer       ,	_SCREAMING_SNAKE_CASE					)
                   for worker in self.retrieval_workers
               ]					)






