




"""simple docstring"""







import enum
import warnings

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
		import tensorflow as tf

		from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

if is_torch_available():
		from ..models.auto.modeling_auto import MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

a      =							logging.get_logger(__name__)






class       SCREAMING_SNAKE_CASE__					(					enum.Enum					):
						_a    						=						0
						_a    						=						1






@add_end_docstrings(_a					)
class       SCREAMING_SNAKE_CASE__					(					_a					):
						_a    						=						'generated'


						def __init__(							self				:			Dict						,						*lowerCAmelCase				:			Optional[Any]						,						**lowerCAmelCase				:			int					):
										super().__init__(*lowerCAmelCase						,						**lowerCAmelCase					)

										self.check_model_type(
										    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
										    if self.framework == """tf"""
										    else MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING					)


						def   __lowercase (							self				:			List[str]						,						lowerCAmelCase				:			Any=None						,						lowerCAmelCase				:			Dict=None						,						lowerCAmelCase				:			List[str]=None						,						lowerCAmelCase				:			Optional[int]=None						,						lowerCAmelCase				:			Optional[Any]=None						,						lowerCAmelCase				:			Tuple=None						,						**lowerCAmelCase				:			str						,						):
										lowerCAmelCase  =				{}
										if truncation is not None:
														lowerCAmelCase  =				truncation

										lowerCAmelCase  =				generate_kwargs

										lowerCAmelCase  =				{}
										if return_tensors is not None and return_type is None:
														lowerCAmelCase  =				ReturnType.TENSORS if return_tensors else ReturnType.TEXT
										if return_type is not None:
														lowerCAmelCase  =				return_type

										if clean_up_tokenization_spaces is not None:
														lowerCAmelCase  =				clean_up_tokenization_spaces

										if stop_sequence is not None:
														lowerCAmelCase  =				self.tokenizer.encode(lowerCAmelCase						,						add_special_tokens=lowerCAmelCase					)
														if len(lowerCAmelCase					) > 1:
																		warnings.warn(
																		    """Stopping on a multiple token sequence is not yet supported on transformers. The first token of"""
																		    """ the stop sequence will be used as the stop sequence string in the interim."""					)
														lowerCAmelCase  =				stop_sequence_ids[0]

										return preprocess_params, forward_params, postprocess_params


						def   __lowercase (							self				:			Optional[Any]						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			int					):
										return True


						def   __lowercase (							self				:			Optional[Any]						,						*lowerCAmelCase				:			Any						,						lowerCAmelCase				:			Union[str, Any]					):
										lowerCAmelCase  =				self.model.config.prefix if self.model.config.prefix is not None else """"""
										if isinstance(args[0]						,						lowerCAmelCase					):
														if self.tokenizer.pad_token_id is None:
																		raise ValueError("""Please make sure that the tokenizer has a pad_token_id when using a batch input"""					)
														lowerCAmelCase  =				([prefix + arg for arg in args[0]],)
														lowerCAmelCase  =				True

										elif isinstance(args[0]						,						lowerCAmelCase					):
														lowerCAmelCase  =				(prefix + args[0],)
														lowerCAmelCase  =				False
										else:
														raise ValueError(
														    f''' `args[0]`: {args[0]} have the wrong format. The should be either of type `str` or type `list`'''					)
										lowerCAmelCase  =				self.tokenizer(*lowerCAmelCase						,						padding=lowerCAmelCase						,						truncation=lowerCAmelCase						,						return_tensors=self.framework					)
										# This is produced by tokenizers but is an invalid generate kwargs
										if "token_type_ids" in inputs:
														del inputs["token_type_ids"]
										return inputs


						def __call__(							self				:			Union[str, Any]						,						*lowerCAmelCase				:			Dict						,						**lowerCAmelCase				:			List[Any]					):
										lowerCAmelCase  =				super().__call__(*lowerCAmelCase						,						**lowerCAmelCase					)
										if (
										    isinstance(args[0]						,						lowerCAmelCase					)
										    and all(isinstance(lowerCAmelCase						,						lowerCAmelCase					) for el in args[0]					)
										    and all(len(lowerCAmelCase					) == 1 for res in result					)
										):
														return [res[0] for res in result]
										return result


						def   __lowercase (							self				:			Optional[int]						,						lowerCAmelCase				:			Any						,						lowerCAmelCase				:			List[Any]=TruncationStrategy.DO_NOT_TRUNCATE						,						**lowerCAmelCase				:			int					):
										lowerCAmelCase  =				self._parse_and_tokenize(lowerCAmelCase						,						truncation=lowerCAmelCase						,						**lowerCAmelCase					)
										return inputs


						def   __lowercase (							self				:			str						,						lowerCAmelCase				:			Optional[Any]						,						**lowerCAmelCase				:			Tuple					):
										if self.framework == "pt":
														lowerCAmelCase ,						lowerCAmelCase  =				model_inputs["""input_ids"""].shape
										elif self.framework == "tf":
														lowerCAmelCase ,						lowerCAmelCase  =				tf.shape(model_inputs["""input_ids"""]					).numpy()

										lowerCAmelCase  =				generate_kwargs.get("""min_length"""						,						self.model.config.min_length					)
										lowerCAmelCase  =				generate_kwargs.get("""max_length"""						,						self.model.config.max_length					)
										self.check_inputs(lowerCAmelCase						,						generate_kwargs["""min_length"""]						,						generate_kwargs["""max_length"""]					)
										lowerCAmelCase  =				self.model.generate(**lowerCAmelCase						,						**lowerCAmelCase					)
										lowerCAmelCase  =				output_ids.shape[0]
										if self.framework == "pt":
														lowerCAmelCase  =				output_ids.reshape(lowerCAmelCase						,						out_b // in_b						,						*output_ids.shape[1:]					)
										elif self.framework == "tf":
														lowerCAmelCase  =				tf.reshape(lowerCAmelCase						,						(in_b, out_b // in_b, *output_ids.shape[1:])					)
										return {"output_ids": output_ids}




						def   __lowercase (							self				:			Tuple						,						lowerCAmelCase				:			List[str]						,						lowerCAmelCase				:			int=ReturnType.TEXT						,						lowerCAmelCase				:			Dict=False					):
										lowerCAmelCase  =				[]
										for output_ids in model_outputs["output_ids"][0]:
														if return_type == ReturnType.TENSORS:
																		lowerCAmelCase  =				{f'''{self.return_name}_token_ids''': output_ids}
														elif return_type == ReturnType.TEXT:
																		lowerCAmelCase  =				{
																		    f'''{self.return_name}_text''': self.tokenizer.decode(
																		        lowerCAmelCase						,						skip_special_tokens=lowerCAmelCase						,						clean_up_tokenization_spaces=lowerCAmelCase						,						)
																		}
														records.append(lowerCAmelCase					)
										return records






@add_end_docstrings(_a					)
class       SCREAMING_SNAKE_CASE__					(					_a					):
						_a    						=						'summary'


						def __call__(							self				:			Union[str, Any]						,						*lowerCAmelCase				:			List[str]						,						**lowerCAmelCase				:			Optional[int]					):
										return super().__call__(*lowerCAmelCase						,						**lowerCAmelCase					)




						def   __lowercase (							self				:			List[Any]						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			int					):
										if max_length < min_length:
														logger.warning(f'''Your min_length={min_length} must be inferior than your max_length={max_length}.'''					)

										if input_length < max_length:
														logger.warning(
														    f'''Your max_length is set to {max_length}, but your input_length is only {input_length}. Since this is '''
														    """a summarization task, where outputs shorter than the input are typically wanted, you might """
														    f'''consider decreasing max_length manually, e.g. summarizer(\'...\', max_length={input_length//2})'''					)






@add_end_docstrings(_a					)
class       SCREAMING_SNAKE_CASE__					(					_a					):
						_a    						=						'translation'


						def   __lowercase (							self				:			Tuple						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			int					):
										if input_length > 0.9 * max_length:
														logger.warning(
														    f'''Your input_length: {input_length} is bigger than 0.9 * max_length: {max_length}. You might consider '''
														    """increasing your max_length manually, e.g. translator('...', max_length=400)"""					)
										return True


						def   __lowercase (							self				:			Dict						,						*lowerCAmelCase				:			Optional[Any]						,						lowerCAmelCase				:			List[str]=TruncationStrategy.DO_NOT_TRUNCATE						,						lowerCAmelCase				:			Dict=None						,						lowerCAmelCase				:			int=None					):
										if getattr(self.tokenizer						,						"""_build_translation_inputs"""						,						lowerCAmelCase					):
														return self.tokenizer._build_translation_inputs(
														    *lowerCAmelCase						,						return_tensors=self.framework						,						truncation=lowerCAmelCase						,						src_lang=lowerCAmelCase						,						tgt_lang=lowerCAmelCase					)
										else:
														return super()._parse_and_tokenize(*lowerCAmelCase						,						truncation=lowerCAmelCase					)


						def   __lowercase (							self				:			Union[str, Any]						,						lowerCAmelCase				:			int=None						,						lowerCAmelCase				:			Union[str, Any]=None						,						**lowerCAmelCase				:			Dict					):
										lowerCAmelCase ,						lowerCAmelCase ,						lowerCAmelCase  =				super()._sanitize_parameters(**lowerCAmelCase					)
										if src_lang is not None:
														lowerCAmelCase  =				src_lang
										if tgt_lang is not None:
														lowerCAmelCase  =				tgt_lang
										if src_lang is None and tgt_lang is None:
														# Backward compatibility, direct arguments use is preferred.
														lowerCAmelCase  =				kwargs.get("""task"""						,						self.task					)
														lowerCAmelCase  =				task.split("""_"""					)
														if task and len(lowerCAmelCase					) == 4:
																		# translation, XX, to YY
																		lowerCAmelCase  =				items[1]
																		lowerCAmelCase  =				items[3]
										return preprocess_params, forward_params, postprocess_params




						def __call__(							self				:			List[str]						,						*lowerCAmelCase				:			str						,						**lowerCAmelCase				:			List[str]					):
										return super().__call__(*lowerCAmelCase						,						**lowerCAmelCase					)











"""simple docstring"""







import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
		import tensorflow as tf
		from tensorflow.python.framework.errors_impl import ResourceExhaustedError

		from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
		import pyanvml.pyanvml as nvml

a      =							logging.get_logger(__name__)





def 							lowercase					(snake_case__   : bool      ,       snake_case__   : bool )     ->  Tuple:


				'''simple docstring'''






				def run_func(snake_case__   : Any ):
								@wraps(snake_case__ )
								def run_in_eager_mode(*snake_case__   : Optional[Any]      ,       **snake_case__   : int ):
												return func(*snake_case__      ,       **snake_case__ )

								@wraps(snake_case__ )
								@tf.function(experimental_compile=snake_case__ )
								def run_in_graph_mode(*snake_case__   : int      ,       **snake_case__   : Tuple ):
												return func(*snake_case__      ,       **snake_case__ )

								if do_eager_mode is True:
												if use_xla is not False:
																raise ValueError(
																    """Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`.""" )
												return run_in_eager_mode
								else:
												return run_in_graph_mode

				return run_func





def 							lowercase					(snake_case__   : int      ,       snake_case__   : int      ,       snake_case__   : int )     ->  ["tf.Tensor"]:


				'''simple docstring'''






				lowerCAmelCase  =				random.Random()
				lowerCAmelCase  =				[rng.randint(0      ,       vocab_size - 1 ) for i in range(batch_size * sequence_length )]
				return tf.constant(snake_case__      ,       shape=(batch_size, sequence_length)      ,       dtype=tf.intaa )






class       SCREAMING_SNAKE_CASE__					(					_a					):
						_a    						=						42
						_a    						=						42
						_a    						=						"TensorFlow"


						@property
						def   __lowercase (							self				:			Optional[int]					):
										return tf.__version__


						def   __lowercase (							self				:			Optional[Any]						,						lowerCAmelCase				:			str						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			int					):
										# initialize GPU on separate process
										lowerCAmelCase  =				self.args.strategy
										if strategy is None:
														raise ValueError("""A device strategy has to be initialized before using TensorFlow."""					)
										lowerCAmelCase  =				self._prepare_inference_func(lowerCAmelCase						,						lowerCAmelCase						,						lowerCAmelCase					)
										return self._measure_speed(_inference					)


						def   __lowercase (							self				:			str						,						lowerCAmelCase				:			str						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			int					):
										lowerCAmelCase  =				self.args.strategy
										if strategy is None:
														raise ValueError("""A device strategy has to be initialized before using TensorFlow."""					)
										lowerCAmelCase  =				self._prepare_train_func(lowerCAmelCase						,						lowerCAmelCase						,						lowerCAmelCase					)
										return self._measure_speed(_train					)


						def   __lowercase (							self				:			str						,						lowerCAmelCase				:			str						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			int					):
										# initialize GPU on separate process
										if self.args.is_gpu:
														tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]						,						lowerCAmelCase					)
										lowerCAmelCase  =				self.args.strategy
										if strategy is None:
														raise ValueError("""A device strategy has to be initialized before using TensorFlow."""					)
										lowerCAmelCase  =				self._prepare_inference_func(lowerCAmelCase						,						lowerCAmelCase						,						lowerCAmelCase					)
										return self._measure_memory(_inference					)


						def   __lowercase (							self				:			int						,						lowerCAmelCase				:			str						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			int					):
										if self.args.is_gpu:
														tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]						,						lowerCAmelCase					)
										lowerCAmelCase  =				self.args.strategy
										if strategy is None:
														raise ValueError("""A device strategy has to be initialized before using TensorFlow."""					)

										lowerCAmelCase  =				self._prepare_train_func(lowerCAmelCase						,						lowerCAmelCase						,						lowerCAmelCase					)
										return self._measure_memory(_train					)


						def   __lowercase (							self				:			List[str]						,						lowerCAmelCase				:			str						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			int					):
										lowerCAmelCase  =				self.config_dict[model_name]

										if self.args.fpaa:
														raise NotImplementedError("""Mixed precision is currently not supported."""					)

										lowerCAmelCase  =				(
										    hasattr(lowerCAmelCase						,						"""architectures"""					)
										    and isinstance(config.architectures						,						lowerCAmelCase					)
										    and len(config.architectures					) > 0
										)
										if not self.args.only_pretrain_model and has_model_class_in_config:
														try:
																		lowerCAmelCase  =				"""TF""" + config.architectures[0]  # prepend 'TF' for tensorflow model
																		lowerCAmelCase  =				__import__("""transformers"""						,						fromlist=[model_class]					)
																		lowerCAmelCase  =				getattr(lowerCAmelCase						,						lowerCAmelCase					)
																		lowerCAmelCase  =				model_cls(lowerCAmelCase					)
														except ImportError:
																		raise ImportError(
																		    f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
																		    """ set `--only_pretrain_model` or `args.only_pretrain_model=True`."""					)
										else:
														lowerCAmelCase  =				TF_MODEL_MAPPING[config.__class__](lowerCAmelCase					)

										# encoder-decoder has vocab size saved differently
										lowerCAmelCase  =				config.vocab_size if hasattr(lowerCAmelCase						,						"""vocab_size"""					) else config.encoder.vocab_size
										lowerCAmelCase  =				random_input_ids(lowerCAmelCase						,						lowerCAmelCase						,						lowerCAmelCase					)

										@run_with_tf_optimizations(self.args.eager_mode						,						self.args.use_xla					)
										def encoder_decoder_forward():
														return model(lowerCAmelCase						,						decoder_input_ids=lowerCAmelCase						,						training=lowerCAmelCase					)

										@run_with_tf_optimizations(self.args.eager_mode						,						self.args.use_xla					)
										def encoder_forward():
														return model(lowerCAmelCase						,						training=lowerCAmelCase					)

										lowerCAmelCase  =				encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

										return _inference


						def   __lowercase (							self				:			Optional[int]						,						lowerCAmelCase				:			str						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			int					):
										lowerCAmelCase  =				self.config_dict[model_name]

										if self.args.eager_mode is not False:
														raise ValueError("""Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`."""					)

										if self.args.fpaa:
														raise NotImplementedError("""Mixed precision is currently not supported."""					)

										lowerCAmelCase  =				(
										    hasattr(lowerCAmelCase						,						"""architectures"""					)
										    and isinstance(config.architectures						,						lowerCAmelCase					)
										    and len(config.architectures					) > 0
										)
										if not self.args.only_pretrain_model and has_model_class_in_config:
														try:
																		lowerCAmelCase  =				"""TF""" + config.architectures[0]  # prepend 'TF' for tensorflow model
																		lowerCAmelCase  =				__import__("""transformers"""						,						fromlist=[model_class]					)
																		lowerCAmelCase  =				getattr(lowerCAmelCase						,						lowerCAmelCase					)
																		lowerCAmelCase  =				model_cls(lowerCAmelCase					)
														except ImportError:
																		raise ImportError(
																		    f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
																		    """ set `--only_pretrain_model` or `args.only_pretrain_model=True`."""					)
										else:
														lowerCAmelCase  =				TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](lowerCAmelCase					)

										# encoder-decoder has vocab size saved differently
										lowerCAmelCase  =				config.vocab_size if hasattr(lowerCAmelCase						,						"""vocab_size"""					) else config.encoder.vocab_size
										lowerCAmelCase  =				random_input_ids(lowerCAmelCase						,						lowerCAmelCase						,						lowerCAmelCase					)

										@run_with_tf_optimizations(self.args.eager_mode						,						self.args.use_xla					)
										def encoder_decoder_train():
														lowerCAmelCase  =				model(lowerCAmelCase						,						decoder_input_ids=lowerCAmelCase						,						labels=lowerCAmelCase						,						training=lowerCAmelCase					)[0]
														lowerCAmelCase  =				tf.gradients(lowerCAmelCase						,						model.trainable_variables					)
														return gradients

										@run_with_tf_optimizations(self.args.eager_mode						,						self.args.use_xla					)
										def encoder_train():
														lowerCAmelCase  =				model(lowerCAmelCase						,						labels=lowerCAmelCase						,						training=lowerCAmelCase					)[0]
														lowerCAmelCase  =				tf.gradients(lowerCAmelCase						,						model.trainable_variables					)
														return gradients

										lowerCAmelCase  =				encoder_decoder_train if config.is_encoder_decoder else encoder_train

										return _train


						def   __lowercase (							self				:			Optional[int]						,						lowerCAmelCase				:			List[str]					):
										with self.args.strategy.scope():
														try:
																		if self.args.is_tpu or self.args.use_xla:
																						# run additional 10 times to stabilize compilation for tpu
																						logger.info("""Do inference on TPU. Running model 5 times to stabilize compilation"""					)
																						timeit.repeat(lowerCAmelCase						,						repeat=1						,						number=5					)

																		# as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
																		lowerCAmelCase  =				timeit.repeat(
																		    lowerCAmelCase						,						repeat=self.args.repeat						,						number=10						,						)

																		return min(lowerCAmelCase					) / 10.0
														except ResourceExhaustedError as e:
																		self.print_fn(f'''Doesn\'t fit on GPU. {e}'''					)




						def   __lowercase (							self				:			Optional[Any]						,						lowerCAmelCase				:			Callable[[], None]					):
										logger.info(
										    """Note that TensorFlow allocates more memory than """
										    """it might need to speed up computation. """
										    """The memory reported here corresponds to the memory """
										    """reported by `nvidia-smi`, which can vary depending """
										    """on total available memory on the GPU that is used."""					)
										with self.args.strategy.scope():
														try:
																		if self.args.trace_memory_line_by_line:
																						if not self.args.eager_mode:
																										raise ValueError(
																										    """`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory"""
																										    """ consumption line by line."""					)
																						lowerCAmelCase  =				start_memory_tracing("""transformers"""					)

																		if self.args.is_tpu:
																						# tpu
																						raise NotImplementedError(
																						    """Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking"""
																						    """ with `args.memory=False`"""					)
																		elif self.args.is_gpu:
																						# gpu
																						if not is_pyanvml_available():
																										logger.warning(
																										    """py3nvml not installed, we won't log GPU memory usage. """
																										    """Install py3nvml (pip install py3nvml) to log information about GPU."""					)
																										lowerCAmelCase  =				"""N/A"""
																						else:
																										logger.info(
																										    """Measuring total GPU usage on GPU device. Make sure to not have additional processes"""
																										    """ running on the same GPU."""					)
																										# init nvml
																										nvml.nvmlInit()
																										func()
																										lowerCAmelCase  =				nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx					)
																										lowerCAmelCase  =				nvml.nvmlDeviceGetMemoryInfo(lowerCAmelCase					)
																										lowerCAmelCase  =				meminfo.used
																										lowerCAmelCase  =				Memory(lowerCAmelCase					)
																										# shutdown nvml
																										nvml.nvmlShutdown()
																		else:
																						# cpu
																						if self.args.trace_memory_line_by_line:
																										logger.info(
																										    """When enabling line by line tracing, the max peak memory for CPU is inaccurate in"""
																										    """ TensorFlow."""					)
																										lowerCAmelCase  =				None
																						else:
																										lowerCAmelCase  =				measure_peak_memory_cpu(lowerCAmelCase					)
																										lowerCAmelCase  =				Memory(lowerCAmelCase					) if isinstance(lowerCAmelCase						,						lowerCAmelCase					) else memory_bytes
																		if self.args.trace_memory_line_by_line:
																						lowerCAmelCase  =				stop_memory_tracing(lowerCAmelCase					)
																						if memory is None:
																										lowerCAmelCase  =				summary.total
																		else:
																						lowerCAmelCase  =				None

																		return memory, summary
														except ResourceExhaustedError as e:
																		self.print_fn(f'''Doesn\'t fit on GPU. {e}'''					)
																		return "N/A", None





