


"""simple docstring"""



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 (
    BertTokenizer,
    ViltConfig,
    ViltForImageAndTextRetrieval,
    ViltForImagesAndTextClassification,
    ViltForMaskedLM,
    ViltForQuestionAnswering,
    ViltImageProcessor,
    ViltProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
__snake_case   			=    logging.get_logger(__name__)







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


    """simple docstring"""





    snake_case    :      int      					=    []
    for i in range(config.num_hidden_layers ):
        # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
        rename_keys.append((F'transformer.blocks.{i}.norm1.weight', F'vilt.encoder.layer.{i}.layernorm_before.weight') )
        rename_keys.append((F'transformer.blocks.{i}.norm1.bias', F'vilt.encoder.layer.{i}.layernorm_before.bias') )
        rename_keys.append(
            (F'transformer.blocks.{i}.attn.proj.weight', F'vilt.encoder.layer.{i}.attention.output.dense.weight') )
        rename_keys.append(
            (F'transformer.blocks.{i}.attn.proj.bias', F'vilt.encoder.layer.{i}.attention.output.dense.bias') )
        rename_keys.append((F'transformer.blocks.{i}.norm2.weight', F'vilt.encoder.layer.{i}.layernorm_after.weight') )
        rename_keys.append((F'transformer.blocks.{i}.norm2.bias', F'vilt.encoder.layer.{i}.layernorm_after.bias') )
        rename_keys.append(
            (F'transformer.blocks.{i}.mlp.fc1.weight', F'vilt.encoder.layer.{i}.intermediate.dense.weight') )
        rename_keys.append((F'transformer.blocks.{i}.mlp.fc1.bias', F'vilt.encoder.layer.{i}.intermediate.dense.bias') )
        rename_keys.append((F'transformer.blocks.{i}.mlp.fc2.weight', F'vilt.encoder.layer.{i}.output.dense.weight') )
        rename_keys.append((F'transformer.blocks.{i}.mlp.fc2.bias', F'vilt.encoder.layer.{i}.output.dense.bias') )

    # embeddings
    rename_keys.extend(
        [
            # text embeddings
            ("text_embeddings.word_embeddings.weight", "vilt.embeddings.text_embeddings.word_embeddings.weight"),
            (
                "text_embeddings.position_embeddings.weight",
                "vilt.embeddings.text_embeddings.position_embeddings.weight",
            ),
            ("text_embeddings.position_ids", "vilt.embeddings.text_embeddings.position_ids"),
            (
                "text_embeddings.token_type_embeddings.weight",
                "vilt.embeddings.text_embeddings.token_type_embeddings.weight",
            ),
            ("text_embeddings.LayerNorm.weight", "vilt.embeddings.text_embeddings.LayerNorm.weight"),
            ("text_embeddings.LayerNorm.bias", "vilt.embeddings.text_embeddings.LayerNorm.bias"),
            # patch embeddings
            ("transformer.cls_token", "vilt.embeddings.cls_token"),
            ("transformer.patch_embed.proj.weight", "vilt.embeddings.patch_embeddings.projection.weight"),
            ("transformer.patch_embed.proj.bias", "vilt.embeddings.patch_embeddings.projection.bias"),
            ("transformer.pos_embed", "vilt.embeddings.position_embeddings"),
            # token type embeddings
            ("token_type_embeddings.weight", "vilt.embeddings.token_type_embeddings.weight"),
        ] )

    # final layernorm + pooler
    rename_keys.extend(
        [
            ("transformer.norm.weight", "vilt.layernorm.weight"),
            ("transformer.norm.bias", "vilt.layernorm.bias"),
            ("pooler.dense.weight", "vilt.pooler.dense.weight"),
            ("pooler.dense.bias", "vilt.pooler.dense.bias"),
        ] )

    # classifier head(s)
    if vqa_model:
        # classification head
        rename_keys.extend(
            [
                ("vqa_classifier.0.weight", "classifier.0.weight"),
                ("vqa_classifier.0.bias", "classifier.0.bias"),
                ("vqa_classifier.1.weight", "classifier.1.weight"),
                ("vqa_classifier.1.bias", "classifier.1.bias"),
                ("vqa_classifier.3.weight", "classifier.3.weight"),
                ("vqa_classifier.3.bias", "classifier.3.bias"),
            ] )
    elif nlvr_model:
        # classification head
        rename_keys.extend(
            [
                ("nlvr2_classifier.0.weight", "classifier.0.weight"),
                ("nlvr2_classifier.0.bias", "classifier.0.bias"),
                ("nlvr2_classifier.1.weight", "classifier.1.weight"),
                ("nlvr2_classifier.1.bias", "classifier.1.bias"),
                ("nlvr2_classifier.3.weight", "classifier.3.weight"),
                ("nlvr2_classifier.3.bias", "classifier.3.bias"),
            ] )
    else:
        pass

    return rename_keys







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


    """simple docstring"""





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







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


    """simple docstring"""





    snake_case    :      Tuple      					=    ["head.weight", "head.bias"]
    for k in ignore_keys:
        state_dict.pop(lowercase						,							lowercase )







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


    """simple docstring"""





    snake_case    :      Tuple      					=    dct.pop(lowercase )
    snake_case    :      Any      					=    val







@torch.no_grad()
def   __lowerCAmelCase						(						lowercase  :  Dict						,							lowercase  :  Optional[int] )	->						Optional[Any]:


    """simple docstring"""





    snake_case    :      Tuple      					=    ViltConfig(image_size=384						,							patch_size=32						,							tie_word_embeddings=lowercase )
    snake_case    :      Tuple      					=    False
    snake_case    :      Optional[Any]      					=    False
    snake_case    :      Optional[Any]      					=    False
    snake_case    :      Optional[Any]      					=    False
    if "vqa" in checkpoint_url:
        snake_case    :      Tuple      					=    True
        snake_case    :      List[Any]      					=    3129
        snake_case    :      Any      					=    "huggingface/label-files"
        snake_case    :      List[str]      					=    "vqa2-id2label.json"
        snake_case    :      Optional[int]      					=    json.load(open(hf_hub_download(lowercase						,							lowercase						,							repo_type="dataset" )						,							"r" ) )
        snake_case    :      Dict      					=    {int(lowercase ): v for k, v in idalabel.items()}
        snake_case    :      Union[str, Any]      					=    idalabel
        snake_case    :      Tuple      					=    {v: k for k, v in idalabel.items()}
        snake_case    :      List[str]      					=    ViltForQuestionAnswering(lowercase )
    elif "nlvr" in checkpoint_url:
        snake_case    :      str      					=    True
        snake_case    :      Dict      					=    2
        snake_case    :      Dict      					=    {0: "False", 1: "True"}
        snake_case    :      Any      					=    {v: k for k, v in config.idalabel.items()}
        snake_case    :      Any      					=    3
        snake_case    :      Optional[Any]      					=    ViltForImagesAndTextClassification(lowercase )
    elif "irtr" in checkpoint_url:
        snake_case    :      List[str]      					=    True
        snake_case    :      List[str]      					=    ViltForImageAndTextRetrieval(lowercase )
    elif "mlm_itm" in checkpoint_url:
        snake_case    :      Optional[int]      					=    True
        snake_case    :      List[Any]      					=    ViltForMaskedLM(lowercase )
    else:
        raise ValueError("Unknown model type" )

    # load state_dict of original model, remove and rename some keys
    snake_case    :      Dict      					=    torch.hub.load_state_dict_from_url(lowercase						,							map_location="cpu" )["state_dict"]
    snake_case    :      Union[str, Any]      					=    create_rename_keys(lowercase						,							lowercase						,							lowercase						,							lowercase )
    for src, dest in rename_keys:
        rename_key(lowercase						,							lowercase						,							lowercase )
    read_in_q_k_v(lowercase						,							lowercase )
    if mlm_model or irtr_model:
        snake_case    :      List[str]      					=    ["itm_score.fc.weight", "itm_score.fc.bias"]
        for k in ignore_keys:
            state_dict.pop(lowercase						,							lowercase )

    # load state dict into HuggingFace model
    model.eval()
    if mlm_model:
        snake_case					,snake_case    :      Dict      					=    model.load_state_dict(lowercase						,							strict=lowercase )
        assert missing_keys == ["mlm_score.decoder.bias"]
    else:
        model.load_state_dict(lowercase )

    # Define processor
    snake_case    :      List[str]      					=    ViltImageProcessor(size=384 )
    snake_case    :      Optional[Any]      					=    BertTokenizer.from_pretrained("bert-base-uncased" )
    snake_case    :      List[str]      					=    ViltProcessor(lowercase						,							lowercase )

    # Forward pass on example inputs (image + text)
    if nlvr_model:
        snake_case    :      Union[str, Any]      					=    Image.open(requests.get("https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg"						,							stream=lowercase ).raw )
        snake_case    :      Dict      					=    Image.open(requests.get("https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg"						,							stream=lowercase ).raw )
        snake_case    :      Dict      					=    (
            "The left image contains twice the number of dogs as the right image, and at least two dogs in total are"
            " standing."
        )
        snake_case    :      str      					=    processor(lowercase						,							lowercase						,							return_tensors="pt" )
        snake_case    :      Optional[int]      					=    processor(lowercase						,							lowercase						,							return_tensors="pt" )
        snake_case    :      int      					=    model(
            input_ids=encoding_a.input_ids						,							pixel_values=encoding_a.pixel_values						,							pixel_values_a=encoding_a.pixel_values						,							)
    else:
        snake_case    :      Optional[int]      					=    Image.open(requests.get("http://images.cocodataset.org/val2017/000000039769.jpg"						,							stream=lowercase ).raw )
        if mlm_model:
            snake_case    :      Optional[int]      					=    "a bunch of [MASK] laying on a [MASK]."
        else:
            snake_case    :      str      					=    "How many cats are there?"
        snake_case    :      Optional[Any]      					=    processor(lowercase						,							lowercase						,							return_tensors="pt" )
        snake_case    :      Any      					=    model(**lowercase )

    # Verify outputs
    if mlm_model:
        snake_case    :      int      					=    torch.Size([1, 11, 3_0522] )
        snake_case    :      int      					=    torch.tensor([-12.5061, -12.5123, -12.5174] )
        assert outputs.logits.shape == expected_shape
        assert torch.allclose(outputs.logits[0, 0, :3]						,							lowercase						,							atol=1e-4 )

        # verify masked token prediction equals "cats"
        snake_case    :      str      					=    outputs.logits[0, 4, :].argmax(-1 ).item()
        assert tokenizer.decode([predicted_id] ) == "cats"
    elif vqa_model:
        snake_case    :      Dict      					=    torch.Size([1, 3129] )
        snake_case    :      Dict      					=    torch.tensor([-15.9495, -18.1472, -10.3041] )
        assert torch.allclose(outputs.logits[0, :3]						,							lowercase						,							atol=1e-4 )
        assert outputs.logits.shape == expected_shape
        assert torch.allclose(outputs.logits[0, 0, :3]						,							lowercase						,							atol=1e-4 )

        # verify vqa prediction equals "2"
        snake_case    :      Any      					=    outputs.logits.argmax(-1 ).item()
        assert model.config.idalabel[predicted_idx] == "2"
    elif nlvr_model:
        snake_case    :      Optional[Any]      					=    torch.Size([1, 2] )
        snake_case    :      List[Any]      					=    torch.tensor([-2.8721, 2.1291] )
        assert torch.allclose(outputs.logits[0, :3]						,							lowercase						,							atol=1e-4 )
        assert outputs.logits.shape == expected_shape

    Path(lowercase ).mkdir(exist_ok=lowercase )
    print(F'Saving model and processor to {pytorch_dump_folder_path}' )
    model.save_pretrained(lowercase )
    processor.save_pretrained(lowercase )


if __name__ == "__main__":
  __snake_case   			=    argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--checkpoint_url""",
      default="""https://github.com/dandelin/ViLT/releases/download/200k/vilt_200k_mlm_itm.ckpt""",
      type=str,
      help="""URL of the checkpoint you'd like to convert.""",
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
  )

  __snake_case   			=    parser.parse_args()
  convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









"""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]						)						)





