"""simple docstring"""



import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
   import os

   import jax.numpy as jnp
   from jax import jit

   from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
   from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

   SCREAMING_SNAKE_CASE__:Union[str, Any]									=			"""0.12"""  # assumed parallelism: 8


if is_torch_available():
   import torch



def 	_lowerCamelCase(	a   ,    a   ,    a=None				):
       if rng is None:
              __a       =    random.Random()

       __a       =    1
       for dim in shape:
              total_dims *= dim

       __a       =    []
       for _ in range(a				):
              values.append(rng.randint(0   ,    vocab_size - 1				)				)

       __a       =    np.array(a   ,    dtype=jnp.intaa				).reshape(a				)

       return output



def 	_lowerCamelCase(	a   ,    a=None				):
       __a       =    ids_tensor(a   ,    vocab_size=2   ,    rng=a				)
       # make sure that at least one token is attended to for each batch
       __a       =    1
       return attn_mask






@require_flax
class  snake_case__      :
 _snake_case       :					int							  =						None
 _snake_case       :					Optional[Any]							  =						()


 def 						a__   ( self						):
        __a ,   __a       =    self.model_tester.prepare_config_and_inputs_for_common()

        # cut to half length & take max batch_size 3
        __a       =    2
        __a       =    inputs["input_ids"].shape[-1] // 2
        __a       =    inputs["input_ids"][:max_batch_size, :sequence_length]

        __a       =    jnp.ones_like(lowerCamelCase						)
        __a       =    attention_mask[:max_batch_size, :sequence_length]

        # generate max 5 tokens
        __a       =    input_ids.shape[-1] + 5
        if config.eos_token_id is not None and config.pad_token_id is None:
               # hack to allow generate for models such as GPT2 as is done in `generate()`
               __a       =    config.eos_token_id
        return config, input_ids, attention_mask, max_length


 @is_pt_flax_cross_test
 def 						a__   ( self						):
        __a ,   __a ,   __a ,   __a       =    self._get_input_ids_and_config()
        __a       =    False
        __a       =    max_length
        __a       =    0

        for model_class in self.all_generative_model_classes:
               __a       =    model_class(lowerCamelCase						)

               __a       =    model_class.__name__[4:]  # Skip the "Flax" at the beginning
               __a       =    getattr(lowerCamelCase    ,		lowerCamelCase						)
               __a       =    pt_model_class(lowerCamelCase						).eval()
               __a       =    load_flax_weights_in_pytorch_model(lowerCamelCase    ,		flax_model.params						)

               __a       =    flax_model.generate(lowerCamelCase						).sequences
               __a       =    pt_model.generate(torch.tensor(lowerCamelCase    ,		dtype=torch.long						)						)

               if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
                      __a       =    flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

               self.assertListEqual(pt_generation_outputs.numpy().tolist()    ,		flax_generation_outputs.tolist()						)


 def 						a__   ( self						):
        __a ,   __a ,   __a ,   __a       =    self._get_input_ids_and_config()
        __a       =    False
        __a       =    max_length

        for model_class in self.all_generative_model_classes:
               __a       =    model_class(lowerCamelCase						)

               __a       =    model.generate(lowerCamelCase						).sequences
               self.assertEqual(generation_outputs.shape[-1]    ,		lowerCamelCase						)

               __a       =    jit(model.generate						)
               __a       =    jit_generate(lowerCamelCase						).sequences

               self.assertListEqual(generation_outputs.tolist()    ,		jit_generation_outputs.tolist()						)


 def 						a__   ( self						):
        __a ,   __a ,   __a ,   __a       =    self._get_input_ids_and_config()
        __a       =    True
        __a       =    max_length

        for model_class in self.all_generative_model_classes:
               __a       =    model_class(lowerCamelCase						)

               __a       =    model.generate(lowerCamelCase						).sequences
               self.assertEqual(generation_outputs.shape[-1]    ,		lowerCamelCase						)

               __a       =    jit(model.generate						)
               __a       =    jit_generate(lowerCamelCase						).sequences

               self.assertListEqual(generation_outputs.tolist()    ,		jit_generation_outputs.tolist()						)


 def 						a__   ( self						):
        __a ,   __a ,   __a ,   __a       =    self._get_input_ids_and_config()
        __a       =    False
        __a       =    max_length
        __a       =    2

        for model_class in self.all_generative_model_classes:
               __a       =    model_class(lowerCamelCase						)

               __a       =    model.generate(lowerCamelCase						).sequences
               self.assertEqual(generation_outputs.shape[-1]    ,		lowerCamelCase						)

               __a       =    jit(model.generate						)
               __a       =    jit_generate(lowerCamelCase						).sequences

               self.assertListEqual(generation_outputs.tolist()    ,		jit_generation_outputs.tolist()						)


 def 						a__   ( self						):
        __a ,   __a ,   __a ,   __a       =    self._get_input_ids_and_config()
        __a       =    False
        __a       =    max_length
        __a       =    2
        __a       =    2

        for model_class in self.all_generative_model_classes:
               __a       =    model_class(lowerCamelCase						)

               __a       =    model.generate(lowerCamelCase						).sequences
               self.assertEqual(generation_outputs.shape[0]    ,		input_ids.shape[0] * config.num_return_sequences						)


 def 						a__   ( self						):
        __a ,   __a ,   __a ,   __a       =    self._get_input_ids_and_config()
        __a       =    True
        __a       =    max_length
        __a       =    0.8
        __a       =    10
        __a       =    0.3
        __a       =    1
        __a       =    8
        __a       =    9

        for model_class in self.all_generative_model_classes:
               __a       =    model_class(lowerCamelCase						)

               __a       =    model.generate(lowerCamelCase						).sequences
               self.assertEqual(generation_outputs.shape[-1]    ,		lowerCamelCase						)

               __a       =    jit(model.generate						)
               __a       =    jit_generate(lowerCamelCase						).sequences

               self.assertListEqual(generation_outputs.tolist()    ,		jit_generation_outputs.tolist()						)


 def 						a__   ( self						):
        __a ,   __a ,   __a ,   __a       =    self._get_input_ids_and_config()
        __a       =    max_length
        __a       =    1
        __a       =    8
        __a       =    9

        for model_class in self.all_generative_model_classes:
               __a       =    model_class(lowerCamelCase						)

               __a       =    model.generate(lowerCamelCase						).sequences
               self.assertEqual(generation_outputs.shape[-1]    ,		lowerCamelCase						)

               __a       =    jit(model.generate						)
               __a       =    jit_generate(lowerCamelCase						).sequences

               self.assertListEqual(generation_outputs.tolist()    ,		jit_generation_outputs.tolist()						)


 def 						a__   ( self						):
        __a ,   __a ,   __a ,   __a       =    self._get_input_ids_and_config()
        __a       =    max_length
        __a       =    2
        __a       =    1
        __a       =    8
        __a       =    9

        for model_class in self.all_generative_model_classes:
               __a       =    model_class(lowerCamelCase						)

               __a       =    model.generate(lowerCamelCase						).sequences
               self.assertEqual(generation_outputs.shape[-1]    ,		lowerCamelCase						)

               __a       =    jit(model.generate						)
               __a       =    jit_generate(lowerCamelCase						).sequences

               self.assertListEqual(generation_outputs.tolist()    ,		jit_generation_outputs.tolist()						)


 def 						a__   ( self						):
        __a ,   __a ,   __a ,   __a       =    self._get_input_ids_and_config()

        # pad attention mask on the left
        __a       =    attention_mask.at[(0, 0)].set(0						)

        __a       =    False
        __a       =    max_length

        for model_class in self.all_generative_model_classes:
               __a       =    model_class(lowerCamelCase						)

               __a       =    model.generate(lowerCamelCase    ,		attention_mask=lowerCamelCase						).sequences
               self.assertEqual(generation_outputs.shape[-1]    ,		lowerCamelCase						)

               __a       =    jit(model.generate						)
               __a       =    jit_generate(lowerCamelCase    ,		attention_mask=lowerCamelCase						).sequences

               self.assertListEqual(generation_outputs.tolist()    ,		jit_generation_outputs.tolist()						)


 def 						a__   ( self						):
        __a ,   __a ,   __a ,   __a       =    self._get_input_ids_and_config()

        # pad attention mask on the left
        __a       =    attention_mask.at[(0, 0)].set(0						)

        __a       =    True
        __a       =    max_length

        for model_class in self.all_generative_model_classes:
               __a       =    model_class(lowerCamelCase						)

               __a       =    model.generate(lowerCamelCase    ,		attention_mask=lowerCamelCase						).sequences
               self.assertEqual(generation_outputs.shape[-1]    ,		lowerCamelCase						)

               __a       =    jit(model.generate						)
               __a       =    jit_generate(lowerCamelCase    ,		attention_mask=lowerCamelCase						).sequences

               self.assertListEqual(generation_outputs.tolist()    ,		jit_generation_outputs.tolist()						)






 def 						a__   ( self						):
        __a ,   __a ,   __a ,   __a       =    self._get_input_ids_and_config()

        # pad attention mask on the left
        __a       =    attention_mask.at[(0, 0)].set(0						)

        __a       =    2
        __a       =    max_length

        for model_class in self.all_generative_model_classes:
               __a       =    model_class(lowerCamelCase						)

               __a       =    model.generate(lowerCamelCase    ,		attention_mask=lowerCamelCase						).sequences
               self.assertEqual(generation_outputs.shape[-1]    ,		lowerCamelCase						)

               __a       =    jit(model.generate						)
               __a       =    jit_generate(lowerCamelCase    ,		attention_mask=lowerCamelCase						).sequences

               self.assertListEqual(generation_outputs.tolist()    ,		jit_generation_outputs.tolist()						)






@require_flax
class  snake_case__      (       unittest.TestCase       ):


 def 						a__   ( self						):
        __a       =    AutoTokenizer.from_pretrained("hf-internal-testing/tiny-bert"						)
        __a       =    FlaxAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-bert-flax-only"						)

        __a       =    "Hello world"
        __a       =    tokenizer(lowerCamelCase    ,		return_tensors="np"						).input_ids

        # typos are quickly detected (the correct argument is `do_sample`)
        with self.assertRaisesRegex(lowerCamelCase    ,		"do_samples"						):
               model.generate(lowerCamelCase    ,		do_samples=lowerCamelCase						)

        # arbitrary arguments that will not be used anywhere are also not accepted
        with self.assertRaisesRegex(lowerCamelCase    ,		"foo"						):
               __a       =    {"foo": "bar"}
               model.generate(lowerCamelCase    ,		**lowerCamelCase						)

"""simple docstring"""



import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible






class  snake_case__      (       unittest.TestCase       ):


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.bin",
            "safety_checker/model.safetensors",
            "vae/diffusion_pytorch_model.bin",
            "vae/diffusion_pytorch_model.safetensors",
            "text_encoder/pytorch_model.bin",
            "text_encoder/model.safetensors",
            "unet/diffusion_pytorch_model.bin",
            "unet/diffusion_pytorch_model.safetensors",
        ]
        self.assertTrue(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "unet/diffusion_pytorch_model.bin",
            "unet/diffusion_pytorch_model.safetensors",
        ]
        self.assertTrue(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.bin",
            "safety_checker/model.safetensors",
            "vae/diffusion_pytorch_model.bin",
            "vae/diffusion_pytorch_model.safetensors",
            "text_encoder/pytorch_model.bin",
            "text_encoder/model.safetensors",
            "unet/diffusion_pytorch_model.bin",
            # Removed: 'unet/diffusion_pytorch_model.safetensors',
        ]
        self.assertFalse(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "text_encoder/pytorch_model.bin",
            "text_encoder/model.safetensors",
        ]
        self.assertTrue(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.bin",
            "safety_checker/model.safetensors",
            "vae/diffusion_pytorch_model.bin",
            "vae/diffusion_pytorch_model.safetensors",
            "text_encoder/pytorch_model.bin",
            # Removed: 'text_encoder/model.safetensors',
            "unet/diffusion_pytorch_model.bin",
            "unet/diffusion_pytorch_model.safetensors",
        ]
        self.assertFalse(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.fp16.bin",
            "safety_checker/model.fp16.safetensors",
            "vae/diffusion_pytorch_model.fp16.bin",
            "vae/diffusion_pytorch_model.fp16.safetensors",
            "text_encoder/pytorch_model.fp16.bin",
            "text_encoder/model.fp16.safetensors",
            "unet/diffusion_pytorch_model.fp16.bin",
            "unet/diffusion_pytorch_model.fp16.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "unet/diffusion_pytorch_model.fp16.bin",
            "unet/diffusion_pytorch_model.fp16.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        # pass variant but use the non-variant filenames
        __a       =    [
            "unet/diffusion_pytorch_model.bin",
            "unet/diffusion_pytorch_model.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.fp16.bin",
            "safety_checker/model.fp16.safetensors",
            "vae/diffusion_pytorch_model.fp16.bin",
            "vae/diffusion_pytorch_model.fp16.safetensors",
            "text_encoder/pytorch_model.fp16.bin",
            "text_encoder/model.fp16.safetensors",
            "unet/diffusion_pytorch_model.fp16.bin",
            # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
        ]
        __a       =    "fp16"
        self.assertFalse(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "text_encoder/pytorch_model.fp16.bin",
            "text_encoder/model.fp16.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        # pass variant but use the non-variant filenames
        __a       =    [
            "text_encoder/pytorch_model.bin",
            "text_encoder/model.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)






 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.fp16.bin",
            "safety_checker/model.fp16.safetensors",
            "vae/diffusion_pytorch_model.fp16.bin",
            "vae/diffusion_pytorch_model.fp16.safetensors",
            "text_encoder/pytorch_model.fp16.bin",
            # 'text_encoder/model.fp16.safetensors',
            "unet/diffusion_pytorch_model.fp16.bin",
            "unet/diffusion_pytorch_model.fp16.safetensors",
        ]
        __a       =    "fp16"
        self.assertFalse(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)
