


import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import YolosConfig, YolosForObjectDetection, YolosImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
SCREAMING_SNAKE_CASE       :      Dict												= logging.get_logger(__name__)



def 				UpperCamelCase_(      lowerCamelCase_ )    ->						YolosConfig:
    _lowercase       :							Optional[Any]      		=		YolosConfig()

    # size of the architecture
    if "yolos_ti" in yolos_name:
        _lowercase       :							Optional[Any]      		=		192
        _lowercase       :							Optional[int]      		=		768
        _lowercase       :							List[Any]      		=		12
        _lowercase       :							List[str]      		=		3
        _lowercase       :							List[Any]      		=		[800, 1333]
        _lowercase       :							int      		=		False
    elif yolos_name == "yolos_s_dWr":
        _lowercase       :							Optional[int]      		=		330
        _lowercase       :							Optional[Any]      		=		14
        _lowercase       :							List[Any]      		=		6
        _lowercase       :							Tuple      		=		1320
    elif "yolos_s" in yolos_name:
        _lowercase       :							int      		=		384
        _lowercase       :							int      		=		1536
        _lowercase       :							List[str]      		=		12
        _lowercase       :							int      		=		6
    elif "yolos_b" in yolos_name:
        _lowercase       :							Dict      		=		[800, 1344]

    _lowercase       :							Tuple      		=		91
    _lowercase       :							str      		=		'huggingface/label-files'
    _lowercase       :							int      		=		'coco-detection-id2label.json'
    _lowercase       :							List[Any]      		=		json.load(open(hf_hub_download(lowerCamelCase_							, lowerCamelCase_							, repo_type='dataset' )							, 'r' ) )
    _lowercase       :							List[str]      		=		{int(lowerCamelCase_ ): v for k, v in idalabel.items()}
    _lowercase       :							Optional[int]      		=		idalabel
    _lowercase       :							int      		=		{v: k for k, v in idalabel.items()}

    return config



def 				UpperCamelCase_(      lowerCamelCase_							, lowerCamelCase_							, lowerCamelCase_ = False )    ->						List[Any]:
    for i in range(config.num_hidden_layers ):
        # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
        _lowercase       :							Any      		=		state_dict.pop(F'''blocks.{i}.attn.qkv.weight''' )
        _lowercase       :							Optional[int]      		=		state_dict.pop(F'''blocks.{i}.attn.qkv.bias''' )
        # next, add query, keys and values (in that order) to the state dict
        _lowercase       :							Any      		=		in_proj_weight[: config.hidden_size, :]
        _lowercase       :							Optional[int]      		=		in_proj_bias[: config.hidden_size]
        _lowercase       :							List[str]      		=		in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        _lowercase       :							Tuple      		=		in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        _lowercase       :							Optional[Any]      		=		in_proj_weight[-config.hidden_size :, :]
        _lowercase       :							Union[str, Any]      		=		in_proj_bias[-config.hidden_size :]



def 				UpperCamelCase_(      lowerCamelCase_ )    ->						str:
    if "backbone" in name:
        _lowercase       :							str      		=		name.replace('backbone'							, 'vit' )
    if "cls_token" in name:
        _lowercase       :							Union[str, Any]      		=		name.replace('cls_token'							, 'embeddings.cls_token' )
    if "det_token" in name:
        _lowercase       :							Tuple      		=		name.replace('det_token'							, 'embeddings.detection_tokens' )
    if "mid_pos_embed" in name:
        _lowercase       :							Dict      		=		name.replace('mid_pos_embed'							, 'encoder.mid_position_embeddings' )
    if "pos_embed" in name:
        _lowercase       :							Optional[Any]      		=		name.replace('pos_embed'							, 'embeddings.position_embeddings' )
    if "patch_embed.proj" in name:
        _lowercase       :							Any      		=		name.replace('patch_embed.proj'							, 'embeddings.patch_embeddings.projection' )
    if "blocks" in name:
        _lowercase       :							List[Any]      		=		name.replace('blocks'							, 'encoder.layer' )
    if "attn.proj" in name:
        _lowercase       :							Union[str, Any]      		=		name.replace('attn.proj'							, 'attention.output.dense' )
    if "attn" in name:
        _lowercase       :							List[Any]      		=		name.replace('attn'							, 'attention.self' )
    if "norm1" in name:
        _lowercase       :							List[str]      		=		name.replace('norm1'							, 'layernorm_before' )
    if "norm2" in name:
        _lowercase       :							Union[str, Any]      		=		name.replace('norm2'							, 'layernorm_after' )
    if "mlp.fc1" in name:
        _lowercase       :							Union[str, Any]      		=		name.replace('mlp.fc1'							, 'intermediate.dense' )
    if "mlp.fc2" in name:
        _lowercase       :							Optional[Any]      		=		name.replace('mlp.fc2'							, 'output.dense' )
    if "class_embed" in name:
        _lowercase       :							Dict      		=		name.replace('class_embed'							, 'class_labels_classifier' )
    if "bbox_embed" in name:
        _lowercase       :							Optional[Any]      		=		name.replace('bbox_embed'							, 'bbox_predictor' )
    if "vit.norm" in name:
        _lowercase       :							Tuple      		=		name.replace('vit.norm'							, 'vit.layernorm' )

    return name



def 				UpperCamelCase_(      lowerCamelCase_							, lowerCamelCase_ )    ->						dict:
    for key in orig_state_dict.copy().keys():
        _lowercase       :							Dict      		=		orig_state_dict.pop(lowerCamelCase_ )

        if "qkv" in key:
            _lowercase       :							List[Any]      		=		key.split('.' )
            _lowercase       :							Any      		=		int(key_split[2] )
            _lowercase       :							Union[str, Any]      		=		model.vit.encoder.layer[layer_num].attention.attention.all_head_size
            if "weight" in key:
                _lowercase       :							str      		=		val[:dim, :]
                _lowercase       :							List[Any]      		=		val[
                    dim : dim * 2, :
                ]
                _lowercase       :							int      		=		val[-dim:, :]
            else:
                _lowercase       :							Optional[int]      		=		val[:dim]
                _lowercase       :							Optional[Any]      		=		val[dim : dim * 2]
                _lowercase       :							List[Any]      		=		val[-dim:]
        else:
            _lowercase       :							str      		=		val

    return orig_state_dict



def 				UpperCamelCase_(      )    ->						torch.Tensor:
    _lowercase       :							List[str]      		=		'http://images.cocodataset.org/val2017/000000039769.jpg'
    _lowercase       :							Optional[int]      		=		Image.open(requests.get(lowerCamelCase_							, stream=lowerCamelCase_ ).raw )
    return im



@torch.no_grad()
def 				UpperCamelCase_(      lowerCamelCase_							, lowerCamelCase_							, lowerCamelCase_							, lowerCamelCase_ = False )    ->						str:
    _lowercase       :							int      		=		get_yolos_config(lowerCamelCase_ )

    # load original state_dict
    _lowercase       :							Dict      		=		torch.load(lowerCamelCase_							, map_location='cpu' )['model']

    # load 🤗 model
    _lowercase       :							Optional[Any]      		=		YolosForObjectDetection(lowerCamelCase_ )
    model.eval()
    _lowercase       :							Optional[int]      		=		convert_state_dict(lowerCamelCase_							, lowerCamelCase_ )
    model.load_state_dict(lowerCamelCase_ )

    # Check outputs on an image, prepared by YolosImageProcessor
    _lowercase       :							List[str]      		=		800 if yolos_name != 'yolos_ti' else 512
    _lowercase       :							Any      		=		YolosImageProcessor(format='coco_detection'							, size=lowerCamelCase_ )
    _lowercase       :							int      		=		image_processor(images=prepare_img()							, return_tensors='pt' )
    _lowercase       :							Union[str, Any]      		=		model(**lowerCamelCase_ )
    _lowercase		,				_lowercase       :							Any      		=		outputs.logits, outputs.pred_boxes

    _lowercase		,				_lowercase       :							List[str]      		=		None, None
    if yolos_name == "yolos_ti":
        _lowercase       :							Optional[int]      		=		torch.tensor(
            [[-39.50_22, -11.98_20, -17.68_88], [-29.95_74, -9.97_69, -17.76_91], [-42.32_81, -20.72_00, -30.62_94]] )
        _lowercase       :							Optional[int]      		=		torch.tensor(
            [[0.40_21, 0.08_36, 0.79_79], [0.01_84, 0.26_09, 0.03_64], [0.17_81, 0.20_04, 0.20_95]] )
    elif yolos_name == "yolos_s_200_pre":
        _lowercase       :							Dict      		=		torch.tensor(
            [[-24.02_48, -10.30_24, -14.82_90], [-42.03_92, -16.82_00, -27.43_34], [-27.27_43, -11.81_54, -18.71_48]] )
        _lowercase       :							Tuple      		=		torch.tensor(
            [[0.25_59, 0.54_55, 0.47_06], [0.29_89, 0.72_79, 0.18_75], [0.77_32, 0.40_17, 0.44_62]] )
    elif yolos_name == "yolos_s_300_pre":
        _lowercase       :							Any      		=		torch.tensor(
            [[-36.22_20, -14.43_85, -23.54_57], [-35.69_70, -14.75_83, -21.39_35], [-31.59_39, -13.60_42, -16.80_49]] )
        _lowercase       :							Optional[Any]      		=		torch.tensor(
            [[0.76_14, 0.23_16, 0.47_28], [0.71_68, 0.44_95, 0.38_55], [0.49_96, 0.14_66, 0.99_96]] )
    elif yolos_name == "yolos_s_dWr":
        _lowercase       :							Optional[Any]      		=		torch.tensor(
            [[-42.86_68, -24.10_49, -41.16_90], [-34.74_56, -14.12_74, -24.91_94], [-33.78_98, -12.19_46, -25.64_95]] )
        _lowercase       :							Optional[int]      		=		torch.tensor(
            [[0.55_87, 0.27_73, 0.06_05], [0.50_04, 0.30_14, 0.99_94], [0.49_99, 0.15_48, 0.99_94]] )
    elif yolos_name == "yolos_base":
        _lowercase       :							Optional[int]      		=		torch.tensor(
            [[-40.60_64, -24.30_84, -32.64_47], [-55.19_90, -30.77_19, -35.58_77], [-51.43_11, -33.35_07, -35.64_62]] )
        _lowercase       :							Dict      		=		torch.tensor(
            [[0.55_55, 0.27_94, 0.06_55], [0.90_49, 0.26_64, 0.18_94], [0.91_83, 0.19_84, 0.16_35]] )
    else:
        raise ValueError(F'''Unknown yolos_name: {yolos_name}''' )

    assert torch.allclose(logits[0, :3, :3]							, lowerCamelCase_							, atol=1e-4 )
    assert torch.allclose(pred_boxes[0, :3, :3]							, lowerCamelCase_							, atol=1e-4 )

    Path(lowerCamelCase_ ).mkdir(exist_ok=lowerCamelCase_ )
    print(F'''Saving model {yolos_name} to {pytorch_dump_folder_path}''' )
    model.save_pretrained(lowerCamelCase_ )
    print(F'''Saving image processor to {pytorch_dump_folder_path}''' )
    image_processor.save_pretrained(lowerCamelCase_ )

    if push_to_hub:
        _lowercase       :							str      		=		{
            'yolos_ti': 'yolos-tiny',
            'yolos_s_200_pre': 'yolos-small',
            'yolos_s_300_pre': 'yolos-small-300',
            'yolos_s_dWr': 'yolos-small-dwr',
            'yolos_base': 'yolos-base',
        }

        print('Pushing to the hub...' )
        _lowercase       :							Any      		=		model_mapping[yolos_name]
        image_processor.push_to_hub(lowerCamelCase_							, organization='hustvl' )
        model.push_to_hub(lowerCamelCase_							, organization='hustvl' )


if __name__ == "__main__":
 SCREAMING_SNAKE_CASE       :      Optional[int]												= argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     "--yolos_name",
     default="yolos_s_200_pre",
     type=str,
     help=(
         "Name of the YOLOS model you'd like to convert. Should be one of 'yolos_ti', 'yolos_s_200_pre',"
         " 'yolos_s_300_pre', 'yolos_s_dWr', 'yolos_base'."
     ),
 )
 parser.add_argument(
     "--checkpoint_path", default=None, type=str, help="Path to the original state dict (.pth file)."
 )
 parser.add_argument(
     "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
 )
 parser.add_argument(
     "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
 )

 SCREAMING_SNAKE_CASE       :      Dict												= parser.parse_args()
 convert_yolos_checkpoint(args.yolos_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)






import unittest

from transformers import JukeboxTokenizer
from transformers.testing_utils import require_torch






class _lowerCamelCase(						unittest.TestCase   ):
       lowercase_ :	Dict     =			JukeboxTokenizer
       lowercase_ :	Dict     =			{
           """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 			UpperCamelCase     (					self)       ->   Optional[int]:




           """simple docstring"""

           import torch

           _lowercase       :							str      		=		JukeboxTokenizer.from_pretrained('openai/jukebox-1b-lyrics')
           _lowercase       :							Optional[Any]      		=		tokenizer(**self.metas)['input_ids']
           # fmt: off
           _lowercase       :							Optional[int]      		=		[
               torch.tensor([[
                   0, 0, 0, 71_69, 5_07, 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, 10_69, 11]]),
               torch.tensor([[0, 0, 0, 10_69, 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 			UpperCamelCase     (					self)       ->   int:




           """simple docstring"""

           import torch

           _lowercase       :							List[str]      		=		JukeboxTokenizer.from_pretrained('openai/jukebox-5b-lyrics')
           _lowercase       :							List[str]      		=		tokenizer(**self.metas)['input_ids']
           # fmt: off
           _lowercase       :							Optional[int]      		=		[
               torch.tensor([[
                   0, 0, 0, 10_69, 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, 10_69, 11, -1, -1, -1, -1]]),
               torch.tensor([[0, 0, 0, 10_69, 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]))


