






import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import PaddingStrategy, logging
from .tokenization_realm import RealmTokenizer


UpperCAmelCase__							  =						logging.get_logger(__name__)

UpperCAmelCase__							  =						{"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}

UpperCAmelCase__							  =						{
    "vocab_file": {
        "google/realm-cc-news-pretrained-embedder": (
            "https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/vocab.txt"
        ),
        "google/realm-cc-news-pretrained-encoder": (
            "https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/vocab.txt"
        ),
        "google/realm-cc-news-pretrained-scorer": (
            "https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/vocab.txt"
        ),
        "google/realm-cc-news-pretrained-openqa": (
            "https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/vocab.txt"
        ),
        "google/realm-orqa-nq-openqa": "https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/vocab.txt",
        "google/realm-orqa-nq-reader": "https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/vocab.txt",
        "google/realm-orqa-wq-openqa": "https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/vocab.txt",
        "google/realm-orqa-wq-reader": "https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/vocab.txt",
    },
    "tokenizer_file": {
        "google/realm-cc-news-pretrained-embedder": (
            "https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/tokenizer.jsont"
        ),
        "google/realm-cc-news-pretrained-encoder": (
            "https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/tokenizer.json"
        ),
        "google/realm-cc-news-pretrained-scorer": (
            "https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/tokenizer.json"
        ),
        "google/realm-cc-news-pretrained-openqa": (
            "https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/tokenizer.json"
        ),
        "google/realm-orqa-nq-openqa": (
            "https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/tokenizer.json"
        ),
        "google/realm-orqa-nq-reader": (
            "https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/tokenizer.json"
        ),
        "google/realm-orqa-wq-openqa": (
            "https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/tokenizer.json"
        ),
        "google/realm-orqa-wq-reader": (
            "https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/tokenizer.json"
        ),
    },
}

UpperCAmelCase__							  =						{
    "google/realm-cc-news-pretrained-embedder": 512,
    "google/realm-cc-news-pretrained-encoder": 512,
    "google/realm-cc-news-pretrained-scorer": 512,
    "google/realm-cc-news-pretrained-openqa": 512,
    "google/realm-orqa-nq-openqa": 512,
    "google/realm-orqa-nq-reader": 512,
    "google/realm-orqa-wq-openqa": 512,
    "google/realm-orqa-wq-reader": 512,
}

UpperCAmelCase__							  =						{
    "google/realm-cc-news-pretrained-embedder": {"do_lower_case": True},
    "google/realm-cc-news-pretrained-encoder": {"do_lower_case": True},
    "google/realm-cc-news-pretrained-scorer": {"do_lower_case": True},
    "google/realm-cc-news-pretrained-openqa": {"do_lower_case": True},
    "google/realm-orqa-nq-openqa": {"do_lower_case": True},
    "google/realm-orqa-nq-reader": {"do_lower_case": True},
    "google/realm-orqa-wq-openqa": {"do_lower_case": True},
    "google/realm-orqa-wq-reader": {"do_lower_case": True},
}






class 						lowercase_   (      lowercase  ):






      '''simple docstring'''




      __snake_case	    =    VOCAB_FILES_NAMES
      __snake_case	    =    PRETRAINED_VOCAB_FILES_MAP
      __snake_case	    =    PRETRAINED_INIT_CONFIGURATION
      __snake_case	    =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __snake_case	    =    RealmTokenizer







      def __init__( self :		List[Any]					,   __UpperCAmelCase :		Union[str, Any]=None					,   __UpperCAmelCase :		Tuple=None					,   __UpperCAmelCase :		List[Any]=True					,   __UpperCAmelCase :		Any="[UNK]"					,   __UpperCAmelCase :		Dict="[SEP]"					,   __UpperCAmelCase :		List[str]="[PAD]"					,   __UpperCAmelCase :		str="[CLS]"					,   __UpperCAmelCase :		Optional[int]="[MASK]"					,   __UpperCAmelCase :		str=True					,   __UpperCAmelCase :		List[str]=None					,   **__UpperCAmelCase :		str					,   )       ->Optional[int]:




       """simple docstring"""


       super().__init__(
           __UpperCAmelCase					,   tokenizer_file=__UpperCAmelCase					,   do_lower_case=__UpperCAmelCase					,   unk_token=__UpperCAmelCase					,   sep_token=__UpperCAmelCase					,   pad_token=__UpperCAmelCase					,   cls_token=__UpperCAmelCase					,   mask_token=__UpperCAmelCase					,   tokenize_chinese_chars=__UpperCAmelCase					,   strip_accents=__UpperCAmelCase					,   **__UpperCAmelCase					,   )

       a						    = json.loads(self.backend_tokenizer.normalizer.__getstate__()    )
       if (
           normalizer_state.get('''lowercase'''					,   __UpperCAmelCase    ) != do_lower_case
           or normalizer_state.get('''strip_accents'''					,   __UpperCAmelCase    ) != strip_accents
           or normalizer_state.get('''handle_chinese_chars'''					,   __UpperCAmelCase    ) != tokenize_chinese_chars
       ):
        a						    = getattr(__UpperCAmelCase					,   normalizer_state.pop('''type'''    )    )
        a						    = do_lower_case
        a						    = strip_accents
        a						    = tokenize_chinese_chars
        a						    = normalizer_class(**__UpperCAmelCase    )

       a						    = do_lower_case







      def  __lowerCAmelCase   ( self :		str					,   __UpperCAmelCase :		List[Any]					,   **__UpperCAmelCase :		List[str]    )       ->List[Any]:




       """simple docstring"""


       a						    = PaddingStrategy.MAX_LENGTH

       a						    = text
       a						    = kwargs.pop('''text_pair'''					,   __UpperCAmelCase    )
       a						    = kwargs.pop('''return_tensors'''					,   __UpperCAmelCase    )

       a						    = {
           '''input_ids''': [],
           '''attention_mask''': [],
           '''token_type_ids''': [],
       }

       for idx, candidate_text in enumerate(__UpperCAmelCase    ):
        if batch_text_pair is not None:
         a						    = batch_text_pair[idx]
        else:
         a						    = None

        a						    = super().__call__(__UpperCAmelCase					,   __UpperCAmelCase					,   return_tensors=__UpperCAmelCase					,   **__UpperCAmelCase    )

        a						    = encoded_candidates.get('''input_ids'''    )
        a						    = encoded_candidates.get('''attention_mask'''    )
        a						    = encoded_candidates.get('''token_type_ids'''    )

        if encoded_input_ids is not None:
         output_data["input_ids"].append(__UpperCAmelCase    )
        if encoded_attention_mask is not None:
         output_data["attention_mask"].append(__UpperCAmelCase    )
        if encoded_token_type_ids is not None:
         output_data["token_type_ids"].append(__UpperCAmelCase    )

       a						    = {key: item for key, item in output_data.items() if len(__UpperCAmelCase    ) != 0}

       return BatchEncoding(__UpperCAmelCase					,   tensor_type=__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Tuple					,   __UpperCAmelCase :		Optional[int]					,   __UpperCAmelCase :		Dict=None    )       ->str:




       """simple docstring"""


       a						    = [self.cls_token_id] + token_ids_a + [self.sep_token_id]

       if token_ids_a:
        output += token_ids_a + [self.sep_token_id]

       return output







      def  __lowerCAmelCase   ( self :		str					,   __UpperCAmelCase :		List[int]					,   __UpperCAmelCase :		Optional[List[int]] = None    )       ->List[int]:




       """simple docstring"""


       a						    = [self.sep_token_id]
       a						    = [self.cls_token_id]
       if token_ids_a is None:
        return len(cls + token_ids_a + sep    ) * [0]
       return len(cls + token_ids_a + sep    ) * [0] + len(token_ids_a + sep    ) * [1]







      def  __lowerCAmelCase   ( self :		Dict					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		Optional[str] = None    )       ->Tuple[str]:




       """simple docstring"""


       a						    = self._tokenizer.model.save(__UpperCAmelCase					,   name=__UpperCAmelCase    )
       return tuple(__UpperCAmelCase    )










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 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




      def  __lowerCAmelCase   ( self :		Dict    )       ->str:




       """simple docstring"""


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()







      @property
      def  __lowerCAmelCase   ( self :		List[str]    )       ->Any:




       """simple docstring"""


       return 12







      @property
      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->Dict:




       """simple docstring"""


       return 12







      @property
      def  __lowerCAmelCase   ( self :		Any    )       ->Any:




       """simple docstring"""


       return 32







      @property
      def  __lowerCAmelCase   ( self :		int    )       ->List[str]:




       """simple docstring"""


       torch.manual_seed(0    )
       a						    = 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 :		str    )       ->str:




       """simple docstring"""


       a						    = CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    )
       return tokenizer







      @property
      def  __lowerCAmelCase   ( self :		str    )       ->Union[str, Any]:




       """simple docstring"""


       torch.manual_seed(0    )
       a						    = CLIPTextConfig(
           bos_token_id=0					,   eos_token_id=2					,   hidden_size=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=1_000					,   )
       return CLIPTextModel(__UpperCAmelCase    )







      @property
      def  __lowerCAmelCase   ( self :		Any    )       ->Tuple:




       """simple docstring"""


       torch.manual_seed(0    )

       a						    = 12
       a						    = 12

       a						    = {
           '''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''',
       }

       a						    = TransformeraDModel(**__UpperCAmelCase    )
       return model







      def  __lowerCAmelCase   ( self :		Dict    )       ->Any:




       """simple docstring"""


       a						    = '''cpu'''

       a						    = self.dummy_vqvae
       a						    = self.dummy_text_encoder
       a						    = self.dummy_tokenizer
       a						    = self.dummy_transformer
       a						    = VQDiffusionScheduler(self.num_embed    )
       a						    = LearnedClassifierFreeSamplingEmbeddings(learnable=__UpperCAmelCase    )

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

       a						    = '''teddy bear playing in the pool'''

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe([prompt]					,   generator=__UpperCAmelCase					,   num_inference_steps=2					,   output_type='''np'''    )
       a						    = output.images

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe(
           [prompt]					,   generator=__UpperCAmelCase					,   output_type='''np'''					,   return_dict=__UpperCAmelCase					,   num_inference_steps=2    )[0]

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

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

       a						    = 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 :		Tuple    )       ->str:




       """simple docstring"""


       a						    = '''cpu'''

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

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

       a						    = '''teddy bear playing in the pool'''

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe([prompt]					,   generator=__UpperCAmelCase					,   num_inference_steps=2					,   output_type='''np'''    )
       a						    = output.images

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe(
           [prompt]					,   generator=__UpperCAmelCase					,   output_type='''np'''					,   return_dict=__UpperCAmelCase					,   num_inference_steps=2    )[0]

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

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

       a						    = 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 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




      def  __lowerCAmelCase   ( self :		Union[str, Any]    )       ->Union[str, Any]:




       """simple docstring"""


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()







      def  __lowerCAmelCase   ( self :		Dict    )       ->int:




       """simple docstring"""


       a						    = load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''    )

       a						    = VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''    )
       a						    = pipeline.to(__UpperCAmelCase    )
       pipeline.set_progress_bar_config(disable=__UpperCAmelCase    )

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

       a						    = output.images[0]

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


