
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     __snake_case		(   enum.Enum):
					_lowerCAmelCase			     =	0
					_lowerCAmelCase			     =	1





@add_end_docstrings(a__)
class     __snake_case		(   a__):
					_lowerCAmelCase			     =	'''generated'''
					def __init__(     self,    *A,    **A      ):

									"""simple docstring"""


									super().__init__(*UpperCAmelCase__,    **UpperCAmelCase__      )

									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,    A=None,    A=None,    A=None,    A=None,    A=None,    A=None,    **A,    ):

									"""simple docstring"""


									lowerCamelCase					:			int									= {}
									if truncation is not None:
													lowerCamelCase					:			Union[str, Any]									= truncation

									lowerCamelCase					:			Any									= generate_kwargs

									lowerCamelCase					:			Optional[int]									= {}
									if return_tensors is not None and return_type is None:
													lowerCamelCase					:			Optional[Any]									= ReturnType.TENSORS if return_tensors else ReturnType.TEXT
									if return_type is not None:
													lowerCamelCase					:			Dict									= return_type

									if clean_up_tokenization_spaces is not None:
													lowerCamelCase					:			str									= clean_up_tokenization_spaces

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

									return preprocess_params, forward_params, postprocess_params
					def     UpperCAmelCase_							(     self,    A,    A,    A      ):

									"""simple docstring"""


									return True
					def     UpperCAmelCase_							(     self,    *A,    A      ):

									"""simple docstring"""


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

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

									"""simple docstring"""


									lowerCamelCase					:			Optional[Any]									= super().__call__(*UpperCAmelCase__,    **UpperCAmelCase__      )
									if (
									    isinstance(args[0],    UpperCAmelCase__      )
									    and all(isinstance(UpperCAmelCase__,    UpperCAmelCase__      ) for el in args[0]      )
									    and all(len(UpperCAmelCase__      ) == 1 for res in result      )
									):
													return [res[0] for res in result]
									return result
					def     UpperCAmelCase_							(     self,    A,    A=TruncationStrategy.DO_NOT_TRUNCATE,    **A      ):

									"""simple docstring"""


									lowerCamelCase					:			List[str]									= self._parse_and_tokenize(UpperCAmelCase__,    truncation=UpperCAmelCase__,    **UpperCAmelCase__      )
									return inputs
					def     UpperCAmelCase_							(     self,    A,    **A      ):

									"""simple docstring"""


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

									lowerCamelCase					:			List[Any]									= generate_kwargs.get('min_length',    self.model.config.min_length      )
									lowerCamelCase					:			Optional[Any]									= generate_kwargs.get('max_length',    self.model.config.max_length      )
									self.check_inputs(UpperCAmelCase__,    generate_kwargs['min_length'],    generate_kwargs['max_length']      )
									lowerCamelCase					:			Union[str, Any]									= self.model.generate(**UpperCAmelCase__,    **UpperCAmelCase__      )
									lowerCamelCase					:			str									= output_ids.shape[0]
									if self.framework == "pt":
													lowerCamelCase					:			Optional[int]									= output_ids.reshape(UpperCAmelCase__,    out_b // in_b,    *output_ids.shape[1:]      )
									elif self.framework == "tf":
													lowerCamelCase					:			List[Any]									= tf.reshape(UpperCAmelCase__,    (in_b, out_b // in_b, *output_ids.shape[1:])      )
									return {"output_ids": output_ids}




					def     UpperCAmelCase_							(     self,    A,    A=ReturnType.TEXT,    A=False      ):

									"""simple docstring"""


									lowerCamelCase					:			Any									= []
									for output_ids in model_outputs["output_ids"][0]:
													if return_type == ReturnType.TENSORS:
																	lowerCamelCase					:			List[Any]									= {F'''{self.return_name}_token_ids''': output_ids}
													elif return_type == ReturnType.TEXT:
																	lowerCamelCase					:			Dict									= {
																	    F'''{self.return_name}_text''': self.tokenizer.decode(
																	        UpperCAmelCase__,    skip_special_tokens=UpperCAmelCase__,    clean_up_tokenization_spaces=UpperCAmelCase__,    )
																	}
													records.append(UpperCAmelCase__      )
									return records





@add_end_docstrings(a__)
class     __snake_case		(   a__):
					_lowerCAmelCase			     =	'''summary'''
					def __call__(     self,    *A,    **A      ):

									"""simple docstring"""


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




					def     UpperCAmelCase_							(     self,    A,    A,    A      ):

									"""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(a__)
class     __snake_case		(   a__):
					_lowerCAmelCase			     =	'''translation'''
					def     UpperCAmelCase_							(     self,    A,    A,    A      ):

									"""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,    *A,    A=TruncationStrategy.DO_NOT_TRUNCATE,    A=None,    A=None      ):

									"""simple docstring"""


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

									"""simple docstring"""


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




					def __call__(     self,    *A,    **A      ):

									"""simple docstring"""


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


'''simple docstring'''






import unittest
from queue import Empty
from threading import Thread

from transformers import AutoTokenizer, TextIteratorStreamer, TextStreamer, is_torch_available
from transformers.testing_utils import CaptureStdout, require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
					import torch

					from transformers import AutoModelForCausalLM




@require_torch
class     __snake_case		(   unittest.TestCase):
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Any									= AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'      )
									lowerCamelCase					:			Optional[int]									= AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'      ).to(A      )
									lowerCamelCase					:			Union[str, Any]									= -1

									lowerCamelCase					:			Optional[Any]									= ids_tensor((1, 5),    vocab_size=model.config.vocab_size      ).to(A      )
									lowerCamelCase					:			List[Any]									= model.generate(A,    max_new_tokens=10,    do_sample=A      )
									lowerCamelCase					:			int									= tokenizer.decode(greedy_ids[0]      )

									with CaptureStdout() as cs:
													lowerCamelCase					:			Optional[int]									= TextStreamer(A      )
													model.generate(A,    max_new_tokens=10,    do_sample=A,    streamer=A      )
									# The greedy text should be printed to stdout, except for the final "\n" in the streamer
									lowerCamelCase					:			int									= cs.out[:-1]

									self.assertEqual(A,    A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[str]									= AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'      )
									lowerCamelCase					:			Dict									= AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'      ).to(A      )
									lowerCamelCase					:			List[Any]									= -1

									lowerCamelCase					:			Optional[int]									= ids_tensor((1, 5),    vocab_size=model.config.vocab_size      ).to(A      )
									lowerCamelCase					:			Any									= model.generate(A,    max_new_tokens=10,    do_sample=A      )
									lowerCamelCase					:			Optional[int]									= tokenizer.decode(greedy_ids[0]      )

									lowerCamelCase					:			Union[str, Any]									= TextIteratorStreamer(A      )
									lowerCamelCase					:			Tuple									= {'input_ids': input_ids, 'max_new_tokens': 10, 'do_sample': False, 'streamer': streamer}
									lowerCamelCase					:			List[Any]									= Thread(target=model.generate,    kwargs=A      )
									thread.start()
									lowerCamelCase					:			List[Any]									= ''
									for new_text in streamer:
													streamer_text += new_text

									self.assertEqual(A,    A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[int]									= AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'      )
									lowerCamelCase					:			int									= AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'      ).to(A      )
									lowerCamelCase					:			Optional[Any]									= -1

									lowerCamelCase					:			Union[str, Any]									= ids_tensor((1, 5),    vocab_size=model.config.vocab_size      ).to(A      )
									lowerCamelCase					:			Union[str, Any]									= model.generate(A,    max_new_tokens=10,    do_sample=A      )
									lowerCamelCase					:			int									= greedy_ids[:, input_ids.shape[1] :]
									lowerCamelCase					:			Dict									= tokenizer.decode(new_greedy_ids[0]      )

									with CaptureStdout() as cs:
													lowerCamelCase					:			Tuple									= TextStreamer(A,    skip_prompt=A      )
													model.generate(A,    max_new_tokens=10,    do_sample=A,    streamer=A      )
									# The greedy text should be printed to stdout, except for the final "\n" in the streamer
									lowerCamelCase					:			Optional[int]									= cs.out[:-1]

									self.assertEqual(A,    A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[Any]									= AutoTokenizer.from_pretrained('distilgpt2'      )
									lowerCamelCase					:			Dict									= AutoModelForCausalLM.from_pretrained('distilgpt2'      ).to(A      )
									lowerCamelCase					:			Dict									= -1

									lowerCamelCase					:			Any									= torch.ones((1, 5),    device=A      ).long() * model.config.bos_token_id
									with CaptureStdout() as cs:
													lowerCamelCase					:			int									= TextStreamer(A,    skip_special_tokens=A      )
													model.generate(A,    max_new_tokens=1,    do_sample=A,    streamer=A      )

									# The prompt contains a special token, so the streamer should not print it. As such, the output text, when
									# re-tokenized, must only contain one token
									lowerCamelCase					:			List[str]									= cs.out[:-1]  # Remove the final "\n"
									lowerCamelCase					:			Optional[Any]									= tokenizer(A,    return_tensors='pt'      )
									self.assertEqual(streamer_text_tokenized.input_ids.shape,    (1, 1)      )




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'      )
									lowerCamelCase					:			List[str]									= AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'      ).to(A      )
									lowerCamelCase					:			Any									= -1

									lowerCamelCase					:			int									= ids_tensor((1, 5),    vocab_size=model.config.vocab_size      ).to(A      )
									lowerCamelCase					:			int									= TextIteratorStreamer(A,    timeout=0.001      )
									lowerCamelCase					:			Any									= {'input_ids': input_ids, 'max_new_tokens': 10, 'do_sample': False, 'streamer': streamer}
									lowerCamelCase					:			Optional[Any]									= Thread(target=model.generate,    kwargs=A      )
									thread.start()

									# The streamer will timeout after 0.001 seconds, so an exception will be raised
									with self.assertRaises(A      ):
													lowerCamelCase					:			Dict									= ''
													for new_text in streamer:
																	streamer_text += new_text
