


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

lowerCAmelCase__    :str       =				logging.get_logger(__name__)

class 						__a (							enum.Enum     ):
		_a      :    List[Any]           =      0
		_a      :    Union[str, Any]           =      1



@add_end_docstrings(UpperCAmelCase     )
class 						__a (							UpperCAmelCase     ):
		_a      :    List[Any]           =      'generated'
		def __init__(     self     ,       *_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)					->			int:




									"""simple docstring"""





									super().__init__(*_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)

									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       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE=None     ,       _SCREAMING_SNAKE_CASE=None     ,       _SCREAMING_SNAKE_CASE=None     ,       _SCREAMING_SNAKE_CASE=None     ,       _SCREAMING_SNAKE_CASE=None     ,       _SCREAMING_SNAKE_CASE=None     ,       **_SCREAMING_SNAKE_CASE     ,       )					->			Tuple:




									"""simple docstring"""





									_UpperCAmelCase											=     {}
									if truncation is not None:
																_UpperCAmelCase											=     truncation

									_UpperCAmelCase											=     generate_kwargs

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

									if clean_up_tokenization_spaces is not None:
																_UpperCAmelCase											=     clean_up_tokenization_spaces

									if stop_sequence is not None:
																_UpperCAmelCase											=     self.tokenizer.encode(_SCREAMING_SNAKE_CASE     ,       add_special_tokens=_SCREAMING_SNAKE_CASE		)
																if len(_SCREAMING_SNAKE_CASE		) > 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.'		)
																_UpperCAmelCase											=     stop_sequence_ids[0]

									return preprocess_params, forward_params, postprocess_params
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			List[Any]:




									"""simple docstring"""





									return True
		def       UpperCAmelCase__							(     self     ,       *_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     self.model.config.prefix if self.model.config.prefix is not None else ''
									if isinstance(args[0]     ,       _SCREAMING_SNAKE_CASE		):
																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'		)
																_UpperCAmelCase											=     ([prefix + arg for arg in args[0]],)
																_UpperCAmelCase											=     True

									elif isinstance(args[0]     ,       _SCREAMING_SNAKE_CASE		):
																_UpperCAmelCase											=     (prefix + args[0],)
																_UpperCAmelCase											=     False
									else:
																raise ValueError(
																    f''' `args[0]`: {args[0]} have the wrong format. The should be either of type `str` or type `list`'''		)
									_UpperCAmelCase											=     self.tokenizer(*_SCREAMING_SNAKE_CASE     ,       padding=_SCREAMING_SNAKE_CASE     ,       truncation=_SCREAMING_SNAKE_CASE     ,       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     ,       *_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)					->			List[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     super().__call__(*_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)
									if (
									    isinstance(args[0]     ,       _SCREAMING_SNAKE_CASE		)
									    and all(isinstance(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		) for el in args[0]		)
									    and all(len(_SCREAMING_SNAKE_CASE		) == 1 for res in result		)
									):
																return [res[0] for res in result]
									return result
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=TruncationStrategy.DO_NOT_TRUNCATE     ,       **_SCREAMING_SNAKE_CASE		)					->			Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase											=     self._parse_and_tokenize(_SCREAMING_SNAKE_CASE     ,       truncation=_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)
									return inputs
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)					->			List[str]:




									"""simple docstring"""





									if self.framework == "pt":
																_UpperCAmelCase						,			_UpperCAmelCase											=     model_inputs['input_ids'].shape
									elif self.framework == "tf":
																_UpperCAmelCase						,			_UpperCAmelCase											=     tf.shape(model_inputs['input_ids']		).numpy()

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






		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=ReturnType.TEXT     ,       _SCREAMING_SNAKE_CASE=False		)					->			List[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     []
									for output_ids in model_outputs["output_ids"][0]:
																if return_type == ReturnType.TENSORS:
																							_UpperCAmelCase											=     {f'''{self.return_name}_token_ids''': output_ids}
																elif return_type == ReturnType.TEXT:
																							_UpperCAmelCase											=     {
																							    f'''{self.return_name}_text''': self.tokenizer.decode(
																							        _SCREAMING_SNAKE_CASE     ,       skip_special_tokens=_SCREAMING_SNAKE_CASE     ,       clean_up_tokenization_spaces=_SCREAMING_SNAKE_CASE     ,       )
																							}
																records.append(_SCREAMING_SNAKE_CASE		)
									return records



@add_end_docstrings(UpperCAmelCase     )
class 						__a (							UpperCAmelCase     ):
		_a      :    Any           =      'summary'
		def __call__(     self     ,       *_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)					->			Tuple:




									"""simple docstring"""





									return super().__call__(*_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)






		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			bool:




									"""simple docstring"""





									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(UpperCAmelCase     )
class 						__a (							UpperCAmelCase     ):
		_a      :    str           =      'translation'
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Tuple:




									"""simple docstring"""





									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       UpperCAmelCase__							(     self     ,       *_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=TruncationStrategy.DO_NOT_TRUNCATE     ,       _SCREAMING_SNAKE_CASE=None     ,       _SCREAMING_SNAKE_CASE=None		)					->			Tuple:




									"""simple docstring"""





									if getattr(self.tokenizer     ,       '_build_translation_inputs'     ,       _SCREAMING_SNAKE_CASE		):
																return self.tokenizer._build_translation_inputs(
																    *_SCREAMING_SNAKE_CASE     ,       return_tensors=self.framework     ,       truncation=_SCREAMING_SNAKE_CASE     ,       src_lang=_SCREAMING_SNAKE_CASE     ,       tgt_lang=_SCREAMING_SNAKE_CASE		)
									else:
																return super()._parse_and_tokenize(*_SCREAMING_SNAKE_CASE     ,       truncation=_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE=None     ,       _SCREAMING_SNAKE_CASE=None     ,       **_SCREAMING_SNAKE_CASE		)					->			Tuple:




									"""simple docstring"""





									_UpperCAmelCase						,			_UpperCAmelCase						,			_UpperCAmelCase											=     super()._sanitize_parameters(**_SCREAMING_SNAKE_CASE		)
									if src_lang is not None:
																_UpperCAmelCase											=     src_lang
									if tgt_lang is not None:
																_UpperCAmelCase											=     tgt_lang
									if src_lang is None and tgt_lang is None:
																# Backward compatibility, direct arguments use is preferred.
																_UpperCAmelCase											=     kwargs.get('task'     ,       self.task		)
																_UpperCAmelCase											=     task.split('_'		)
																if task and len(_SCREAMING_SNAKE_CASE		) == 4:
																							# translation, XX, to YY
																							_UpperCAmelCase											=     items[1]
																							_UpperCAmelCase											=     items[3]
									return preprocess_params, forward_params, postprocess_params






		def __call__(     self     ,       *_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)					->			Dict:




									"""simple docstring"""





									return super().__call__(*_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)




import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
		import jax

		from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel

class 						__a (							unittest.TestCase     ):
		def __init__(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=13     ,       _SCREAMING_SNAKE_CASE=30     ,       _SCREAMING_SNAKE_CASE=2     ,       _SCREAMING_SNAKE_CASE=3     ,       _SCREAMING_SNAKE_CASE=True     ,       _SCREAMING_SNAKE_CASE=True     ,       _SCREAMING_SNAKE_CASE=32     ,       _SCREAMING_SNAKE_CASE=5     ,       _SCREAMING_SNAKE_CASE=4     ,       _SCREAMING_SNAKE_CASE=37     ,       _SCREAMING_SNAKE_CASE="gelu"     ,       _SCREAMING_SNAKE_CASE=0.1     ,       _SCREAMING_SNAKE_CASE=0.1     ,       _SCREAMING_SNAKE_CASE=10     ,       _SCREAMING_SNAKE_CASE=0.02     ,       )					->			int:




									"""simple docstring"""





									_UpperCAmelCase											=     parent
									_UpperCAmelCase											=     batch_size
									_UpperCAmelCase											=     image_size
									_UpperCAmelCase											=     patch_size
									_UpperCAmelCase											=     num_channels
									_UpperCAmelCase											=     is_training
									_UpperCAmelCase											=     use_labels
									_UpperCAmelCase											=     hidden_size
									_UpperCAmelCase											=     num_hidden_layers
									_UpperCAmelCase											=     num_attention_heads
									_UpperCAmelCase											=     intermediate_size
									_UpperCAmelCase											=     hidden_act
									_UpperCAmelCase											=     hidden_dropout_prob
									_UpperCAmelCase											=     attention_probs_dropout_prob
									_UpperCAmelCase											=     type_sequence_label_size
									_UpperCAmelCase											=     initializer_range

									# in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
									_UpperCAmelCase											=     (image_size // patch_size) ** 2
									_UpperCAmelCase											=     num_patches + 1
		def       UpperCAmelCase__							(     self		)					->			Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase											=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

									_UpperCAmelCase											=     ViTConfig(
									    image_size=self.image_size     ,       patch_size=self.patch_size     ,       num_channels=self.num_channels     ,       hidden_size=self.hidden_size     ,       num_hidden_layers=self.num_hidden_layers     ,       num_attention_heads=self.num_attention_heads     ,       intermediate_size=self.intermediate_size     ,       hidden_act=self.hidden_act     ,       hidden_dropout_prob=self.hidden_dropout_prob     ,       attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,       is_decoder=_SCREAMING_SNAKE_CASE     ,       initializer_range=self.initializer_range     ,       )

									return config, pixel_values
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			str:




									"""simple docstring"""





									_UpperCAmelCase											=     FlaxViTModel(config=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     model(_SCREAMING_SNAKE_CASE		)
									# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
									_UpperCAmelCase											=     (self.image_size, self.image_size)
									_UpperCAmelCase											=     (self.patch_size, self.patch_size)
									_UpperCAmelCase											=     (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
									self.parent.assertEqual(result.last_hidden_state.shape     ,       (self.batch_size, num_patches + 1, self.hidden_size)		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Any:




									"""simple docstring"""





									_UpperCAmelCase											=     self.type_sequence_label_size
									_UpperCAmelCase											=     FlaxViTForImageClassification(config=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     model(_SCREAMING_SNAKE_CASE		)
									self.parent.assertEqual(result.logits.shape     ,       (self.batch_size, self.type_sequence_label_size)		)

									# test greyscale images
									_UpperCAmelCase											=     1
									_UpperCAmelCase											=     FlaxViTForImageClassification(_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
									_UpperCAmelCase											=     model(_SCREAMING_SNAKE_CASE		)






		def       UpperCAmelCase__							(     self		)					->			str:




									"""simple docstring"""





									_UpperCAmelCase											=     self.prepare_config_and_inputs()
									(
									    (
									    _UpperCAmelCase
									)						,			(
									    _UpperCAmelCase
									)						,			
									)											=     config_and_inputs
									_UpperCAmelCase											=     {'pixel_values': pixel_values}
									return config, inputs_dict



@require_flax
class 						__a (							UpperCAmelCase       , unittest.TestCase     ):
		_a      :    List[str]           =      (FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()
		def       UpperCAmelCase__							(     self		)					->			None:




									"""simple docstring"""





									_UpperCAmelCase											=     FlaxViTModelTester(self		)
									_UpperCAmelCase											=     ConfigTester(self     ,       config_class=_SCREAMING_SNAKE_CASE     ,       has_text_modality=_SCREAMING_SNAKE_CASE     ,       hidden_size=37		)
		def       UpperCAmelCase__							(     self		)					->			Dict:




									"""simple docstring"""





									self.config_tester.run_common_tests()
		def       UpperCAmelCase__							(     self		)					->			Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			List[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			str:




									"""simple docstring"""





									_UpperCAmelCase						,			_UpperCAmelCase											=     self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
																_UpperCAmelCase											=     model_class(_SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase											=     inspect.signature(model.__call__		)
																# signature.parameters is an OrderedDict => so arg_names order is deterministic
																_UpperCAmelCase											=     [*signature.parameters.keys()]

																_UpperCAmelCase											=     ['pixel_values']
																self.assertListEqual(arg_names[:1]     ,       _SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			int:




									"""simple docstring"""





									_UpperCAmelCase						,			_UpperCAmelCase											=     self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
																with self.subTest(model_class.__name__		):
																							_UpperCAmelCase											=     self._prepare_for_class(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
																							_UpperCAmelCase											=     model_class(_SCREAMING_SNAKE_CASE		)

																							@jax.jit
																							def model_jitted(_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		):
																														return model(pixel_values=_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)

																							with self.subTest('JIT Enabled'		):
																														_UpperCAmelCase											=     model_jitted(**_SCREAMING_SNAKE_CASE		).to_tuple()

																							with self.subTest('JIT Disabled'		):
																														with jax.disable_jit():
																																					_UpperCAmelCase											=     model_jitted(**_SCREAMING_SNAKE_CASE		).to_tuple()

																							self.assertEqual(len(_SCREAMING_SNAKE_CASE		)     ,       len(_SCREAMING_SNAKE_CASE		)		)
																							for jitted_output, output in zip(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		):
																														self.assertEqual(jitted_output.shape     ,       output.shape		)






		@slow
		def       UpperCAmelCase__							(     self		)					->			Dict:




									"""simple docstring"""





									for model_class_name in self.all_model_classes:
																_UpperCAmelCase											=     model_class_name.from_pretrained('google/vit-base-patch16-224'		)
																_UpperCAmelCase											=     model(np.ones((1, 3, 224, 224)		)		)
																self.assertIsNotNone(_SCREAMING_SNAKE_CASE		)
