





"""simple docstring"""





import json
import os
import re
import unittest

from transformers import CodeGenTokenizer, CodeGenTokenizerFast
from transformers.models.codegen.tokenization_codegen import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin




@require_tokenizers
class 							snake_case							(			__UpperCAmelCase      ,  unittest.TestCase ):

      '''simple docstring'''



      _A     :       List[Any]         =       CodeGenTokenizer
      _A     :       str         =       CodeGenTokenizerFast
      _A     :       Any         =       True
      _A     :       List[str]         =       {'add_prefix_space': True}
      _A     :       List[Any]         =       False

      def       A_					(							self       :      Optional[Any] ):


        '''simple docstring'''




        super().setUp()

        # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
        __UpperCAmelCase    :						int       =						[
            '''l''',
            '''o''',
            '''w''',
            '''e''',
            '''r''',
            '''s''',
            '''t''',
            '''i''',
            '''d''',
            '''n''',
            '''\u0120''',
            '''\u0120l''',
            '''\u0120n''',
            '''\u0120lo''',
            '''\u0120low''',
            '''er''',
            '''\u0120lowest''',
            '''\u0120newer''',
            '''\u0120wider''',
            '''<unk>''',
            '''<|endoftext|>''',
        ]
        __UpperCAmelCase    :						Optional[int]       =						dict(zip(__lowercase						,				range(len(__lowercase ) ) ) )
        __UpperCAmelCase    :						Dict       =						['''#version: 0.2''', '''\u0120 l''', '''\u0120l o''', '''\u0120lo w''', '''e r''', '''''']
        __UpperCAmelCase    :						Dict       =						{'''unk_token''': '''<unk>'''}

        __UpperCAmelCase    :						Optional[int]       =						os.path.join(self.tmpdirname						,				VOCAB_FILES_NAMES['''vocab_file'''] )
        __UpperCAmelCase    :						List[Any]       =						os.path.join(self.tmpdirname						,				VOCAB_FILES_NAMES['''merges_file'''] )
        with open(self.vocab_file						,				'''w'''						,				encoding='''utf-8''' ) as fp:
          fp.write(json.dumps(__lowercase ) + '''\n''' )
        with open(self.merges_file						,				'''w'''						,				encoding='''utf-8''' ) as fp:
          fp.write('''\n'''.join(__lowercase ) )

      def       A_					(							self       :      List[str]						,				**__lowercase       :      Optional[int] ):


        '''simple docstring'''




        kwargs.update(self.special_tokens_map )
        return CodeGenTokenizer.from_pretrained(self.tmpdirname						,				**__lowercase )

      def       A_					(							self       :      Tuple						,				**__lowercase       :      Optional[Any] ):


        '''simple docstring'''




        kwargs.update(self.special_tokens_map )
        return CodeGenTokenizerFast.from_pretrained(self.tmpdirname						,				**__lowercase )

      def       A_					(							self       :      Tuple						,				__lowercase       :      Optional[Any] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[Any]       =						'''lower newer'''
        __UpperCAmelCase    :						Tuple       =						'''lower newer'''
        return input_text, output_text

      def       A_					(							self       :      Tuple ):


        '''simple docstring'''




        __UpperCAmelCase    :						Dict       =						CodeGenTokenizer(self.vocab_file						,				self.merges_file						,				**self.special_tokens_map )
        __UpperCAmelCase    :						int       =						'''lower newer'''
        __UpperCAmelCase    :						Dict       =						['''\u0120low''', '''er''', '''\u0120''', '''n''', '''e''', '''w''', '''er''']
        __UpperCAmelCase    :						Optional[Any]       =						tokenizer.tokenize(__lowercase						,				add_prefix_space=__lowercase )
        self.assertListEqual(__lowercase						,				__lowercase )

        __UpperCAmelCase    :						Optional[int]       =						tokens + [tokenizer.unk_token]
        __UpperCAmelCase    :						Any       =						[14, 15, 10, 9, 3, 2, 15, 19]
        self.assertListEqual(tokenizer.convert_tokens_to_ids(__lowercase )						,				__lowercase )

      def       A_					(							self       :      str ):


        '''simple docstring'''




        if not self.test_rust_tokenizer:
          return

        __UpperCAmelCase    :						Optional[Any]       =						self.get_tokenizer()
        __UpperCAmelCase    :						Any       =						self.get_rust_tokenizer(add_prefix_space=__lowercase )

        __UpperCAmelCase    :						Any       =						'''lower newer'''

        # Testing tokenization
        __UpperCAmelCase    :						int       =						tokenizer.tokenize(__lowercase						,				add_prefix_space=__lowercase )
        __UpperCAmelCase    :						Optional[Any]       =						rust_tokenizer.tokenize(__lowercase )
        self.assertListEqual(__lowercase						,				__lowercase )

        # Testing conversion to ids without special tokens
        __UpperCAmelCase    :						int       =						tokenizer.encode(__lowercase						,				add_special_tokens=__lowercase						,				add_prefix_space=__lowercase )
        __UpperCAmelCase    :						Optional[Any]       =						rust_tokenizer.encode(__lowercase						,				add_special_tokens=__lowercase )
        self.assertListEqual(__lowercase						,				__lowercase )

        # Testing conversion to ids with special tokens
        __UpperCAmelCase    :						Optional[int]       =						self.get_rust_tokenizer(add_prefix_space=__lowercase )
        __UpperCAmelCase    :						str       =						tokenizer.encode(__lowercase						,				add_prefix_space=__lowercase )
        __UpperCAmelCase    :						Dict       =						rust_tokenizer.encode(__lowercase )
        self.assertListEqual(__lowercase						,				__lowercase )

        # Testing the unknown token
        __UpperCAmelCase    :						Optional[int]       =						tokens + [rust_tokenizer.unk_token]
        __UpperCAmelCase    :						Dict       =						[14, 15, 10, 9, 3, 2, 15, 19]
        self.assertListEqual(rust_tokenizer.convert_tokens_to_ids(__lowercase )						,				__lowercase )

      def       A_					(							self       :      Any						,				*__lowercase       :      str						,				**__lowercase       :      str ):


        '''simple docstring'''




        pass

      def       A_					(							self       :      Dict						,				__lowercase       :      int=15 ):


        '''simple docstring'''




        for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
          with self.subTest(f'''{tokenizer.__class__.__name__} ({pretrained_name})''' ):
            __UpperCAmelCase    :						Any       =						self.rust_tokenizer_class.from_pretrained(__lowercase						,				**__lowercase )

            # Simple input
            __UpperCAmelCase    :						Union[str, Any]       =						'''This is a simple input'''
            __UpperCAmelCase    :						List[str]       =						['''This is a simple input 1''', '''This is a simple input 2''']
            __UpperCAmelCase    :						int       =						('''This is a simple input''', '''This is a pair''')
            __UpperCAmelCase    :						str       =						[
                ('''This is a simple input 1''', '''This is a simple input 2'''),
                ('''This is a simple pair 1''', '''This is a simple pair 2'''),
            ]

            # Simple input tests
            self.assertRaises(__lowercase						,				tokenizer_r.encode						,				__lowercase						,				max_length=__lowercase						,				padding='''max_length''' )

            # Simple input
            self.assertRaises(__lowercase						,				tokenizer_r.encode_plus						,				__lowercase						,				max_length=__lowercase						,				padding='''max_length''' )

            # Simple input
            self.assertRaises(
                __lowercase						,				tokenizer_r.batch_encode_plus						,				__lowercase						,				max_length=__lowercase						,				padding='''max_length'''						,				)

            # Pair input
            self.assertRaises(__lowercase						,				tokenizer_r.encode						,				__lowercase						,				max_length=__lowercase						,				padding='''max_length''' )

            # Pair input
            self.assertRaises(__lowercase						,				tokenizer_r.encode_plus						,				__lowercase						,				max_length=__lowercase						,				padding='''max_length''' )

            # Pair input
            self.assertRaises(
                __lowercase						,				tokenizer_r.batch_encode_plus						,				__lowercase						,				max_length=__lowercase						,				padding='''max_length'''						,				)

      def       A_					(							self       :      int ):


        '''simple docstring'''




        __UpperCAmelCase    :						Union[str, Any]       =						CodeGenTokenizer.from_pretrained(self.tmpdirname						,				pad_token='''<pad>''' )

        # Simple input
        __UpperCAmelCase    :						int       =						'''This is a simple input'''
        __UpperCAmelCase    :						List[str]       =						['''This is a simple input looooooooong''', '''This is a simple input''']
        __UpperCAmelCase    :						List[Any]       =						('''This is a simple input''', '''This is a pair''')
        __UpperCAmelCase    :						str       =						[
            ('''This is a simple input loooooong''', '''This is a simple input'''),
            ('''This is a simple pair loooooong''', '''This is a simple pair'''),
        ]

        __UpperCAmelCase    :						Any       =						tokenizer.pad_token_id

        __UpperCAmelCase    :						int       =						tokenizer(__lowercase						,				padding='''max_length'''						,				max_length=30						,				return_tensors='''np''' )
        __UpperCAmelCase    :						List[Any]       =						tokenizer(__lowercase						,				padding=__lowercase						,				truncate=__lowercase						,				return_tensors='''np''' )
        __UpperCAmelCase    :						Optional[int]       =						tokenizer(*__lowercase						,				padding='''max_length'''						,				max_length=60						,				return_tensors='''np''' )
        __UpperCAmelCase    :						Optional[Any]       =						tokenizer(__lowercase						,				padding=__lowercase						,				truncate=__lowercase						,				return_tensors='''np''' )

        # s
        # test single string max_length padding
        self.assertEqual(out_s['''input_ids'''].shape[-1]						,				30 )
        self.assertTrue(pad_token_id in out_s['''input_ids'''] )
        self.assertTrue(0 in out_s['''attention_mask'''] )

        # s2
        # test automatic padding
        self.assertEqual(out_sa['''input_ids'''].shape[-1]						,				33 )
        # long slice doesn't have padding
        self.assertFalse(pad_token_id in out_sa['''input_ids'''][0] )
        self.assertFalse(0 in out_sa['''attention_mask'''][0] )
        # short slice does have padding
        self.assertTrue(pad_token_id in out_sa['''input_ids'''][1] )
        self.assertTrue(0 in out_sa['''attention_mask'''][1] )

        # p
        # test single pair max_length padding
        self.assertEqual(out_p['''input_ids'''].shape[-1]						,				60 )
        self.assertTrue(pad_token_id in out_p['''input_ids'''] )
        self.assertTrue(0 in out_p['''attention_mask'''] )

        # p2
        # test automatic padding pair
        self.assertEqual(out_pa['''input_ids'''].shape[-1]						,				52 )
        # long slice pair doesn't have padding
        self.assertFalse(pad_token_id in out_pa['''input_ids'''][0] )
        self.assertFalse(0 in out_pa['''attention_mask'''][0] )
        # short slice pair does have padding
        self.assertTrue(pad_token_id in out_pa['''input_ids'''][1] )
        self.assertTrue(0 in out_pa['''attention_mask'''][1] )

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


        '''simple docstring'''




        __UpperCAmelCase    :						int       =						'''$$$'''
        __UpperCAmelCase    :						Optional[Any]       =						CodeGenTokenizer.from_pretrained(self.tmpdirname						,				bos_token=__lowercase						,				add_bos_token=__lowercase )

        __UpperCAmelCase    :						str       =						'''This is a simple input'''
        __UpperCAmelCase    :						Dict       =						['''This is a simple input 1''', '''This is a simple input 2''']

        __UpperCAmelCase    :						Any       =						tokenizer.bos_token_id

        __UpperCAmelCase    :						Any       =						tokenizer(__lowercase )
        __UpperCAmelCase    :						str       =						tokenizer(__lowercase )

        self.assertEqual(out_s.input_ids[0]						,				__lowercase )
        self.assertTrue(all(o[0] == bos_token_id for o in out_sa.input_ids ) )

        __UpperCAmelCase    :						Optional[int]       =						tokenizer.decode(out_s.input_ids )
        __UpperCAmelCase    :						Dict       =						tokenizer.batch_decode(out_sa.input_ids )

        self.assertEqual(decode_s.split()[0]						,				__lowercase )
        self.assertTrue(all(d.split()[0] == bos_token for d in decode_sa ) )

      @slow
      def       A_					(							self       :      int ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[Any]       =						CodeGenTokenizer.from_pretrained('''Salesforce/codegen-350M-mono''' )

        __UpperCAmelCase    :						Optional[int]       =						'''\nif len_a > len_b:\n    result = a\nelse:\n    result = b\n\n\n\n#'''
        __UpperCAmelCase    :						int       =						'''\nif len_a > len_b:      result = a\nelse:      result = b'''

        __UpperCAmelCase    :						Dict       =						tokenizer.encode(__lowercase )
        __UpperCAmelCase    :						Optional[Any]       =						['''^#''', re.escape('''<|endoftext|>''' ), '''^\'\'\'''', '''^"""''', '''\n\n\n''']
        __UpperCAmelCase    :						Any       =						tokenizer.decode(__lowercase						,				truncate_before_pattern=__lowercase )
        self.assertEqual(__lowercase						,				__lowercase )







      def       A_					(							self       :      List[Any] ):


        '''simple docstring'''




        pass






"""simple docstring"""





import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
  import jax
  import jax.numpy as jnp
  from flax.jax_utils import replicate
  from flax.training.common_utils import shard




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

      '''simple docstring'''



      def       A_					(							self       :      str ):


        '''simple docstring'''




        super().tearDown()
        gc.collect()

      def       A_					(							self       :      List[Any] ):


        '''simple docstring'''




        __UpperCAmelCase       ,     __UpperCAmelCase    :						List[Any]       =						FlaxStableDiffusionPipeline.from_pretrained(
            '''stabilityai/stable-diffusion-2'''						,				revision='''bf16'''						,				dtype=jnp.bfloataa						,				)

        __UpperCAmelCase    :						Dict       =						'''A painting of a squirrel eating a burger'''
        __UpperCAmelCase    :						Any       =						jax.device_count()
        __UpperCAmelCase    :						Optional[int]       =						num_samples * [prompt]
        __UpperCAmelCase    :						Tuple       =						sd_pipe.prepare_inputs(__lowercase )

        __UpperCAmelCase    :						List[Any]       =						replicate(__lowercase )
        __UpperCAmelCase    :						Optional[Any]       =						shard(__lowercase )

        __UpperCAmelCase    :						str       =						jax.random.PRNGKey(0 )
        __UpperCAmelCase    :						Dict       =						jax.random.split(__lowercase						,				jax.device_count() )

        __UpperCAmelCase    :						Tuple       =						sd_pipe(__lowercase						,				__lowercase						,				__lowercase						,				num_inference_steps=25						,				jit=__lowercase )[0]
        assert images.shape == (jax.device_count(), 1, 768, 768, 3)

        __UpperCAmelCase    :						Any       =						images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:] )
        __UpperCAmelCase    :						Union[str, Any]       =						images[0, 253:256, 253:256, -1]

        __UpperCAmelCase    :						str       =						jnp.asarray(jax.device_get(image_slice.flatten() ) )
        __UpperCAmelCase    :						List[Any]       =						jnp.array([0.4_2_3_8, 0.4_4_1_4, 0.4_3_9_5, 0.4_4_5_3, 0.4_6_2_9, 0.4_5_9_0, 0.4_5_3_1, 0.4_5_5_0_8, 0.4_5_1_2] )
        print(f'''output_slice: {output_slice}''' )
        assert jnp.abs(output_slice - expected_slice ).max() < 1e-2







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


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[int]       =						'''stabilityai/stable-diffusion-2'''
        __UpperCAmelCase       ,     __UpperCAmelCase    :						Any       =						FlaxDPMSolverMultistepScheduler.from_pretrained(__lowercase						,				subfolder='''scheduler''' )
        __UpperCAmelCase       ,     __UpperCAmelCase    :						Optional[int]       =						FlaxStableDiffusionPipeline.from_pretrained(
            __lowercase						,				scheduler=__lowercase						,				revision='''bf16'''						,				dtype=jnp.bfloataa						,				)
        __UpperCAmelCase    :						Dict       =						scheduler_params

        __UpperCAmelCase    :						int       =						'''A painting of a squirrel eating a burger'''
        __UpperCAmelCase    :						List[Any]       =						jax.device_count()
        __UpperCAmelCase    :						Optional[int]       =						num_samples * [prompt]
        __UpperCAmelCase    :						Tuple       =						sd_pipe.prepare_inputs(__lowercase )

        __UpperCAmelCase    :						Optional[int]       =						replicate(__lowercase )
        __UpperCAmelCase    :						List[Any]       =						shard(__lowercase )

        __UpperCAmelCase    :						Union[str, Any]       =						jax.random.PRNGKey(0 )
        __UpperCAmelCase    :						List[Any]       =						jax.random.split(__lowercase						,				jax.device_count() )

        __UpperCAmelCase    :						List[Any]       =						sd_pipe(__lowercase						,				__lowercase						,				__lowercase						,				num_inference_steps=25						,				jit=__lowercase )[0]
        assert images.shape == (jax.device_count(), 1, 768, 768, 3)

        __UpperCAmelCase    :						Tuple       =						images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:] )
        __UpperCAmelCase    :						Any       =						images[0, 253:256, 253:256, -1]

        __UpperCAmelCase    :						List[str]       =						jnp.asarray(jax.device_get(image_slice.flatten() ) )
        __UpperCAmelCase    :						Any       =						jnp.array([0.4_3_3_6, 0.4_2_9_6_9, 0.4_4_5_3, 0.4_1_9_9, 0.4_2_9_7, 0.4_5_3_1, 0.4_4_3_4, 0.4_4_3_4, 0.4_2_9_7] )
        print(f'''output_slice: {output_slice}''' )
        assert jnp.abs(output_slice - expected_slice ).max() < 1e-2