codingutf8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert trocr checkpoints from the unilm repository import argparse from pathlib import path import requests import torch from pil import image from transformers import robertatokenizer trocrconfig trocrforcausallm trocrprocessor visionencoderdecodermodel vitconfig vitimageprocessor vitmodel from transformers utils import logging logging setverbosityinfo logger logging getloggername here we list all keys to be renamed original name on the left our name on the right def createrenamekeysencoderconfig decoderconfig renamekeys for i in rangeencoderconfig numhiddenlayers encoder layers output projection 2 feedforward neural networks and 2 layernorms renamekeys append fencoder deit blocks i norm1 weight fencoder encoder layer i layernormbefore weight renamekeys appendfencoder deit blocks i norm1 bias fencoder encoder layer i layernormbefore bias renamekeys append fencoder deit blocks i attn proj weight fencoder encoder layer i attention output dense weight renamekeys append fencoder deit blocks i attn proj bias fencoder encoder layer i attention output dense bias renamekeys append fencoder deit blocks i norm2 weight fencoder encoder layer i layernormafter weight renamekeys appendfencoder deit blocks i norm2 bias fencoder encoder layer i layernormafter bias renamekeys append fencoder deit blocks i mlp fc1 weight fencoder encoder layer i intermediate dense weight renamekeys append fencoder deit blocks i mlp fc1 bias fencoder encoder layer i intermediate dense bias renamekeys append fencoder deit blocks i mlp fc2 weight fencoder encoder layer i output dense weight renamekeys appendfencoder deit blocks i mlp fc2 bias fencoder encoder layer i output dense bias cls token position embeddings and patch embeddings of encoder renamekeys extend encoder deit clstoken encoder embeddings clstoken encoder deit posembed encoder embeddings positionembeddings encoder deit patchembed proj weight encoder embeddings patchembeddings projection weight encoder deit patchembed proj bias encoder embeddings patchembeddings projection bias encoder deit norm weight encoder layernorm weight encoder deit norm bias encoder layernorm bias return renamekeys we split up the matrix of each encoder layer into queries keys and values def readinqkvstatedict encoderconfig for i in rangeencoderconfig numhiddenlayers queries keys and values only weights no biases inprojweight statedict popfencoder deit blocks i attn qkv weight statedictfencoder encoder layer i attention attention query weight inprojweight encoderconfig hiddensize statedictfencoder encoder layer i attention attention key weight inprojweight encoderconfig hiddensize encoderconfig hiddensize 2 statedictfencoder encoder layer i attention attention value weight inprojweight encoderconfig hiddensize def renamekeydct old new val dct popold dctnew val we will verify our results on an image of the iam handwriting database def prepareimgcheckpointurl if handwritten in checkpointurl url https fki tic heiafr chstaticimga011220200 jpg industry url https fki tic heiafr chstaticimga011220212 jpg have url https fki tic heiafr chstaticimga011220210 jpg let url https fki tic heiafr chstaticimga0112202 jpg url https fki tic heiafr chstaticimga01122 jpg elif printed in checkpointurl or stage1 in checkpointurl url https www researchgate netprofiledinhsangpublication338099565figurefig8as 8404132293509221577381536857anreceiptexampleinthesroie2019datasetq640 jpg im image openrequests geturl streamtrue raw convertrgb return im torch nograd def converttrocrcheckpointcheckpointurl pytorchdumpfolderpath define encoder and decoder configs based on checkpointurl encoderconfig vitconfigimagesize384 qkvbiasfalse decoderconfig trocrconfig size of the architecture if base in checkpointurl decoderconfig encoderhiddensize 768 elif large in checkpointurl use vitlarge encoder encoderconfig hiddensize 1024 encoderconfig intermediatesize 4096 encoderconfig numhiddenlayers 24 encoderconfig numattentionheads 16 decoderconfig encoderhiddensize 1024 else raise valueerrorshould either find base or large in checkpoint url the largeprinted stage1 checkpoints uses sinusoidal position embeddings no layernorm afterwards if largeprinted in checkpointurl or stage1 in checkpointurl decoderconfig tiewordembeddings false decoderconfig activationfunction relu decoderconfig maxpositionembeddings 1024 decoderconfig scaleembedding true decoderconfig uselearnedpositionembeddings false decoderconfig layernormembedding false load huggingface model encoder vitmodelencoderconfig addpoolinglayerfalse decoder trocrforcausallmdecoderconfig model visionencoderdecodermodelencoderencoder decoderdecoder model eval load statedict of original model rename some keys statedict torch hub loadstatedictfromurlcheckpointurl maplocationcpu checkhashtruemodel renamekeys createrenamekeysencoderconfig decoderconfig for src dest in renamekeys renamekeystatedict src dest readinqkvstatedict encoderconfig remove parameters we don t need del statedictencoder deit head weight del statedictencoder deit head bias del statedictdecoder version add prefix to decoder keys for key val in statedict copy items val statedict popkey if key startswithdecoder and outputprojection not in key statedictdecoder model key val else statedictkey val load state dict model loadstatedictstatedict check outputs on an image imageprocessor vitimageprocessorsizeencoderconfig imagesize tokenizer robertatokenizer frompretrainedrobertalarge processor trocrprocessorimageprocessor tokenizer pixelvalues processorimagesprepareimgcheckpointurl returntensorspt pixelvalues verify logits decoderinputids torch tensormodel config decoder decoderstarttokenid outputs modelpixelvaluespixelvalues decoderinputidsdecoderinputids logits outputs logits expectedshape torch size1 1 50265 if trocrbasehandwritten in checkpointurl expectedslice torch tensor 1 4502 4 6683 0 5347 2 9291 9 1435 3 0571 8 9764 1 7560 8 7358 1 5311 elif trocrlargehandwritten in checkpointurl expectedslice torch tensor 2 6437 1 3129 2 2596 5 3455 6 3539 1 7604 5 4991 1 4702 5 6113 2 0170 elif trocrbaseprinted in checkpointurl expectedslice torch tensor 5 6816 5 8388 1 1398 6 9034 6 8505 2 4393 1 2284 1 0232 1 9661 3 9210 elif trocrlargeprinted in checkpointurl expectedslice torch tensor 6 0162 7 0959 4 4155 5 1063 7 0468 3 1631 2 6466 0 3081 0 8106 1 7535 if stage1 not in checkpointurl assert logits shape expectedshape shape of logits not as expected assert torch allcloselogits0 0 10 expectedslice atol1e3 first elements of logits not as expected pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath printfsaving processor to pytorchdumpfolderpath processor savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser parser addargument checkpointurl defaulthttps layoutlm blob core windows nettrocrmodelzoofairseqtrocrbasehandwritten pt typestr helpurl to the original pytorch checkpoint pth file parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the folder to output pytorch model args parser parseargs converttrocrcheckpointargs checkpointurl args pytorchdumpfolderpath coding utf 8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert trocr checkpoints from the unilm repository here we list all keys to be renamed original name on the left our name on the right encoder layers output projection 2 feedforward neural networks and 2 layernorms cls token position embeddings and patch embeddings of encoder we split up the matrix of each encoder layer into queries keys and values queries keys and values only weights no biases we will verify our results on an image of the iam handwriting database industry url https fki tic heia fr ch static img a01 122 02 12 jpg have url https fki tic heia fr ch static img a01 122 02 10 jpg let url https fki tic heia fr ch static img a01 122 02 jpg url https fki tic heia fr ch static img a01 122 jpg copy paste tweak model s weights to our visionencoderdecodermodel structure define encoder and decoder configs based on checkpoint_url size of the architecture use vit large encoder the large printed stage1 checkpoints uses sinusoidal position embeddings no layernorm afterwards load huggingface model load state_dict of original model rename some keys remove parameters we don t need add prefix to decoder keys load state dict check outputs on an image verify logits
import argparse
from pathlib import Path

import requests
import torch
from PIL import Image

from transformers import (
    RobertaTokenizer,
    TrOCRConfig,
    TrOCRForCausalLM,
    TrOCRProcessor,
    VisionEncoderDecoderModel,
    ViTConfig,
    ViTImageProcessor,
    ViTModel,
)
from transformers.utils import logging


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



def create_rename_keys(encoder_config, decoder_config):
    rename_keys = []
    for i in range(encoder_config.num_hidden_layers):
        
        rename_keys.append(
            (f"encoder.deit.blocks.{i}.norm1.weight", f"encoder.encoder.layer.{i}.layernorm_before.weight")
        )
        rename_keys.append((f"encoder.deit.blocks.{i}.norm1.bias", f"encoder.encoder.layer.{i}.layernorm_before.bias"))
        rename_keys.append(
            (f"encoder.deit.blocks.{i}.attn.proj.weight", f"encoder.encoder.layer.{i}.attention.output.dense.weight")
        )
        rename_keys.append(
            (f"encoder.deit.blocks.{i}.attn.proj.bias", f"encoder.encoder.layer.{i}.attention.output.dense.bias")
        )
        rename_keys.append(
            (f"encoder.deit.blocks.{i}.norm2.weight", f"encoder.encoder.layer.{i}.layernorm_after.weight")
        )
        rename_keys.append((f"encoder.deit.blocks.{i}.norm2.bias", f"encoder.encoder.layer.{i}.layernorm_after.bias"))
        rename_keys.append(
            (f"encoder.deit.blocks.{i}.mlp.fc1.weight", f"encoder.encoder.layer.{i}.intermediate.dense.weight")
        )
        rename_keys.append(
            (f"encoder.deit.blocks.{i}.mlp.fc1.bias", f"encoder.encoder.layer.{i}.intermediate.dense.bias")
        )
        rename_keys.append(
            (f"encoder.deit.blocks.{i}.mlp.fc2.weight", f"encoder.encoder.layer.{i}.output.dense.weight")
        )
        rename_keys.append((f"encoder.deit.blocks.{i}.mlp.fc2.bias", f"encoder.encoder.layer.{i}.output.dense.bias"))

    
    rename_keys.extend(
        [
            ("encoder.deit.cls_token", "encoder.embeddings.cls_token"),
            ("encoder.deit.pos_embed", "encoder.embeddings.position_embeddings"),
            ("encoder.deit.patch_embed.proj.weight", "encoder.embeddings.patch_embeddings.projection.weight"),
            ("encoder.deit.patch_embed.proj.bias", "encoder.embeddings.patch_embeddings.projection.bias"),
            ("encoder.deit.norm.weight", "encoder.layernorm.weight"),
            ("encoder.deit.norm.bias", "encoder.layernorm.bias"),
        ]
    )

    return rename_keys



def read_in_q_k_v(state_dict, encoder_config):
    for i in range(encoder_config.num_hidden_layers):
        
        in_proj_weight = state_dict.pop(f"encoder.deit.blocks.{i}.attn.qkv.weight")

        state_dict[f"encoder.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[
            : encoder_config.hidden_size, :
        ]
        state_dict[f"encoder.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
            encoder_config.hidden_size : encoder_config.hidden_size * 2, :
        ]
        state_dict[f"encoder.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[
            -encoder_config.hidden_size :, :
        ]


def rename_key(dct, old, new):
    val = dct.pop(old)
    dct[new] = val



def prepare_img(checkpoint_url):
    if "handwritten" in checkpoint_url:
        url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-00.jpg"  
        
        
        
        
    elif "printed" in checkpoint_url or "stage1" in checkpoint_url:
        url = "https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg"
    im = Image.open(requests.get(url, stream=True).raw).convert("RGB")
    return im


@torch.no_grad()
def convert_tr_ocr_checkpoint(checkpoint_url, pytorch_dump_folder_path):
    
    
    encoder_config = ViTConfig(image_size=384, qkv_bias=False)
    decoder_config = TrOCRConfig()

    
    if "base" in checkpoint_url:
        decoder_config.encoder_hidden_size = 768
    elif "large" in checkpoint_url:
        
        encoder_config.hidden_size = 1024
        encoder_config.intermediate_size = 4096
        encoder_config.num_hidden_layers = 24
        encoder_config.num_attention_heads = 16
        decoder_config.encoder_hidden_size = 1024
    else:
        raise ValueError("Should either find 'base' or 'large' in checkpoint URL")

    
    if "large-printed" in checkpoint_url or "stage1" in checkpoint_url:
        decoder_config.tie_word_embeddings = False
        decoder_config.activation_function = "relu"
        decoder_config.max_position_embeddings = 1024
        decoder_config.scale_embedding = True
        decoder_config.use_learned_position_embeddings = False
        decoder_config.layernorm_embedding = False

    
    encoder = ViTModel(encoder_config, add_pooling_layer=False)
    decoder = TrOCRForCausalLM(decoder_config)
    model = VisionEncoderDecoderModel(encoder=encoder, decoder=decoder)
    model.eval()

    
    state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu", check_hash=True)["model"]

    rename_keys = create_rename_keys(encoder_config, decoder_config)
    for src, dest in rename_keys:
        rename_key(state_dict, src, dest)
    read_in_q_k_v(state_dict, encoder_config)

    
    del state_dict["encoder.deit.head.weight"]
    del state_dict["encoder.deit.head.bias"]
    del state_dict["decoder.version"]

    
    for key, val in state_dict.copy().items():
        val = state_dict.pop(key)
        if key.startswith("decoder") and "output_projection" not in key:
            state_dict["decoder.model." + key] = val
        else:
            state_dict[key] = val

    
    model.load_state_dict(state_dict)

    
    image_processor = ViTImageProcessor(size=encoder_config.image_size)
    tokenizer = RobertaTokenizer.from_pretrained("roberta-large")
    processor = TrOCRProcessor(image_processor, tokenizer)

    pixel_values = processor(images=prepare_img(checkpoint_url), return_tensors="pt").pixel_values

    
    decoder_input_ids = torch.tensor([[model.config.decoder.decoder_start_token_id]])
    outputs = model(pixel_values=pixel_values, decoder_input_ids=decoder_input_ids)
    logits = outputs.logits

    expected_shape = torch.Size([1, 1, 50265])
    if "trocr-base-handwritten" in checkpoint_url:
        expected_slice = torch.tensor(
            [-1.4502, -4.6683, -0.5347, -2.9291, 9.1435, -3.0571, 8.9764, 1.7560, 8.7358, -1.5311]
        )
    elif "trocr-large-handwritten" in checkpoint_url:
        expected_slice = torch.tensor(
            [-2.6437, -1.3129, -2.2596, -5.3455, 6.3539, 1.7604, 5.4991, 1.4702, 5.6113, 2.0170]
        )
    elif "trocr-base-printed" in checkpoint_url:
        expected_slice = torch.tensor(
            [-5.6816, -5.8388, 1.1398, -6.9034, 6.8505, -2.4393, 1.2284, -1.0232, -1.9661, -3.9210]
        )
    elif "trocr-large-printed" in checkpoint_url:
        expected_slice = torch.tensor(
            [-6.0162, -7.0959, 4.4155, -5.1063, 7.0468, -3.1631, 2.6466, -0.3081, -0.8106, -1.7535]
        )

    if "stage1" not in checkpoint_url:
        assert logits.shape == expected_shape, "Shape of logits not as expected"
        assert torch.allclose(logits[0, 0, :10], expected_slice, atol=1e-3), "First elements of logits not as expected"

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)
    print(f"Saving processor to {pytorch_dump_folder_path}")
    processor.save_pretrained(pytorch_dump_folder_path)


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

    parser.add_argument(
        "--checkpoint_url",
        default="https://layoutlm.blob.core.windows.net/trocr/model_zoo/fairseq/trocr-base-handwritten.pt",
        type=str,
        help="URL to the original PyTorch checkpoint (.pth file).",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
    )
    args = parser.parse_args()
    convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)