





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
  _snake_case				= """platform"""

  import jax
  import jax.numpy as jnp

  from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer


def _A   (    __magic_name__  ,					__magic_name__  ,					__magic_name__=None  ,					__magic_name__=None       ):
   if attention_mask is None:
      lowercase__    =  np.where(input_ids != config.pad_token_id  ,					1  ,					0       )
   return {
       "input_ids": input_ids,
       "attention_mask": attention_mask,
   }





@require_flax
class 					lowerCAmelCase							:







  def __init__(			self  :int    ,							_lowercase  :Tuple    ,							_lowercase  :int=13    ,							_lowercase  :Union[str, Any]=7    ,							_lowercase  :Any=True    ,							_lowercase  :Optional[int]=False    ,							_lowercase  :Optional[Any]=99    ,							_lowercase  :List[str]=16    ,							_lowercase  :Any=2    ,							_lowercase  :str=4    ,							_lowercase  :Any=4    ,							_lowercase  :Optional[Any]="gelu"    ,							_lowercase  :List[Any]=0.1    ,							_lowercase  :Tuple=0.1    ,							_lowercase  :Dict=20    ,							_lowercase  :List[str]=2    ,							_lowercase  :Any=1    ,							_lowercase  :Union[str, Any]=0    ,							_lowercase  :Optional[int]=16    ,							_lowercase  :int=16    ,							_lowercase  :Dict=0.02    ,							):






     '''simple docstring'''





     lowercase__    =  parent
     lowercase__    =  batch_size
     lowercase__    =  seq_length
     lowercase__    =  is_training
     lowercase__    =  use_labels
     lowercase__    =  vocab_size
     lowercase__    =  hidden_size
     lowercase__    =  num_hidden_layers
     lowercase__    =  num_attention_heads
     lowercase__    =  intermediate_size
     lowercase__    =  hidden_act
     lowercase__    =  hidden_dropout_prob
     lowercase__    =  attention_probs_dropout_prob
     lowercase__    =  max_position_embeddings
     lowercase__    =  eos_token_id
     lowercase__    =  pad_token_id
     lowercase__    =  bos_token_id
     lowercase__    =  embed_dim
     lowercase__    =  word_embed_proj_dim
     lowercase__    =  initializer_range
     lowercase__    =  False







  def   UpperCAmelCase     (			self  :Optional[Any]					):






     '''simple docstring'''





     lowercase__    =  np.clip(ids_tensor([self.batch_size, self.seq_length - 1]    ,							self.vocab_size					)    ,							3    ,							self.vocab_size					)
     lowercase__    =  np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)    ,							dtype=np.intaa					))    ,							-1					)

     lowercase__    =  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=_lowercase    ,							word_embed_proj_dim=self.word_embed_proj_dim    ,							initializer_range=self.initializer_range    ,							use_cache=_lowercase    ,							)
     lowercase__    =  prepare_opt_inputs_dict(_lowercase    ,							_lowercase					)
     return config, inputs_dict







  def   UpperCAmelCase     (			self  :Union[str, Any]					):






     '''simple docstring'''





     lowercase__   ,							lowercase__    =  self.prepare_config_and_inputs()
     return config, inputs_dict







  def   UpperCAmelCase     (			self  :Any    ,							_lowercase  :List[str]    ,							_lowercase  :str    ,							_lowercase  :Union[str, Any]					):






     '''simple docstring'''





     lowercase__    =  20
     lowercase__    =  model_class_name(_lowercase					)

     lowercase__    =  inputs_dict["input_ids"]
     lowercase__    =  inputs_dict["attention_mask"]

     lowercase__    =  model.init_cache(input_ids.shape[0]    ,							_lowercase					)
     lowercase__    =  jnp.ones((input_ids.shape[0], max_length)    ,							dtype="i4"					)

     lowercase__    =  jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1					)[None, :]    ,							(input_ids.shape[0], input_ids.shape[-1] - 1)    ,							)
     lowercase__    =  model(
         input_ids[:, :-1]    ,							attention_mask=_lowercase    ,							past_key_values=_lowercase    ,							position_ids=_lowercase    ,							)

     lowercase__    =  jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]    ,							dtype="i4"					)
     lowercase__    =  model(
         input_ids[:, -1:]    ,							attention_mask=_lowercase    ,							past_key_values=outputs_cache.past_key_values    ,							position_ids=_lowercase    ,							)

     lowercase__    =  model(_lowercase					)

     lowercase__    =  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   UpperCAmelCase     (			self  :List[str]    ,							_lowercase  :Optional[Any]    ,							_lowercase  :List[Any]    ,							_lowercase  :List[str]					):






     '''simple docstring'''





     lowercase__    =  20
     lowercase__    =  model_class_name(_lowercase					)

     lowercase__   ,							lowercase__    =  (
         inputs_dict["input_ids"],
         inputs_dict["attention_mask"],
     )

     lowercase__    =  jnp.concatenate(
         [
             attention_mask,
             jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])					),
         ]    ,							axis=-1    ,							)

     lowercase__    =  model.init_cache(input_ids.shape[0]    ,							_lowercase					)
     lowercase__    =  jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1					)[None, :]    ,							(input_ids.shape[0], input_ids.shape[-1] - 1)    ,							)

     lowercase__    =  model(
         input_ids[:, :-1]    ,							attention_mask=_lowercase    ,							past_key_values=_lowercase    ,							position_ids=_lowercase    ,							)
     lowercase__    =  jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]    ,							dtype="i4"					)
     lowercase__    =  model(
         input_ids[:, -1:]    ,							past_key_values=outputs_cache.past_key_values    ,							attention_mask=_lowercase    ,							position_ids=_lowercase    ,							)

     lowercase__    =  model(_lowercase    ,							attention_mask=_lowercase					)

     lowercase__    =  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 					lowerCAmelCase							(   lowercase_							,							unittest.TestCase							,							lowercase_					):
  __lowerCamelCase		=			(FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
  __lowerCamelCase		=			() if is_flax_available() else ()







  def   UpperCAmelCase     (			self  :List[Any]					):






     '''simple docstring'''





     lowercase__    =  FlaxOPTModelTester(self					)







  def   UpperCAmelCase     (			self  :List[Any]					):






     '''simple docstring'''





     lowercase__   ,							lowercase__    =  self.model_tester.prepare_config_and_inputs()
     for model_class in self.all_model_classes:
        self.model_tester.check_use_cache_forward(_lowercase    ,							_lowercase    ,							_lowercase					)







  def   UpperCAmelCase     (			self  :str					):






     '''simple docstring'''





     lowercase__   ,							lowercase__    =  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(_lowercase    ,							_lowercase    ,							_lowercase					)







  @slow
  def   UpperCAmelCase     (			self  :Optional[int]					):






     '''simple docstring'''





     for model_class_name in self.all_model_classes:
        lowercase__    =  model_class_name.from_pretrained("facebook/opt-125m"					)
        lowercase__    =  np.ones((1, 1)					) * model.config.eos_token_id
        lowercase__    =  model(_lowercase					)
        self.assertIsNotNone(_lowercase					)





@require_sentencepiece
@require_flax
class 					lowerCAmelCase							(   unittest.TestCase					):







  @slow
  def   UpperCAmelCase     (			self  :Union[str, Any]					):






     '''simple docstring'''





     lowercase__    =  FlaxOPTModel.from_pretrained("facebook/opt-350m"					)
     lowercase__    =  jnp.array([[0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 4_60_78, 15_88, 2]]					)
     lowercase__    =  model(input_ids=_lowercase					).last_hidden_state
     lowercase__    =  (1, 11, 5_12)
     self.assertEqual(output.shape    ,							_lowercase					)
     lowercase__    =  jnp.array(
         [[-0.2867, -1.9256, -0.3062], [-1.2711, -0.1337, -0.1897], [0.4109, 0.1187, -1.3142]]					)
     self.assertTrue(jnp.allclose(output[:, :3, :3]    ,							_lowercase    ,							atol=4e-2					)					)





@require_flax
@slow
class 					lowerCAmelCase							(   unittest.TestCase					):







  def   UpperCAmelCase     (			self  :Optional[int]					):






     '''simple docstring'''





     super().setUp()
     lowercase__    =  "facebook/opt-350m"







  def   UpperCAmelCase     (			self  :List[Any]					):






     '''simple docstring'''





     lowercase__    =  FlaxOPTForCausalLM.from_pretrained(self.path_model					)
     lowercase__    =  GPTaTokenizer.from_pretrained(self.path_model					)

     lowercase__    =  [
         "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
     lowercase__    =  tokenizer(_lowercase    ,							return_tensors="jax"    ,							padding=_lowercase    ,							add_special_tokens=_lowercase					)
     lowercase__    =  model(inputs.input_ids    ,							attention_mask=inputs.attention_mask					)[0].mean(axis=-1					)
     lowercase__    =  jnp.array(
         [
             [1.3851, -13.8923, -10.5229, -10.7533, -0.2309, -10.2384, -0.5365, -9.0947, -5.1670],
             [-4.7073, -10.6276, -3.9415, -21.5242, -0.2822, -0.2822, -0.2822, -0.2822, -0.2822],
             [0.6247, -3.4229, -8.9179, -1.4297, -14.1650, 1.4146, -9.0218, -0.2703, -0.2703],
             [6.4783, -1.9913, -10.7926, -2.3336, 1.5092, -0.9974, -6.8213, 1.3477, 1.3477],
         ]					)
     self.assertTrue(jnp.allclose(_lowercase    ,							_lowercase    ,							atol=4e-2					)					)

     lowercase__    =  jax.jit(_lowercase					)
     lowercase__    =  model(inputs.input_ids    ,							attention_mask=inputs.attention_mask					)[0].mean(axis=-1					)
     self.assertTrue(jnp.allclose(_lowercase    ,							_lowercase    ,							atol=4e-2					)					)





@require_flax
@slow
class 					lowerCAmelCase							(   unittest.TestCase					):







  @property
  def   UpperCAmelCase     (			self  :Optional[int]					):






     '''simple docstring'''





     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   UpperCAmelCase     (			self  :str					):






     '''simple docstring'''





     lowercase__    =  "facebook/opt-125m"

     lowercase__    =  [
         "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",
     ]

     lowercase__    =  []

     lowercase__    =  FlaxOPTForCausalLM.from_pretrained(_lowercase					)
     lowercase__    =  GPTaTokenizer.from_pretrained(_lowercase					)

     for prompt in self.prompts:
        lowercase__    =  tokenizer(_lowercase    ,							return_tensors="jax"					).input_ids

        lowercase__    =  model.generate(_lowercase    ,							max_length=10					)
        lowercase__    =  generated_ids[0]

        lowercase__    =  tokenizer.batch_decode(_lowercase    ,							skip_special_tokens=_lowercase					)
        predicted_outputs += generated_string

     self.assertListEqual(_lowercase    ,							_lowercase					)







  def   UpperCAmelCase     (			self  :str					):






     '''simple docstring'''





     lowercase__    =  "facebook/opt-350m"

     lowercase__    =  [
         "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",
     ]

     lowercase__    =  []
     lowercase__    =  FlaxOPTForCausalLM.from_pretrained(_lowercase					)
     lowercase__    =  GPTaTokenizer.from_pretrained(_lowercase					)

     for prompt in self.prompts:
        lowercase__    =  tokenizer(_lowercase    ,							return_tensors="jax"					).input_ids

        lowercase__    =  model.generate(_lowercase    ,							max_length=10					)
        lowercase__    =  generated_ids[0]

        lowercase__    =  tokenizer.batch_decode(_lowercase    ,							skip_special_tokens=_lowercase					)
        predicted_outputs += generated_string

     self.assertListEqual(_lowercase    ,							_lowercase					)







  def   UpperCAmelCase     (			self  :List[str]					):






     '''simple docstring'''





     lowercase__    =  "facebook/opt-125m"
     lowercase__    =  [
         "Today is a beautiful day and I want to thank",
         "In the city of Rome Canaver Canaver Canaver Canaver",
     ]
     lowercase__    =  FlaxOPTForCausalLM.from_pretrained(_lowercase					)
     lowercase__    =  GPTaTokenizer.from_pretrained(_lowercase					)
     lowercase__    =  tokenizer(
         [
             "Today is a beautiful day and I want to",
             "In the city of",
         ]    ,							return_tensors="jax"    ,							padding=_lowercase    ,							)

     lowercase__    =  jax.jit(model.generate					)

     lowercase__    =  jit_generate(inputs["input_ids"]    ,							attention_mask=inputs["attention_mask"]					).sequences

     lowercase__    =  tokenizer.batch_decode(_lowercase    ,							skip_special_tokens=_lowercase					)

     self.assertIsNotNone(_lowercase    ,							_lowercase					)







  def   UpperCAmelCase     (			self  :Dict					):






     '''simple docstring'''





     lowercase__    =  "facebook/opt-350m"

     lowercase__    =  GPTaTokenizer.from_pretrained(_lowercase					)
     lowercase__    =  FlaxOPTForCausalLM.from_pretrained(_lowercase					)

     lowercase__    =  "left"

     # use different length sentences to test batching
     lowercase__    =  [
         "Hello, my dog is a little",
         "Today, I",
     ]

     lowercase__    =  tokenizer(_lowercase    ,							return_tensors="jax"    ,							padding=_lowercase					)
     lowercase__    =  inputs["input_ids"]

     lowercase__    =  model.generate(input_ids=_lowercase    ,							attention_mask=inputs["attention_mask"]    ,							trace=_lowercase					)

     lowercase__    =  tokenizer(sentences[0]    ,							return_tensors="jax"					).input_ids
     lowercase__    =  model.generate(input_ids=_lowercase					)

     lowercase__    =  inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].sum()
     lowercase__    =  tokenizer(sentences[1]    ,							return_tensors="jax"					).input_ids
     lowercase__    =  model.generate(input_ids=_lowercase    ,							max_length=model.config.max_length - num_paddings					)

     lowercase__    =  tokenizer.batch_decode(outputs[0]    ,							skip_special_tokens=_lowercase					)
     lowercase__    =  tokenizer.decode(output_non_padded[0][0]    ,							skip_special_tokens=_lowercase					)
     lowercase__    =  tokenizer.decode(output_padded[0][0]    ,							skip_special_tokens=_lowercase					)

     lowercase__    =  [
         "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(_lowercase    ,							_lowercase					)
     self.assertListEqual(_lowercase    ,							[non_padded_sentence, padded_sentence]					)













from typing import List, Optional, Tuple, Union

import torch

from ...models import UNetaDModel
from ...schedulers import PNDMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput





class 					lowerCAmelCase							(   lowercase_					):
  __lowerCamelCase		=			42
  __lowerCamelCase		=			42







  def __init__(			self  :List[Any]    ,							_lowercase  :UNetaDModel    ,							_lowercase  :PNDMScheduler					):






     '''simple docstring'''





     super().__init__()

     lowercase__    =  PNDMScheduler.from_config(scheduler.config					)

     self.register_modules(unet=_lowercase    ,							scheduler=_lowercase					)







  @torch.no_grad()
  def __call__(			self  :Union[str, Any]    ,							_lowercase  :int = 1    ,							_lowercase  :int = 50    ,							_lowercase  :Optional[Union[torch.Generator, List[torch.Generator]]] = None    ,							_lowercase  :Optional[str] = "pil"    ,							_lowercase  :bool = True    ,							**_lowercase  :List[Any]    ,							):






     '''simple docstring'''





     lowercase__    =  randn_tensor(
         (batch_size, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size)    ,							generator=_lowercase    ,							device=self.device    ,							)

     self.scheduler.set_timesteps(_lowercase					)
     for t in self.progress_bar(self.scheduler.timesteps					):
        lowercase__    =  self.unet(_lowercase    ,							_lowercase					).sample

        lowercase__    =  self.scheduler.step(_lowercase    ,							_lowercase    ,							_lowercase					).prev_sample

     lowercase__    =  (image / 2 + 0.5).clamp(0    ,							1					)
     lowercase__    =  image.cpu().permute(0    ,							2    ,							3    ,							1					).numpy()
     if output_type == "pil":
        lowercase__    =  self.numpy_to_pil(_lowercase					)

     if not return_dict:
        return (image,)

     return ImagePipelineOutput(images=_lowercase					)






