import argparse

import torch
from internvl.model.internvl_chat import InternVLChatModel
from transformers import AutoModel, AutoTokenizer

argparse = argparse.ArgumentParser()
argparse.add_argument('model_path', type=str, default='')
argparse.add_argument('llm_path', type=str, default='')

args = argparse.parse_args()

if args.model_path[-1] == '/':
    args.model_path = args.model_path[:-1]

model = InternVLChatModel.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)

llm = AutoModel.from_pretrained(
    args.llm_path, trust_remote_code=True, torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained(
    args.llm_path, trust_remote_code=True)
model.language_model = llm
model.config.llm_config = llm.config
model.to(torch.bfloat16)

output_path = args.model_path + '_replace_llm'
model.save_pretrained(output_path)
tokenizer.save_pretrained(output_path)
print('finished')