





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


UpperCamelCase__           =			False

class 							__SCREAMING_SNAKE_CASE       (  unittest.TestCase							):






							def 					_lowerCamelCase		(    self	):
														# clean up the VRAM after each test
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()






							@property
							def 					_lowerCamelCase		(    self	):
														return 12






							@property
							def 					_lowerCamelCase		(    self	):
														return 12






							@property
							def 					_lowerCamelCase		(    self	):
														return 32






							@property
							def 					_lowerCamelCase		(    self	):
														torch.manual_seed(0	)
														UpperCamelCase__					  =      VQModel(
														    block_out_channels=[32, 64]	,			in_channels=3	,			out_channels=3	,			down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]	,			up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]	,			latent_channels=3	,			num_vq_embeddings=self.num_embed	,			vq_embed_dim=3	,			)
														return model






							@property
							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)
														return tokenizer






							@property
							def 					_lowerCamelCase		(    self	):
														torch.manual_seed(0	)
														UpperCamelCase__					  =      CLIPTextConfig(
														    bos_token_id=0	,			eos_token_id=2	,			hidden_size=self.text_embedder_hidden_size	,			intermediate_size=37	,			layer_norm_eps=1E-05	,			num_attention_heads=4	,			num_hidden_layers=5	,			pad_token_id=1	,			vocab_size=1000	,			)
														return CLIPTextModel(__lowerCAmelCase	)






							@property
							def 					_lowerCamelCase		(    self	):
														torch.manual_seed(0	)

														UpperCamelCase__					  =      12
														UpperCamelCase__					  =      12

														UpperCamelCase__					  =      {
														    """attention_bias""": True,
														    """cross_attention_dim""": 32,
														    """attention_head_dim""": height * width,
														    """num_attention_heads""": 1,
														    """num_vector_embeds""": self.num_embed,
														    """num_embeds_ada_norm""": self.num_embeds_ada_norm,
														    """norm_num_groups""": 32,
														    """sample_size""": width,
														    """activation_fn""": """geglu-approximate""",
														}

														UpperCamelCase__					  =      TransformeraDModel(**__lowerCAmelCase	)
														return model






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      """cpu"""

														UpperCamelCase__					  =      self.dummy_vqvae
														UpperCamelCase__					  =      self.dummy_text_encoder
														UpperCamelCase__					  =      self.dummy_tokenizer
														UpperCamelCase__					  =      self.dummy_transformer
														UpperCamelCase__					  =      VQDiffusionScheduler(self.num_embed	)
														UpperCamelCase__					  =      LearnedClassifierFreeSamplingEmbeddings(learnable=__lowerCAmelCase	)

														UpperCamelCase__					  =      VQDiffusionPipeline(
														    vqvae=__lowerCAmelCase	,			text_encoder=__lowerCAmelCase	,			tokenizer=__lowerCAmelCase	,			transformer=__lowerCAmelCase	,			scheduler=__lowerCAmelCase	,			learned_classifier_free_sampling_embeddings=__lowerCAmelCase	,			)
														UpperCamelCase__					  =      pipe.to(__lowerCAmelCase	)
														pipe.set_progress_bar_config(disable=__lowerCAmelCase	)

														UpperCamelCase__					  =      """teddy bear playing in the pool"""

														UpperCamelCase__					  =      torch.Generator(device=__lowerCAmelCase	).manual_seed(0	)
														UpperCamelCase__					  =      pipe([prompt]	,			generator=__lowerCAmelCase	,			num_inference_steps=2	,			output_type="""np"""	)
														UpperCamelCase__					  =      output.images

														UpperCamelCase__					  =      torch.Generator(device=__lowerCAmelCase	).manual_seed(0	)
														UpperCamelCase__					  =      pipe(
														    [prompt]	,			generator=__lowerCAmelCase	,			output_type="""np"""	,			return_dict=__lowerCAmelCase	,			num_inference_steps=2	)[0]

														UpperCamelCase__					  =      image[0, -3:, -3:, -1]
														UpperCamelCase__					  =      image_from_tuple[0, -3:, -3:, -1]

														assert image.shape == (1, 24, 24, 3)

														UpperCamelCase__					  =      np.array([0.6551, 0.6168, 0.5008, 0.5676, 0.5659, 0.4295, 0.6073, 0.5599, 0.4992]	)

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







							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      """cpu"""

														UpperCamelCase__					  =      self.dummy_vqvae
														UpperCamelCase__					  =      self.dummy_text_encoder
														UpperCamelCase__					  =      self.dummy_tokenizer
														UpperCamelCase__					  =      self.dummy_transformer
														UpperCamelCase__					  =      VQDiffusionScheduler(self.num_embed	)
														UpperCamelCase__					  =      LearnedClassifierFreeSamplingEmbeddings(
														    learnable=__lowerCAmelCase	,			hidden_size=self.text_embedder_hidden_size	,			length=tokenizer.model_max_length	)

														UpperCamelCase__					  =      VQDiffusionPipeline(
														    vqvae=__lowerCAmelCase	,			text_encoder=__lowerCAmelCase	,			tokenizer=__lowerCAmelCase	,			transformer=__lowerCAmelCase	,			scheduler=__lowerCAmelCase	,			learned_classifier_free_sampling_embeddings=__lowerCAmelCase	,			)
														UpperCamelCase__					  =      pipe.to(__lowerCAmelCase	)
														pipe.set_progress_bar_config(disable=__lowerCAmelCase	)

														UpperCamelCase__					  =      """teddy bear playing in the pool"""

														UpperCamelCase__					  =      torch.Generator(device=__lowerCAmelCase	).manual_seed(0	)
														UpperCamelCase__					  =      pipe([prompt]	,			generator=__lowerCAmelCase	,			num_inference_steps=2	,			output_type="""np"""	)
														UpperCamelCase__					  =      output.images

														UpperCamelCase__					  =      torch.Generator(device=__lowerCAmelCase	).manual_seed(0	)
														UpperCamelCase__					  =      pipe(
														    [prompt]	,			generator=__lowerCAmelCase	,			output_type="""np"""	,			return_dict=__lowerCAmelCase	,			num_inference_steps=2	)[0]

														UpperCamelCase__					  =      image[0, -3:, -3:, -1]
														UpperCamelCase__					  =      image_from_tuple[0, -3:, -3:, -1]

														assert image.shape == (1, 24, 24, 3)

														UpperCamelCase__					  =      np.array([0.6693, 0.6075, 0.4959, 0.5701, 0.5583, 0.4333, 0.6171, 0.5684, 0.4988]	)

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






@slow
@require_torch_gpu
class 							__SCREAMING_SNAKE_CASE       (  unittest.TestCase							):






							def 					_lowerCamelCase		(    self	):
														# clean up the VRAM after each test
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()







							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      load_numpy(
														    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
														    """/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy"""	)

														UpperCamelCase__					  =      VQDiffusionPipeline.from_pretrained("""microsoft/vq-diffusion-ithq"""	)
														UpperCamelCase__					  =      pipeline.to(__lowerCAmelCase	)
														pipeline.set_progress_bar_config(disable=__lowerCAmelCase	)

														# requires GPU generator for gumbel softmax
														# don't use GPU generator in tests though
														UpperCamelCase__					  =      torch.Generator(device=__lowerCAmelCase	).manual_seed(0	)
														UpperCamelCase__					  =      pipeline(
														    """teddy bear playing in the pool"""	,			num_images_per_prompt=1	,			generator=__lowerCAmelCase	,			output_type="""np"""	,			)

														UpperCamelCase__					  =      output.images[0]

														assert image.shape == (256, 256, 3)
														assert np.abs(expected_image - image	).max() < 2.0











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 							__SCREAMING_SNAKE_CASE       (  _a ,		unittest.TestCase							):
							snake_case  :   Dict      =							CodeGenTokenizer
							snake_case  :   Dict      =							CodeGenTokenizerFast
							snake_case  :   Tuple      =							True
							snake_case  :   Optional[int]      =							{"""add_prefix_space""": True}
							snake_case  :   int      =							False






							def 					_lowerCamelCase		(    self	):
														super().setUp()

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

														UpperCamelCase__					  =      os.path.join(self.tmpdirname	,			VOCAB_FILES_NAMES["""vocab_file"""]	)
														UpperCamelCase__					  =      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(__lowerCAmelCase	) + """\n"""	)
														with open(self.merges_file	,			"""w"""	,			encoding="""utf-8"""	) as fp:
																					fp.write("""\n""".join(__lowerCAmelCase	)	)






							def 					_lowerCamelCase		(    self	,			**__lowerCAmelCase	):
														kwargs.update(self.special_tokens_map	)
														return CodeGenTokenizer.from_pretrained(self.tmpdirname	,			**__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	,			**__lowerCAmelCase	):
														kwargs.update(self.special_tokens_map	)
														return CodeGenTokenizerFast.from_pretrained(self.tmpdirname	,			**__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	):
														UpperCamelCase__					  =      """lower newer"""
														UpperCamelCase__					  =      """lower newer"""
														return input_text, output_text






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      CodeGenTokenizer(self.vocab_file	,			self.merges_file	,			**self.special_tokens_map	)
														UpperCamelCase__					  =      """lower newer"""
														UpperCamelCase__					  =      ["""\u0120low""", """er""", """\u0120""", """n""", """e""", """w""", """er"""]
														UpperCamelCase__					  =      tokenizer.tokenize(__lowerCAmelCase	,			add_prefix_space=__lowerCAmelCase	)
														self.assertListEqual(__lowerCAmelCase	,			__lowerCAmelCase	)

														UpperCamelCase__					  =      tokens + [tokenizer.unk_token]
														UpperCamelCase__					  =      [14, 15, 10, 9, 3, 2, 15, 19]
														self.assertListEqual(tokenizer.convert_tokens_to_ids(__lowerCAmelCase	)	,			__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	):
														if not self.test_rust_tokenizer:
																					return

														UpperCamelCase__					  =      self.get_tokenizer()
														UpperCamelCase__					  =      self.get_rust_tokenizer(add_prefix_space=__lowerCAmelCase	)

														UpperCamelCase__					  =      """lower newer"""

														# Testing tokenization
														UpperCamelCase__					  =      tokenizer.tokenize(__lowerCAmelCase	,			add_prefix_space=__lowerCAmelCase	)
														UpperCamelCase__					  =      rust_tokenizer.tokenize(__lowerCAmelCase	)
														self.assertListEqual(__lowerCAmelCase	,			__lowerCAmelCase	)

														# Testing conversion to ids without special tokens
														UpperCamelCase__					  =      tokenizer.encode(__lowerCAmelCase	,			add_special_tokens=__lowerCAmelCase	,			add_prefix_space=__lowerCAmelCase	)
														UpperCamelCase__					  =      rust_tokenizer.encode(__lowerCAmelCase	,			add_special_tokens=__lowerCAmelCase	)
														self.assertListEqual(__lowerCAmelCase	,			__lowerCAmelCase	)

														# Testing conversion to ids with special tokens
														UpperCamelCase__					  =      self.get_rust_tokenizer(add_prefix_space=__lowerCAmelCase	)
														UpperCamelCase__					  =      tokenizer.encode(__lowerCAmelCase	,			add_prefix_space=__lowerCAmelCase	)
														UpperCamelCase__					  =      rust_tokenizer.encode(__lowerCAmelCase	)
														self.assertListEqual(__lowerCAmelCase	,			__lowerCAmelCase	)

														# Testing the unknown token
														UpperCamelCase__					  =      tokens + [rust_tokenizer.unk_token]
														UpperCamelCase__					  =      [14, 15, 10, 9, 3, 2, 15, 19]
														self.assertListEqual(rust_tokenizer.convert_tokens_to_ids(__lowerCAmelCase	)	,			__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	,			*__lowerCAmelCase	,			**__lowerCAmelCase	):
														# It's very difficult to mix/test pretokenization with byte-level
														# And get both CodeGen and Roberta to work at the same time (mostly an issue of adding a space before the string)
														pass






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase=15	):
														for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
																					with self.subTest(f"""{tokenizer.__class__.__name__} ({pretrained_name})"""	):
																												UpperCamelCase__					  =      self.rust_tokenizer_class.from_pretrained(__lowerCAmelCase	,			**__lowerCAmelCase	)

																												# Simple input
																												UpperCamelCase__					  =      """This is a simple input"""
																												UpperCamelCase__					  =      ["""This is a simple input 1""", """This is a simple input 2"""]
																												UpperCamelCase__					  =      ("""This is a simple input""", """This is a pair""")
																												UpperCamelCase__					  =      [
																												    ("""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(__lowerCAmelCase	,			tokenizer_r.encode	,			__lowerCAmelCase	,			max_length=__lowerCAmelCase	,			padding="""max_length"""	)

																												# Simple input
																												self.assertRaises(__lowerCAmelCase	,			tokenizer_r.encode_plus	,			__lowerCAmelCase	,			max_length=__lowerCAmelCase	,			padding="""max_length"""	)

																												# Simple input
																												self.assertRaises(
																												    __lowerCAmelCase	,			tokenizer_r.batch_encode_plus	,			__lowerCAmelCase	,			max_length=__lowerCAmelCase	,			padding="""max_length"""	,			)

																												# Pair input
																												self.assertRaises(__lowerCAmelCase	,			tokenizer_r.encode	,			__lowerCAmelCase	,			max_length=__lowerCAmelCase	,			padding="""max_length"""	)

																												# Pair input
																												self.assertRaises(__lowerCAmelCase	,			tokenizer_r.encode_plus	,			__lowerCAmelCase	,			max_length=__lowerCAmelCase	,			padding="""max_length"""	)

																												# Pair input
																												self.assertRaises(
																												    __lowerCAmelCase	,			tokenizer_r.batch_encode_plus	,			__lowerCAmelCase	,			max_length=__lowerCAmelCase	,			padding="""max_length"""	,			)






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      CodeGenTokenizer.from_pretrained(self.tmpdirname	,			pad_token="""<pad>"""	)

														# Simple input
														UpperCamelCase__					  =      """This is a simple input"""
														UpperCamelCase__					  =      ["""This is a simple input looooooooong""", """This is a simple input"""]
														UpperCamelCase__					  =      ("""This is a simple input""", """This is a pair""")
														UpperCamelCase__					  =      [
														    ("""This is a simple input loooooong""", """This is a simple input"""),
														    ("""This is a simple pair loooooong""", """This is a simple pair"""),
														]

														UpperCamelCase__					  =      tokenizer.pad_token_id

														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	,			padding="""max_length"""	,			max_length=30	,			return_tensors="""np"""	)
														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	,			padding=__lowerCAmelCase	,			truncate=__lowerCAmelCase	,			return_tensors="""np"""	)
														UpperCamelCase__					  =      tokenizer(*__lowerCAmelCase	,			padding="""max_length"""	,			max_length=60	,			return_tensors="""np"""	)
														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	,			padding=__lowerCAmelCase	,			truncate=__lowerCAmelCase	,			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 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      """$$$"""
														UpperCamelCase__					  =      CodeGenTokenizer.from_pretrained(self.tmpdirname	,			bos_token=__lowerCAmelCase	,			add_bos_token=__lowerCAmelCase	)

														UpperCamelCase__					  =      """This is a simple input"""
														UpperCamelCase__					  =      ["""This is a simple input 1""", """This is a simple input 2"""]

														UpperCamelCase__					  =      tokenizer.bos_token_id

														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	)
														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	)

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

														UpperCamelCase__					  =      tokenizer.decode(out_s.input_ids	)
														UpperCamelCase__					  =      tokenizer.batch_decode(out_sa.input_ids	)

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






							@slow
							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      CodeGenTokenizer.from_pretrained("""Salesforce/codegen-350M-mono"""	)

														UpperCamelCase__					  =      """\nif len_a > len_b:\n    result = a\nelse:\n    result = b\n\n\n\n#"""
														UpperCamelCase__					  =      """\nif len_a > len_b:      result = a\nelse:      result = b"""

														UpperCamelCase__					  =      tokenizer.encode(__lowerCAmelCase	)
														UpperCamelCase__					  =      ["""^#""", re.escape("""<|endoftext|>"""	), """^'''""", """^\"\"\"""", """\n\n\n"""]
														UpperCamelCase__					  =      tokenizer.decode(__lowerCAmelCase	,			truncate_before_pattern=__lowerCAmelCase	)
														self.assertEqual(__lowerCAmelCase	,			__lowerCAmelCase	)







							def 					_lowerCamelCase		(    self	):
														pass




