register_pruned_smolvlm.py

from transformers import AutoModelForVision2Seq, AutoProcessor
from SmolVLM_500M_Instruct_bl_15.modeling_idefics3bl import Idefics3blForConditionalGeneration
from SmolVLM_500M_Instruct_bl_15.configuration_idefics3bl import Idefics3blConfig, Idefics3blVisionConfig, LlamablConfig

import torch
from peft import PeftModel

Idefics3blConfig.register_for_auto_class()
Idefics3blVisionConfig.register_for_auto_class()
LlamablConfig.register_for_auto_class()

Idefics3blForConditionalGeneration.register_for_auto_class('AutoModelForVision2Seq')

# base_path = "SmolVLM_500M_Instruct_bl_0_4"
# base_path = "SmolVLM_500M_Instruct_bl_15"
# lora_path = "/home/kaixin/programs/LLM-Pruner/tune_log/SmolVLM-500M-Instruct-bl-0_2-4-28-tuned_qlora-r16-llava150K"

base_path = "/home/kaixin/programs/LLM-Pruner/SmolVLM-500M-Instruct-bl-15"
lora_path = "/home/kaixin/programs/LLM-Pruner/tune_log/SmolVLM-500M-Instruct-bl-0_2-4-28-tuned_qlora-r16-llava150K+textvqa_freezevision"
stage3_path = "/home/kaixin/programs/LLM-Pruner/tune_log/SmolVLM-500M-Instruct-bl-15_stage3_moredata_new"


hf_name = "kartmannXu"
hf_name = "HuangRT"

config = Idefics3blConfig.from_pretrained(base_path)
base_model = Idefics3blForConditionalGeneration.from_pretrained(base_path, config=config)
tokenizer = AutoProcessor.from_pretrained(base_path)

model = PeftModel.from_pretrained(
                base_model,
                lora_path,
                torch_dtype=torch.float32,
                device_map="cuda"
                # torch_dtype="bfloat16",
            )
model = model.merge_and_unload()

# merge stage3
model = PeftModel.from_pretrained(
                model,
                stage3_path,
                torch_dtype=torch.float32,
                device_map="cuda"
                # torch_dtype="bfloat16",
            )
model = model.merge_and_unload()
print(model)

model.push_to_hub(hf_name + '/SmolVLM-500M-Instruct-bl-15-stage3-new')
tokenizer.push_to_hub(hf_name + '/SmolVLM-500M-Instruct-bl-15-stage3-new')