"""simple docstring"""


import os
import shutil
import sys
import tempfile
import unittest
from pathlib import Path

import pytest

import transformers
from transformers import (
    BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP,
    AutoTokenizer,
    BertConfig,
    BertTokenizer,
    BertTokenizerFast,
    CTRLTokenizer,
    GPTaTokenizer,
    GPTaTokenizerFast,
    PreTrainedTokenizerFast,
    RobertaTokenizer,
    RobertaTokenizerFast,
    is_tokenizers_available,
)
from transformers.models.auto.configuration_auto import CONFIG_MAPPING, AutoConfig
from transformers.models.auto.tokenization_auto import (
    TOKENIZER_MAPPING,
    get_tokenizer_config,
    tokenizer_class_from_name,
)
from transformers.models.roberta.configuration_roberta import RobertaConfig
from transformers.testing_utils import (
    DUMMY_DIFF_TOKENIZER_IDENTIFIER,
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tokenizers,
    slow,
)


sys.path.append(str(Path(__file__).parent.parent.parent.parent / """utils"""))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_tokenization import CustomTokenizer  # noqa E402


if is_tokenizers_available():
    from test_module.custom_tokenization_fast import CustomTokenizerFast


class 							a					(		unittest.TestCase ):


    def   lowerCAmelCase_	(       self						:	List[Any]       ):
      _UpperCAmelCase    			=   0


    @slow
    def   lowerCAmelCase_	(       self						:	List[str]       ):
      for model_name in (x for x in BERT_PRETRAINED_CONFIG_ARCHIVE_MAP.keys() if "japanese" not in x):
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase       )
        self.assertIsNotNone(__lowerCAmelCase       )
        self.assertIsInstance(__lowerCAmelCase ,					(BertTokenizer, BertTokenizerFast)       )
        self.assertGreater(len(__lowerCAmelCase       ) ,					0       )

      for model_name in GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP.keys():
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase       )
        self.assertIsNotNone(__lowerCAmelCase       )
        self.assertIsInstance(__lowerCAmelCase ,					(GPTaTokenizer, GPTaTokenizerFast)       )
        self.assertGreater(len(__lowerCAmelCase       ) ,					0       )


    def   lowerCAmelCase_	(       self						:	Union[str, Any]       ):
      _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase       )
      self.assertIsInstance(__lowerCAmelCase ,					(BertTokenizer, BertTokenizerFast)       )
      self.assertEqual(tokenizer.vocab_size ,					12       )


    def   lowerCAmelCase_	(       self						:	List[str]       ):
      _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase       )
      self.assertIsInstance(__lowerCAmelCase ,					(RobertaTokenizer, RobertaTokenizerFast)       )
      self.assertEqual(tokenizer.vocab_size ,					20       )


    def   lowerCAmelCase_	(       self						:	Any       ):
      _UpperCAmelCase    			=   AutoConfig.from_pretrained(__lowerCAmelCase       )
      self.assertIsInstance(__lowerCAmelCase ,					__lowerCAmelCase       )
      # Check that tokenizer_type ≠ model_type
      _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase ,					config=__lowerCAmelCase       )
      self.assertIsInstance(__lowerCAmelCase ,					(BertTokenizer, BertTokenizerFast)       )
      self.assertEqual(tokenizer.vocab_size ,					12       )


    def   lowerCAmelCase_	(       self						:	List[str]       ):
      with tempfile.TemporaryDirectory() as tmp_dir:
        shutil.copy("""./tests/fixtures/vocab.txt""" ,					os.path.join(__lowerCAmelCase ,					"""vocab.txt"""       )       )

        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase ,					tokenizer_type="""bert""" ,					use_fast=__lowerCAmelCase       )
        self.assertIsInstance(__lowerCAmelCase ,					__lowerCAmelCase       )

      with tempfile.TemporaryDirectory() as tmp_dir:
        shutil.copy("""./tests/fixtures/vocab.json""" ,					os.path.join(__lowerCAmelCase ,					"""vocab.json"""       )       )
        shutil.copy("""./tests/fixtures/merges.txt""" ,					os.path.join(__lowerCAmelCase ,					"""merges.txt"""       )       )

        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase ,					tokenizer_type="""gpt2""" ,					use_fast=__lowerCAmelCase       )
        self.assertIsInstance(__lowerCAmelCase ,					__lowerCAmelCase       )


    @require_tokenizers
    def   lowerCAmelCase_	(       self						:	Tuple       ):
      with tempfile.TemporaryDirectory() as tmp_dir:
        shutil.copy("""./tests/fixtures/vocab.txt""" ,					os.path.join(__lowerCAmelCase ,					"""vocab.txt"""       )       )

        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase ,					tokenizer_type="""bert"""       )
        self.assertIsInstance(__lowerCAmelCase ,					__lowerCAmelCase       )

      with tempfile.TemporaryDirectory() as tmp_dir:
        shutil.copy("""./tests/fixtures/vocab.json""" ,					os.path.join(__lowerCAmelCase ,					"""vocab.json"""       )       )
        shutil.copy("""./tests/fixtures/merges.txt""" ,					os.path.join(__lowerCAmelCase ,					"""merges.txt"""       )       )

        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase ,					tokenizer_type="""gpt2"""       )
        self.assertIsInstance(__lowerCAmelCase ,					__lowerCAmelCase       )


    def   lowerCAmelCase_	(       self						:	List[Any]       ):
      with pytest.raises(__lowerCAmelCase       ):
        AutoTokenizer.from_pretrained("""./""" ,					tokenizer_type="""xxx"""       )


    @require_tokenizers
    def   lowerCAmelCase_	(       self						:	Optional[int]       ):
      for tokenizer_class in [BertTokenizer, BertTokenizerFast, AutoTokenizer]:
        _UpperCAmelCase    			=   tokenizer_class.from_pretrained("""wietsedv/bert-base-dutch-cased"""       )
        self.assertIsInstance(__lowerCAmelCase ,					(BertTokenizer, BertTokenizerFast)       )

        if isinstance(__lowerCAmelCase ,					__lowerCAmelCase       ):
          self.assertEqual(tokenizer.basic_tokenizer.do_lower_case ,					__lowerCAmelCase       )
        else:
          self.assertEqual(tokenizer.do_lower_case ,					__lowerCAmelCase       )

        self.assertEqual(tokenizer.model_max_length ,					512       )


    @require_tokenizers
    def   lowerCAmelCase_	(       self						:	List[str]       ):
      for tokenizer_class in [BertTokenizer, BertTokenizerFast, AutoTokenizer]:
        with self.assertRaisesRegex(
            __lowerCAmelCase ,					"""julien-c/herlolip-not-exists is not a local folder and is not a valid model identifier""" ,					):
          _UpperCAmelCase    			=   tokenizer_class.from_pretrained("""julien-c/herlolip-not-exists"""       )


    def   lowerCAmelCase_	(       self						:	Optional[Any]       ):
      # tests: https://github.com/huggingface/transformers/pull/13251
      # 1. models with `-`, e.g. xlm-roberta -> xlm_roberta
      # 2. models that don't remap 1-1 from model-name to model file, e.g., openai-gpt -> openai
      _UpperCAmelCase    			=   TOKENIZER_MAPPING.values()
      _UpperCAmelCase    			=   []

      for slow_tok, fast_tok in tokenizers:
        if slow_tok is not None:
          tokenizer_names.append(slow_tok.__name__       )

        if fast_tok is not None:
          tokenizer_names.append(fast_tok.__name__       )

      for tokenizer_name in tokenizer_names:
        # must find the right class
        tokenizer_class_from_name(__lowerCAmelCase       )


    @require_tokenizers
    def   lowerCAmelCase_	(       self						:	Optional[Any]       ):
      self.assertIsInstance(AutoTokenizer.from_pretrained("""bert-base-cased""" ,					use_fast=__lowerCAmelCase       ) ,					__lowerCAmelCase       )
      self.assertIsInstance(AutoTokenizer.from_pretrained("""bert-base-cased"""       ) ,					__lowerCAmelCase       )


    @require_tokenizers
    def   lowerCAmelCase_	(       self						:	int       ):
      _UpperCAmelCase    			=   AutoTokenizer.from_pretrained("""distilbert-base-uncased""" ,					do_lower_case=__lowerCAmelCase       )
      _UpperCAmelCase    			=   """Hello, world. How are you?"""
      _UpperCAmelCase    			=   tokenizer.tokenize(__lowerCAmelCase       )
      self.assertEqual("""[UNK]""" ,					tokens[0]       )

      _UpperCAmelCase    			=   AutoTokenizer.from_pretrained("""microsoft/mpnet-base""" ,					do_lower_case=__lowerCAmelCase       )
      _UpperCAmelCase    			=   tokenizer.tokenize(__lowerCAmelCase       )
      self.assertEqual("""[UNK]""" ,					tokens[0]       )


    @require_tokenizers
    def   lowerCAmelCase_	(       self						:	Optional[Any]       ):
      _UpperCAmelCase    			=   AutoTokenizer.from_pretrained("""robot-test/dummy-tokenizer-fast-with-model-config"""       )
      self.assertEqual(type(__lowerCAmelCase       ) ,					__lowerCAmelCase       )
      self.assertEqual(tokenizer.model_max_length ,					512       )
      self.assertEqual(tokenizer.vocab_size ,					3_0000       )
      self.assertEqual(tokenizer.unk_token ,					"""[UNK]"""       )
      self.assertEqual(tokenizer.padding_side ,					"""right"""       )
      self.assertEqual(tokenizer.truncation_side ,					"""right"""       )


    def   lowerCAmelCase_	(       self						:	Optional[int]       ):
      _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase       )
      self.assertIsInstance(__lowerCAmelCase ,					(BertTokenizer, BertTokenizerFast)       )
      with tempfile.TemporaryDirectory() as tmp_dir:
        tokenizer.save_pretrained(__lowerCAmelCase       )
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase       )

      self.assertIsInstance(__lowerCAmelCase ,					tokenizer.__class__       )
      self.assertEqual(tokenizera.vocab_size ,					12       )


    def   lowerCAmelCase_	(       self						:	Optional[int]       ):
      _UpperCAmelCase    			=   AutoTokenizer.from_pretrained("""ctrl"""       )
      # There is no fast CTRL so this always gives us a slow tokenizer.
      self.assertIsInstance(__lowerCAmelCase ,					__lowerCAmelCase       )


    def   lowerCAmelCase_	(       self						:	str       ):
      # Check we can load the tokenizer config of an online model.
      _UpperCAmelCase    			=   get_tokenizer_config("""bert-base-cased"""       )
      _UpperCAmelCase    			=   config.pop("""_commit_hash""" ,					__lowerCAmelCase       )
      # If we ever update bert-base-cased tokenizer config, this dict here will need to be updated.
      self.assertEqual(__lowerCAmelCase ,					{"""do_lower_case""": False}       )

      # This model does not have a tokenizer_config so we get back an empty dict.
      _UpperCAmelCase    			=   get_tokenizer_config(__lowerCAmelCase       )
      self.assertDictEqual(__lowerCAmelCase ,					{}       )

      # A tokenizer saved with `save_pretrained` always creates a tokenizer config.
      _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase       )
      with tempfile.TemporaryDirectory() as tmp_dir:
        tokenizer.save_pretrained(__lowerCAmelCase       )
        _UpperCAmelCase    			=   get_tokenizer_config(__lowerCAmelCase       )

      # Check the class of the tokenizer was properly saved (note that it always saves the slow class).
      self.assertEqual(config["""tokenizer_class"""] ,					"""BertTokenizer"""       )


    def   lowerCAmelCase_	(       self						:	List[Any]       ):
      try:
        AutoConfig.register("""custom""" ,					__lowerCAmelCase       )

        AutoTokenizer.register(__lowerCAmelCase ,					slow_tokenizer_class=__lowerCAmelCase       )
        # Trying to register something existing in the Transformers library will raise an error
        with self.assertRaises(__lowerCAmelCase       ):
          AutoTokenizer.register(__lowerCAmelCase ,					slow_tokenizer_class=__lowerCAmelCase       )

        _UpperCAmelCase    			=   CustomTokenizer.from_pretrained(__lowerCAmelCase       )
        with tempfile.TemporaryDirectory() as tmp_dir:
          tokenizer.save_pretrained(__lowerCAmelCase       )

          _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase       )
          self.assertIsInstance(__lowerCAmelCase ,					__lowerCAmelCase       )

      finally:
        if "custom" in CONFIG_MAPPING._extra_content:
          del CONFIG_MAPPING._extra_content["custom"]
        if CustomConfig in TOKENIZER_MAPPING._extra_content:
          del TOKENIZER_MAPPING._extra_content[CustomConfig]


    @require_tokenizers
    def   lowerCAmelCase_	(       self						:	List[Any]       ):
      try:
        AutoConfig.register("""custom""" ,					__lowerCAmelCase       )

        # Can register in two steps
        AutoTokenizer.register(__lowerCAmelCase ,					slow_tokenizer_class=__lowerCAmelCase       )
        self.assertEqual(TOKENIZER_MAPPING[CustomConfig] ,					(CustomTokenizer, None)       )
        AutoTokenizer.register(__lowerCAmelCase ,					fast_tokenizer_class=__lowerCAmelCase       )
        self.assertEqual(TOKENIZER_MAPPING[CustomConfig] ,					(CustomTokenizer, CustomTokenizerFast)       )

        del TOKENIZER_MAPPING._extra_content[CustomConfig]
        # Can register in one step
        AutoTokenizer.register(
            __lowerCAmelCase ,					slow_tokenizer_class=__lowerCAmelCase ,					fast_tokenizer_class=__lowerCAmelCase       )
        self.assertEqual(TOKENIZER_MAPPING[CustomConfig] ,					(CustomTokenizer, CustomTokenizerFast)       )

        # Trying to register something existing in the Transformers library will raise an error
        with self.assertRaises(__lowerCAmelCase       ):
          AutoTokenizer.register(__lowerCAmelCase ,					fast_tokenizer_class=__lowerCAmelCase       )

        # We pass through a bert tokenizer fast cause there is no converter slow to fast for our new toknizer
        # and that model does not have a tokenizer.json
        with tempfile.TemporaryDirectory() as tmp_dir:
          _UpperCAmelCase    			=   BertTokenizerFast.from_pretrained(__lowerCAmelCase       )
          bert_tokenizer.save_pretrained(__lowerCAmelCase       )
          _UpperCAmelCase    			=   CustomTokenizerFast.from_pretrained(__lowerCAmelCase       )

        with tempfile.TemporaryDirectory() as tmp_dir:
          tokenizer.save_pretrained(__lowerCAmelCase       )

          _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase       )
          self.assertIsInstance(__lowerCAmelCase ,					__lowerCAmelCase       )

          _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase ,					use_fast=__lowerCAmelCase       )
          self.assertIsInstance(__lowerCAmelCase ,					__lowerCAmelCase       )

      finally:
        if "custom" in CONFIG_MAPPING._extra_content:
          del CONFIG_MAPPING._extra_content["custom"]
        if CustomConfig in TOKENIZER_MAPPING._extra_content:
          del TOKENIZER_MAPPING._extra_content[CustomConfig]


    def   lowerCAmelCase_	(       self						:	List[str]       ):
      # If remote code is not set, we will time out when asking whether to load the model.
      with self.assertRaises(__lowerCAmelCase       ):
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained("""hf-internal-testing/test_dynamic_tokenizer"""       )
      # If remote code is disabled, we can't load this config.
      with self.assertRaises(__lowerCAmelCase       ):
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(
            """hf-internal-testing/test_dynamic_tokenizer""" ,					trust_remote_code=__lowerCAmelCase       )

      _UpperCAmelCase    			=   AutoTokenizer.from_pretrained("""hf-internal-testing/test_dynamic_tokenizer""" ,					trust_remote_code=__lowerCAmelCase       )
      self.assertTrue(tokenizer.special_attribute_present       )
      # Test tokenizer can be reloaded.
      with tempfile.TemporaryDirectory() as tmp_dir:
        tokenizer.save_pretrained(__lowerCAmelCase       )
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase ,					trust_remote_code=__lowerCAmelCase       )
      self.assertTrue(reloaded_tokenizer.special_attribute_present       )

      if is_tokenizers_available():
        self.assertEqual(tokenizer.__class__.__name__ ,					"""NewTokenizerFast"""       )
        self.assertEqual(reloaded_tokenizer.__class__.__name__ ,					"""NewTokenizerFast"""       )

        # Test we can also load the slow version
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(
            """hf-internal-testing/test_dynamic_tokenizer""" ,					trust_remote_code=__lowerCAmelCase ,					use_fast=__lowerCAmelCase       )
        self.assertTrue(tokenizer.special_attribute_present       )
        self.assertEqual(tokenizer.__class__.__name__ ,					"""NewTokenizer"""       )
        # Test tokenizer can be reloaded.
        with tempfile.TemporaryDirectory() as tmp_dir:
          tokenizer.save_pretrained(__lowerCAmelCase       )
          _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase ,					trust_remote_code=__lowerCAmelCase ,					use_fast=__lowerCAmelCase       )
        self.assertEqual(reloaded_tokenizer.__class__.__name__ ,					"""NewTokenizer"""       )
        self.assertTrue(reloaded_tokenizer.special_attribute_present       )
      else:
        self.assertEqual(tokenizer.__class__.__name__ ,					"""NewTokenizer"""       )
        self.assertEqual(reloaded_tokenizer.__class__.__name__ ,					"""NewTokenizer"""       )


    @require_tokenizers
    def   lowerCAmelCase_	(       self						:	List[str]       ):


      class 							a					(		lowerCAmelCase_ ):
          _snake_case  :	Tuple  				=				False



      class 							a					(		lowerCAmelCase_ ):
          _snake_case  :	Optional[int]  				=				NewTokenizer
          _snake_case  :	Any  				=				False

      try:
        AutoConfig.register("""custom""" ,					__lowerCAmelCase       )
        AutoTokenizer.register(__lowerCAmelCase ,					slow_tokenizer_class=__lowerCAmelCase       )
        AutoTokenizer.register(__lowerCAmelCase ,					fast_tokenizer_class=__lowerCAmelCase       )
        # If remote code is not set, the default is to use local
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained("""hf-internal-testing/test_dynamic_tokenizer"""       )
        self.assertEqual(tokenizer.__class__.__name__ ,					"""NewTokenizerFast"""       )
        self.assertFalse(tokenizer.special_attribute_present       )
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained("""hf-internal-testing/test_dynamic_tokenizer""" ,					use_fast=__lowerCAmelCase       )
        self.assertEqual(tokenizer.__class__.__name__ ,					"""NewTokenizer"""       )
        self.assertFalse(tokenizer.special_attribute_present       )

        # If remote code is disabled, we load the local one.
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(
            """hf-internal-testing/test_dynamic_tokenizer""" ,					trust_remote_code=__lowerCAmelCase       )
        self.assertEqual(tokenizer.__class__.__name__ ,					"""NewTokenizerFast"""       )
        self.assertFalse(tokenizer.special_attribute_present       )
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(
            """hf-internal-testing/test_dynamic_tokenizer""" ,					trust_remote_code=__lowerCAmelCase ,					use_fast=__lowerCAmelCase       )
        self.assertEqual(tokenizer.__class__.__name__ ,					"""NewTokenizer"""       )
        self.assertFalse(tokenizer.special_attribute_present       )

        # If remote is enabled, we load from the Hub
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(
            """hf-internal-testing/test_dynamic_tokenizer""" ,					trust_remote_code=__lowerCAmelCase       )
        self.assertEqual(tokenizer.__class__.__name__ ,					"""NewTokenizerFast"""       )
        self.assertTrue(tokenizer.special_attribute_present       )
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(
            """hf-internal-testing/test_dynamic_tokenizer""" ,					trust_remote_code=__lowerCAmelCase ,					use_fast=__lowerCAmelCase       )
        self.assertEqual(tokenizer.__class__.__name__ ,					"""NewTokenizer"""       )
        self.assertTrue(tokenizer.special_attribute_present       )

      finally:
        if "custom" in CONFIG_MAPPING._extra_content:
          del CONFIG_MAPPING._extra_content["custom"]
        if CustomConfig in TOKENIZER_MAPPING._extra_content:
          del TOKENIZER_MAPPING._extra_content[CustomConfig]


    def   lowerCAmelCase_	(       self						:	Dict       ):
      _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(
          """hf-internal-testing/test_dynamic_tokenizer_legacy""" ,					trust_remote_code=__lowerCAmelCase       )
      self.assertTrue(tokenizer.special_attribute_present       )
      if is_tokenizers_available():
        self.assertEqual(tokenizer.__class__.__name__ ,					"""NewTokenizerFast"""       )

        # Test we can also load the slow version
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(
            """hf-internal-testing/test_dynamic_tokenizer_legacy""" ,					trust_remote_code=__lowerCAmelCase ,					use_fast=__lowerCAmelCase       )
        self.assertTrue(tokenizer.special_attribute_present       )
        self.assertEqual(tokenizer.__class__.__name__ ,					"""NewTokenizer"""       )
      else:
        self.assertEqual(tokenizer.__class__.__name__ ,					"""NewTokenizer"""       )


    def   lowerCAmelCase_	(       self						:	List[Any]       ):
      with self.assertRaisesRegex(
          __lowerCAmelCase ,					"""bert-base is not a local folder and is not a valid model identifier"""       ):
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained("""bert-base"""       )


    def   lowerCAmelCase_	(       self						:	int       ):
      with self.assertRaisesRegex(
          __lowerCAmelCase ,					R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""       ):
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained(__lowerCAmelCase ,					revision="""aaaaaa"""       )


    def   lowerCAmelCase_	(       self						:	List[str]       ):
      # Make sure we have cached the tokenizer.
      _UpperCAmelCase    			=   AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bert"""       )
      with RequestCounter() as counter:
        _UpperCAmelCase    			=   AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bert"""       )
        self.assertEqual(counter.get_request_count ,					0       )
        self.assertEqual(counter.head_request_count ,					1       )
        self.assertEqual(counter.other_request_count ,					0       )

"""simple docstring"""


from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from transformers import CLIPImageProcessor, CLIPVisionModel

from ...models import PriorTransformer
from ...pipelines import DiffusionPipeline
from ...schedulers import HeunDiscreteScheduler
from ...utils import (
    BaseOutput,
    is_accelerate_available,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from .renderer import ShapERenderer


UpperCAmelCase__			=    logging.get_logger(__name__)  # pylint: disable=invalid-name

UpperCAmelCase__			=    """
    Examples:
        ```py
        >>> from PIL import Image
        >>> import torch
        >>> from diffusers import DiffusionPipeline
        >>> from diffusers.utils import export_to_gif, load_image

        >>> device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")

        >>> repo = \"openai/shap-e-img2img\"
        >>> pipe = DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.float16)
        >>> pipe = pipe.to(device)

        >>> guidance_scale = 3.0
        >>> image_url = \"https://hf.co/datasets/diffusers/docs-images/resolve/main/shap-e/corgi.png\"
        >>> image = load_image(image_url).convert(\"RGB\")

        >>> images = pipe(
        ...     image,
        ...     guidance_scale=guidance_scale,
        ...     num_inference_steps=64,
        ...     frame_size=256,
        ... ).images

        >>> gif_path = export_to_gif(images[0], \"corgi_3d.gif\")
        ```
"""


@dataclass
class 							a					(		lowerCAmelCase_ ):
    _snake_case  :	Union[PIL.Image.Image, np.ndarray]



class 							a					(		lowerCAmelCase_ ):


    def __init__(       self						:	Dict ,					__lowerCAmelCase						:	PriorTransformer ,					__lowerCAmelCase						:	CLIPVisionModel ,					__lowerCAmelCase						:	CLIPImageProcessor ,					__lowerCAmelCase						:	HeunDiscreteScheduler ,					__lowerCAmelCase						:	ShapERenderer ,					):
      super().__init__()

      self.register_modules(
          prior=__lowerCAmelCase ,					image_encoder=__lowerCAmelCase ,					image_processor=__lowerCAmelCase ,					scheduler=__lowerCAmelCase ,					renderer=__lowerCAmelCase ,					)


    def   lowerCAmelCase_	(       self						:	List[Any] ,					__lowerCAmelCase						:	Dict ,					__lowerCAmelCase						:	Union[str, Any] ,					__lowerCAmelCase						:	str ,					__lowerCAmelCase						:	List[Any] ,					__lowerCAmelCase						:	Dict ,					__lowerCAmelCase						:	List[Any]       ):
      if latents is None:
        _UpperCAmelCase    			=   randn_tensor(__lowerCAmelCase ,					generator=__lowerCAmelCase ,					device=__lowerCAmelCase ,					dtype=__lowerCAmelCase       )
      else:
        if latents.shape != shape:
          raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {shape}'''       )
        _UpperCAmelCase    			=   latents.to(__lowerCAmelCase       )

      _UpperCAmelCase    			=   latents * scheduler.init_noise_sigma
      return latents


    def   lowerCAmelCase_	(       self						:	str ,					__lowerCAmelCase						:	List[str]=0       ):
      if is_accelerate_available():
        from accelerate import cpu_offload
      else:
        raise ImportError("""Please install accelerate via `pip install accelerate`"""       )

      _UpperCAmelCase    			=   torch.device(f'''cuda:{gpu_id}'''       )

      _UpperCAmelCase    			=   [self.image_encoder, self.prior]
      for cpu_offloaded_model in models:
        if cpu_offloaded_model is not None:
          cpu_offload(__lowerCAmelCase ,					__lowerCAmelCase       )


    @property
    def   lowerCAmelCase_	(       self						:	List[str]       ):
      if self.device != torch.device("""meta"""       ) or not hasattr(self.image_encoder ,					"""_hf_hook"""       ):
        return self.device
      for module in self.image_encoder.modules():
        if (
            hasattr(__lowerCAmelCase ,					"""_hf_hook"""       )
            and hasattr(module._hf_hook ,					"""execution_device"""       )
            and module._hf_hook.execution_device is not None
        ):
          return torch.device(module._hf_hook.execution_device       )
      return self.device


    def   lowerCAmelCase_	(       self						:	int ,					__lowerCAmelCase						:	Optional[int] ,					__lowerCAmelCase						:	Any ,					__lowerCAmelCase						:	int ,					__lowerCAmelCase						:	Any ,					):
      if isinstance(__lowerCAmelCase ,					__lowerCAmelCase       ) and isinstance(image[0] ,					torch.Tensor       ):
        _UpperCAmelCase    			=   torch.cat(__lowerCAmelCase ,					axis=0       ) if image[0].ndim == 4 else torch.stack(__lowerCAmelCase ,					axis=0       )

      if not isinstance(__lowerCAmelCase ,					torch.Tensor       ):
        _UpperCAmelCase    			=   self.image_processor(__lowerCAmelCase ,					return_tensors="""pt"""       ).pixel_values[0].unsqueeze(0       )

      _UpperCAmelCase    			=   image.to(dtype=self.image_encoder.dtype ,					device=__lowerCAmelCase       )

      _UpperCAmelCase    			=   self.image_encoder(__lowerCAmelCase       )["""last_hidden_state"""]
      _UpperCAmelCase    			=   image_embeds[:, 1:, :].contiguous()  # batch_size, dim, 256

      _UpperCAmelCase    			=   image_embeds.repeat_interleave(__lowerCAmelCase ,					dim=0       )

      if do_classifier_free_guidance:
        _UpperCAmelCase    			=   torch.zeros_like(__lowerCAmelCase       )

        # For classifier free guidance, we need to do two forward passes.
        # Here we concatenate the unconditional and text embeddings into a single batch
        # to avoid doing two forward passes
        _UpperCAmelCase    			=   torch.cat([negative_image_embeds, image_embeds]       )

      return image_embeds


    @torch.no_grad()
    @replace_example_docstring(__lowerCAmelCase       )
    def __call__(       self						:	Union[str, Any] ,					__lowerCAmelCase						:	Union[PIL.Image.Image, List[PIL.Image.Image]] ,					__lowerCAmelCase						:	int = 1 ,					__lowerCAmelCase						:	int = 25 ,					__lowerCAmelCase						:	Optional[Union[torch.Generator, List[torch.Generator]]] = None ,					__lowerCAmelCase						:	Optional[torch.FloatTensor] = None ,					__lowerCAmelCase						:	float = 4.0 ,					__lowerCAmelCase						:	int = 64 ,					__lowerCAmelCase						:	Optional[str] = "pil" ,					__lowerCAmelCase						:	bool = True ,					):

      if isinstance(__lowerCAmelCase ,					PIL.Image.Image       ):
        _UpperCAmelCase    			=   1
      elif isinstance(__lowerCAmelCase ,					torch.Tensor       ):
        _UpperCAmelCase    			=   image.shape[0]
      elif isinstance(__lowerCAmelCase ,					__lowerCAmelCase       ) and isinstance(image[0] ,					(torch.Tensor, PIL.Image.Image)       ):
        _UpperCAmelCase    			=   len(__lowerCAmelCase       )
      else:
        raise ValueError(
            f'''`image` has to be of type `PIL.Image.Image`, `torch.Tensor`, `List[PIL.Image.Image]` or `List[torch.Tensor]` but is {type(__lowerCAmelCase       )}'''       )

      _UpperCAmelCase    			=   self._execution_device

      _UpperCAmelCase    			=   batch_size * num_images_per_prompt

      _UpperCAmelCase    			=   guidance_scale > 1.0
      _UpperCAmelCase    			=   self._encode_image(__lowerCAmelCase ,					__lowerCAmelCase ,					__lowerCAmelCase ,					__lowerCAmelCase       )

      # prior

      self.scheduler.set_timesteps(__lowerCAmelCase ,					device=__lowerCAmelCase       )
      _UpperCAmelCase    			=   self.scheduler.timesteps

      _UpperCAmelCase    			=   self.prior.config.num_embeddings
      _UpperCAmelCase    			=   self.prior.config.embedding_dim

      _UpperCAmelCase    			=   self.prepare_latents(
          (batch_size, num_embeddings * embedding_dim) ,					image_embeds.dtype ,					__lowerCAmelCase ,					__lowerCAmelCase ,					__lowerCAmelCase ,					self.scheduler ,					)

      # YiYi notes: for testing only to match ldm, we can directly create a latents with desired shape: batch_size, num_embeddings, embedding_dim
      _UpperCAmelCase    			=   latents.reshape(latents.shape[0] ,					__lowerCAmelCase ,					__lowerCAmelCase       )

      for i, t in enumerate(self.progress_bar(__lowerCAmelCase       )       ):
        # expand the latents if we are doing classifier free guidance
        _UpperCAmelCase    			=   torch.cat([latents] * 2       ) if do_classifier_free_guidance else latents
        _UpperCAmelCase    			=   self.scheduler.scale_model_input(__lowerCAmelCase ,					__lowerCAmelCase       )

        _UpperCAmelCase    			=   self.prior(
            __lowerCAmelCase ,					timestep=__lowerCAmelCase ,					proj_embedding=__lowerCAmelCase ,					).predicted_image_embedding

        # remove the variance
        _UpperCAmelCase			,					_UpperCAmelCase    			=   noise_pred.split(
            scaled_model_input.shape[2] ,					dim=2       )  # batch_size, num_embeddings, embedding_dim

        if do_classifier_free_guidance is not None:
          _UpperCAmelCase			,					_UpperCAmelCase    			=   noise_pred.chunk(2       )
          _UpperCAmelCase    			=   noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)

        _UpperCAmelCase    			=   self.scheduler.step(
            __lowerCAmelCase ,					timestep=__lowerCAmelCase ,					sample=__lowerCAmelCase ,					).prev_sample

      if output_type == "latent":
        return ShapEPipelineOutput(images=__lowerCAmelCase       )

      _UpperCAmelCase    			=   []
      for i, latent in enumerate(__lowerCAmelCase       ):
        print()
        _UpperCAmelCase    			=   self.renderer.decode(
            latent[None, :] ,					__lowerCAmelCase ,					size=__lowerCAmelCase ,					ray_batch_size=4096 ,					n_coarse_samples=64 ,					n_fine_samples=128 ,					)

        images.append(__lowerCAmelCase       )

      _UpperCAmelCase    			=   torch.stack(__lowerCAmelCase       )

      if output_type not in ["np", "pil"]:
        raise ValueError(f'''Only the output types `pil` and `np` are supported not output_type={output_type}'''       )

      _UpperCAmelCase    			=   images.cpu().numpy()

      if output_type == "pil":
        _UpperCAmelCase    			=   [self.numpy_to_pil(__lowerCAmelCase       ) for image in images]

      # Offload last model to CPU
      if hasattr(self ,					"""final_offload_hook"""       ) and self.final_offload_hook is not None:
        self.final_offload_hook.offload()

      if not return_dict:
        return (images,)

      return ShapEPipelineOutput(images=__lowerCAmelCase       )
