

"""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      _UpperCAmelCase   (							unittest.TestCase       ):



	def        a  (  self     :       Union[str, Any]						):
		__UpperCAmelCase      		= AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''						)
		__UpperCAmelCase      		= AutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''						).to(_lowercase						)
		__UpperCAmelCase      		= -1

		__UpperCAmelCase      		= ids_tensor((1, 5)			,    vocab_size=model.config.vocab_size						).to(_lowercase						)
		__UpperCAmelCase      		= model.generate(_lowercase			,    max_new_tokens=10			,    do_sample=_lowercase						)
		__UpperCAmelCase      		= tokenizer.decode(greedy_ids[0]						)

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

		self.assertEqual(_lowercase			,    _lowercase						)



	def        a  (  self     :       Optional[Any]						):
		__UpperCAmelCase      		= AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''						)
		__UpperCAmelCase      		= AutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''						).to(_lowercase						)
		__UpperCAmelCase      		= -1

		__UpperCAmelCase      		= ids_tensor((1, 5)			,    vocab_size=model.config.vocab_size						).to(_lowercase						)
		__UpperCAmelCase      		= model.generate(_lowercase			,    max_new_tokens=10			,    do_sample=_lowercase						)
		__UpperCAmelCase      		= tokenizer.decode(greedy_ids[0]						)

		__UpperCAmelCase      		= TextIteratorStreamer(_lowercase						)
		__UpperCAmelCase      		= {'''input_ids''': input_ids, '''max_new_tokens''': 10, '''do_sample''': False, '''streamer''': streamer}
		__UpperCAmelCase      		= Thread(target=model.generate			,    kwargs=_lowercase						)
		thread.start()
		__UpperCAmelCase      		= ''''''
		for new_text in streamer:
			streamer_text += new_text

		self.assertEqual(_lowercase			,    _lowercase						)



	def        a  (  self     :       str						):
		__UpperCAmelCase      		= AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''						)
		__UpperCAmelCase      		= AutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''						).to(_lowercase						)
		__UpperCAmelCase      		= -1

		__UpperCAmelCase      		= ids_tensor((1, 5)			,    vocab_size=model.config.vocab_size						).to(_lowercase						)
		__UpperCAmelCase      		= model.generate(_lowercase			,    max_new_tokens=10			,    do_sample=_lowercase						)
		__UpperCAmelCase      		= greedy_ids[:, input_ids.shape[1] :]
		__UpperCAmelCase      		= tokenizer.decode(new_greedy_ids[0]						)

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

		self.assertEqual(_lowercase			,    _lowercase						)



	def        a  (  self     :       Tuple						):
		# Tests that we can pass `decode_kwargs` to the streamer to control how the tokens are decoded. Must be tested
		# with actual models -- the dummy models' tokenizers are not aligned with their models, and
		# `skip_special_tokens=True` has no effect on them
		__UpperCAmelCase      		= AutoTokenizer.from_pretrained('''distilgpt2'''						)
		__UpperCAmelCase      		= AutoModelForCausalLM.from_pretrained('''distilgpt2'''						).to(_lowercase						)
		__UpperCAmelCase      		= -1

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

		# 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
		__UpperCAmelCase      		= cs.out[:-1]  # Remove the final "\n"
		__UpperCAmelCase      		= tokenizer(_lowercase			,    return_tensors='''pt'''						)
		self.assertEqual(streamer_text_tokenized.input_ids.shape			,    (1, 1)						)



	def        a  (  self     :       Tuple						):
		__UpperCAmelCase      		= AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''						)
		__UpperCAmelCase      		= AutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''						).to(_lowercase						)
		__UpperCAmelCase      		= -1

		__UpperCAmelCase      		= ids_tensor((1, 5)			,    vocab_size=model.config.vocab_size						).to(_lowercase						)
		__UpperCAmelCase      		= TextIteratorStreamer(_lowercase			,    timeout=0.001						)
		__UpperCAmelCase      		= {'''input_ids''': input_ids, '''max_new_tokens''': 10, '''do_sample''': False, '''streamer''': streamer}
		__UpperCAmelCase      		= Thread(target=model.generate			,    kwargs=_lowercase						)
		thread.start()

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



import flax.linen as nn
import jax
import jax.numpy as jnp





class 					__lowerCAmelCase	(   nn.Module				):
  _a		     =    42
  _a		     =    jnp.floataa




  def     A__	(							self   )    ->					int:
        '''simple docstring'''






        _lowercase      			=nn.Conv(
            self.out_channels				,  kernel_size=(3, 3)				,  strides=(1, 1)				,  padding=((1, 1), (1, 1))				,  dtype=self.dtype				,  )




  def __call__(							self				,  lowerCAmelCase   )    ->					Any:
        '''simple docstring'''






        _lowercase			,			_lowercase			,			_lowercase			,			_lowercase      			=hidden_states.shape
        _lowercase      			=jax.image.resize(
            lowerCAmelCase				,  shape=(batch, height * 2, width * 2, channels)				,  method='nearest'				,  )
        _lowercase      			=self.conv(lowerCAmelCase   )
        return hidden_states





class 					__lowerCAmelCase	(   nn.Module				):
  _a		     =    42
  _a		     =    jnp.floataa




  def     A__	(							self   )    ->					Any:
        '''simple docstring'''






        _lowercase      			=nn.Conv(
            self.out_channels				,  kernel_size=(3, 3)				,  strides=(2, 2)				,  padding=((1, 1), (1, 1))				,  dtype=self.dtype				,  )




  def __call__(							self				,  lowerCAmelCase   )    ->					Optional[int]:
        '''simple docstring'''






        _lowercase      			=self.conv(lowerCAmelCase   )
        return hidden_states





class 					__lowerCAmelCase	(   nn.Module				):
  _a		     =    42
  _a		     =    None
  _a		     =    0.0
  _a		     =    None
  _a		     =    jnp.floataa




  def     A__	(							self   )    ->					Optional[int]:
        '''simple docstring'''






        _lowercase      			=self.in_channels if self.out_channels is None else self.out_channels

        _lowercase      			=nn.GroupNorm(num_groups=32				,  epsilon=1e-5   )
        _lowercase      			=nn.Conv(
            lowerCAmelCase				,  kernel_size=(3, 3)				,  strides=(1, 1)				,  padding=((1, 1), (1, 1))				,  dtype=self.dtype				,  )

        _lowercase      			=nn.Dense(lowerCAmelCase				,  dtype=self.dtype   )

        _lowercase      			=nn.GroupNorm(num_groups=32				,  epsilon=1e-5   )
        _lowercase      			=nn.Dropout(self.dropout_prob   )
        _lowercase      			=nn.Conv(
            lowerCAmelCase				,  kernel_size=(3, 3)				,  strides=(1, 1)				,  padding=((1, 1), (1, 1))				,  dtype=self.dtype				,  )

        _lowercase      			=self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

        _lowercase      			=None
        if use_nin_shortcut:
              _lowercase      			=nn.Conv(
                  lowerCAmelCase				,  kernel_size=(1, 1)				,  strides=(1, 1)				,  padding='VALID'				,  dtype=self.dtype				,  )




  def __call__(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase=True   )    ->					Tuple:
        '''simple docstring'''






        _lowercase      			=hidden_states
        _lowercase      			=self.norma(lowerCAmelCase   )
        _lowercase      			=nn.swish(lowerCAmelCase   )
        _lowercase      			=self.conva(lowerCAmelCase   )

        _lowercase      			=self.time_emb_proj(nn.swish(lowerCAmelCase   )   )
        _lowercase      			=jnp.expand_dims(jnp.expand_dims(lowerCAmelCase				,  1   )				,  1   )
        _lowercase      			=hidden_states + temb

        _lowercase      			=self.norma(lowerCAmelCase   )
        _lowercase      			=nn.swish(lowerCAmelCase   )
        _lowercase      			=self.dropout(lowerCAmelCase				,  lowerCAmelCase   )
        _lowercase      			=self.conva(lowerCAmelCase   )

        if self.conv_shortcut is not None:
              _lowercase      			=self.conv_shortcut(lowerCAmelCase   )

        return hidden_states + residual


