codingutf8 2023 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 bros checkpoints import argparse import bros original repo import torch from transformers import brosconfig brosmodel brosprocessor from transformers utils import logging logging setverbosityinfo logger logging getloggername def getconfigsmodelname brosconfig brosconfig frompretrainedmodelname return brosconfig def removeignorekeysstatedict ignorekeys embeddings bboxsinusoidemb invfreq for k in ignorekeys statedict popk none def renamekeyname if name embeddings bboxprojection weight name bboxembeddings bboxprojection weight if name embeddings bboxsinusoidemb xposemb invfreq name bboxembeddings bboxsinusoidemb xposemb invfreq if name embeddings bboxsinusoidemb yposemb invfreq name bboxembeddings bboxsinusoidemb yposemb invfreq return name def convertstatedictorigstatedict model rename keys for key in origstatedict copy keys val origstatedict popkey origstatedictrenamekeykey val remove ignore keys removeignorekeysorigstatedict return origstatedict def convertbroscheckpointmodelname pytorchdumpfolderpathnone pushtohubfalse load original model originalmodel bros brosmodel frompretrainedmodelname eval load huggingface model brosconfig getconfigsmodelname model brosmodel frompretrainedmodelname configbrosconfig model eval statedict originalmodel statedict newstatedict convertstatedictstatedict model model loadstatedictnewstatedict verify results original bros model require 4 points 8 float values for each bbox prepare bbox with batchsize seqlen 8 shape bbox torch tensor 0 0000 0 0000 0 0000 0 0000 0 0000 0 0000 0 0000 0 0000 0 4396 0 6720 0 4659 0 6720 0 4659 0 6850 0 4396 0 6850 0 4698 0 6720 0 4843 0 6720 0 4843 0 6850 0 4698 0 6850 0 4698 0 6720 0 4843 0 6720 0 4843 0 6850 0 4698 0 6850 0 2047 0 6870 0 2730 0 6870 0 2730 0 7000 0 2047 0 7000 0 2047 0 6870 0 2730 0 6870 0 2730 0 7000 0 2047 0 7000 1 0000 1 0000 1 0000 1 0000 1 0000 1 0000 1 0000 1 0000 processor brosprocessor frompretrainedmodelname encoding processorhis name is rocco returntensorspt encodingbbox bbox originalhiddenstates originalmodelencoding lasthiddenstate pixelvalues processorimage returntensorspt pixelvalues lasthiddenstates modelencoding lasthiddenstate assert torch allcloseoriginalhiddenstates lasthiddenstates atol1e4 if pytorchdumpfolderpath is not none printfsaving model and processor to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath processor savepretrainedpytorchdumpfolderpath if pushtohub model pushtohubjinho8345 modelname split1 commitmessageupdate model processor pushtohubjinho8345 modelname split1 commitmessageupdate model if name main parser argparse argumentparser required parameters parser addargument modelname defaultjinho8345brosbaseuncased requiredfalse typestr helpname of the original model you d like to convert parser addargument pytorchdumpfolderpath defaultnone requiredfalse typestr helppath to the output pytorch model directory parser addargument pushtohub actionstoretrue helpwhether or not to push the converted model and processor to the hub args parser parseargs convertbroscheckpointargs modelname args pytorchdumpfolderpath args pushtohub coding utf 8 2023 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 bros checkpoints original repo rename keys remove ignore keys load original model load huggingface model verify results original bros model require 4 points 8 float values for each bbox prepare bbox with batch_size seq_len 8 shape pixel_values processor image return_tensors pt pixel_values required parameters
import argparse

import bros  
import torch

from transformers import BrosConfig, BrosModel, BrosProcessor
from transformers.utils import logging


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


def get_configs(model_name):
    bros_config = BrosConfig.from_pretrained(model_name)
    return bros_config


def remove_ignore_keys_(state_dict):
    ignore_keys = [
        "embeddings.bbox_sinusoid_emb.inv_freq",
    ]
    for k in ignore_keys:
        state_dict.pop(k, None)


def rename_key(name):
    if name == "embeddings.bbox_projection.weight":
        name = "bbox_embeddings.bbox_projection.weight"

    if name == "embeddings.bbox_sinusoid_emb.x_pos_emb.inv_freq":
        name = "bbox_embeddings.bbox_sinusoid_emb.x_pos_emb.inv_freq"

    if name == "embeddings.bbox_sinusoid_emb.y_pos_emb.inv_freq":
        name = "bbox_embeddings.bbox_sinusoid_emb.y_pos_emb.inv_freq"

    return name


def convert_state_dict(orig_state_dict, model):
    
    for key in orig_state_dict.copy().keys():
        val = orig_state_dict.pop(key)
        orig_state_dict[rename_key(key)] = val

    
    remove_ignore_keys_(orig_state_dict)

    return orig_state_dict


def convert_bros_checkpoint(model_name, pytorch_dump_folder_path=None, push_to_hub=False):
    
    original_model = bros.BrosModel.from_pretrained(model_name).eval()

    
    bros_config = get_configs(model_name)
    model = BrosModel.from_pretrained(model_name, config=bros_config)
    model.eval()

    state_dict = original_model.state_dict()
    new_state_dict = convert_state_dict(state_dict, model)
    model.load_state_dict(new_state_dict)

    

    
    bbox = torch.tensor(
        [
            [
                [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
                [0.4396, 0.6720, 0.4659, 0.6720, 0.4659, 0.6850, 0.4396, 0.6850],
                [0.4698, 0.6720, 0.4843, 0.6720, 0.4843, 0.6850, 0.4698, 0.6850],
                [0.4698, 0.6720, 0.4843, 0.6720, 0.4843, 0.6850, 0.4698, 0.6850],
                [0.2047, 0.6870, 0.2730, 0.6870, 0.2730, 0.7000, 0.2047, 0.7000],
                [0.2047, 0.6870, 0.2730, 0.6870, 0.2730, 0.7000, 0.2047, 0.7000],
                [1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
            ]
        ]
    )

    processor = BrosProcessor.from_pretrained(model_name)

    encoding = processor("His name is Rocco.", return_tensors="pt")
    encoding["bbox"] = bbox

    original_hidden_states = original_model(**encoding).last_hidden_state
    

    last_hidden_states = model(**encoding).last_hidden_state

    assert torch.allclose(original_hidden_states, last_hidden_states, atol=1e-4)

    if pytorch_dump_folder_path is not None:
        print(f"Saving model and processor to {pytorch_dump_folder_path}")
        model.save_pretrained(pytorch_dump_folder_path)
        processor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        model.push_to_hub("jinho8345/" + model_name.split("/")[-1], commit_message="Update model")
        processor.push_to_hub("jinho8345/" + model_name.split("/")[-1], commit_message="Update model")


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

    
    parser.add_argument(
        "--model_name",
        default="jinho8345/bros-base-uncased",
        required=False,
        type=str,
        help="Name of the original model you'd like to convert.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default=None,
        required=False,
        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 and processor to the 🤗 hub.",
    )

    args = parser.parse_args()
    convert_bros_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)