



import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest



class  __lowerCAmelCase     (	lowerCAmelCase__   ):
 lowerCAmelCase__       					=  (DEISMultistepScheduler,)
 lowerCAmelCase__       					=  (("""num_inference_steps""", 2_5),)




 def        lowerCamelCase      (						self	,				**__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    {
        '''num_train_timesteps''': 1000,
        '''beta_start''': 0.0_001,
        '''beta_end''': 0.02,
        '''beta_schedule''': '''linear''',
        '''solver_order''': 2,
    }

    config.update(**__UpperCAmelCase				)
    return config




 def        lowerCamelCase      (						self	,				__UpperCAmelCase=0	,				**__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    dict(self.forward_default_kwargs				)
    __lowerCamelCase													=    kwargs.pop('''num_inference_steps'''	,				__UpperCAmelCase				)
    __lowerCamelCase													=    self.dummy_sample
    __lowerCamelCase													=    0.1 * sample
    __lowerCamelCase													=    [residual + 0.2, residual + 0.15, residual + 0.10]

    for scheduler_class in self.scheduler_classes:
       __lowerCamelCase													=    self.get_scheduler_config(**__UpperCAmelCase				)
       __lowerCamelCase													=    scheduler_class(**__UpperCAmelCase				)
       scheduler.set_timesteps(__UpperCAmelCase				)
       # copy over dummy past residuals
       __lowerCamelCase													=    dummy_past_residuals[: scheduler.config.solver_order]

       with tempfile.TemporaryDirectory() as tmpdirname:
          scheduler.save_config(__UpperCAmelCase				)
          __lowerCamelCase													=    scheduler_class.from_pretrained(__UpperCAmelCase				)
          new_scheduler.set_timesteps(__UpperCAmelCase				)
          # copy over dummy past residuals
          __lowerCamelCase													=    dummy_past_residuals[: new_scheduler.config.solver_order]

       __lowerCamelCase							,__lowerCamelCase													=    sample, sample
       for t in range(__UpperCAmelCase	,				time_step + scheduler.config.solver_order + 1				):
          __lowerCamelCase													=    scheduler.step(__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				**__UpperCAmelCase				).prev_sample
          __lowerCamelCase													=    new_scheduler.step(__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				**__UpperCAmelCase				).prev_sample

          assert torch.sum(torch.abs(output - new_output				)				) < 1E-5, "Scheduler outputs are not identical"




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    pass




 def        lowerCamelCase      (						self	,				__UpperCAmelCase=0	,				**__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    dict(self.forward_default_kwargs				)
    __lowerCamelCase													=    kwargs.pop('''num_inference_steps'''	,				__UpperCAmelCase				)
    __lowerCamelCase													=    self.dummy_sample
    __lowerCamelCase													=    0.1 * sample
    __lowerCamelCase													=    [residual + 0.2, residual + 0.15, residual + 0.10]

    for scheduler_class in self.scheduler_classes:
       __lowerCamelCase													=    self.get_scheduler_config()
       __lowerCamelCase													=    scheduler_class(**__UpperCAmelCase				)
       scheduler.set_timesteps(__UpperCAmelCase				)

       # copy over dummy past residuals (must be after setting timesteps)
       __lowerCamelCase													=    dummy_past_residuals[: scheduler.config.solver_order]

       with tempfile.TemporaryDirectory() as tmpdirname:
          scheduler.save_config(__UpperCAmelCase				)
          __lowerCamelCase													=    scheduler_class.from_pretrained(__UpperCAmelCase				)
          # copy over dummy past residuals
          new_scheduler.set_timesteps(__UpperCAmelCase				)

          # copy over dummy past residual (must be after setting timesteps)
          __lowerCamelCase													=    dummy_past_residuals[: new_scheduler.config.solver_order]

       __lowerCamelCase													=    scheduler.step(__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				**__UpperCAmelCase				).prev_sample
       __lowerCamelCase													=    new_scheduler.step(__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				**__UpperCAmelCase				).prev_sample

       assert torch.sum(torch.abs(output - new_output				)				) < 1E-5, "Scheduler outputs are not identical"




 def        lowerCamelCase      (						self	,				__UpperCAmelCase=None	,				**__UpperCAmelCase				):



    '''simple docstring'''
    if scheduler is None:
       __lowerCamelCase													=    self.scheduler_classes[0]
       __lowerCamelCase													=    self.get_scheduler_config(**__UpperCAmelCase				)
       __lowerCamelCase													=    scheduler_class(**__UpperCAmelCase				)

    __lowerCamelCase													=    self.scheduler_classes[0]
    __lowerCamelCase													=    self.get_scheduler_config(**__UpperCAmelCase				)
    __lowerCamelCase													=    scheduler_class(**__UpperCAmelCase				)

    __lowerCamelCase													=    10
    __lowerCamelCase													=    self.dummy_model()
    __lowerCamelCase													=    self.dummy_sample_deter
    scheduler.set_timesteps(__UpperCAmelCase				)

    for i, t in enumerate(scheduler.timesteps				):
       __lowerCamelCase													=    model(__UpperCAmelCase	,				__UpperCAmelCase				)
       __lowerCamelCase													=    scheduler.step(__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				).prev_sample

    return sample




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    dict(self.forward_default_kwargs				)

    __lowerCamelCase													=    kwargs.pop('''num_inference_steps'''	,				__UpperCAmelCase				)

    for scheduler_class in self.scheduler_classes:
       __lowerCamelCase													=    self.get_scheduler_config()
       __lowerCamelCase													=    scheduler_class(**__UpperCAmelCase				)

       __lowerCamelCase													=    self.dummy_sample
       __lowerCamelCase													=    0.1 * sample

       if num_inference_steps is not None and hasattr(__UpperCAmelCase	,				'''set_timesteps'''				):
          scheduler.set_timesteps(__UpperCAmelCase				)
       elif num_inference_steps is not None and not hasattr(__UpperCAmelCase	,				'''set_timesteps'''				):
          __lowerCamelCase													=    num_inference_steps

       # copy over dummy past residuals (must be done after set_timesteps)
       __lowerCamelCase													=    [residual + 0.2, residual + 0.15, residual + 0.10]
       __lowerCamelCase													=    dummy_past_residuals[: scheduler.config.solver_order]

       __lowerCamelCase													=    scheduler.timesteps[5]
       __lowerCamelCase													=    scheduler.timesteps[6]

       __lowerCamelCase													=    scheduler.step(__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				**__UpperCAmelCase				).prev_sample
       __lowerCamelCase													=    scheduler.step(__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				**__UpperCAmelCase				).prev_sample

       self.assertEqual(output_a.shape	,				sample.shape				)
       self.assertEqual(output_a.shape	,				output_a.shape				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    # make sure that iterating over schedulers with same config names gives same results
    # for defaults
    __lowerCamelCase													=    DEISMultistepScheduler(**self.get_scheduler_config()				)
    __lowerCamelCase													=    self.full_loop(scheduler=__UpperCAmelCase				)
    __lowerCamelCase													=    torch.mean(torch.abs(__UpperCAmelCase				)				)

    assert abs(result_mean.item() - 0.23_916				) < 1E-3

    __lowerCamelCase													=    DPMSolverSinglestepScheduler.from_config(scheduler.config				)
    __lowerCamelCase													=    DPMSolverMultistepScheduler.from_config(scheduler.config				)
    __lowerCamelCase													=    UniPCMultistepScheduler.from_config(scheduler.config				)
    __lowerCamelCase													=    DEISMultistepScheduler.from_config(scheduler.config				)

    __lowerCamelCase													=    self.full_loop(scheduler=__UpperCAmelCase				)
    __lowerCamelCase													=    torch.mean(torch.abs(__UpperCAmelCase				)				)

    assert abs(result_mean.item() - 0.23_916				) < 1E-3




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    for timesteps in [25, 50, 100, 999, 1000]:
       self.check_over_configs(num_train_timesteps=__UpperCAmelCase				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    self.check_over_configs(thresholding=__UpperCAmelCase				)
    for order in [1, 2, 3]:
       for solver_type in ["logrho"]:
          for threshold in [0.5, 1.0, 2.0]:
             for prediction_type in ["epsilon", "sample"]:
                self.check_over_configs(
                    thresholding=__UpperCAmelCase	,				prediction_type=__UpperCAmelCase	,				sample_max_value=__UpperCAmelCase	,				algorithm_type='''deis'''	,				solver_order=__UpperCAmelCase	,				solver_type=__UpperCAmelCase	,				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    for prediction_type in ["epsilon", "v_prediction"]:
       self.check_over_configs(prediction_type=__UpperCAmelCase				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    for algorithm_type in ["deis"]:
       for solver_type in ["logrho"]:
          for order in [1, 2, 3]:
             for prediction_type in ["epsilon", "sample"]:
                self.check_over_configs(
                    solver_order=__UpperCAmelCase	,				solver_type=__UpperCAmelCase	,				prediction_type=__UpperCAmelCase	,				algorithm_type=__UpperCAmelCase	,				)
                __lowerCamelCase													=    self.full_loop(
                    solver_order=__UpperCAmelCase	,				solver_type=__UpperCAmelCase	,				prediction_type=__UpperCAmelCase	,				algorithm_type=__UpperCAmelCase	,				)
                assert not torch.isnan(__UpperCAmelCase				).any(), "Samples have nan numbers"




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    self.check_over_configs(lower_order_final=__UpperCAmelCase				)
    self.check_over_configs(lower_order_final=__UpperCAmelCase				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    for num_inference_steps in [1, 2, 3, 5, 10, 50, 100, 999, 1000]:
       self.check_over_forward(num_inference_steps=__UpperCAmelCase	,				time_step=0				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.full_loop()
    __lowerCamelCase													=    torch.mean(torch.abs(__UpperCAmelCase				)				)

    assert abs(result_mean.item() - 0.23_916				) < 1E-3




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.full_loop(prediction_type='''v_prediction'''				)
    __lowerCamelCase													=    torch.mean(torch.abs(__UpperCAmelCase				)				)

    assert abs(result_mean.item() - 0.091				) < 1E-3




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.scheduler_classes[0]
    __lowerCamelCase													=    self.get_scheduler_config(thresholding=__UpperCAmelCase	,				dynamic_thresholding_ratio=0				)
    __lowerCamelCase													=    scheduler_class(**__UpperCAmelCase				)

    __lowerCamelCase													=    10
    __lowerCamelCase													=    self.dummy_model()
    __lowerCamelCase													=    self.dummy_sample_deter.half()
    scheduler.set_timesteps(__UpperCAmelCase				)

    for i, t in enumerate(scheduler.timesteps				):
       __lowerCamelCase													=    model(__UpperCAmelCase	,				__UpperCAmelCase				)
       __lowerCamelCase													=    scheduler.step(__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				).prev_sample

    assert sample.dtype == torch.floataa











import os
import unicodedata
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


a_			=					logging.get_logger(__name__)
a_			=					{"""vocab_file""": """spiece.model"""}

a_			=					{
    """vocab_file""": {
        """albert-base-v1""": """https://huggingface.co/albert-base-v1/resolve/main/spiece.model""",
        """albert-large-v1""": """https://huggingface.co/albert-large-v1/resolve/main/spiece.model""",
        """albert-xlarge-v1""": """https://huggingface.co/albert-xlarge-v1/resolve/main/spiece.model""",
        """albert-xxlarge-v1""": """https://huggingface.co/albert-xxlarge-v1/resolve/main/spiece.model""",
        """albert-base-v2""": """https://huggingface.co/albert-base-v2/resolve/main/spiece.model""",
        """albert-large-v2""": """https://huggingface.co/albert-large-v2/resolve/main/spiece.model""",
        """albert-xlarge-v2""": """https://huggingface.co/albert-xlarge-v2/resolve/main/spiece.model""",
        """albert-xxlarge-v2""": """https://huggingface.co/albert-xxlarge-v2/resolve/main/spiece.model""",
    }
}

a_			=					{
    """albert-base-v1""": 512,
    """albert-large-v1""": 512,
    """albert-xlarge-v1""": 512,
    """albert-xxlarge-v1""": 512,
    """albert-base-v2""": 512,
    """albert-large-v2""": 512,
    """albert-xlarge-v2""": 512,
    """albert-xxlarge-v2""": 512,
}

a_			=					"""▁"""



class  __lowerCAmelCase     (	lowerCAmelCase__   ):
 lowerCAmelCase__       					=  VOCAB_FILES_NAMES
 lowerCAmelCase__       					=  PRETRAINED_VOCAB_FILES_MAP
 lowerCAmelCase__       					=  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES




 def __init__(						self	,				__UpperCAmelCase	,				__UpperCAmelCase=True	,				__UpperCAmelCase=True	,				__UpperCAmelCase=False	,				__UpperCAmelCase="[CLS]"	,				__UpperCAmelCase="[SEP]"	,				__UpperCAmelCase="<unk>"	,				__UpperCAmelCase="[SEP]"	,				__UpperCAmelCase="<pad>"	,				__UpperCAmelCase="[CLS]"	,				__UpperCAmelCase="[MASK]"	,				__UpperCAmelCase = None	,				**__UpperCAmelCase	,				):



    '''simple docstring'''
    # Mask token behave like a normal word, i.e. include the space before it and
    # is included in the raw text, there should be a match in a non-normalized sentence.
    __lowerCamelCase													=    (
        AddedToken(__UpperCAmelCase	,				lstrip=__UpperCAmelCase	,				rstrip=__UpperCAmelCase	,				normalized=__UpperCAmelCase				)
        if isinstance(__UpperCAmelCase	,				__UpperCAmelCase				)
        else mask_token
    )

    __lowerCamelCase													=    {} if sp_model_kwargs is None else sp_model_kwargs

    super().__init__(
        do_lower_case=__UpperCAmelCase	,				remove_space=__UpperCAmelCase	,				keep_accents=__UpperCAmelCase	,				bos_token=__UpperCAmelCase	,				eos_token=__UpperCAmelCase	,				unk_token=__UpperCAmelCase	,				sep_token=__UpperCAmelCase	,				pad_token=__UpperCAmelCase	,				cls_token=__UpperCAmelCase	,				mask_token=__UpperCAmelCase	,				sp_model_kwargs=self.sp_model_kwargs	,				**__UpperCAmelCase	,				)

    __lowerCamelCase													=    do_lower_case
    __lowerCamelCase													=    remove_space
    __lowerCamelCase													=    keep_accents
    __lowerCamelCase													=    vocab_file

    __lowerCamelCase													=    spm.SentencePieceProcessor(**self.sp_model_kwargs				)
    self.sp_model.Load(__UpperCAmelCase				)




 @property
 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    return len(self.sp_model				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    {self.convert_ids_to_tokens(__UpperCAmelCase				): i for i in range(self.vocab_size				)}
    vocab.update(self.added_tokens_encoder				)
    return vocab




 def __getstate__(						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.__dict__.copy()
    __lowerCamelCase													=    None
    return state




 def __setstate__(						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    d

    # for backward compatibility
    if not hasattr(self	,				'''sp_model_kwargs'''				):
       __lowerCamelCase													=    {}

    __lowerCamelCase													=    spm.SentencePieceProcessor(**self.sp_model_kwargs				)
    self.sp_model.Load(self.vocab_file				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    if self.remove_space:
       __lowerCamelCase													=    ''' '''.join(inputs.strip().split()				)
    else:
       __lowerCamelCase													=    inputs
    __lowerCamelCase													=    outputs.replace('''``'''	,				'''"'''				).replace('''\'\''''	,				'''"'''				)

    if not self.keep_accents:
       __lowerCamelCase													=    unicodedata.normalize('''NFKD'''	,				__UpperCAmelCase				)
       __lowerCamelCase													=    ''''''.join([c for c in outputs if not unicodedata.combining(__UpperCAmelCase				)]				)
    if self.do_lower_case:
       __lowerCamelCase													=    outputs.lower()

    return outputs




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    self.preprocess_text(__UpperCAmelCase				)
    __lowerCamelCase													=    self.sp_model.encode(__UpperCAmelCase	,				out_type=__UpperCAmelCase				)
    __lowerCamelCase													=    []
    for piece in pieces:
       if len(__UpperCAmelCase				) > 1 and piece[-1] == str(''','''				) and piece[-2].isdigit():
          __lowerCamelCase													=    self.sp_model.EncodeAsPieces(piece[:-1].replace(__UpperCAmelCase	,				''''''				)				)
          if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
             if len(cur_pieces[0]				) == 1:
                __lowerCamelCase													=    cur_pieces[1:]
             else:
                __lowerCamelCase													=    cur_pieces[0][1:]
          cur_pieces.append(piece[-1]				)
          new_pieces.extend(__UpperCAmelCase				)
       else:
          new_pieces.append(__UpperCAmelCase				)

    return new_pieces




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    return self.sp_model.PieceToId(__UpperCAmelCase				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    return self.sp_model.IdToPiece(__UpperCAmelCase				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    []
    __lowerCamelCase													=    ''''''
    __lowerCamelCase													=    False
    for token in tokens:
       # make sure that special tokens are not decoded using sentencepiece model
       if token in self.all_special_tokens:
          if not prev_is_special:
             out_string += " "
          out_string += self.sp_model.decode(__UpperCAmelCase				) + token
          __lowerCamelCase													=    True
          __lowerCamelCase													=    []
       else:
          current_sub_tokens.append(__UpperCAmelCase				)
          __lowerCamelCase													=    False
    out_string += self.sp_model.decode(__UpperCAmelCase				)
    return out_string.strip()




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase = None				):



    '''simple docstring'''
    __lowerCamelCase													=    [self.sep_token_id]
    __lowerCamelCase													=    [self.cls_token_id]
    if token_ids_a is None:
       return cls + token_ids_a + sep
    return cls + token_ids_a + sep + token_ids_a + sep




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase = None	,				__UpperCAmelCase = False				):



    '''simple docstring'''

    if already_has_special_tokens:
       return super().get_special_tokens_mask(
           token_ids_a=__UpperCAmelCase	,				token_ids_a=__UpperCAmelCase	,				already_has_special_tokens=__UpperCAmelCase				)

    if token_ids_a is not None:
       return [1] + ([0] * len(__UpperCAmelCase				)) + [1] + ([0] * len(__UpperCAmelCase				)) + [1]
    return [1] + ([0] * len(__UpperCAmelCase				)) + [1]




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase = None				):



    '''simple docstring'''
    __lowerCamelCase													=    [self.sep_token_id]
    __lowerCamelCase													=    [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	,				__UpperCAmelCase	,				__UpperCAmelCase = None				):



    '''simple docstring'''
    if not os.path.isdir(__UpperCAmelCase				):
       logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""				)
       return
    __lowerCamelCase													=    os.path.join(
        __UpperCAmelCase	,				(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']				)

    if os.path.abspath(self.vocab_file				) != os.path.abspath(__UpperCAmelCase				) and os.path.isfile(self.vocab_file				):
       copyfile(self.vocab_file	,				__UpperCAmelCase				)
    elif not os.path.isfile(self.vocab_file				):
       with open(__UpperCAmelCase	,				'''wb'''				) as fi:
          __lowerCamelCase													=    self.sp_model.serialized_model_proto()
          fi.write(__UpperCAmelCase				)

    return (out_vocab_file,)






