OEvortex/HelpingAI-Vision

1# Copyright 2023 The HuggingFace Inc. team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14import argparse15 16import torch17 18from transformers import (19    AddedToken,20    AutoConfig,21    AutoTokenizer,22)23from configuration_llava import LlavaConfig24from modeling_llava import LlavaForConditionalGeneration25 26 27KEYS_TO_MODIFY_MAPPING = {28    "transformer.vision_tower.vision_tower": "vision_model",29    "transformer.mm_projector": "multi_modal_projector",30    "transformer": "language_model.transformer",31    "lm_head": "language_model.lm_head",32    "model.model": "language_model.transformer",33    "multi_modal_projector.0": "multi_modal_projector.linear_1",34    "multi_modal_projector.2": "multi_modal_projector.linear_2",35}36 37 38def convert_state_dict_to_hf(state_dict):39    new_state_dict = {}40    for key, value in state_dict.items():41        for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():42            if key_to_modify in key:43                key = key.replace(key_to_modify, new_key)44 45        new_state_dict[key] = value46    return new_state_dict47 48 49def convert_llava_llama_to_hf(text_model_id, vision_model_id, projector_tokens_num, output_path, old_state_dict_path):50    torch.set_default_dtype(torch.float16)51    text_config = AutoConfig.from_pretrained(text_model_id, trust_remote_code=True)52 53    tokenizer = AutoTokenizer.from_pretrained(text_model_id)54    tokenizer.add_tokens(AddedToken("<image>", special=True, normalized=False), special_tokens=True)55    tokenizer.add_special_tokens({"pad_token": "<pad>"})56 57    config = LlavaConfig(text_config=text_config, vocab_size=51200, vision_tower_name=vision_model_id, projector_tokens_num=projector_tokens_num)58    config.text_config.vocab_size = config.vocab_size59 60    with torch.device("cuda"):61        model = LlavaForConditionalGeneration(config)62    63    state_dict = torch.load(old_state_dict_path, map_location="cpu")64    state_dict = convert_state_dict_to_hf(state_dict)65    model.load_state_dict(state_dict, strict=True, assign=True)66 67    model.config.vocab_size = model.config.vocab_size68    model.config.text_config.vocab_size = model.config.text_config.vocab_size69 70    model.save_pretrained(output_path)71    tokenizer.save_pretrained(output_path)72 73 74def main():75    parser = argparse.ArgumentParser()76    parser.add_argument(77        "--text_model_id",78        help="Hub location of the text model",79    )80    parser.add_argument(81        "--vision_model_id",82        help="Hub location of the vision model",83    )84    parser.add_argument(85        "--output_path",86        help="Location of the converted model",87    )88    parser.add_argument(89        "--old_state_dict_path",90        help="Location on the hub of the raw state dict of the original model. The filename needs to be `model_state_dict.bin`",91    )92    parser.add_argument(93        "--tokens_num",94        type=int,95        default=196    )97    args = parser.parse_args()98    convert_llava_llama_to_hf(args.text_model_id, args.vision_model_id, args.tokens_num, args.output_path, args.old_state_dict_path)99 100 101if __name__ == "__main__":102    main()