





'''simple docstring'''

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

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD
class 				__lowerCAmelCase     (    nn.Module ):






       """simple docstring"""
       _snake_case	:     Optional[int]										=    4_2
       _snake_case	:     Dict										=    4_2
       _snake_case	:     Dict										=    0.0
       _snake_case	:     Tuple										=    1
       _snake_case	:     Tuple										=    1
       _snake_case	:     int										=    True
       _snake_case	:     Dict										=    False
       _snake_case	:     Optional[Any]										=    False
       _snake_case	:     List[Any]										=    False
       _snake_case	:     List[str]										=    jnp.floataa





       def snake_case__						(							self   :		Optional[int]     )    ->     str:

            '''simple docstring'''




            _UpperCamelCase				    =       []
            _UpperCamelCase				    =       []

            for i in range(self.num_layers     ):
                 _UpperCamelCase				    =       self.in_channels if i == 0 else self.out_channels

                 _UpperCamelCase				    =       FlaxResnetBlockaD(
                     in_channels=_a       , out_channels=self.out_channels       , dropout_prob=self.dropout       , dtype=self.dtype       , )
                 resnets.append(_a     )

                 _UpperCamelCase				    =       FlaxTransformeraDModel(
                     in_channels=self.out_channels       , n_heads=self.num_attention_heads       , d_head=self.out_channels // self.num_attention_heads       , depth=1       , use_linear_projection=self.use_linear_projection       , only_cross_attention=self.only_cross_attention       , use_memory_efficient_attention=self.use_memory_efficient_attention       , dtype=self.dtype       , )
                 attentions.append(_a     )

            _UpperCamelCase				    =       resnets
            _UpperCamelCase				    =       attentions

            if self.add_downsample:
                 _UpperCamelCase				    =       FlaxDownsampleaD(self.out_channels       , dtype=self.dtype     )







       def __call__(							self   :		List[Any]       , lowerCAmelCase__   :		Dict       , lowerCAmelCase__   :		Tuple       , lowerCAmelCase__   :		Any       , lowerCAmelCase__   :		List[Any]=True     )    ->     List[Any]:

            '''simple docstring'''




            _UpperCamelCase				    =       ()

            for resnet, attn in zip(self.resnets       , self.attentions     ):
                 _UpperCamelCase				    =       resnet(_a       , _a       , deterministic=_a     )
                 _UpperCamelCase				    =       attn(_a       , _a       , deterministic=_a     )
                 output_states += (hidden_states,)

            if self.add_downsample:
                 _UpperCamelCase				    =       self.downsamplers_a(_a     )
                 output_states += (hidden_states,)

            return hidden_states, output_states

class 				__lowerCAmelCase     (    nn.Module ):






       """simple docstring"""
       _snake_case	:     Tuple										=    4_2
       _snake_case	:     List[Any]										=    4_2
       _snake_case	:     Tuple										=    0.0
       _snake_case	:     Optional[int]										=    1
       _snake_case	:     List[str]										=    True
       _snake_case	:     str										=    jnp.floataa





       def snake_case__						(							self   :		Union[str, Any]     )    ->     Optional[Any]:

            '''simple docstring'''




            _UpperCamelCase				    =       []

            for i in range(self.num_layers     ):
                 _UpperCamelCase				    =       self.in_channels if i == 0 else self.out_channels

                 _UpperCamelCase				    =       FlaxResnetBlockaD(
                     in_channels=_a       , out_channels=self.out_channels       , dropout_prob=self.dropout       , dtype=self.dtype       , )
                 resnets.append(_a     )
            _UpperCamelCase				    =       resnets

            if self.add_downsample:
                 _UpperCamelCase				    =       FlaxDownsampleaD(self.out_channels       , dtype=self.dtype     )







       def __call__(							self   :		Tuple       , lowerCAmelCase__   :		Dict       , lowerCAmelCase__   :		Union[str, Any]       , lowerCAmelCase__   :		int=True     )    ->     List[str]:

            '''simple docstring'''




            _UpperCamelCase				    =       ()

            for resnet in self.resnets:
                 _UpperCamelCase				    =       resnet(_a       , _a       , deterministic=_a     )
                 output_states += (hidden_states,)

            if self.add_downsample:
                 _UpperCamelCase				    =       self.downsamplers_a(_a     )
                 output_states += (hidden_states,)

            return hidden_states, output_states

class 				__lowerCAmelCase     (    nn.Module ):






       """simple docstring"""
       _snake_case	:     Dict										=    4_2
       _snake_case	:     Union[str, Any]										=    4_2
       _snake_case	:     List[str]										=    4_2
       _snake_case	:     List[Any]										=    0.0
       _snake_case	:     int										=    1
       _snake_case	:     Tuple										=    1
       _snake_case	:     List[str]										=    True
       _snake_case	:     Dict										=    False
       _snake_case	:     Optional[Any]										=    False
       _snake_case	:     int										=    False
       _snake_case	:     int										=    jnp.floataa





       def snake_case__						(							self   :		Union[str, Any]     )    ->     Union[str, Any]:

            '''simple docstring'''




            _UpperCamelCase				    =       []
            _UpperCamelCase				    =       []

            for i in range(self.num_layers     ):
                 _UpperCamelCase				    =       self.in_channels if (i == self.num_layers - 1) else self.out_channels
                 _UpperCamelCase				    =       self.prev_output_channel if i == 0 else self.out_channels

                 _UpperCamelCase				    =       FlaxResnetBlockaD(
                     in_channels=resnet_in_channels + res_skip_channels       , out_channels=self.out_channels       , dropout_prob=self.dropout       , dtype=self.dtype       , )
                 resnets.append(_a     )

                 _UpperCamelCase				    =       FlaxTransformeraDModel(
                     in_channels=self.out_channels       , n_heads=self.num_attention_heads       , d_head=self.out_channels // self.num_attention_heads       , depth=1       , use_linear_projection=self.use_linear_projection       , only_cross_attention=self.only_cross_attention       , use_memory_efficient_attention=self.use_memory_efficient_attention       , dtype=self.dtype       , )
                 attentions.append(_a     )

            _UpperCamelCase				    =       resnets
            _UpperCamelCase				    =       attentions

            if self.add_upsample:
                 _UpperCamelCase				    =       FlaxUpsampleaD(self.out_channels       , dtype=self.dtype     )







       def __call__(							self   :		str       , lowerCAmelCase__   :		int       , lowerCAmelCase__   :		Tuple       , lowerCAmelCase__   :		Any       , lowerCAmelCase__   :		Optional[Any]       , lowerCAmelCase__   :		str=True     )    ->     Tuple:

            '''simple docstring'''




            for resnet, attn in zip(self.resnets       , self.attentions     ):
                 # pop res hidden states
                 _UpperCamelCase				    =       res_hidden_states_tuple[-1]
                 _UpperCamelCase				    =       res_hidden_states_tuple[:-1]
                 _UpperCamelCase				    =       jnp.concatenate((hidden_states, res_hidden_states)       , axis=-1     )

                 _UpperCamelCase				    =       resnet(_a       , _a       , deterministic=_a     )
                 _UpperCamelCase				    =       attn(_a       , _a       , deterministic=_a     )

            if self.add_upsample:
                 _UpperCamelCase				    =       self.upsamplers_a(_a     )

            return hidden_states

class 				__lowerCAmelCase     (    nn.Module ):






       """simple docstring"""
       _snake_case	:     Tuple										=    4_2
       _snake_case	:     str										=    4_2
       _snake_case	:     Dict										=    4_2
       _snake_case	:     Optional[Any]										=    0.0
       _snake_case	:     Any										=    1
       _snake_case	:     Tuple										=    True
       _snake_case	:     Any										=    jnp.floataa





       def snake_case__						(							self   :		Tuple     )    ->     str:

            '''simple docstring'''




            _UpperCamelCase				    =       []

            for i in range(self.num_layers     ):
                 _UpperCamelCase				    =       self.in_channels if (i == self.num_layers - 1) else self.out_channels
                 _UpperCamelCase				    =       self.prev_output_channel if i == 0 else self.out_channels

                 _UpperCamelCase				    =       FlaxResnetBlockaD(
                     in_channels=resnet_in_channels + res_skip_channels       , out_channels=self.out_channels       , dropout_prob=self.dropout       , dtype=self.dtype       , )
                 resnets.append(_a     )

            _UpperCamelCase				    =       resnets

            if self.add_upsample:
                 _UpperCamelCase				    =       FlaxUpsampleaD(self.out_channels       , dtype=self.dtype     )







       def __call__(							self   :		int       , lowerCAmelCase__   :		List[str]       , lowerCAmelCase__   :		Any       , lowerCAmelCase__   :		List[Any]       , lowerCAmelCase__   :		str=True     )    ->     str:

            '''simple docstring'''




            for resnet in self.resnets:
                 # pop res hidden states
                 _UpperCamelCase				    =       res_hidden_states_tuple[-1]
                 _UpperCamelCase				    =       res_hidden_states_tuple[:-1]
                 _UpperCamelCase				    =       jnp.concatenate((hidden_states, res_hidden_states)       , axis=-1     )

                 _UpperCamelCase				    =       resnet(_a       , _a       , deterministic=_a     )

            if self.add_upsample:
                 _UpperCamelCase				    =       self.upsamplers_a(_a     )

            return hidden_states

class 				__lowerCAmelCase     (    nn.Module ):






       """simple docstring"""
       _snake_case	:     Dict										=    4_2
       _snake_case	:     Union[str, Any]										=    0.0
       _snake_case	:     Any										=    1
       _snake_case	:     List[Any]										=    1
       _snake_case	:     Any										=    False
       _snake_case	:     Tuple										=    False
       _snake_case	:     Union[str, Any]										=    jnp.floataa





       def snake_case__						(							self   :		Union[str, Any]     )    ->     Union[str, Any]:

            '''simple docstring'''




            _UpperCamelCase				    =       [
                FlaxResnetBlockaD(
                    in_channels=self.in_channels       , out_channels=self.in_channels       , dropout_prob=self.dropout       , dtype=self.dtype       , )
            ]

            _UpperCamelCase				    =       []

            for _ in range(self.num_layers     ):
                 _UpperCamelCase				    =       FlaxTransformeraDModel(
                     in_channels=self.in_channels       , n_heads=self.num_attention_heads       , d_head=self.in_channels // self.num_attention_heads       , depth=1       , use_linear_projection=self.use_linear_projection       , use_memory_efficient_attention=self.use_memory_efficient_attention       , dtype=self.dtype       , )
                 attentions.append(_a     )

                 _UpperCamelCase				    =       FlaxResnetBlockaD(
                     in_channels=self.in_channels       , out_channels=self.in_channels       , dropout_prob=self.dropout       , dtype=self.dtype       , )
                 resnets.append(_a     )

            _UpperCamelCase				    =       resnets
            _UpperCamelCase				    =       attentions







       def __call__(							self   :		Optional[int]       , lowerCAmelCase__   :		List[Any]       , lowerCAmelCase__   :		int       , lowerCAmelCase__   :		List[str]       , lowerCAmelCase__   :		Dict=True     )    ->     int:

            '''simple docstring'''




            _UpperCamelCase				    =       self.resnets[0](_a       , _a     )
            for attn, resnet in zip(self.attentions       , self.resnets[1:]     ):
                 _UpperCamelCase				    =       attn(_a       , _a       , deterministic=_a     )
                 _UpperCamelCase				    =       resnet(_a       , _a       , deterministic=_a     )

            return hidden_states









"""simple docstring"""




import unittest

from transformers import (
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TextaTextGenerationPipeline,
    pipeline,
)
from transformers.testing_utils import is_pipeline_test, require_tf, require_torch
from transformers.utils import is_torch_available

from .test_pipelines_common import ANY


if is_torch_available():
					import torch






@is_pipeline_test
class     __magic_name__							(       unittest.TestCase				):






				'''simple docstring'''
				__UpperCamelCase			  =						MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
				__UpperCamelCase			  =						TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING





				def    _lowerCAmelCase				(       self		, _a		, _a		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					TextaTextGenerationPipeline(model=_a		, tokenizer=_a )
							return generator, ["Something to write", "Something else"]





				def    _lowerCAmelCase				(       self		, _a		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					generator("""Something there""" )
							self.assertEqual(_a		, [{"""generated_text""": ANY(_a )}] )
							# These are encoder decoder, they don't just append to incoming string
							self.assertFalse(outputs[0]["""generated_text"""].startswith("""Something there""" ) )

							lowerCamelCase			      =					generator(["""This is great !""", """Something else"""]		, num_return_sequences=2		, do_sample=_a )
							self.assertEqual(
							    _a		, [
							        [{"""generated_text""": ANY(_a )}, {"""generated_text""": ANY(_a )}],
							        [{"""generated_text""": ANY(_a )}, {"""generated_text""": ANY(_a )}],
							    ]		, )

							lowerCamelCase			      =					generator(
							    ["""This is great !""", """Something else"""]		, num_return_sequences=2		, batch_size=2		, do_sample=_a )
							self.assertEqual(
							    _a		, [
							        [{"""generated_text""": ANY(_a )}, {"""generated_text""": ANY(_a )}],
							        [{"""generated_text""": ANY(_a )}, {"""generated_text""": ANY(_a )}],
							    ]		, )

							with self.assertRaises(_a ):
										generator(4 )





				@require_torch
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					pipeline("""text2text-generation"""		, model="""patrickvonplaten/t5-tiny-random"""		, framework="""pt""" )
							# do_sample=False necessary for reproducibility
							lowerCamelCase			      =					generator("""Something there"""		, do_sample=_a )
							self.assertEqual(_a		, [{"""generated_text""": """"""}] )

							lowerCamelCase			      =					3
							lowerCamelCase			      =					generator(
							    """Something there"""		, num_return_sequences=_a		, num_beams=_a		, )
							lowerCamelCase			      =					[
							    {"""generated_text""": """Beide Beide Beide Beide Beide Beide Beide Beide Beide"""},
							    {"""generated_text""": """Beide Beide Beide Beide Beide Beide Beide Beide"""},
							    {"""generated_text""": """"""},
							]
							self.assertEqual(_a		, _a )

							lowerCamelCase			      =					generator("""This is a test"""		, do_sample=_a		, num_return_sequences=2		, return_tensors=_a )
							self.assertEqual(
							    _a		, [
							        {"""generated_token_ids""": ANY(torch.Tensor )},
							        {"""generated_token_ids""": ANY(torch.Tensor )},
							    ]		, )
							lowerCamelCase			      =					generator.model.config.eos_token_id
							lowerCamelCase			      =					"""<pad>"""
							lowerCamelCase			      =					generator(
							    ["""This is a test""", """This is a second test"""]		, do_sample=_a		, num_return_sequences=2		, batch_size=2		, return_tensors=_a		, )
							self.assertEqual(
							    _a		, [
							        [
							            {"""generated_token_ids""": ANY(torch.Tensor )},
							            {"""generated_token_ids""": ANY(torch.Tensor )},
							        ],
							        [
							            {"""generated_token_ids""": ANY(torch.Tensor )},
							            {"""generated_token_ids""": ANY(torch.Tensor )},
							        ],
							    ]		, )





				@require_tf
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					pipeline("""text2text-generation"""		, model="""patrickvonplaten/t5-tiny-random"""		, framework="""tf""" )
							# do_sample=False necessary for reproducibility
							lowerCamelCase			      =					generator("""Something there"""		, do_sample=_a )
							self.assertEqual(_a		, [{"""generated_text""": """"""}] )



