

'''simple docstring'''






import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import MBartConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
from transformers.utils import cached_property

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 AutoTokenizer
      from transformers.models.mbart.modeling_flax_mbart import (
          FlaxMBartForConditionalGeneration,
          FlaxMBartForQuestionAnswering,
          FlaxMBartForSequenceClassification,
          FlaxMBartModel,
          shift_tokens_right,
      )


def 					__lowerCamelCase			(      SCREAMING_SNAKE_CASE_     :	Any  ,  SCREAMING_SNAKE_CASE_     :	List[str]  ,  SCREAMING_SNAKE_CASE_     :	Optional[Any]=None  ,  SCREAMING_SNAKE_CASE_     :	Tuple=None  ,  SCREAMING_SNAKE_CASE_     :	Dict=None  ,  SCREAMING_SNAKE_CASE_     :	Tuple=None  ,  SCREAMING_SNAKE_CASE_     :	Any=None  ,  SCREAMING_SNAKE_CASE_     :	List[Any]=None  ,  )							->   Tuple:
      """simple docstring"""

      if attention_mask is None:
            SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       np.where(input_ids != config.pad_token_id  ,  1  ,  0    )
      if decoder_attention_mask is None:
            SCREAMING_SNAKE_CASE_      :			str     =       np.where(decoder_input_ids != config.pad_token_id  ,  1  ,  0    )
      if head_mask is None:
            SCREAMING_SNAKE_CASE_      :			Dict     =       np.ones((config.encoder_layers, config.encoder_attention_heads)    )
      if decoder_head_mask is None:
            SCREAMING_SNAKE_CASE_      :			int     =       np.ones((config.decoder_layers, config.decoder_attention_heads)    )
      if cross_attn_head_mask is None:
            SCREAMING_SNAKE_CASE_      :			Dict     =       np.ones((config.decoder_layers, config.decoder_attention_heads)    )
      return {
          "input_ids": input_ids,
          "decoder_input_ids": decoder_input_ids,
          "attention_mask": attention_mask,
          "decoder_attention_mask": decoder_attention_mask,
      }
class 	SCREAMING_SNAKE_CASE__			(				unittest.TestCase      ):




    def __init__(					self				,     lowercase__				,     lowercase__=13				,     lowercase__=7				,     lowercase__=True				,     lowercase__=False				,     lowercase__=99				,     lowercase__=16				,     lowercase__=2				,     lowercase__=4				,     lowercase__=4				,     lowercase__="gelu"				,     lowercase__=0.1				,     lowercase__=0.1				,     lowercase__=32				,     lowercase__=2				,     lowercase__=1				,     lowercase__=0				,     lowercase__=2				,     lowercase__=0.02				,     ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       parent
          SCREAMING_SNAKE_CASE_      :			Dict     =       batch_size
          SCREAMING_SNAKE_CASE_      :			List[str]     =       seq_length
          SCREAMING_SNAKE_CASE_      :			List[str]     =       is_training
          SCREAMING_SNAKE_CASE_      :			Any     =       use_labels
          SCREAMING_SNAKE_CASE_      :			Dict     =       vocab_size
          SCREAMING_SNAKE_CASE_      :			str     =       hidden_size
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       num_hidden_layers
          SCREAMING_SNAKE_CASE_      :			Tuple     =       num_attention_heads
          SCREAMING_SNAKE_CASE_      :			Tuple     =       intermediate_size
          SCREAMING_SNAKE_CASE_      :			int     =       hidden_act
          SCREAMING_SNAKE_CASE_      :			int     =       hidden_dropout_prob
          SCREAMING_SNAKE_CASE_      :			str     =       attention_probs_dropout_prob
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       max_position_embeddings
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       eos_token_id
          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       pad_token_id
          SCREAMING_SNAKE_CASE_      :			Any     =       bos_token_id
          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       decoder_start_token_id
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       initializer_range




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






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

          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       shift_tokens_right(lowercase__				,     1			)

          SCREAMING_SNAKE_CASE_      :			Any     =       MBartConfig(
              vocab_size=self.vocab_size				,     d_model=self.hidden_size				,     encoder_layers=self.num_hidden_layers				,     decoder_layers=self.num_hidden_layers				,     encoder_attention_heads=self.num_attention_heads				,     decoder_attention_heads=self.num_attention_heads				,     encoder_ffn_dim=self.intermediate_size				,     decoder_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				,     decoder_start_token_id=self.decoder_start_token_id				,     initializer_range=self.initializer_range				,     use_cache=lowercase__				,     )
          SCREAMING_SNAKE_CASE_      :			str     =       prepare_mbart_inputs_dict(lowercase__				,     lowercase__				,     lowercase__			)
          return config, inputs_dict




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_      :			int     =       self.prepare_config_and_inputs()
          return config, inputs_dict




    def   __lowerCamelCase      (					self				,     lowercase__				,     lowercase__				,     lowercase__			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Tuple     =       20
          SCREAMING_SNAKE_CASE_      :			List[str]     =       model_class_name(lowercase__			)

          SCREAMING_SNAKE_CASE_      :			Dict     =       model.encode(inputs_dict["input_ids"]			)

          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       (
              inputs_dict["decoder_input_ids"],
              inputs_dict["decoder_attention_mask"],
          )

          SCREAMING_SNAKE_CASE_      :			Tuple     =       model.init_cache(decoder_input_ids.shape[0]				,     lowercase__				,     lowercase__			)
          SCREAMING_SNAKE_CASE_      :			int     =       jnp.ones((decoder_input_ids.shape[0], max_decoder_length)				,     dtype="i4"			)

          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       jnp.broadcast_to(
              jnp.arange(decoder_input_ids.shape[-1] - 1			)[None, :]				,     (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)				,     )
          SCREAMING_SNAKE_CASE_      :			Tuple     =       model.decode(
              decoder_input_ids[:, :-1]				,     lowercase__				,     decoder_attention_mask=lowercase__				,     past_key_values=lowercase__				,     decoder_position_ids=lowercase__				,     )

          SCREAMING_SNAKE_CASE_      :			int     =       jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,     dtype="i4"			)
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       model.decode(
              decoder_input_ids[:, -1:]				,     lowercase__				,     decoder_attention_mask=lowercase__				,     past_key_values=outputs_cache.past_key_values				,     decoder_position_ids=lowercase__				,     )

          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       model.decode(lowercase__				,     lowercase__			)

          SCREAMING_SNAKE_CASE_      :			Any     =       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   __lowerCamelCase      (					self				,     lowercase__				,     lowercase__				,     lowercase__			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       20
          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       model_class_name(lowercase__			)

          SCREAMING_SNAKE_CASE_      :			List[str]     =       model.encode(inputs_dict["input_ids"]			)

          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_      :			Optional[int]     =       (
              inputs_dict["decoder_input_ids"],
              inputs_dict["decoder_attention_mask"],
          )

          SCREAMING_SNAKE_CASE_      :			List[str]     =       jnp.concatenate(
              [
                  decoder_attention_mask,
                  jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])			),
              ]				,     axis=-1				,     )

          SCREAMING_SNAKE_CASE_      :			Tuple     =       model.init_cache(decoder_input_ids.shape[0]				,     lowercase__				,     lowercase__			)
          SCREAMING_SNAKE_CASE_      :			List[str]     =       jnp.broadcast_to(
              jnp.arange(decoder_input_ids.shape[-1] - 1			)[None, :]				,     (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)				,     )

          SCREAMING_SNAKE_CASE_      :			Tuple     =       model.decode(
              decoder_input_ids[:, :-1]				,     lowercase__				,     decoder_attention_mask=lowercase__				,     past_key_values=lowercase__				,     decoder_position_ids=lowercase__				,     )
          SCREAMING_SNAKE_CASE_      :			Dict     =       jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,     dtype="i4"			)
          SCREAMING_SNAKE_CASE_      :			Tuple     =       model.decode(
              decoder_input_ids[:, -1:]				,     lowercase__				,     past_key_values=outputs_cache.past_key_values				,     decoder_attention_mask=lowercase__				,     decoder_position_ids=lowercase__				,     )

          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       model.decode(lowercase__				,     lowercase__				,     decoder_attention_mask=lowercase__			)

          SCREAMING_SNAKE_CASE_      :			Any     =       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 	SCREAMING_SNAKE_CASE__			(				unittest.TestCase      ):
    _A			    =							99




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       np.array(
              [
                  [71, 82, 18, 33, 46, 91, 2],
                  [68, 34, 26, 58, 30, 82, 2],
                  [5, 97, 17, 39, 94, 40, 2],
                  [76, 83, 94, 25, 70, 78, 2],
                  [87, 59, 41, 35, 48, 66, 2],
                  [55, 13, 16, 58, 5, 2, 1],  # note padding
                  [64, 27, 31, 51, 12, 75, 2],
                  [52, 64, 86, 17, 83, 39, 2],
                  [48, 61, 9, 24, 71, 82, 2],
                  [26, 1, 60, 48, 22, 13, 2],
                  [21, 5, 62, 28, 14, 76, 2],
                  [45, 98, 37, 86, 59, 48, 2],
                  [70, 70, 50, 9, 28, 0, 2],
              ]				,     dtype=np.intaa				,     )

          SCREAMING_SNAKE_CASE_      :			Dict     =       input_ids.shape[0]
          SCREAMING_SNAKE_CASE_      :			str     =       MBartConfig(
              vocab_size=self.vocab_size				,     d_model=24				,     encoder_layers=2				,     decoder_layers=2				,     encoder_attention_heads=2				,     decoder_attention_heads=2				,     encoder_ffn_dim=32				,     decoder_ffn_dim=32				,     max_position_embeddings=48				,     eos_token_id=2				,     pad_token_id=1				,     bos_token_id=0				,     )
          return config, input_ids, batch_size




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_      :			List[Any]     =       self._get_config_and_data()
          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       FlaxMBartForSequenceClassification(lowercase__			)
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       model(input_ids=lowercase__				,     decoder_input_ids=lowercase__			)
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       (batch_size, config.num_labels)
          self.assertEqual(outputs["logits"].shape				,     lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_      :			int     =       self._get_config_and_data()
          SCREAMING_SNAKE_CASE_      :			int     =       FlaxMBartForQuestionAnswering(lowercase__			)
          SCREAMING_SNAKE_CASE_      :			List[str]     =       model(input_ids=lowercase__			)

          self.assertEqual(outputs["start_logits"].shape				,     input_ids.shape			)
          self.assertEqual(outputs["end_logits"].shape				,     input_ids.shape			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_      :			Any     =       self._get_config_and_data()
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       FlaxMBartForConditionalGeneration(lowercase__			)
          SCREAMING_SNAKE_CASE_      :			str     =       lm_model(input_ids=lowercase__			)
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       (batch_size, input_ids.shape[1], config.vocab_size)
          self.assertEqual(outputs["logits"].shape				,     lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			str     =       MBartConfig(
              vocab_size=self.vocab_size				,     d_model=14				,     encoder_layers=2				,     decoder_layers=2				,     encoder_attention_heads=2				,     decoder_attention_heads=2				,     encoder_ffn_dim=8				,     decoder_ffn_dim=8				,     max_position_embeddings=48				,     )
          SCREAMING_SNAKE_CASE_      :			Tuple     =       FlaxMBartForConditionalGeneration(lowercase__			)
          SCREAMING_SNAKE_CASE_      :			Any     =       np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]				,     dtype=np.intaa			)
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]				,     dtype=np.intaa			)
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       lm_model(input_ids=lowercase__				,     decoder_input_ids=lowercase__			)
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       (*summary.shape, config.vocab_size)
          self.assertEqual(outputs["logits"].shape				,     lowercase__			)







    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			int     =       np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]				,     dtype=np.intaa			)
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       shift_tokens_right(lowercase__				,     1			)
          SCREAMING_SNAKE_CASE_      :			List[str]     =       np.equal(lowercase__				,     1			).astype(np.floataa			).sum()
          SCREAMING_SNAKE_CASE_      :			Tuple     =       np.equal(lowercase__				,     1			).astype(np.floataa			).sum()
          self.assertEqual(shifted.shape				,     input_ids.shape			)
          self.assertEqual(lowercase__				,     n_pad_before - 1			)
          self.assertTrue(np.equal(shifted[:, 0]				,     2			).all()			)






@require_flax
class 	SCREAMING_SNAKE_CASE__			(				_UpperCAmelCase,unittest.TestCase,_UpperCAmelCase      ):
    _A			    =							True
    _A			    =							(
        (
            FlaxMBartModel,
            FlaxMBartForConditionalGeneration,
            FlaxMBartForSequenceClassification,
            FlaxMBartForQuestionAnswering,
        )
        if is_flax_available()
        else ()
    )
    _A			    =							(FlaxMBartForConditionalGeneration,) if is_flax_available() else ()




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			List[str]     =       FlaxMBartModelTester(self			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_      :			Any     =       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   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_      :			str     =       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__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_      :			Dict     =       self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
                with self.subTest(model_class.__name__			):
                      SCREAMING_SNAKE_CASE_      :			int     =       self._prepare_for_class(lowercase__				,     lowercase__			)
                      SCREAMING_SNAKE_CASE_      :			Optional[int]     =       model_class(lowercase__			)

                      @jax.jit
                      def encode_jitted(lowercase__				,     lowercase__=None				,     **lowercase__			):
                            return model.encode(input_ids=lowercase__				,     attention_mask=lowercase__			)

                      with self.subTest("JIT Enabled"			):
                            SCREAMING_SNAKE_CASE_      :			Tuple     =       encode_jitted(**lowercase__			).to_tuple()

                      with self.subTest("JIT Disabled"			):
                            with jax.disable_jit():
                                  SCREAMING_SNAKE_CASE_      :			Dict     =       encode_jitted(**lowercase__			).to_tuple()

                      self.assertEqual(len(lowercase__			)				,     len(lowercase__			)			)
                      for jitted_output, output in zip(lowercase__				,     lowercase__			):
                            self.assertEqual(jitted_output.shape				,     output.shape			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_      :			Any     =       self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
                with self.subTest(model_class.__name__			):
                      SCREAMING_SNAKE_CASE_      :			Tuple     =       model_class(lowercase__			)
                      SCREAMING_SNAKE_CASE_      :			Dict     =       model.encode(inputs_dict["input_ids"]				,     inputs_dict["attention_mask"]			)

                      SCREAMING_SNAKE_CASE_      :			int     =       {
                          "decoder_input_ids": inputs_dict["decoder_input_ids"],
                          "decoder_attention_mask": inputs_dict["decoder_attention_mask"],
                          "encoder_outputs": encoder_outputs,
                      }

                      @jax.jit
                      def decode_jitted(lowercase__				,     lowercase__				,     lowercase__			):
                            return model.decode(
                                decoder_input_ids=lowercase__				,     decoder_attention_mask=lowercase__				,     encoder_outputs=lowercase__				,     )

                      with self.subTest("JIT Enabled"			):
                            SCREAMING_SNAKE_CASE_      :			Dict     =       decode_jitted(**lowercase__			).to_tuple()

                      with self.subTest("JIT Disabled"			):
                            with jax.disable_jit():
                                  SCREAMING_SNAKE_CASE_      :			str     =       decode_jitted(**lowercase__			).to_tuple()

                      self.assertEqual(len(lowercase__			)				,     len(lowercase__			)			)
                      for jitted_output, output in zip(lowercase__				,     lowercase__			):
                            self.assertEqual(jitted_output.shape				,     output.shape			)







    @slow
    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          for model_class_name in self.all_model_classes:
                SCREAMING_SNAKE_CASE_      :			Tuple     =       model_class_name.from_pretrained("facebook/mbart-large-cc25"				,     from_pt=lowercase__			)
                # FlaxMBartForSequenceClassification expects eos token in input_ids
                SCREAMING_SNAKE_CASE_      :			Any     =       np.ones((1, 1)			) * model.config.eos_token_id
                SCREAMING_SNAKE_CASE_      :			Any     =       model(lowercase__			)
                self.assertIsNotNone(lowercase__			)






@require_flax
@require_sentencepiece
@require_tokenizers
class 	SCREAMING_SNAKE_CASE__			(				unittest.TestCase      ):
    _A			    =							[
        " UN Chief Says There Is No Military Solution in Syria",
    ]
    _A			    =							[
        "Şeful ONU declară că nu există o soluţie militară în Siria",
    ]
    _A			    =							"facebook/mbart-large-en-ro"




    @cached_property
    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          return AutoTokenizer.from_pretrained(self.model_name			)




    @cached_property
    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Tuple     =       FlaxMBartForConditionalGeneration.from_pretrained(self.model_name				,     from_pt=lowercase__			)
          return model




    def   __lowerCamelCase      (					self				,     **lowercase__			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Any     =       self.translate_src_text(**lowercase__			)
          self.assertListEqual(self.expected_text				,     lowercase__			)




    def   __lowerCamelCase      (					self				,     **lowercase__			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			List[str]     =       self.tokenizer(self.src_text				,     **lowercase__				,     return_tensors="np"			)
          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       self.model.generate(
              model_inputs.input_ids				,     attention_mask=model_inputs.attention_mask				,     decoder_start_token_id=self.tokenizer.lang_code_to_id["ro_RO"]				,     early_stopping=lowercase__				,     num_beams=2				,     ).sequences
          SCREAMING_SNAKE_CASE_      :			Any     =       self.tokenizer.batch_decode(lowercase__				,     skip_special_tokens=lowercase__			)
          return generated_words







    @slow
    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          self._assert_generated_batch_equal_expected()




'''simple docstring'''






from __future__ import annotations

from PIL import Image

# Define the first generation of cells
# fmt: off
snake_case_       =       [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]


def 					__lowerCamelCase			(      SCREAMING_SNAKE_CASE_     :	int    )							->   list[int]:
      """simple docstring"""

      return [int(SCREAMING_SNAKE_CASE_    ) for c in F"{ruleset:08}"[:8]]


def 					__lowerCamelCase			(      SCREAMING_SNAKE_CASE_     :	list[list[int]]  ,  SCREAMING_SNAKE_CASE_     :	list[int]  ,  SCREAMING_SNAKE_CASE_     :	int    )							->   list[int]:
      """simple docstring"""

      SCREAMING_SNAKE_CASE_      :			List[str]     =       len(cells[0]    )  # 31
      SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       []
      for i in range(SCREAMING_SNAKE_CASE_    ):
            # Get the neighbors of each cell
            # Handle neighbours outside bounds by using 0 as their value
            SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       0 if i == 0 else cells[time][i - 1]
            SCREAMING_SNAKE_CASE_      :			Tuple     =       0 if i == population - 1 else cells[time][i + 1]
            # Define a new cell and add it to the new generation
            SCREAMING_SNAKE_CASE_      :			Dict     =       7 - int(F"{left_neighbor}{cells[time][i]}{right_neighbor}"  ,  2    )
            next_generation.append(rule[situation]    )
      return next_generation


def 					__lowerCamelCase			(      SCREAMING_SNAKE_CASE_     :	list[list[int]]    )							->   Image.Image:
      """simple docstring"""

      SCREAMING_SNAKE_CASE_      :			List[Any]     =       Image.new("RGB"  ,  (len(cells[0]    ), len(SCREAMING_SNAKE_CASE_    ))    )
      SCREAMING_SNAKE_CASE_      :			Optional[int]     =       img.load()
      # Generates image
      for w in range(img.width    ):
            for h in range(img.height    ):
                  SCREAMING_SNAKE_CASE_      :			List[str]     =       2_5_5 - int(2_5_5 * cells[h][w]    )
                  SCREAMING_SNAKE_CASE_      :			str     =       (color, color, color)
      return img


if __name__ == "__main__":
      snake_case_       =       bin(int(input('Rule:\n').strip()))[2:]
      snake_case_       =       format_ruleset(int(rule_num))
      for time in range(1_6):
            CELLS.append(new_generation(CELLS, rule, time))
      snake_case_       =       generate_image(CELLS)
      # Uncomment to save the image
      # img.save(f"rule_{rule_num}.png")
      img.show()

