code: import unittest
  
  import numpy as np
  import timeout_decorator  # noqa
  
  from transformers import OPTConfig, is_flax_available
  from transformers.testing_utils import require_flax, require_sentencepiece, slow
  
  from ...generation.test_flax_utils import FlaxGenerationTesterMixin
  from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor
  
  
  if is_flax_available():
  	import os
  
  	# The slow tests are often failing with OOM error on GPU
  	# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
  	# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
  	lowercase__    :       int						=			'''platform'''
  
  	import jax
  	import jax.numpy as jnp
  
  	from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer
  
  def      SCREAMING_SNAKE_CASE_  (		snake_case__     ,						snake_case__     ,						snake_case__=None     ,						snake_case__=None							)  ->			Tuple:
  							if attention_mask is None:
  														lowerCAmelCase				=	np.where(input_ids != config.pad_token_id     ,						1     ,						0							)
  							return {
  							    "input_ids": input_ids,
  							    "attention_mask": attention_mask,
  							}
  
  
  
  
  @require_flax
  class      lowercase_						:
  
  
  
  	"""simple docstring"""
  
  
  
  
  
  	def __init__(							self	,    __SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE=13	,    __SCREAMING_SNAKE_CASE=7	,    __SCREAMING_SNAKE_CASE=True	,    __SCREAMING_SNAKE_CASE=False	,    __SCREAMING_SNAKE_CASE=99	,    __SCREAMING_SNAKE_CASE=16	,    __SCREAMING_SNAKE_CASE=2	,    __SCREAMING_SNAKE_CASE=4	,    __SCREAMING_SNAKE_CASE=4	,    __SCREAMING_SNAKE_CASE="gelu"	,    __SCREAMING_SNAKE_CASE=0.1	,    __SCREAMING_SNAKE_CASE=0.1	,    __SCREAMING_SNAKE_CASE=20	,    __SCREAMING_SNAKE_CASE=2	,    __SCREAMING_SNAKE_CASE=1	,    __SCREAMING_SNAKE_CASE=0	,    __SCREAMING_SNAKE_CASE=16	,    __SCREAMING_SNAKE_CASE=16	,    __SCREAMING_SNAKE_CASE=0.0_2	,    )		->Optional[int]:
  								lowerCAmelCase				=	parent
  								lowerCAmelCase				=	batch_size
  								lowerCAmelCase				=	seq_length
  								lowerCAmelCase				=	is_training
  								lowerCAmelCase				=	use_labels
  								lowerCAmelCase				=	vocab_size
  								lowerCAmelCase				=	hidden_size
  								lowerCAmelCase				=	num_hidden_layers
  								lowerCAmelCase				=	num_attention_heads
  								lowerCAmelCase				=	intermediate_size
  								lowerCAmelCase				=	hidden_act
  								lowerCAmelCase				=	hidden_dropout_prob
  								lowerCAmelCase				=	attention_probs_dropout_prob
  								lowerCAmelCase				=	max_position_embeddings
  								lowerCAmelCase				=	eos_token_id
  								lowerCAmelCase				=	pad_token_id
  								lowerCAmelCase				=	bos_token_id
  								lowerCAmelCase				=	embed_dim
  								lowerCAmelCase				=	word_embed_proj_dim
  								lowerCAmelCase				=	initializer_range
  								lowerCAmelCase				=	False
  
  
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->List[Any]:
  								lowerCAmelCase				=	np.clip(ids_tensor([self.batch_size, self.seq_length - 1]	,    self.vocab_size       )	,    3	,    self.vocab_size       )
  								lowerCAmelCase				=	np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)	,    dtype=np.intaa       ))	,    -1       )
  
  								lowerCAmelCase				=	OPTConfig(
  								    vocab_size=self.vocab_size	,    hidden_size=self.hidden_size	,    num_hidden_layers=self.num_hidden_layers	,    num_attention_heads=self.num_attention_heads	,    ffn_dim=self.intermediate_size	,    dropout=self.hidden_dropout_prob	,    attention_dropout=self.attention_probs_dropout_prob	,    max_position_embeddings=self.max_position_embeddings	,    eos_token_id=self.eos_token_id	,    bos_token_id=self.bos_token_id	,    pad_token_id=self.pad_token_id	,    embed_dim=self.embed_dim	,    is_encoder_decoder=__SCREAMING_SNAKE_CASE	,    word_embed_proj_dim=self.word_embed_proj_dim	,    initializer_range=self.initializer_range	,    use_cache=__SCREAMING_SNAKE_CASE	,    )
  								lowerCAmelCase				=	prepare_opt_inputs_dict(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )
  								return config, inputs_dict
  
  
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->Optional[int]:
  								lowerCAmelCase	,   lowerCAmelCase				=	self.prepare_config_and_inputs()
  								return config, inputs_dict
  
  
  	def     SCREAMING_SNAKE_CASE_       (							self	,    __SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )		->Optional[int]:
  								lowerCAmelCase				=	20
  								lowerCAmelCase				=	model_class_name(__SCREAMING_SNAKE_CASE       )
  
  								lowerCAmelCase				=	inputs_dict['''input_ids''']
  								lowerCAmelCase				=	inputs_dict['''attention_mask''']
  
  								lowerCAmelCase				=	model.init_cache(input_ids.shape[0]	,    __SCREAMING_SNAKE_CASE       )
  								lowerCAmelCase				=	jnp.ones((input_ids.shape[0], max_length)	,    dtype='''i4'''       )
  
  								lowerCAmelCase				=	jnp.broadcast_to(
  								    jnp.arange(input_ids.shape[-1] - 1       )[None, :]	,    (input_ids.shape[0], input_ids.shape[-1] - 1)	,    )
  								lowerCAmelCase				=	model(
  								    input_ids[:, :-1]	,    attention_mask=__SCREAMING_SNAKE_CASE	,    past_key_values=__SCREAMING_SNAKE_CASE	,    position_ids=__SCREAMING_SNAKE_CASE	,    )
  
  								lowerCAmelCase				=	jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]	,    dtype='''i4'''       )
  								lowerCAmelCase				=	model(
  								    input_ids[:, -1:]	,    attention_mask=__SCREAMING_SNAKE_CASE	,    past_key_values=outputs_cache.past_key_values	,    position_ids=__SCREAMING_SNAKE_CASE	,    )
  
  								lowerCAmelCase				=	model(__SCREAMING_SNAKE_CASE       )
  
  								lowerCAmelCase				=	np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])       )       )
  								self.parent.assertTrue(diff < 1e-3	,    msg=F"Max diff is {diff}"       )
  
  
  
  
  	def     SCREAMING_SNAKE_CASE_       (							self	,    __SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )		->Optional[Any]:
  								lowerCAmelCase				=	20
  								lowerCAmelCase				=	model_class_name(__SCREAMING_SNAKE_CASE       )
  
  								lowerCAmelCase	,   lowerCAmelCase				=	(
  								    inputs_dict['''input_ids'''],
  								    inputs_dict['''attention_mask'''],
  								)
  
  								lowerCAmelCase				=	jnp.concatenate(
  								    [
  								        attention_mask,
  								        jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])       ),
  								    ]	,    axis=-1	,    )
  
  								lowerCAmelCase				=	model.init_cache(input_ids.shape[0]	,    __SCREAMING_SNAKE_CASE       )
  								lowerCAmelCase				=	jnp.broadcast_to(
  								    jnp.arange(input_ids.shape[-1] - 1       )[None, :]	,    (input_ids.shape[0], input_ids.shape[-1] - 1)	,    )
  
  								lowerCAmelCase				=	model(
  								    input_ids[:, :-1]	,    attention_mask=__SCREAMING_SNAKE_CASE	,    past_key_values=__SCREAMING_SNAKE_CASE	,    position_ids=__SCREAMING_SNAKE_CASE	,    )
  								lowerCAmelCase				=	jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]	,    dtype='''i4'''       )
  								lowerCAmelCase				=	model(
  								    input_ids[:, -1:]	,    past_key_values=outputs_cache.past_key_values	,    attention_mask=__SCREAMING_SNAKE_CASE	,    position_ids=__SCREAMING_SNAKE_CASE	,    )
  
  								lowerCAmelCase				=	model(__SCREAMING_SNAKE_CASE	,    attention_mask=__SCREAMING_SNAKE_CASE       )
  
  								lowerCAmelCase				=	np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])       )       )
  								self.parent.assertTrue(diff < 1e-3	,    msg=F"Max diff is {diff}"       )
  
  
  
  
  @require_flax
  class      lowercase_						(						UpperCamelCase_   ,  unittest.TestCase   ,  UpperCamelCase_		):
  
  
  
  	"""simple docstring"""
  
  
  
  
  
  	UpperCAmelCase_						:      List[str]     =       (FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
  	UpperCAmelCase_						:      int     =       () if is_flax_available() else ()
  
  
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->int:
  								lowerCAmelCase				=	FlaxOPTModelTester(self       )
  
  
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->Union[str, Any]:
  								lowerCAmelCase	,   lowerCAmelCase				=	self.model_tester.prepare_config_and_inputs()
  								for model_class in self.all_model_classes:
  															self.model_tester.check_use_cache_forward(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )
  
  
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->Optional[Any]:
  								lowerCAmelCase	,   lowerCAmelCase				=	self.model_tester.prepare_config_and_inputs()
  								for model_class in self.all_model_classes:
  															self.model_tester.check_use_cache_forward_with_attn_mask(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )
  
  
  
  
  	@slow
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->Any:
  								for model_class_name in self.all_model_classes:
  															lowerCAmelCase				=	model_class_name.from_pretrained('''facebook/opt-125m'''       )
  															lowerCAmelCase				=	np.ones((1, 1)       ) * model.config.eos_token_id
  															lowerCAmelCase				=	model(__SCREAMING_SNAKE_CASE       )
  															self.assertIsNotNone(__SCREAMING_SNAKE_CASE       )
  
  
  
  
  @require_sentencepiece
  @require_flax
  class      lowercase_						(						unittest.TestCase		):
  
  
  
  	"""simple docstring"""
  
  
  
  
  
  	@slow
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->Any:
  								lowerCAmelCase				=	FlaxOPTModel.from_pretrained('''facebook/opt-350m'''       )
  								lowerCAmelCase				=	jnp.array([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]]       )
  								lowerCAmelCase				=	model(input_ids=__SCREAMING_SNAKE_CASE       ).last_hidden_state
  								lowerCAmelCase				=	(1, 11, 512)
  								self.assertEqual(output.shape	,    __SCREAMING_SNAKE_CASE       )
  								lowerCAmelCase				=	jnp.array(
  								    [[-0.2_8_6_7, -1.9_2_5_6, -0.3_0_6_2], [-1.2_7_1_1, -0.1_3_3_7, -0.1_8_9_7], [0.4_1_0_9, 0.1_1_8_7, -1.3_1_4_2]]       )
  								self.assertTrue(jnp.allclose(output[:, :3, :3]	,    __SCREAMING_SNAKE_CASE	,    atol=4e-2       )       )
  
  
  
  
  @require_flax
  @slow
  class      lowercase_						(						unittest.TestCase		):
  
  
  
  	"""simple docstring"""
  
  
  
  
  
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->Union[str, Any]:
  								super().setUp()
  								lowerCAmelCase				=	'''facebook/opt-350m'''
  
  
  
  
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->Optional[Any]:
  								lowerCAmelCase				=	FlaxOPTForCausalLM.from_pretrained(self.path_model       )
  								lowerCAmelCase				=	GPTaTokenizer.from_pretrained(self.path_model       )
  
  								lowerCAmelCase				=	[
  								    '''Today is a beautiful day and I want to''',
  								    '''In the city of''',
  								    '''Paris is the capital of France and''',
  								    '''Computers and mobile phones have taken''',
  								]
  								# verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
  								lowerCAmelCase				=	tokenizer(__SCREAMING_SNAKE_CASE	,    return_tensors='''jax'''	,    padding=__SCREAMING_SNAKE_CASE	,    add_special_tokens=__SCREAMING_SNAKE_CASE       )
  								lowerCAmelCase				=	model(inputs.input_ids	,    attention_mask=inputs.attention_mask       )[0].mean(axis=-1       )
  								lowerCAmelCase				=	jnp.array(
  								    [
  								        [1.3_8_5_1, -1_3.8_9_2_3, -1_0.5_2_2_9, -1_0.7_5_3_3, -0.2_3_0_9, -1_0.2_3_8_4, -0.5_3_6_5, -9.0_9_4_7, -5.1_6_7_0],
  								        [-4.7_0_7_3, -1_0.6_2_7_6, -3.9_4_1_5, -2_1.5_2_4_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2, -0.2_8_2_2],
  								        [0.6_2_4_7, -3.4_2_2_9, -8.9_1_7_9, -1.4_2_9_7, -1_4.1_6_5_0, 1.4_1_4_6, -9.0_2_1_8, -0.2_7_0_3, -0.2_7_0_3],
  								        [6.4_7_8_3, -1.9_9_1_3, -1_0.7_9_2_6, -2.3_3_3_6, 1.5_0_9_2, -0.9_9_7_4, -6.8_2_1_3, 1.3_4_7_7, 1.3_4_7_7],
  								    ]       )
  								self.assertTrue(jnp.allclose(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    atol=4e-2       )       )
  
  								lowerCAmelCase				=	jax.jit(__SCREAMING_SNAKE_CASE       )
  								lowerCAmelCase				=	model(inputs.input_ids	,    attention_mask=inputs.attention_mask       )[0].mean(axis=-1       )
  								self.assertTrue(jnp.allclose(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    atol=4e-2       )       )
  
  
  
  
  @require_flax
  @slow
  class      lowercase_						(						unittest.TestCase		):
  
  
  
  	"""simple docstring"""
  
  
  
  
  
  	@property
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->Dict:
  								return [
  								    "Today is a beautiful day and I want",
  								    "In the city of",
  								    "Paris is the capital of France and",
  								    "Computers and mobile phones have taken",
  								]
  
  
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->Optional[int]:
  								lowerCAmelCase				=	'''facebook/opt-125m'''
  
  								lowerCAmelCase				=	[
  								    '''Today is a beautiful day and I want to''',
  								    '''In the city of New York, the city''',
  								    '''Paris is the capital of France and the capital''',
  								    '''Computers and mobile phones have taken over the''',
  								]
  
  								lowerCAmelCase				=	[]
  
  								lowerCAmelCase				=	FlaxOPTForCausalLM.from_pretrained(__SCREAMING_SNAKE_CASE       )
  								lowerCAmelCase				=	GPTaTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE       )
  
  								for prompt in self.prompts:
  															lowerCAmelCase				=	tokenizer(__SCREAMING_SNAKE_CASE	,    return_tensors='''jax'''       ).input_ids
  
  															lowerCAmelCase				=	model.generate(__SCREAMING_SNAKE_CASE	,    max_length=10       )
  															lowerCAmelCase				=	generated_ids[0]
  
  															lowerCAmelCase				=	tokenizer.batch_decode(__SCREAMING_SNAKE_CASE	,    skip_special_tokens=__SCREAMING_SNAKE_CASE       )
  															predicted_outputs += generated_string
  
  								self.assertListEqual(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )
  
  
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->str:
  								lowerCAmelCase				=	'''facebook/opt-350m'''
  
  								lowerCAmelCase				=	[
  								    '''Today is a beautiful day and I want to''',
  								    '''In the city of San Francisco, the city''',
  								    '''Paris is the capital of France and the capital''',
  								    '''Computers and mobile phones have taken over the''',
  								]
  
  								lowerCAmelCase				=	[]
  								lowerCAmelCase				=	FlaxOPTForCausalLM.from_pretrained(__SCREAMING_SNAKE_CASE       )
  								lowerCAmelCase				=	GPTaTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE       )
  
  								for prompt in self.prompts:
  															lowerCAmelCase				=	tokenizer(__SCREAMING_SNAKE_CASE	,    return_tensors='''jax'''       ).input_ids
  
  															lowerCAmelCase				=	model.generate(__SCREAMING_SNAKE_CASE	,    max_length=10       )
  															lowerCAmelCase				=	generated_ids[0]
  
  															lowerCAmelCase				=	tokenizer.batch_decode(__SCREAMING_SNAKE_CASE	,    skip_special_tokens=__SCREAMING_SNAKE_CASE       )
  															predicted_outputs += generated_string
  
  								self.assertListEqual(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )
  
  
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->Any:
  								lowerCAmelCase				=	'''facebook/opt-125m'''
  								lowerCAmelCase				=	[
  								    '''Today is a beautiful day and I want to thank''',
  								    '''In the city of Rome Canaver Canaver Canaver Canaver''',
  								]
  								lowerCAmelCase				=	FlaxOPTForCausalLM.from_pretrained(__SCREAMING_SNAKE_CASE       )
  								lowerCAmelCase				=	GPTaTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE       )
  								lowerCAmelCase				=	tokenizer(
  								    [
  								        '''Today is a beautiful day and I want to''',
  								        '''In the city of''',
  								    ]	,    return_tensors='''jax'''	,    padding=__SCREAMING_SNAKE_CASE	,    )
  
  								lowerCAmelCase				=	jax.jit(model.generate       )
  
  								lowerCAmelCase				=	jit_generate(inputs['''input_ids''']	,    attention_mask=inputs['''attention_mask''']       ).sequences
  
  								lowerCAmelCase				=	tokenizer.batch_decode(__SCREAMING_SNAKE_CASE	,    skip_special_tokens=__SCREAMING_SNAKE_CASE       )
  
  								self.assertIsNotNone(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )
  
  
  
  
  	def     SCREAMING_SNAKE_CASE_       (							self       )		->Any:
  								lowerCAmelCase				=	'''facebook/opt-350m'''
  
  								lowerCAmelCase				=	GPTaTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE       )
  								lowerCAmelCase				=	FlaxOPTForCausalLM.from_pretrained(__SCREAMING_SNAKE_CASE       )
  
  								lowerCAmelCase				=	'''left'''
  
  								# use different length sentences to test batching
  								lowerCAmelCase				=	[
  								    '''Hello, my dog is a little''',
  								    '''Today, I''',
  								]
  
  								lowerCAmelCase				=	tokenizer(__SCREAMING_SNAKE_CASE	,    return_tensors='''jax'''	,    padding=__SCREAMING_SNAKE_CASE       )
  								lowerCAmelCase				=	inputs['''input_ids''']
  
  								lowerCAmelCase				=	model.generate(input_ids=__SCREAMING_SNAKE_CASE	,    attention_mask=inputs['''attention_mask''']	,    trace=__SCREAMING_SNAKE_CASE       )
  
  								lowerCAmelCase				=	tokenizer(sentences[0]	,    return_tensors='''jax'''       ).input_ids
  								lowerCAmelCase				=	model.generate(input_ids=__SCREAMING_SNAKE_CASE       )
  
  								lowerCAmelCase				=	inputs_non_padded.shape[-1] - inputs['''attention_mask'''][-1].sum()
  								lowerCAmelCase				=	tokenizer(sentences[1]	,    return_tensors='''jax'''       ).input_ids
  								lowerCAmelCase				=	model.generate(input_ids=__SCREAMING_SNAKE_CASE	,    max_length=model.config.max_length - num_paddings       )
  
  								lowerCAmelCase				=	tokenizer.batch_decode(outputs[0]	,    skip_special_tokens=__SCREAMING_SNAKE_CASE       )
  								lowerCAmelCase				=	tokenizer.decode(output_non_padded[0][0]	,    skip_special_tokens=__SCREAMING_SNAKE_CASE       )
  								lowerCAmelCase				=	tokenizer.decode(output_padded[0][0]	,    skip_special_tokens=__SCREAMING_SNAKE_CASE       )
  
  								lowerCAmelCase				=	[
  								    '''Hello, my dog is a little bit of a dork.\nI\'m a little bit''',
  								    '''Today, I was in the middle of a conversation with a friend about the''',
  								]
  								self.assertListEqual(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )
  								self.assertListEqual(__SCREAMING_SNAKE_CASE	,    [non_padded_sentence, padded_sentence]       )
  
  
  
  
style_context: from dataclasses import dataclass
  from typing import List, Optional, Union
  
  import numpy as np
  import PIL
  
  from ...utils import BaseOutput, is_invisible_watermark_available, is_torch_available, is_transformers_available
  
  
  
  
  @dataclass
  class      lowercase_						(						UpperCamelCase_		):
  
  
  
  	"""simple docstring"""
  
  
  
  
  
  	UpperCAmelCase_						:      Union[List[PIL.Image.Image], np.ndarray]
  
  
  if is_transformers_available() and is_torch_available() and is_invisible_watermark_available():
  	from .pipeline_stable_diffusion_xl import StableDiffusionXLPipeline
  	from .pipeline_stable_diffusion_xl_imgaimg import StableDiffusionXLImgaImgPipeline
  
  
  
  