




import argparse

import torch

from diffusers import UNetaDConditionModel



def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase=None,    __UpperCamelCase=None,    __UpperCamelCase=None     ):
    assert isinstance(__UpperCamelCase,    __UpperCamelCase     ), "Paths should be a list of dicts containing 'old' and 'new' keys."

    # Splits the attention layers into three variables.
    if attention_paths_to_split is not None:
        for path, path_map in attention_paths_to_split.items():
            SCREAMING_SNAKE_CASE__       =old_checkpoint[path]
            SCREAMING_SNAKE_CASE__       =old_tensor.shape[0] // 3

            SCREAMING_SNAKE_CASE__       =(-1, channels) if len(old_tensor.shape     ) == 3 else (-1)

            SCREAMING_SNAKE_CASE__       =old_tensor.shape[0] // config["""num_head_channels"""] // 3

            SCREAMING_SNAKE_CASE__       =old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]     )
            SCREAMING_SNAKE_CASE__      ,		SCREAMING_SNAKE_CASE__      ,		SCREAMING_SNAKE_CASE__       =old_tensor.split(channels // num_heads,    dim=1     )

            SCREAMING_SNAKE_CASE__       =query.reshape(__UpperCamelCase     )
            SCREAMING_SNAKE_CASE__       =key.reshape(__UpperCamelCase     )
            SCREAMING_SNAKE_CASE__       =value.reshape(__UpperCamelCase     )

    for path in paths:
        SCREAMING_SNAKE_CASE__       =path["""new"""]

        # These have already been assigned
        if attention_paths_to_split is not None and new_path in attention_paths_to_split:
            continue

        if additional_replacements is not None:
            for replacement in additional_replacements:
                SCREAMING_SNAKE_CASE__       =new_path.replace(replacement["""old"""],    replacement["""new"""]     )

        # proj_attn.weight has to be converted from conv 1D to linear
        SCREAMING_SNAKE_CASE__       =old_checkpoint[path["""old"""]]
        SCREAMING_SNAKE_CASE__       =["""proj_attn.weight"""]
        SCREAMING_SNAKE_CASE__       =["""proj_out.weight""", """proj_in.weight"""]
        if any(k in new_path for k in names     ):
            SCREAMING_SNAKE_CASE__       =weight[:, :, 0]
        elif any(k in new_path for k in names_a     ) and len(weight.shape     ) > 2 and ".attentions." not in new_path:
            SCREAMING_SNAKE_CASE__       =weight[:, :, 0]
        else:
            SCREAMING_SNAKE_CASE__       =weight



def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase=0     ):
    SCREAMING_SNAKE_CASE__       =[]
    for old_item in old_list:
        SCREAMING_SNAKE_CASE__       =old_item

        #         new_item = new_item.replace('norm.weight', 'group_norm.weight')
        #         new_item = new_item.replace('norm.bias', 'group_norm.bias')

        #         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
        #         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')

        #         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)

        mapping.append({"""old""": old_item, """new""": new_item}     )

    return mapping



def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase=1     ):
    if n_shave_prefix_segments >= 0:
        return ".".join(path.split("""."""     )[n_shave_prefix_segments:]     )
    else:
        return ".".join(path.split("""."""     )[:n_shave_prefix_segments]     )



def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase=0     ):
    SCREAMING_SNAKE_CASE__       =[]
    for old_item in old_list:
        mapping.append({"""old""": old_item, """new""": old_item}     )

    return mapping



def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase=0     ):
    SCREAMING_SNAKE_CASE__       =[]
    for old_item in old_list:
        SCREAMING_SNAKE_CASE__       =old_item.replace("""in_layers.0""",    """norm1"""     )
        SCREAMING_SNAKE_CASE__       =new_item.replace("""in_layers.2""",    """conv1"""     )

        SCREAMING_SNAKE_CASE__       =new_item.replace("""out_layers.0""",    """norm2"""     )
        SCREAMING_SNAKE_CASE__       =new_item.replace("""out_layers.3""",    """conv2"""     )

        SCREAMING_SNAKE_CASE__       =new_item.replace("""emb_layers.1""",    """time_emb_proj"""     )
        SCREAMING_SNAKE_CASE__       =new_item.replace("""skip_connection""",    """conv_shortcut"""     )

        SCREAMING_SNAKE_CASE__       =shave_segments(__UpperCamelCase,    n_shave_prefix_segments=__UpperCamelCase     )

        if "temopral_conv" not in old_item:
            mapping.append({"""old""": old_item, """new""": new_item}     )

    return mapping



def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase=None,    __UpperCamelCase=False     ):
    SCREAMING_SNAKE_CASE__       ={}
    SCREAMING_SNAKE_CASE__       =list(checkpoint.keys()     )

    SCREAMING_SNAKE_CASE__       ="""model.diffusion_model."""

    # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA
    if sum(k.startswith("""model_ema"""     ) for k in keys     ) > 100 and extract_ema:
        print(f"""Checkpoint {path} has both EMA and non-EMA weights."""     )
        print(
            """In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA"""
            """ weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag."""     )
        for key in keys:
            if key.startswith("""model.diffusion_model"""     ):
                SCREAMING_SNAKE_CASE__       ="""model_ema.""" + """""".join(key.split("""."""     )[1:]     )
                SCREAMING_SNAKE_CASE__       =checkpoint.pop(__UpperCamelCase     )
    else:
        if sum(k.startswith("""model_ema"""     ) for k in keys     ) > 100:
            print(
                """In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA"""
                """ weights (usually better for inference), please make sure to add the `--extract_ema` flag."""     )

        for key in keys:
            SCREAMING_SNAKE_CASE__       =checkpoint.pop(__UpperCamelCase     )

    SCREAMING_SNAKE_CASE__       ={}

    SCREAMING_SNAKE_CASE__       =unet_state_dict["""time_embed.0.weight"""]
    SCREAMING_SNAKE_CASE__       =unet_state_dict["""time_embed.0.bias"""]
    SCREAMING_SNAKE_CASE__       =unet_state_dict["""time_embed.2.weight"""]
    SCREAMING_SNAKE_CASE__       =unet_state_dict["""time_embed.2.bias"""]

    if config["class_embed_type"] is None:
        # No parameters to port
        ...
    elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection":
        SCREAMING_SNAKE_CASE__       =unet_state_dict["""label_emb.0.0.weight"""]
        SCREAMING_SNAKE_CASE__       =unet_state_dict["""label_emb.0.0.bias"""]
        SCREAMING_SNAKE_CASE__       =unet_state_dict["""label_emb.0.2.weight"""]
        SCREAMING_SNAKE_CASE__       =unet_state_dict["""label_emb.0.2.bias"""]
    else:
        raise NotImplementedError(f"""Not implemented `class_embed_type`: {config["class_embed_type"]}"""     )

    SCREAMING_SNAKE_CASE__       =unet_state_dict["""input_blocks.0.0.weight"""]
    SCREAMING_SNAKE_CASE__       =unet_state_dict["""input_blocks.0.0.bias"""]

    SCREAMING_SNAKE_CASE__       =[v for v in unet_state_dict if v.startswith("""input_blocks.0.1"""     )]
    SCREAMING_SNAKE_CASE__       =renew_attention_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__       ={"""old""": """input_blocks.0.1""", """new""": """transformer_in"""}
    assign_to_checkpoint(__UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    additional_replacements=[meta_path],    config=__UpperCamelCase     )

    SCREAMING_SNAKE_CASE__       =unet_state_dict["""out.0.weight"""]
    SCREAMING_SNAKE_CASE__       =unet_state_dict["""out.0.bias"""]
    SCREAMING_SNAKE_CASE__       =unet_state_dict["""out.2.weight"""]
    SCREAMING_SNAKE_CASE__       =unet_state_dict["""out.2.bias"""]

    # Retrieves the keys for the input blocks only
    SCREAMING_SNAKE_CASE__       =len({""".""".join(layer.split("""."""     )[:2]     ) for layer in unet_state_dict if """input_blocks""" in layer}     )
    SCREAMING_SNAKE_CASE__       ={
        layer_id: [key for key in unet_state_dict if f"""input_blocks.{layer_id}""" in key]
        for layer_id in range(__UpperCamelCase     )
    }

    # Retrieves the keys for the middle blocks only
    SCREAMING_SNAKE_CASE__       =len({""".""".join(layer.split("""."""     )[:2]     ) for layer in unet_state_dict if """middle_block""" in layer}     )
    SCREAMING_SNAKE_CASE__       ={
        layer_id: [key for key in unet_state_dict if f"""middle_block.{layer_id}""" in key]
        for layer_id in range(__UpperCamelCase     )
    }

    # Retrieves the keys for the output blocks only
    SCREAMING_SNAKE_CASE__       =len({""".""".join(layer.split("""."""     )[:2]     ) for layer in unet_state_dict if """output_blocks""" in layer}     )
    SCREAMING_SNAKE_CASE__       ={
        layer_id: [key for key in unet_state_dict if f"""output_blocks.{layer_id}""" in key]
        for layer_id in range(__UpperCamelCase     )
    }

    for i in range(1,    __UpperCamelCase     ):
        SCREAMING_SNAKE_CASE__       =(i - 1) // (config["""layers_per_block"""] + 1)
        SCREAMING_SNAKE_CASE__       =(i - 1) % (config["""layers_per_block"""] + 1)

        SCREAMING_SNAKE_CASE__       =[
            key for key in input_blocks[i] if f"""input_blocks.{i}.0""" in key and f"""input_blocks.{i}.0.op""" not in key
        ]
        SCREAMING_SNAKE_CASE__       =[key for key in input_blocks[i] if f"""input_blocks.{i}.1""" in key]
        SCREAMING_SNAKE_CASE__       =[key for key in input_blocks[i] if f"""input_blocks.{i}.2""" in key]

        if f"""input_blocks.{i}.op.weight""" in unet_state_dict:
            SCREAMING_SNAKE_CASE__       =unet_state_dict.pop(
                f"""input_blocks.{i}.op.weight"""     )
            SCREAMING_SNAKE_CASE__       =unet_state_dict.pop(
                f"""input_blocks.{i}.op.bias"""     )

        SCREAMING_SNAKE_CASE__       =renew_resnet_paths(__UpperCamelCase     )
        SCREAMING_SNAKE_CASE__       ={"""old""": f"""input_blocks.{i}.0""", """new""": f"""down_blocks.{block_id}.resnets.{layer_in_block_id}"""}
        assign_to_checkpoint(
            __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    additional_replacements=[meta_path],    config=__UpperCamelCase     )

        SCREAMING_SNAKE_CASE__       =[key for key in resnets if """temopral_conv""" in key]
        SCREAMING_SNAKE_CASE__       =renew_temp_conv_paths(__UpperCamelCase     )
        SCREAMING_SNAKE_CASE__       ={
            """old""": f"""input_blocks.{i}.0.temopral_conv""",
            """new""": f"""down_blocks.{block_id}.temp_convs.{layer_in_block_id}""",
        }
        assign_to_checkpoint(
            __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    additional_replacements=[meta_path],    config=__UpperCamelCase     )

        if len(__UpperCamelCase     ):
            SCREAMING_SNAKE_CASE__       =renew_attention_paths(__UpperCamelCase     )
            SCREAMING_SNAKE_CASE__       ={"""old""": f"""input_blocks.{i}.1""", """new""": f"""down_blocks.{block_id}.attentions.{layer_in_block_id}"""}
            assign_to_checkpoint(
                __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    additional_replacements=[meta_path],    config=__UpperCamelCase     )

        if len(__UpperCamelCase     ):
            SCREAMING_SNAKE_CASE__       =renew_attention_paths(__UpperCamelCase     )
            SCREAMING_SNAKE_CASE__       ={
                """old""": f"""input_blocks.{i}.2""",
                """new""": f"""down_blocks.{block_id}.temp_attentions.{layer_in_block_id}""",
            }
            assign_to_checkpoint(
                __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    additional_replacements=[meta_path],    config=__UpperCamelCase     )

    SCREAMING_SNAKE_CASE__       =middle_blocks[0]
    SCREAMING_SNAKE_CASE__       =[key for key in resnet_a if """temopral_conv""" in key]
    SCREAMING_SNAKE_CASE__       =middle_blocks[1]
    SCREAMING_SNAKE_CASE__       =middle_blocks[2]
    SCREAMING_SNAKE_CASE__       =middle_blocks[3]
    SCREAMING_SNAKE_CASE__       =[key for key in resnet_a if """temopral_conv""" in key]

    SCREAMING_SNAKE_CASE__       =renew_resnet_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__       ={"""old""": """middle_block.0""", """new""": """mid_block.resnets.0"""}
    assign_to_checkpoint(
        __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    config=__UpperCamelCase,    additional_replacements=[meta_path]     )

    SCREAMING_SNAKE_CASE__       =renew_temp_conv_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__       ={"""old""": """middle_block.0.temopral_conv""", """new""": """mid_block.temp_convs.0"""}
    assign_to_checkpoint(
        __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    config=__UpperCamelCase,    additional_replacements=[meta_path]     )

    SCREAMING_SNAKE_CASE__       =renew_resnet_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__       ={"""old""": """middle_block.3""", """new""": """mid_block.resnets.1"""}
    assign_to_checkpoint(
        __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    config=__UpperCamelCase,    additional_replacements=[meta_path]     )

    SCREAMING_SNAKE_CASE__       =renew_temp_conv_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__       ={"""old""": """middle_block.3.temopral_conv""", """new""": """mid_block.temp_convs.1"""}
    assign_to_checkpoint(
        __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    config=__UpperCamelCase,    additional_replacements=[meta_path]     )

    SCREAMING_SNAKE_CASE__       =renew_attention_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__       ={"""old""": """middle_block.1""", """new""": """mid_block.attentions.0"""}
    assign_to_checkpoint(
        __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    additional_replacements=[meta_path],    config=__UpperCamelCase     )

    SCREAMING_SNAKE_CASE__       =renew_attention_paths(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__       ={"""old""": """middle_block.2""", """new""": """mid_block.temp_attentions.0"""}
    assign_to_checkpoint(
        __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    additional_replacements=[meta_path],    config=__UpperCamelCase     )

    for i in range(__UpperCamelCase     ):
        SCREAMING_SNAKE_CASE__       =i // (config["""layers_per_block"""] + 1)
        SCREAMING_SNAKE_CASE__       =i % (config["""layers_per_block"""] + 1)
        SCREAMING_SNAKE_CASE__       =[shave_segments(__UpperCamelCase,    2     ) for name in output_blocks[i]]
        SCREAMING_SNAKE_CASE__       ={}

        for layer in output_block_layers:
            SCREAMING_SNAKE_CASE__      ,		SCREAMING_SNAKE_CASE__       =layer.split("""."""     )[0], shave_segments(__UpperCamelCase,    1     )
            if layer_id in output_block_list:
                output_block_list[layer_id].append(__UpperCamelCase     )
            else:
                SCREAMING_SNAKE_CASE__       =[layer_name]

        if len(__UpperCamelCase     ) > 1:
            SCREAMING_SNAKE_CASE__       =[key for key in output_blocks[i] if f"""output_blocks.{i}.0""" in key]
            SCREAMING_SNAKE_CASE__       =[key for key in output_blocks[i] if f"""output_blocks.{i}.1""" in key]
            SCREAMING_SNAKE_CASE__       =[key for key in output_blocks[i] if f"""output_blocks.{i}.2""" in key]

            SCREAMING_SNAKE_CASE__       =renew_resnet_paths(__UpperCamelCase     )
            SCREAMING_SNAKE_CASE__       =renew_resnet_paths(__UpperCamelCase     )

            SCREAMING_SNAKE_CASE__       ={"""old""": f"""output_blocks.{i}.0""", """new""": f"""up_blocks.{block_id}.resnets.{layer_in_block_id}"""}
            assign_to_checkpoint(
                __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    additional_replacements=[meta_path],    config=__UpperCamelCase     )

            SCREAMING_SNAKE_CASE__       =[key for key in resnets if """temopral_conv""" in key]
            SCREAMING_SNAKE_CASE__       =renew_temp_conv_paths(__UpperCamelCase     )
            SCREAMING_SNAKE_CASE__       ={
                """old""": f"""output_blocks.{i}.0.temopral_conv""",
                """new""": f"""up_blocks.{block_id}.temp_convs.{layer_in_block_id}""",
            }
            assign_to_checkpoint(
                __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    additional_replacements=[meta_path],    config=__UpperCamelCase     )

            SCREAMING_SNAKE_CASE__       ={k: sorted(__UpperCamelCase     ) for k, v in output_block_list.items()}
            if ["conv.bias", "conv.weight"] in output_block_list.values():
                SCREAMING_SNAKE_CASE__       =list(output_block_list.values()     ).index(["""conv.bias""", """conv.weight"""]     )
                SCREAMING_SNAKE_CASE__       =unet_state_dict[
                    f"""output_blocks.{i}.{index}.conv.weight"""
                ]
                SCREAMING_SNAKE_CASE__       =unet_state_dict[
                    f"""output_blocks.{i}.{index}.conv.bias"""
                ]

                # Clear attentions as they have been attributed above.
                if len(__UpperCamelCase     ) == 2:
                    SCREAMING_SNAKE_CASE__       =[]

            if len(__UpperCamelCase     ):
                SCREAMING_SNAKE_CASE__       =renew_attention_paths(__UpperCamelCase     )
                SCREAMING_SNAKE_CASE__       ={
                    """old""": f"""output_blocks.{i}.1""",
                    """new""": f"""up_blocks.{block_id}.attentions.{layer_in_block_id}""",
                }
                assign_to_checkpoint(
                    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    additional_replacements=[meta_path],    config=__UpperCamelCase     )

            if len(__UpperCamelCase     ):
                SCREAMING_SNAKE_CASE__       =renew_attention_paths(__UpperCamelCase     )
                SCREAMING_SNAKE_CASE__       ={
                    """old""": f"""output_blocks.{i}.2""",
                    """new""": f"""up_blocks.{block_id}.temp_attentions.{layer_in_block_id}""",
                }
                assign_to_checkpoint(
                    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    additional_replacements=[meta_path],    config=__UpperCamelCase     )
        else:
            SCREAMING_SNAKE_CASE__       =renew_resnet_paths(__UpperCamelCase,    n_shave_prefix_segments=1     )
            for path in resnet_0_paths:
                SCREAMING_SNAKE_CASE__       =""".""".join(["""output_blocks""", str(__UpperCamelCase     ), path["""old"""]]     )
                SCREAMING_SNAKE_CASE__       =""".""".join(["""up_blocks""", str(__UpperCamelCase     ), """resnets""", str(__UpperCamelCase     ), path["""new"""]]     )
                SCREAMING_SNAKE_CASE__       =unet_state_dict[old_path]

            SCREAMING_SNAKE_CASE__       =[l for l in output_block_layers if """temopral_conv""" in l]
            for path in temopral_conv_paths:
                SCREAMING_SNAKE_CASE__       =path.split("""temopral_conv."""     )[-1]
                SCREAMING_SNAKE_CASE__       =""".""".join(["""output_blocks""", str(__UpperCamelCase     ), str(__UpperCamelCase     ), """temopral_conv""", pruned_path]     )
                SCREAMING_SNAKE_CASE__       =""".""".join(["""up_blocks""", str(__UpperCamelCase     ), """temp_convs""", str(__UpperCamelCase     ), pruned_path]     )
                SCREAMING_SNAKE_CASE__       =unet_state_dict[old_path]

    return new_checkpoint


if __name__ == "__main__":
    lowerCamelCase_  						=       argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert."
    )
    parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.")
    lowerCamelCase_  						=       parser.parse_args()

    lowerCamelCase_  						=       torch.load(args.checkpoint_path, map_location="cpu")
    lowerCamelCase_  						=       UNetaDConditionModel()

    lowerCamelCase_  						=       convert_ldm_unet_checkpoint(unet_checkpoint, unet.config)

    lowerCamelCase_  						=       set(unet.state_dict().keys()) - set(converted_ckpt.keys())
    lowerCamelCase_  						=       set(converted_ckpt.keys()) - set(unet.state_dict().keys())

    assert len(diff_a) == len(diff_a) == 0, "Converted weights don't match"

    # load state_dict
    unet.load_state_dict(converted_ckpt)

    unet.save_pretrained(args.dump_path)

    # -- finish converting the unet --












from collections import Counter
import datasets
import transformers
from transformers.convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS

from transformers.utils import logging

logging.set_verbosity_info()

lowerCamelCase_  						=       {
    name: (getattr(transformers, name), getattr(transformers, name + "Fast")) for name in SLOW_TO_FAST_CONVERTERS
}

lowerCamelCase_  						=       datasets.load_dataset("xnli", split="test+validation")

lowerCamelCase_  						=       0
lowerCamelCase_  						=       0
lowerCamelCase_  						=       0
lowerCamelCase_  						=       0



def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase     ):
    if spm_diff == list(reversed(__UpperCamelCase     )     ):
        # AAA -> AA+A vs A+AA case.
        return True
    elif len(__UpperCamelCase     ) == len(__UpperCamelCase     ) and fast.decode(__UpperCamelCase     ) == fast.decode(__UpperCamelCase     ):
        # Second order OK
        # Barrich -> Barr + ich vs Bar + rich
        return True
    SCREAMING_SNAKE_CASE__       =slow.encode(slow.decode(__UpperCamelCase     )     )
    SCREAMING_SNAKE_CASE__       =fast.encode(fast.decode(__UpperCamelCase     )     )
    if spm_reencoded != spm_diff and spm_reencoded == tok_reencoded:
        # Type 3 error.
        # Snehagatha ->
        #       Sne, h, aga, th, a
        #       Sne, ha, gat, ha
        # Encoding the wrong with sp does not even recover what spm gave us
        # It fits tokenizer however...
        return True
    return False



def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase     ):
    SCREAMING_SNAKE_CASE__       =fast.encode_plus(__UpperCamelCase     )[0]
    SCREAMING_SNAKE_CASE__       =enc.offsets
    SCREAMING_SNAKE_CASE__      ,		SCREAMING_SNAKE_CASE__       =offsets[idx], offsets[idx - 1]
    if curr is not None and line[curr[0] : curr[1]] == "\u200f":
        return True
    if prev is not None and line[prev[0] : prev[1]] == "\u200f":
        return True



def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase     ):
    # Encoding can be the same with same result AAA -> A + AA vs AA + A
    # We can check that we use at least exactly the same number of tokens.
    for i, (spm_id, tok_id) in enumerate(zip(__UpperCamelCase,    __UpperCamelCase     )     ):
        if spm_id != tok_id:
            break
    SCREAMING_SNAKE_CASE__       =i
    for i, (spm_id, tok_id) in enumerate(zip(reversed(__UpperCamelCase     ),    reversed(__UpperCamelCase     )     )     ):
        if spm_id != tok_id:
            break
    SCREAMING_SNAKE_CASE__       =len(__UpperCamelCase     ) - i

    SCREAMING_SNAKE_CASE__       =spm_ids[first:last]
    SCREAMING_SNAKE_CASE__       =tok_ids[first:last]

    if check_diff(__UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase     ):
        return True

    if check_LTR_mark(__UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase     ):
        return True

    if last - first > 5:
        # We might have twice a single problem, attempt to subdivide the disjointed tokens into smaller problems
        SCREAMING_SNAKE_CASE__       =Counter(spm_ids[first:last]     )
        SCREAMING_SNAKE_CASE__       =Counter(tok_ids[first:last]     )

        SCREAMING_SNAKE_CASE__       ={spm_ for (spm_, si) in spms.items() if toks.get(spm_,    0     ) == si}
        SCREAMING_SNAKE_CASE__       =3
        for i in range(last - first - min_width     ):
            if all(spm_ids[first + i + j] in removable_tokens for j in range(__UpperCamelCase     )     ):
                SCREAMING_SNAKE_CASE__       =[
                    k
                    for k in range(last - first - min_width     )
                    if tok_ids[first + k : first + k + min_width] == spm_ids[first + i : first + i + min_width]
                ]
                for j in possible_matches:
                    if check_diff(spm_ids[first : first + i],    tok_ids[first : first + j],    __UpperCamelCase,    __UpperCamelCase     ) and check_details(
                        __UpperCamelCase,    spm_ids[first + i : last],    tok_ids[first + j : last],    __UpperCamelCase,    __UpperCamelCase,    ):
                        return True

    print(f"""Spm: {[fast.decode([spm_ids[i]]     ) for i in range(__UpperCamelCase,    __UpperCamelCase     )]}"""     )
    try:
        print(f"""Tok: {[fast.decode([tok_ids[i]]     ) for i in range(__UpperCamelCase,    __UpperCamelCase     )]}"""     )
    except Exception:
        pass

    SCREAMING_SNAKE_CASE__       =fast.decode(spm_ids[:first]     )
    SCREAMING_SNAKE_CASE__       =fast.decode(spm_ids[last:]     )
    SCREAMING_SNAKE_CASE__       =fast.decode(spm_ids[first:last]     )
    print()
    print(__UpperCamelCase     )
    return False



def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase     ):
    global perfect
    global imperfect
    global wrong
    global total

    SCREAMING_SNAKE_CASE__       =slow.encode(__UpperCamelCase     )
    SCREAMING_SNAKE_CASE__       =fast.encode(__UpperCamelCase     )

    SCREAMING_SNAKE_CASE__       =False
    total += 1

    if slow_ids != fast_ids:
        if check_details(__UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase     ):
            SCREAMING_SNAKE_CASE__       =True
            imperfect += 1
        else:
            wrong += 1
    else:
        perfect += 1

    if total % 10_000 == 0:
        print(f"""({perfect} / {imperfect} / {wrong} ----- {perfect + imperfect + wrong})"""     )

    if skip_assert:
        return

    assert (
        slow_ids == fast_ids
    ), f"""line {text} : \n\n{slow_ids}\n{fast_ids}\n\n{slow.tokenize(__UpperCamelCase     )}\n{fast.tokenize(__UpperCamelCase     )}"""



def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase     ):
    global batch_total
    for i in range(len(__UpperCamelCase     )     ):
        # premise, all languages
        for text in dataset[i]["premise"].values():
            test_string(__UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase     )

        # hypothesis, all languages
        for text in dataset[i]["hypothesis"]["translation"]:
            test_string(__UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase     )


if __name__ == "__main__":
    for name, (slow_class, fast_class) in TOKENIZER_CLASSES.items():
        lowerCamelCase_  						=       list(slow_class.max_model_input_sizes.keys())
        for checkpoint in checkpoint_names:
            lowerCamelCase_  						=       0
            lowerCamelCase_  						=       0
            lowerCamelCase_  						=       0
            lowerCamelCase_  						=       0

            print(f"""========================== Checking {name}: {checkpoint} ==========================""")
            lowerCamelCase_  						=       slow_class.from_pretrained(checkpoint, force_download=True)
            lowerCamelCase_  						=       fast_class.from_pretrained(checkpoint, force_download=True)
            test_tokenizer(slow, fast)
            print(f"""Accuracy {perfect * 100 / total:.2f}""")






