






"""simple docstring"""






from queue import Queue
from typing import TYPE_CHECKING, Optional


if TYPE_CHECKING:
  from ..models.auto import AutoTokenizer




class 							lowerCAmelCase__       :






      '''simple docstring'''

      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]		,			lowercase_  :						List[str]):




          '''simple docstring'''


          raise NotImplementedError()



      def _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


          raise NotImplementedError()




class 							lowerCAmelCase__       (    UpperCAmelCase__				):






      '''simple docstring'''

      def __init__(		self  :						Dict		,			lowercase_  :						"AutoTokenizer"		,			lowercase_  :						bool = False		,			**lowercase_  :						Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[str]							=	tokenizer
          SCREAMING_SNAKE_CASE_   :						str							=	skip_prompt
          SCREAMING_SNAKE_CASE_   :						List[str]							=	decode_kwargs

          # variables used in the streaming process
          SCREAMING_SNAKE_CASE_   :						Dict							=	[]
          SCREAMING_SNAKE_CASE_   :						int							=	0
          SCREAMING_SNAKE_CASE_   :						Tuple							=	True
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple		,			lowercase_  :						Tuple):




          '''simple docstring'''


          if len(value.shape) > 1 and value.shape[0] > 1:
              raise ValueError('''TextStreamer only supports batch size 1''')
          elif len(value.shape) > 1:
              SCREAMING_SNAKE_CASE_   :						List[str]							=	value[0]

          if self.skip_prompt and self.next_tokens_are_prompt:
              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	False
              return

          # Add the new token to the cache and decodes the entire thing.
          self.token_cache.extend(value.tolist())
          SCREAMING_SNAKE_CASE_   :						Any							=	self.tokenizer.decode(self.token_cache		,			**self.decode_kwargs)

          # After the symbol for a new line, we flush the cache.
          if text.endswith('''\n'''):
              SCREAMING_SNAKE_CASE_   :						Any							=	text[self.print_len :]
              SCREAMING_SNAKE_CASE_   :						Any							=	[]
              SCREAMING_SNAKE_CASE_   :						Tuple							=	0
          # If the last token is a CJK character, we print the characters.
          elif len(lowercase_) > 0 and self._is_chinese_char(ord(text[-1])):
              SCREAMING_SNAKE_CASE_   :						Optional[int]							=	text[self.print_len :]
              self.print_len += len(lowercase_)
          # Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words,
          # which may change with the subsequent token -- there are probably smarter ways to do this!)
          else:
              SCREAMING_SNAKE_CASE_   :						str							=	text[self.print_len : text.rfind(''' ''') + 1]
              self.print_len += len(lowercase_)

          self.on_finalized_text(lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          if len(self.token_cache) > 0:
              SCREAMING_SNAKE_CASE_   :						str							=	self.tokenizer.decode(self.token_cache		,			**self.decode_kwargs)
              SCREAMING_SNAKE_CASE_   :						str							=	text[self.print_len :]
              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	[]
              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	0
          else:
              SCREAMING_SNAKE_CASE_   :						Dict							=	''''''

          SCREAMING_SNAKE_CASE_   :						str							=	True
          self.on_finalized_text(lowercase_		,			stream_end=lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[Any]		,			lowercase_  :						str		,			lowercase_  :						bool = False):




          '''simple docstring'''


          print(lowercase_		,			flush=lowercase_		,			end='''''' if not stream_end else None)



      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]		,			lowercase_  :						Tuple):




          '''simple docstring'''


          if (
              (cp >= 0X4E00 and cp <= 0X9FFF)
              or (cp >= 0X3400 and cp <= 0X4DBF)  #
              or (cp >= 0X2_0000 and cp <= 0X2_A6DF)  #
              or (cp >= 0X2_A700 and cp <= 0X2_B73F)  #
              or (cp >= 0X2_B740 and cp <= 0X2_B81F)  #
              or (cp >= 0X2_B820 and cp <= 0X2_CEAF)  #
              or (cp >= 0XF900 and cp <= 0XFAFF)
              or (cp >= 0X2_F800 and cp <= 0X2_FA1F)  #
          ):  #
              return True

          return False




class 							lowerCAmelCase__       (    UpperCAmelCase__				):






      '''simple docstring'''

      def __init__(		self  :						str		,			lowercase_  :						"AutoTokenizer"		,			lowercase_  :						bool = False		,			lowercase_  :						Optional[float] = None		,			**lowercase_  :						str):




          '''simple docstring'''


          super().__init__(lowercase_		,			lowercase_		,			**lowercase_)
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	Queue()
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	None
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	timeout
      def _SCREAMING_SNAKE_CASE				(		self  :						Any		,			lowercase_  :						str		,			lowercase_  :						bool = False):




          '''simple docstring'''


          self.text_queue.put(lowercase_		,			timeout=self.timeout)
          if stream_end:
              self.text_queue.put(self.stop_signal		,			timeout=self.timeout)
      def __iter__(		self  :						Optional[int]):




          '''simple docstring'''


          return self



      def _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	self.text_queue.get(timeout=self.timeout)
          if value == self.stop_signal:
              raise StopIteration()
          else:
              return value










"""simple docstring"""






import unittest

import numpy as np

from transformers.testing_utils import require_flax, require_tf, require_torch
from transformers.utils import (
    expand_dims,
    flatten_dict,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    reshape,
    squeeze,
    transpose,
)


if is_flax_available():
  import jax.numpy as jnp

if is_tf_available():
  import tensorflow as tf

if is_torch_available():
  import torch




class 							lowerCAmelCase__       (    unittest.TestCase				):






      '''simple docstring'''

      def _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						str							=	{
              '''task_specific_params''': {
                  '''summarization''': {'''length_penalty''': 1.0, '''max_length''': 128, '''min_length''': 12, '''num_beams''': 4},
                  '''summarization_cnn''': {'''length_penalty''': 2.0, '''max_length''': 142, '''min_length''': 56, '''num_beams''': 4},
                  '''summarization_xsum''': {'''length_penalty''': 1.0, '''max_length''': 62, '''min_length''': 11, '''num_beams''': 6},
              }
          }
          SCREAMING_SNAKE_CASE_   :						List[str]							=	{
              '''task_specific_params.summarization.length_penalty''': 1.0,
              '''task_specific_params.summarization.max_length''': 128,
              '''task_specific_params.summarization.min_length''': 12,
              '''task_specific_params.summarization.num_beams''': 4,
              '''task_specific_params.summarization_cnn.length_penalty''': 2.0,
              '''task_specific_params.summarization_cnn.max_length''': 142,
              '''task_specific_params.summarization_cnn.min_length''': 56,
              '''task_specific_params.summarization_cnn.num_beams''': 4,
              '''task_specific_params.summarization_xsum.length_penalty''': 1.0,
              '''task_specific_params.summarization_xsum.max_length''': 62,
              '''task_specific_params.summarization_xsum.min_length''': 11,
              '''task_specific_params.summarization_xsum.num_beams''': 6,
          }

          self.assertEqual(flatten_dict(lowercase_)		,			lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	np.random.randn(3		,			4)
          self.assertTrue(np.allclose(transpose(lowercase_)		,			x.transpose()))

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	np.random.randn(3		,			4		,			5)
          self.assertTrue(np.allclose(transpose(lowercase_		,			axes=(1, 2, 0))		,			x.transpose((1, 2, 0))))
      @require_torch
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						str							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						str							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(transpose(lowercase_)		,			transpose(lowercase_).numpy()))

          SCREAMING_SNAKE_CASE_   :						List[Any]							=	np.random.randn(3		,			4		,			5)
          SCREAMING_SNAKE_CASE_   :						Dict							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(transpose(lowercase_		,			axes=(1, 2, 0))		,			transpose(lowercase_		,			axes=(1, 2, 0)).numpy()))
      @require_tf
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						str							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(transpose(lowercase_)		,			transpose(lowercase_).numpy()))

          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	np.random.randn(3		,			4		,			5)
          SCREAMING_SNAKE_CASE_   :						int							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(transpose(lowercase_		,			axes=(1, 2, 0))		,			transpose(lowercase_		,			axes=(1, 2, 0)).numpy()))
      @require_flax
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Tuple							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						Any							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(transpose(lowercase_)		,			np.asarray(transpose(lowercase_))))

          SCREAMING_SNAKE_CASE_   :						Dict							=	np.random.randn(3		,			4		,			5)
          SCREAMING_SNAKE_CASE_   :						Any							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(transpose(lowercase_		,			axes=(1, 2, 0))		,			np.asarray(transpose(lowercase_		,			axes=(1, 2, 0)))))
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Dict							=	np.random.randn(3		,			4)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(4, 3))		,			np.reshape(lowercase_		,			(4, 3))))

          SCREAMING_SNAKE_CASE_   :						Dict							=	np.random.randn(3		,			4		,			5)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(12, 5))		,			np.reshape(lowercase_		,			(12, 5))))
      @require_torch
      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						List[str]							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(4, 3))		,			reshape(lowercase_		,			(4, 3)).numpy()))

          SCREAMING_SNAKE_CASE_   :						int							=	np.random.randn(3		,			4		,			5)
          SCREAMING_SNAKE_CASE_   :						int							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(12, 5))		,			reshape(lowercase_		,			(12, 5)).numpy()))
      @require_tf
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						int							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(4, 3))		,			reshape(lowercase_		,			(4, 3)).numpy()))

          SCREAMING_SNAKE_CASE_   :						List[Any]							=	np.random.randn(3		,			4		,			5)
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(12, 5))		,			reshape(lowercase_		,			(12, 5)).numpy()))
      @require_flax
      def _SCREAMING_SNAKE_CASE				(		self  :						Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Any							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(4, 3))		,			np.asarray(reshape(lowercase_		,			(4, 3)))))

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	np.random.randn(3		,			4		,			5)
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(12, 5))		,			np.asarray(reshape(lowercase_		,			(12, 5)))))
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[Any]							=	np.random.randn(1		,			3		,			4)
          self.assertTrue(np.allclose(squeeze(lowercase_)		,			np.squeeze(lowercase_)))

          SCREAMING_SNAKE_CASE_   :						str							=	np.random.randn(1		,			4		,			1		,			5)
          self.assertTrue(np.allclose(squeeze(lowercase_		,			axis=2)		,			np.squeeze(lowercase_		,			axis=2)))
      @require_torch
      def _SCREAMING_SNAKE_CASE				(		self  :						Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	np.random.randn(1		,			3		,			4)
          SCREAMING_SNAKE_CASE_   :						int							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(squeeze(lowercase_)		,			squeeze(lowercase_).numpy()))

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	np.random.randn(1		,			4		,			1		,			5)
          SCREAMING_SNAKE_CASE_   :						str							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(squeeze(lowercase_		,			axis=2)		,			squeeze(lowercase_		,			axis=2).numpy()))
      @require_tf
      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						int							=	np.random.randn(1		,			3		,			4)
          SCREAMING_SNAKE_CASE_   :						Dict							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(squeeze(lowercase_)		,			squeeze(lowercase_).numpy()))

          SCREAMING_SNAKE_CASE_   :						List[Any]							=	np.random.randn(1		,			4		,			1		,			5)
          SCREAMING_SNAKE_CASE_   :						Any							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(squeeze(lowercase_		,			axis=2)		,			squeeze(lowercase_		,			axis=2).numpy()))
      @require_flax
      def _SCREAMING_SNAKE_CASE				(		self  :						int):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	np.random.randn(1		,			3		,			4)
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(squeeze(lowercase_)		,			np.asarray(squeeze(lowercase_))))

          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	np.random.randn(1		,			4		,			1		,			5)
          SCREAMING_SNAKE_CASE_   :						str							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(squeeze(lowercase_		,			axis=2)		,			np.asarray(squeeze(lowercase_		,			axis=2))))
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						int							=	np.random.randn(3		,			4)
          self.assertTrue(np.allclose(expand_dims(lowercase_		,			axis=1)		,			np.expand_dims(lowercase_		,			axis=1)))
      @require_torch
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						int							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						int							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(expand_dims(lowercase_		,			axis=1)		,			expand_dims(lowercase_		,			axis=1).numpy()))
      @require_tf
      def _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[str]							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						Any							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(expand_dims(lowercase_		,			axis=1)		,			expand_dims(lowercase_		,			axis=1).numpy()))



      @require_flax
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Tuple							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						int							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(expand_dims(lowercase_		,			axis=1)		,			np.asarray(expand_dims(lowercase_		,			axis=1))))


