


"""simple docstring"""



import argparse
import math
import os
from copy import deepcopy

import torch
from audio_diffusion.models import DiffusionAttnUnetaD
from diffusion import sampling
from torch import nn

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel


__snake_case   			=    {
    """gwf-440k""": {
        """url""": """https://model-server.zqevans2.workers.dev/gwf-440k.ckpt""",
        """sample_rate""": 48000,
        """sample_size""": 65536,
    },
    """jmann-small-190k""": {
        """url""": """https://model-server.zqevans2.workers.dev/jmann-small-190k.ckpt""",
        """sample_rate""": 48000,
        """sample_size""": 65536,
    },
    """jmann-large-580k""": {
        """url""": """https://model-server.zqevans2.workers.dev/jmann-large-580k.ckpt""",
        """sample_rate""": 48000,
        """sample_size""": 131072,
    },
    """maestro-uncond-150k""": {
        """url""": """https://model-server.zqevans2.workers.dev/maestro-uncond-150k.ckpt""",
        """sample_rate""": 16000,
        """sample_size""": 65536,
    },
    """unlocked-uncond-250k""": {
        """url""": """https://model-server.zqevans2.workers.dev/unlocked-uncond-250k.ckpt""",
        """sample_rate""": 16000,
        """sample_size""": 65536,
    },
    """honk-140k""": {
        """url""": """https://model-server.zqevans2.workers.dev/honk-140k.ckpt""",
        """sample_rate""": 16000,
        """sample_size""": 65536,
    },
}







def   __lowerCAmelCase						(						lowercase  :  Any						,							lowercase  :  Dict )	->						Tuple:


    """simple docstring"""





    return torch.atana(lowercase						,							lowercase ) / math.pi * 2







def   __lowerCAmelCase						(						lowercase  :  str )	->						List[str]:


    """simple docstring"""





    snake_case    :      str      					=    torch.sin(t * math.pi / 2 ) ** 2
    snake_case    :      Tuple      					=    (1 - sigma**2) ** 0.5
    return alpha_sigma_to_t(lowercase						,							lowercase )







class _lowerCAmelCase				(      snake_case_	):
   pass







class _lowerCAmelCase				(      nn.Module	):







   def __init__(	self						,	UpperCamelCase__						)	->		Optional[int]:

       '''simple docstring'''




       super().__init__()

       snake_case    :      Optional[Any]      					=    DiffusionAttnUnetaD(UpperCamelCase__						,	n_attn_layers=4						)
       snake_case    :      Union[str, Any]      					=    deepcopy(self.diffusion						)
       snake_case    :      Tuple      					=    torch.quasirandom.SobolEngine(1						,	scramble=UpperCamelCase__						)







def   __lowerCAmelCase						(						lowercase  :  Tuple )	->						int:


    """simple docstring"""





    snake_case    :      List[str]      					=    MODELS_MAP[model_name]["url"]
    os.system(F'wget {url} ./' )

    return F'./{model_name}.ckpt'


__snake_case   			=    {
    """1""": """resnets.0""",
    """2""": """attentions.0""",
    """3""": """resnets.1""",
    """4""": """attentions.1""",
    """5""": """resnets.2""",
    """6""": """attentions.2""",
}
__snake_case   			=    {
    """8""": """resnets.0""",
    """9""": """attentions.0""",
    """10""": """resnets.1""",
    """11""": """attentions.1""",
    """12""": """resnets.2""",
    """13""": """attentions.2""",
}
__snake_case   			=    {
    """1""": """resnets.0""",
    """2""": """attentions.0""",
    """3""": """resnets.1""",
    """4""": """attentions.1""",
    """5""": """resnets.2""",
    """6""": """attentions.2""",
    """8""": """resnets.3""",
    """9""": """attentions.3""",
    """10""": """resnets.4""",
    """11""": """attentions.4""",
    """12""": """resnets.5""",
    """13""": """attentions.5""",
}
__snake_case   			=    {
    """0""": """resnets.0""",
    """1""": """resnets.1""",
    """2""": """resnets.2""",
    """4""": """resnets.0""",
    """5""": """resnets.1""",
    """6""": """resnets.2""",
}

__snake_case   			=    {
    """skip""": """conv_skip""",
    """main.0""": """conv_1""",
    """main.1""": """group_norm_1""",
    """main.3""": """conv_2""",
    """main.4""": """group_norm_2""",
}

__snake_case   			=    {
    """norm""": """group_norm""",
    """qkv_proj""": ["""query""", """key""", """value"""],
    """out_proj""": ["""proj_attn"""],
}







def   __lowerCAmelCase						(						lowercase  :  Tuple )	->						List[str]:


    """simple docstring"""





    if name.startswith("skip" ):
        return name.replace("skip"						,							RES_CONV_MAP["skip"] )

    # name has to be of format main.{digit}
    if not name.startswith("main." ):
        raise ValueError(F'ResConvBlock error with {name}' )

    return name.replace(name[:6]						,							RES_CONV_MAP[name[:6]] )







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


    """simple docstring"""





    for key, value in ATTN_MAP.items():
        if name.startswith(lowercase ) and not isinstance(lowercase						,							lowercase ):
            return name.replace(lowercase						,							lowercase )
        elif name.startswith(lowercase ):
            return [name.replace(lowercase						,							lowercase ) for v in value]
    raise ValueError(F'Attn error with {name}' )







def   __lowerCAmelCase						(						lowercase  :  str						,							lowercase  :  Optional[Any]=13 )	->						Optional[int]:


    """simple docstring"""





    snake_case    :      int      					=    input_string

    if string.split("." )[0] == "timestep_embed":
        return string.replace("timestep_embed"						,							"time_proj" )

    snake_case    :      int      					=    0
    if string.startswith("net.3." ):
        depth += 1
        snake_case    :      List[str]      					=    string[6:]
    elif string.startswith("net." ):
        snake_case    :      Optional[Any]      					=    string[4:]

    while string.startswith("main.7." ):
        depth += 1
        snake_case    :      Tuple      					=    string[7:]

    if string.startswith("main." ):
        snake_case    :      int      					=    string[5:]

    # mid block
    if string[:2].isdigit():
        snake_case    :      Optional[Any]      					=    string[:2]
        snake_case    :      List[str]      					=    string[2:]
    else:
        snake_case    :      Optional[int]      					=    string[0]
        snake_case    :      Union[str, Any]      					=    string[1:]

    if depth == max_depth:
        snake_case    :      Optional[Any]      					=    MID_NUM_TO_LAYER[layer_num]
        snake_case    :      Tuple      					=    "mid_block"
    elif depth > 0 and int(lowercase ) < 7:
        snake_case    :      List[str]      					=    DOWN_NUM_TO_LAYER[layer_num]
        snake_case    :      Dict      					=    F'down_blocks.{depth}'
    elif depth > 0 and int(lowercase ) > 7:
        snake_case    :      Union[str, Any]      					=    UP_NUM_TO_LAYER[layer_num]
        snake_case    :      List[Any]      					=    F'up_blocks.{max_depth - depth - 1}'
    elif depth == 0:
        snake_case    :      Optional[int]      					=    DEPTH_0_TO_LAYER[layer_num]
        snake_case    :      List[Any]      					=    F'up_blocks.{max_depth - 1}' if int(lowercase ) > 3 else "down_blocks.0"

    if not string_left.startswith("." ):
        raise ValueError(F'Naming error with {input_string} and string_left: {string_left}.' )

    snake_case    :      Tuple      					=    string_left[1:]

    if "resnets" in new_layer:
        snake_case    :      Union[str, Any]      					=    convert_resconv_naming(lowercase )
    elif "attentions" in new_layer:
        snake_case    :      Union[str, Any]      					=    convert_attn_naming(lowercase )
        snake_case    :      List[str]      					=    new_string_left

    if not isinstance(lowercase						,							lowercase ):
        snake_case    :      List[Any]      					=    prefix + "." + new_layer + "." + string_left
    else:
        snake_case    :      Any      					=    [prefix + "." + new_layer + "." + s for s in string_left]
    return new_string







def   __lowerCAmelCase						(						lowercase  :  List[Any] )	->						Any:


    """simple docstring"""





    snake_case    :      int      					=    {}
    for k, v in state_dict.items():
        if k.endswith("kernel" ):
            # up- and downsample layers, don't have trainable weights
            continue

        snake_case    :      int      					=    rename(lowercase )

        # check if we need to transform from Conv => Linear for attention
        if isinstance(lowercase						,							lowercase ):
            snake_case    :      str      					=    transform_conv_attns(lowercase						,							lowercase						,							lowercase )
        else:
            snake_case    :      str      					=    v

    return new_state_dict







def   __lowerCAmelCase						(						lowercase  :  Optional[Any]						,							lowercase  :  Any						,							lowercase  :  Union[str, Any] )	->						List[str]:


    """simple docstring"""





    if len(lowercase ) == 1:
        if len(v.shape ) == 3:
            # weight
            snake_case    :      Dict      					=    v[:, :, 0]
        else:
            # bias
            snake_case    :      Dict      					=    v
    else:
        # qkv matrices
        snake_case    :      str      					=    v.shape[0]
        snake_case    :      int      					=    trippled_shape // 3
        for i in range(3 ):
            if len(v.shape ) == 3:
                snake_case    :      int      					=    v[i * single_shape : (i + 1) * single_shape, :, 0]
            else:
                snake_case    :      List[str]      					=    v[i * single_shape : (i + 1) * single_shape]
    return new_state_dict







def   __lowerCAmelCase						(						lowercase  :  Tuple )	->						List[str]:


    """simple docstring"""





    snake_case    :      int      					=    torch.device("cuda" if torch.cuda.is_available() else "cpu" )

    snake_case    :      Optional[int]      					=    args.model_path.split("/" )[-1].split("." )[0]
    if not os.path.isfile(args.model_path ):
        assert (
            model_name == args.model_path
        ), F'Make sure to provide one of the official model names {MODELS_MAP.keys()}'
        snake_case    :      int      					=    download(lowercase )

    snake_case    :      Optional[int]      					=    MODELS_MAP[model_name]["sample_rate"]
    snake_case    :      List[str]      					=    MODELS_MAP[model_name]["sample_size"]

    snake_case    :      List[str]      					=    Object()
    snake_case    :      Union[str, Any]      					=    sample_size
    snake_case    :      List[str]      					=    sample_rate
    snake_case    :      int      					=    0

    snake_case    :      List[str]      					=    UNetaDModel(sample_size=lowercase						,							sample_rate=lowercase )
    snake_case    :      Optional[int]      					=    diffusers_model.state_dict()

    snake_case    :      List[Any]      					=    DiffusionUncond(lowercase )
    orig_model.load_state_dict(torch.load(args.model_path						,							map_location=lowercase )["state_dict"] )
    snake_case    :      Any      					=    orig_model.diffusion_ema.eval()
    snake_case    :      str      					=    orig_model.state_dict()
    snake_case    :      str      					=    rename_orig_weights(lowercase )

    snake_case    :      List[Any]      					=    set(renamed_state_dict.keys() ) - set(diffusers_state_dict.keys() )
    snake_case    :      Dict      					=    set(diffusers_state_dict.keys() ) - set(renamed_state_dict.keys() )

    assert len(lowercase ) == 0, F'Problem with {renamed_minus_diffusers}'
    assert all(k.endswith("kernel" ) for k in list(lowercase ) ), F'Problem with {diffusers_minus_renamed}'

    for key, value in renamed_state_dict.items():
        assert (
            diffusers_state_dict[key].squeeze().shape == value.squeeze().shape
        ), F'Shape for {key} doesn\'t match. Diffusers: {diffusers_state_dict[key].shape} vs. {value.shape}'
        if key == "time_proj.weight":
            snake_case    :      int      					=    value.squeeze()

        snake_case    :      Dict      					=    value

    diffusers_model.load_state_dict(lowercase )

    snake_case    :      str      					=    100
    snake_case    :      Dict      					=    33

    snake_case    :      Any      					=    IPNDMScheduler(num_train_timesteps=lowercase )

    snake_case    :      Optional[int]      					=    torch.manual_seed(lowercase )
    snake_case    :      List[Any]      					=    torch.randn([1, 2, config.sample_size]						,							generator=lowercase ).to(lowercase )

    snake_case    :      Union[str, Any]      					=    torch.linspace(1						,							0						,							steps + 1						,							device=lowercase )[:-1]
    snake_case    :      str      					=    get_crash_schedule(lowercase )

    snake_case    :      str      					=    DanceDiffusionPipeline(unet=lowercase						,							scheduler=lowercase )

    snake_case    :      Tuple      					=    torch.manual_seed(33 )
    snake_case    :      List[str]      					=    pipe(num_inference_steps=lowercase						,							generator=lowercase ).audios

    snake_case    :      Optional[int]      					=    sampling.iplms_sample(lowercase						,							lowercase						,							lowercase						,							{} )
    snake_case    :      Any      					=    generated.clamp(-1						,							1 )

    snake_case    :      List[Any]      					=    (generated - audio).abs().sum()
    snake_case    :      Optional[Any]      					=    (generated - audio).abs().max()

    if args.save:
        pipe.save_pretrained(args.checkpoint_path )

    print("Diff sum"						,							lowercase )
    print("Diff max"						,							lowercase )

    assert diff_max < 1e-3, F'Diff max: {diff_max} is too much :-/'

    print(F'Conversion for {model_name} successful!' )


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

  parser.add_argument("""--model_path""", default=None, type=str, required=True, help="""Path to the model to convert.""")
  parser.add_argument(
      """--save""", default=True, type=bool, required=False, help="""Whether to save the converted model or not."""
  )
  parser.add_argument("""--checkpoint_path""", default=None, type=str, required=True, help="""Path to the output model.""")
  __snake_case   			=    parser.parse_args()

  main(args)









"""simple docstring"""



import unittest

from transformers import JukeboxTokenizer
from transformers.testing_utils import require_torch







class _lowerCAmelCase				(      unittest.TestCase	):
   __UpperCAmelCase :       Union[str, Any]							  =  JukeboxTokenizer
   __UpperCAmelCase :       Union[str, Any]							  =  {
       '''artist''': '''Zac Brown Band''',
       '''genres''': '''Country''',
       '''lyrics''': '''I met a traveller from an antique land,
        Who said "Two vast and trunkless legs of stone
        Stand in the desert. . . . Near them, on the sand,
        Half sunk a shattered visage lies, whose frown,
        And wrinkled lip, and sneer of cold command,
        Tell that its sculptor well those passions read
        Which yet survive, stamped on these lifeless things,
        The hand that mocked them, and the heart that fed;
        And on the pedestal, these words appear:
        My name is Ozymandias, King of Kings;
        Look on my Works, ye Mighty, and despair!
        Nothing beside remains. Round the decay
        Of that colossal Wreck, boundless and bare
        The lone and level sands stretch far away
        ''',
   }







   @require_torch
   def  lowerCamelCase		(	self						)	->		int:

       '''simple docstring'''




       import torch

       snake_case    :      Any      					=    JukeboxTokenizer.from_pretrained("openai/jukebox-1b-lyrics"						)
       snake_case    :      Optional[Any]      					=    tokenizer(**self.metas						)["input_ids"]
       # fmt: off
       snake_case    :      Optional[int]      					=    [
           torch.tensor([[
               0, 0, 0, 7169, 507, 9, 76, 39, 31, 46, 76, 27,
               76, 46, 44, 27, 48, 31, 38, 38, 31, 44, 76, 32,
               44, 41, 39, 76, 27, 40, 76, 27, 40, 46, 35, 43,
               47, 31, 76, 38, 27, 40, 30, 64, 78, 76, 76, 76,
               76, 76, 76, 76, 76, 23, 34, 41, 76, 45, 27, 35,
               30, 76, 71, 20, 49, 41, 76, 48, 27, 45, 46, 76,
               27, 40, 30, 76, 46, 44, 47, 40, 37, 38, 31, 45,
               45, 76, 38, 31, 33, 45, 76, 41, 32, 76, 45, 46,
               41, 40, 31, 78, 76, 76, 76, 76, 76, 76, 76, 76,
               19, 46, 27, 40, 30, 76, 35, 40, 76, 46, 34, 31,
               76, 30, 31, 45, 31, 44, 46, 63, 76, 63, 76, 63,
               76, 63, 76, 14, 31, 27, 44, 76, 46, 34, 31, 39,
               64, 76, 41, 40, 76, 46, 34, 31, 76, 45, 27, 40,
               30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76, 8,
               27, 38, 32, 76, 45, 47, 40, 37, 76, 27, 76, 45,
               34, 27, 46, 46, 31, 44, 31, 30, 76, 48, 35, 45,
               27, 33, 31, 76, 38, 35, 31, 45, 64, 76, 49, 34,
               41, 45, 31, 76, 32, 44, 41, 49, 40, 64, 78, 76,
               76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76, 49,
               44, 35, 40, 37, 38, 31, 30, 76, 38, 35, 42, 64,
               76, 27, 40, 30, 76, 45, 40, 31, 31, 44, 76, 41,
               32, 76, 29, 41, 38, 30, 76, 29, 41, 39, 39, 27,
               40, 30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76,
               20, 31, 38, 38, 76, 46, 34, 27, 46, 76, 35, 46,
               45, 76, 45, 29, 47, 38, 42, 46, 41, 44, 76, 49,
               31, 38, 38, 76, 46, 34, 41, 45, 31, 76, 42, 27,
               45, 45, 35, 41, 40, 45, 76, 44, 31, 27, 30, 78,
               76, 76, 76, 76, 76, 76, 76, 76, 23, 34, 35, 29,
               34, 76, 51, 31, 46, 76, 45, 47, 44, 48, 35, 48,
               31, 64, 76, 45, 46, 27, 39, 42, 31, 30, 76, 41,
               40, 76, 46, 34, 31, 45, 31, 76, 38, 35, 32, 31,
               38, 31, 45, 45, 76, 46, 34, 35, 40, 33, 45, 64,
               78, 76, 76, 76, 76, 76, 76, 76, 76, 20, 34, 31,
               76, 34, 27, 40, 30, 76, 46, 34, 27, 46, 76, 39,
               41, 29, 37, 31, 30, 76, 46, 34, 31, 39, 64, 76,
               27, 40, 30, 76, 46, 34, 31, 76, 34, 31, 27, 44,
               46, 76, 46, 34, 27, 46, 76, 32, 31, 30, 66, 78,
               76, 76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76,
               41, 40, 76, 46, 34, 31, 76, 42, 31, 30, 31, 45,
               46, 27, 38, 64, 76, 46, 34, 31, 45, 31, 76, 49,
               41, 44, 30, 45, 76, 27, 42, 42, 31, 27, 44, 65,
               78, 76, 76, 76, 76, 76, 76, 76, 76, 13, 51, 76,
               40, 27, 39, 31, 76, 35, 45, 76, 15, 52, 51, 39,
               27, 40, 30, 35, 27, 45, 64, 76, 11, 35, 40, 33,
               76, 41, 32, 76, 11, 35, 40, 33, 45, 66, 78, 76,
               76, 76, 76, 76, 76, 76, 76, 12, 41, 41, 37, 76,
               41, 40, 76, 39, 51, 76, 23, 41, 44, 37, 45, 64,
               76, 51, 31, 76, 13, 35, 33, 34, 46, 51, 64, 76,
               27, 40, 30, 76, 30, 31, 45, 42, 27, 35, 44, 67,
               78, 76, 76, 76, 76, 76, 76, 76, 76, 14, 41, 46,
               34, 35, 40, 33, 76, 28, 31, 45, 35, 30, 31, 76,
               44, 31, 39, 27, 35, 40, 45, 63, 76, 18, 41, 47,
               40, 30, 76, 46, 34, 31, 76, 30, 31, 29, 27, 51,
               78, 76, 76, 76, 76, 76, 76, 76, 76, 15, 32, 76,
               46, 34, 27, 46, 76, 29, 41, 38, 41, 45, 45, 27,
               38, 76, 23, 44, 31, 29, 37, 64, 76, 28, 41, 47,
               40, 30, 38, 31, 45, 45, 76, 27, 40, 30, 76, 28,
               27, 44, 31, 78, 76, 76, 76, 76, 76, 76, 76, 76,
               20, 34, 31, 76, 38, 41, 40, 31, 76, 27, 40, 30,
               76, 38, 31, 48, 31, 38, 76, 45, 27, 40, 30, 45,
               76, 45, 46, 44, 31, 46, 29, 34, 76, 32, 27, 44,
               76, 27, 49, 27, 51, 78, 76, 76, 76, 76, 76, 76,
               76, 76]]						),
           torch.tensor([[0, 0, 0, 1069, 11]]						),
           torch.tensor([[0, 0, 0, 1069, 11]]						),
       ]
       # fmt: on
       self.assertTrue(torch.allclose(tokens[0]						,	EXPECTED_OUTPUT[0]						)						)
       self.assertTrue(torch.allclose(tokens[1]						,	EXPECTED_OUTPUT[1]						)						)
       self.assertTrue(torch.allclose(tokens[2]						,	EXPECTED_OUTPUT[2]						)						)







   @require_torch
   def  lowerCamelCase		(	self						)	->		Any:

       '''simple docstring'''




       import torch

       snake_case    :      Tuple      					=    JukeboxTokenizer.from_pretrained("openai/jukebox-5b-lyrics"						)
       snake_case    :      Optional[Any]      					=    tokenizer(**self.metas						)["input_ids"]
       # fmt: off
       snake_case    :      List[Any]      					=    [
           torch.tensor([[
               0, 0, 0, 1069, 11, -1, -1, -1, -1, 9, 77, 39,
               31, 46, 77, 27, 77, 46, 44, 27, 48, 31, 38, 38,
               31, 44, 77, 32, 44, 41, 39, 77, 27, 40, 77, 27,
               40, 46, 35, 43, 47, 31, 77, 38, 27, 40, 30, 64,
               79, 77, 77, 77, 77, 77, 77, 77, 77, 23, 34, 41,
               77, 45, 27, 35, 30, 77, 72, 20, 49, 41, 77, 48,
               27, 45, 46, 77, 27, 40, 30, 77, 46, 44, 47, 40,
               37, 38, 31, 45, 45, 77, 38, 31, 33, 45, 77, 41,
               32, 77, 45, 46, 41, 40, 31, 79, 77, 77, 77, 77,
               77, 77, 77, 77, 19, 46, 27, 40, 30, 77, 35, 40,
               77, 46, 34, 31, 77, 30, 31, 45, 31, 44, 46, 63,
               77, 63, 77, 63, 77, 63, 77, 14, 31, 27, 44, 77,
               46, 34, 31, 39, 64, 77, 41, 40, 77, 46, 34, 31,
               77, 45, 27, 40, 30, 64, 79, 77, 77, 77, 77, 77,
               77, 77, 77, 8, 27, 38, 32, 77, 45, 47, 40, 37,
               77, 27, 77, 45, 34, 27, 46, 46, 31, 44, 31, 30,
               77, 48, 35, 45, 27, 33, 31, 77, 38, 35, 31, 45,
               64, 77, 49, 34, 41, 45, 31, 77, 32, 44, 41, 49,
               40, 64, 79, 77, 77, 77, 77, 77, 77, 77, 77, 1,
               40, 30, 77, 49, 44, 35, 40, 37, 38, 31, 30, 77,
               38, 35, 42, 64, 77, 27, 40, 30, 77, 45, 40, 31,
               31, 44, 77, 41, 32, 77, 29, 41, 38, 30, 77, 29,
               41, 39, 39, 27, 40, 30, 64, 79, 77, 77, 77, 77,
               77, 77, 77, 77, 20, 31, 38, 38, 77, 46, 34, 27,
               46, 77, 35, 46, 45, 77, 45, 29, 47, 38, 42, 46,
               41, 44, 77, 49, 31, 38, 38, 77, 46, 34, 41, 45,
               31, 77, 42, 27, 45, 45, 35, 41, 40, 45, 77, 44,
               31, 27, 30, 79, 77, 77, 77, 77, 77, 77, 77, 77,
               23, 34, 35, 29, 34, 77, 51, 31, 46, 77, 45, 47,
               44, 48, 35, 48, 31, 64, 77, 45, 46, 27, 39, 42,
               31, 30, 77, 41, 40, 77, 46, 34, 31, 45, 31, 77,
               38, 35, 32, 31, 38, 31, 45, 45, 77, 46, 34, 35,
               40, 33, 45, 64, 79, 77, 77, 77, 77, 77, 77, 77,
               77, 20, 34, 31, 77, 34, 27, 40, 30, 77, 46, 34,
               27, 46, 77, 39, 41, 29, 37, 31, 30, 77, 46, 34,
               31, 39, 64, 77, 27, 40, 30, 77, 46, 34, 31, 77,
               34, 31, 27, 44, 46, 77, 46, 34, 27, 46, 77, 32,
               31, 30, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77,
               1, 40, 30, 77, 41, 40, 77, 46, 34, 31, 77, 42,
               31, 30, 31, 45, 46, 27, 38, 64, 77, 46, 34, 31,
               45, 31, 77, 49, 41, 44, 30, 45, 77, 27, 42, 42,
               31, 27, 44, 65, 79, 77, 77, 77, 77, 77, 77, 77,
               77, 13, 51, 77, 40, 27, 39, 31, 77, 35, 45, 77,
               15, 52, 51, 39, 27, 40, 30, 35, 27, 45, 64, 77,
               11, 35, 40, 33, 77, 41, 32, 77, 11, 35, 40, 33,
               45, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77, 12,
               41, 41, 37, 77, 41, 40, 77, 39, 51, 77, 23, 41,
               44, 37, 45, 64, 77, 51, 31, 77, 13, 35, 33, 34,
               46, 51, 64, 77, 27, 40, 30, 77, 30, 31, 45, 42,
               27, 35, 44, 67, 79, 77, 77, 77, 77, 77, 77, 77,
               77, 14, 41, 46, 34, 35, 40, 33, 77, 28, 31, 45,
               35, 30, 31, 77, 44, 31, 39, 27, 35, 40, 45, 63,
               77, 18, 41, 47, 40, 30, 77, 46, 34, 31, 77, 30,
               31, 29, 27, 51, 79, 77, 77, 77, 77, 77, 77, 77,
               77, 15, 32, 77, 46, 34, 27, 46, 77, 29, 41, 38,
               41, 45, 45, 27, 38, 77, 23, 44, 31, 29, 37, 64,
               77, 28, 41, 47, 40, 30, 38, 31, 45, 45, 77, 27,
               40, 30, 77, 28, 27, 44, 31, 79, 77, 77, 77, 77,
               77, 77, 77, 77, 20, 34, 31, 77, 38, 41, 40, 31,
               77, 27, 40, 30, 77, 38, 31, 48, 31, 38, 77, 45,
               27, 40, 30, 45, 77, 45, 46, 44, 31, 46, 29, 34,
               77, 32, 27, 44, 77, 27, 49, 27, 51, 79, 77, 77,
               77, 77, 77, 77, 77, 77]]						),
           torch.tensor([[0, 0, 0, 1069, 11, -1, -1, -1, -1]]						),
           torch.tensor([[0, 0, 0, 1069, 11, -1, -1, -1, -1]]						),
       ]
       # fmt: on
       self.assertTrue(torch.allclose(tokens[0]						,	EXPECTED_OUTPUT[0]						)						)
       self.assertTrue(torch.allclose(tokens[1]						,	EXPECTED_OUTPUT[1]						)						)
       self.assertTrue(torch.allclose(tokens[2]						,	EXPECTED_OUTPUT[2]						)						)





