
"""simple docstring"""





import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class A_      (lowercase__							,unittest.TestCase ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__     :		str    	=						KandinskyPriorPipeline
  SCREAMING_SNAKE_CASE__     :		Any    	=						["""prompt"""]
  SCREAMING_SNAKE_CASE__     :		List[str]    	=						["""prompt""", """negative_prompt"""]
  SCREAMING_SNAKE_CASE__     :		int    	=						[
      """num_images_per_prompt""",
      """generator""",
      """num_inference_steps""",
      """latents""",
      """negative_prompt""",
      """guidance_scale""",
      """output_type""",
      """return_dict""",
  ]
  SCREAMING_SNAKE_CASE__     :		Tuple    	=						False







  @property
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     return 32







  @property
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     return 32







  @property
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     return self.time_input_dim







  @property
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     return self.time_input_dim * 4







  @property
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     return 100







  @property
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   List[str]     	=       CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"    )
     return tokenizer







  @property
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     torch.manual_seed(0    )
     UpperCAmelCase_   :   Any     	=       CLIPTextConfig(
         bos_token_id=0							,  eos_token_id=2							,  hidden_size=self.text_embedder_hidden_size							,  projection_dim=self.text_embedder_hidden_size							,  intermediate_size=37							,  layer_norm_eps=1E-0_5							,  num_attention_heads=4							,  num_hidden_layers=5							,  pad_token_id=1							,  vocab_size=1000							,  )
     return CLIPTextModelWithProjection(lowercase_    )







  @property
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     torch.manual_seed(0    )

     UpperCAmelCase_   :   Dict     	=       {
         "num_attention_heads": 2,
         "attention_head_dim": 12,
         "embedding_dim": self.text_embedder_hidden_size,
         "num_layers": 1,
     }

     UpperCAmelCase_   :   Dict     	=       PriorTransformer(**lowercase_    )
     # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
     UpperCAmelCase_   :   List[Any]     	=       nn.Parameter(torch.ones(model.clip_std.shape    )    )
     return model







  @property
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     torch.manual_seed(0    )
     UpperCAmelCase_   :   Tuple     	=       CLIPVisionConfig(
         hidden_size=self.text_embedder_hidden_size							,  image_size=224							,  projection_dim=self.text_embedder_hidden_size							,  intermediate_size=37							,  num_attention_heads=4							,  num_channels=3							,  num_hidden_layers=5							,  patch_size=14							,  )

     UpperCAmelCase_   :   List[Any]     	=       CLIPVisionModelWithProjection(lowercase_    )
     return model







  @property
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   str     	=       CLIPImageProcessor(
         crop_size=224							,  do_center_crop=lowercase_							,  do_normalize=lowercase_							,  do_resize=lowercase_							,  image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]							,  image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]							,  resample=3							,  size=224							,  )

     return image_processor







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[Any]     	=       self.dummy_prior
     UpperCAmelCase_   :   int     	=       self.dummy_image_encoder
     UpperCAmelCase_   :   Union[str, Any]     	=       self.dummy_text_encoder
     UpperCAmelCase_   :   List[Any]     	=       self.dummy_tokenizer
     UpperCAmelCase_   :   List[Any]     	=       self.dummy_image_processor

     UpperCAmelCase_   :   Dict     	=       UnCLIPScheduler(
         variance_type="fixed_small_log"							,  prediction_type="sample"							,  num_train_timesteps=1000							,  clip_sample=lowercase_							,  clip_sample_range=10.0							,  )

     UpperCAmelCase_   :   Optional[int]     	=       {
         "prior": prior,
         "image_encoder": image_encoder,
         "text_encoder": text_encoder,
         "tokenizer": tokenizer,
         "scheduler": scheduler,
         "image_processor": image_processor,
     }

     return components







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_=0    ):

     """simple docstring"""
     if str(lowercase_    ).startswith("mps"    ):
        UpperCAmelCase_   :   List[str]     	=       torch.manual_seed(lowercase_    )
     else:
        UpperCAmelCase_   :   Dict     	=       torch.Generator(device=lowercase_    ).manual_seed(lowercase_    )
     UpperCAmelCase_   :   Any     	=       {
         "prompt": "horse",
         "generator": generator,
         "guidance_scale": 4.0,
         "num_inference_steps": 2,
         "output_type": "np",
     }
     return inputs







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[int]     	=       "cpu"

     UpperCAmelCase_   :   str     	=       self.get_dummy_components()

     UpperCAmelCase_   :   Optional[int]     	=       self.pipeline_class(**lowercase_    )
     UpperCAmelCase_   :   Union[str, Any]     	=       pipe.to(lowercase_    )

     pipe.set_progress_bar_config(disable=lowercase_    )

     UpperCAmelCase_   :   Optional[int]     	=       pipe(**self.get_dummy_inputs(lowercase_    )    )
     UpperCAmelCase_   :   int     	=       output.image_embeds

     UpperCAmelCase_   :   Union[str, Any]     	=       pipe(
         **self.get_dummy_inputs(lowercase_    )							,  return_dict=lowercase_							,  )[0]

     UpperCAmelCase_   :   Optional[Any]     	=       image[0, -10:]
     UpperCAmelCase_   :   List[str]     	=       image_from_tuple[0, -10:]

     assert image.shape == (1, 32)

     UpperCAmelCase_   :   Optional[int]     	=       np.array(
         [-0.05_32, 1.71_20, 0.36_56, -1.08_52, -0.89_46, -1.17_56, 0.43_48, 0.24_82, 0.51_46, -0.11_56]    )

     assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2
     assert np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() < 1E-2







  @skip_mps
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Union[str, Any]     	=       torch_device == "cpu"
     UpperCAmelCase_   :   Union[str, Any]     	=       True
     UpperCAmelCase_   :   Optional[Any]     	=       False

     self._test_inference_batch_single_identical(
         test_max_difference=lowercase_							,  relax_max_difference=lowercase_							,  test_mean_pixel_difference=lowercase_							,  )







  @skip_mps
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[int]     	=       torch_device == "cpu"
     UpperCAmelCase_   :   str     	=       False

     self._test_attention_slicing_forward_pass(
         test_max_difference=lowercase_							,  test_mean_pixel_difference=lowercase_							,  )




"""simple docstring"""





from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import (
        TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFXLMForMultipleChoice,
        TFXLMForQuestionAnsweringSimple,
        TFXLMForSequenceClassification,
        TFXLMForTokenClassification,
        TFXLMModel,
        TFXLMWithLMHeadModel,
        XLMConfig,
    )


class A_      :




  '''simple docstring'''







  def __init__(						self							,  lowercase_							,  ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[Any]     	=       parent
     UpperCAmelCase_   :   Any     	=       13
     UpperCAmelCase_   :   Dict     	=       7
     UpperCAmelCase_   :   str     	=       True
     UpperCAmelCase_   :   Optional[int]     	=       True
     UpperCAmelCase_   :   Optional[Any]     	=       True
     UpperCAmelCase_   :   Tuple     	=       True
     UpperCAmelCase_   :   Tuple     	=       True
     UpperCAmelCase_   :   str     	=       False
     UpperCAmelCase_   :   Optional[Any]     	=       False
     UpperCAmelCase_   :   Any     	=       False
     UpperCAmelCase_   :   Tuple     	=       2
     UpperCAmelCase_   :   str     	=       99
     UpperCAmelCase_   :   str     	=       0
     UpperCAmelCase_   :   Dict     	=       32
     UpperCAmelCase_   :   Dict     	=       2
     UpperCAmelCase_   :   Union[str, Any]     	=       4
     UpperCAmelCase_   :   Union[str, Any]     	=       0.1
     UpperCAmelCase_   :   Dict     	=       0.1
     UpperCAmelCase_   :   Optional[Any]     	=       512
     UpperCAmelCase_   :   Any     	=       16
     UpperCAmelCase_   :   Tuple     	=       2
     UpperCAmelCase_   :   int     	=       0.02
     UpperCAmelCase_   :   Union[str, Any]     	=       3
     UpperCAmelCase_   :   str     	=       4
     UpperCAmelCase_   :   Union[str, Any]     	=       "last"
     UpperCAmelCase_   :   Optional[int]     	=       True
     UpperCAmelCase_   :   str     	=       None
     UpperCAmelCase_   :   Dict     	=       0







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[int]     	=       ids_tensor([self.batch_size, self.seq_length]							,  self.vocab_size    )
     UpperCAmelCase_   :   Any     	=       random_attention_mask([self.batch_size, self.seq_length]							,  dtype=tf.floataa    )

     UpperCAmelCase_   :   Union[str, Any]     	=       None
     if self.use_input_lengths:
        UpperCAmelCase_   :   Dict     	=       (
            ids_tensor([self.batch_size]							,  vocab_size=2    ) + self.seq_length - 2
        )  # small variation of seq_length

     UpperCAmelCase_   :   Union[str, Any]     	=       None
     if self.use_token_type_ids:
        UpperCAmelCase_   :   int     	=       ids_tensor([self.batch_size, self.seq_length]							,  self.n_langs    )

     UpperCAmelCase_   :   Dict     	=       None
     UpperCAmelCase_   :   int     	=       None
     UpperCAmelCase_   :   str     	=       None
     if self.use_labels:
        UpperCAmelCase_   :   Optional[Any]     	=       ids_tensor([self.batch_size]							,  self.type_sequence_label_size    )
        UpperCAmelCase_   :   Optional[int]     	=       ids_tensor([self.batch_size, self.seq_length]							,  self.num_labels    )
        UpperCAmelCase_   :   Optional[int]     	=       ids_tensor([self.batch_size]							,  2							,  dtype=tf.floataa    )
        UpperCAmelCase_   :   Dict     	=       ids_tensor([self.batch_size]							,  self.num_choices    )

     UpperCAmelCase_   :   Optional[Any]     	=       XLMConfig(
         vocab_size=self.vocab_size							,  n_special=self.n_special							,  emb_dim=self.hidden_size							,  n_layers=self.num_hidden_layers							,  n_heads=self.num_attention_heads							,  dropout=self.hidden_dropout_prob							,  attention_dropout=self.attention_probs_dropout_prob							,  gelu_activation=self.gelu_activation							,  sinusoidal_embeddings=self.sinusoidal_embeddings							,  asm=self.asm							,  causal=self.causal							,  n_langs=self.n_langs							,  max_position_embeddings=self.max_position_embeddings							,  initializer_range=self.initializer_range							,  summary_type=self.summary_type							,  use_proj=self.use_proj							,  bos_token_id=self.bos_token_id							,  )

     return (
         config,
         input_ids,
         token_type_ids,
         input_lengths,
         sequence_labels,
         token_labels,
         is_impossible_labels,
         choice_labels,
         input_mask,
     )







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[Any]     	=       TFXLMModel(config=lowercase_    )
     UpperCAmelCase_   :   Any     	=       {"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
     UpperCAmelCase_   :   Tuple     	=       model(lowercase_    )

     UpperCAmelCase_   :   Any     	=       [input_ids, input_mask]
     UpperCAmelCase_   :   Union[str, Any]     	=       model(lowercase_    )
     self.parent.assertEqual(result.last_hidden_state.shape							,  (self.batch_size, self.seq_length, self.hidden_size)    )







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  ):

     """simple docstring"""
     UpperCAmelCase_   :   Tuple     	=       TFXLMWithLMHeadModel(lowercase_    )

     UpperCAmelCase_   :   List[Any]     	=       {"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
     UpperCAmelCase_   :   Dict     	=       model(lowercase_    )

     UpperCAmelCase_   :   Optional[int]     	=       outputs

     self.parent.assertEqual(result.logits.shape							,  (self.batch_size, self.seq_length, self.vocab_size)    )







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  ):

     """simple docstring"""
     UpperCAmelCase_   :   int     	=       TFXLMForQuestionAnsweringSimple(lowercase_    )

     UpperCAmelCase_   :   Any     	=       {"input_ids": input_ids, "lengths": input_lengths}

     UpperCAmelCase_   :   Union[str, Any]     	=       model(lowercase_    )

     self.parent.assertEqual(result.start_logits.shape							,  (self.batch_size, self.seq_length)    )
     self.parent.assertEqual(result.end_logits.shape							,  (self.batch_size, self.seq_length)    )







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  ):

     """simple docstring"""
     UpperCAmelCase_   :   Union[str, Any]     	=       TFXLMForSequenceClassification(lowercase_    )

     UpperCAmelCase_   :   Optional[int]     	=       {"input_ids": input_ids, "lengths": input_lengths}

     UpperCAmelCase_   :   Any     	=       model(lowercase_    )

     self.parent.assertEqual(result.logits.shape							,  (self.batch_size, self.type_sequence_label_size)    )







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  ):

     """simple docstring"""
     UpperCAmelCase_   :   Dict     	=       self.num_labels
     UpperCAmelCase_   :   Any     	=       TFXLMForTokenClassification(config=lowercase_    )
     UpperCAmelCase_   :   Any     	=       {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
     UpperCAmelCase_   :   List[Any]     	=       model(lowercase_    )
     self.parent.assertEqual(result.logits.shape							,  (self.batch_size, self.seq_length, self.num_labels)    )







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  ):

     """simple docstring"""
     UpperCAmelCase_   :   Any     	=       self.num_choices
     UpperCAmelCase_   :   Any     	=       TFXLMForMultipleChoice(config=lowercase_    )
     UpperCAmelCase_   :   Union[str, Any]     	=       tf.tile(tf.expand_dims(lowercase_							,  1    )							,  (1, self.num_choices, 1)    )
     UpperCAmelCase_   :   Any     	=       tf.tile(tf.expand_dims(lowercase_							,  1    )							,  (1, self.num_choices, 1)    )
     UpperCAmelCase_   :   int     	=       tf.tile(tf.expand_dims(lowercase_							,  1    )							,  (1, self.num_choices, 1)    )
     UpperCAmelCase_   :   Any     	=       {
         "input_ids": multiple_choice_inputs_ids,
         "attention_mask": multiple_choice_input_mask,
         "token_type_ids": multiple_choice_token_type_ids,
     }
     UpperCAmelCase_   :   int     	=       model(lowercase_    )
     self.parent.assertEqual(result.logits.shape							,  (self.batch_size, self.num_choices)    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   List[str]     	=       self.prepare_config_and_inputs()
     (
         (
         UpperCAmelCase_
     )	,  (
         UpperCAmelCase_
     )	,  (
         UpperCAmelCase_
     )	,  (
         UpperCAmelCase_
     )	,  (
         UpperCAmelCase_
     )	,  (
         UpperCAmelCase_
     )	,  (
         UpperCAmelCase_
     )	,  (
         UpperCAmelCase_
     )	,  (
         UpperCAmelCase_
     )	,  
     )   :   Dict     	=       config_and_inputs
     UpperCAmelCase_   :   Optional[int]     	=       {
         "input_ids": input_ids,
         "token_type_ids": token_type_ids,
         "langs": token_type_ids,
         "lengths": input_lengths,
     }
     return config, inputs_dict




@require_tf
class A_      (lowercase__							,lowercase__							,unittest.TestCase ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__     :		Any    	=						(
      (
          TFXLMModel,
          TFXLMWithLMHeadModel,
          TFXLMForSequenceClassification,
          TFXLMForQuestionAnsweringSimple,
          TFXLMForTokenClassification,
          TFXLMForMultipleChoice,
      )
      if is_tf_available()
      else ()
  )
  SCREAMING_SNAKE_CASE__     :		List[Any]    	=						(
      (TFXLMWithLMHeadModel,) if is_tf_available() else ()
  )  # TODO (PVP): Check other models whether language generation is also applicable
  SCREAMING_SNAKE_CASE__     :		Any    	=						(
      {
          """feature-extraction""": TFXLMModel,
          """fill-mask""": TFXLMWithLMHeadModel,
          """question-answering""": TFXLMForQuestionAnsweringSimple,
          """text-classification""": TFXLMForSequenceClassification,
          """text-generation""": TFXLMWithLMHeadModel,
          """token-classification""": TFXLMForTokenClassification,
          """zero-shot""": TFXLMForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  SCREAMING_SNAKE_CASE__     :		List[str]    	=						False
  SCREAMING_SNAKE_CASE__     :		Optional[Any]    	=						False







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_							,  lowercase_    ):

     """simple docstring"""
     if (
         pipeline_test_casse_name == "QAPipelineTests"
         and tokenizer_name is not None
         and not tokenizer_name.endswith("Fast"    )
     ):
        # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
        # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
        # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
        return True

     return False







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   List[Any]     	=       TFXLMModelTester(self    )
     UpperCAmelCase_   :   List[str]     	=       ConfigTester(self							,  config_class=lowercase_							,  emb_dim=37    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     self.config_tester.run_common_tests()







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[int]     	=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_model(*lowercase_    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   int     	=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_lm_head(*lowercase_    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[Any]     	=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_qa(*lowercase_    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[Any]     	=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_sequence_classif(*lowercase_    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[int]     	=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_for_token_classification(*lowercase_    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Dict     	=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_for_multiple_choice(*lowercase_    )







  @slow
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCAmelCase_   :   Tuple     	=       TFXLMModel.from_pretrained(lowercase_    )
        self.assertIsNotNone(lowercase_    )




@require_tf
class A_      (unittest.TestCase ):




  '''simple docstring'''







  @slow
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Union[str, Any]     	=       TFXLMWithLMHeadModel.from_pretrained("xlm-mlm-en-2048"    )
     UpperCAmelCase_   :   Optional[Any]     	=       tf.convert_to_tensor([[14, 447]]							,  dtype=tf.intaa    )  # the president
     UpperCAmelCase_   :   Optional[int]     	=       [
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
     ]  # the president the president the president the president the president the president the president the president the president the president
     # TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
     UpperCAmelCase_   :   str     	=       model.generate(lowercase_							,  do_sample=lowercase_    )
     self.assertListEqual(output_ids[0].numpy().tolist()							,  lowercase_    )


