Unslothを​用いた​VLMの​ファインチューニング方​法

from unsloth.trainer import UnslothVisionDataCollator
from trl import SFTTrainer, SFTConfig
from unsloth import is_bf16_supported, FastVisionModel
import torch
from datasets import load_dataset

from data_prep import convert_to_conversation

# --- データセットのロードと整形 ---
dataset = load_dataset("unsloth/LaTeX_OCR", split = "train")
converted_dataset = [convert_to_conversation(sample) for sample in dataset]

instruction = "Write the LaTeX representation for this image."

# ---データセットの「画像」と「正解テキスト」を、対話形式のリストに変換する ---
def convert_to_conversation(sample):
    conversation = [
        { "role": "user",
          "content" : [
            {"type" : "text",  "text"  : instruction},
            {"type" : "image", "image" : sample["image"]} ]
        },
        { "role" : "assistant",
          "content" : [
            {"type" : "text",  "text"  : sample["text"]} ]
        },
    ]
    return { "messages" : conversation }

# 2-1の続き
# --- モデルとトークナイザーのロード ---
model, tokenizer = FastVisionModel.from_pretrained(
    "unsloth/Ministral-3-3B-Instruct-2512",
    load_in_4bit = False, # Use 4bit to reduce memory use. False for 16bit LoRA.
    use_gradient_checkpointing = "unsloth", # True or "unsloth" for long context
)

# --- LoRAアダプターの設定 ---
model = FastVisionModel.get_peft_model(
    model,
    finetune_vision_layers     = True, # False if not finetuning vision layers
    finetune_language_layers   = True, # False if not finetuning language layers
    finetune_attention_modules = True, # False if not finetuning attention layers
    finetune_mlp_modules       = True, # False if not finetuning MLP layers

    r = 32,           # The larger, the higher the accuracy, but might overfit
    lora_alpha = 32,  # Recommended alpha == r at least
    lora_dropout = 0,
    bias = "none",
    random_state = 3407,
    use_rslora = False,  # We support rank stabilized LoRA
    loftq_config = None, # And LoftQ
    # target_modules = "all-linear", # Optional now! Can specify a list if needed
)

# 2-2の続き
# ---トレーナーの設定---
trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    data_collator = UnslothVisionDataCollator(model, tokenizer), # Must use!
    train_dataset = converted_dataset,
    args = SFTConfig(
        per_device_train_batch_size = 4,
        gradient_accumulation_steps = 2,
        warmup_steps = 5,
        max_steps = 30,
        # num_train_epochs = 1, # Set this instead of max_steps for full training runs
        learning_rate = 2e-4,
        logging_steps = 1,
        optim = "adamw_8bit",
        fp16 = not is_bf16_supported(), # Use fp16 if bf16 is not supported
        bf16 = is_bf16_supported(), # Use bf16 if supported
        weight_decay = 0.001,
        lr_scheduler_type = "linear",
        seed = 3407,
        output_dir = "outputs",
        report_to = "tensorboard",     # For Weights and Biases

        # You MUST put the below items for vision finetuning:
        remove_unused_columns = False,
        dataset_text_field = "",
        dataset_kwargs = {"skip_prepare_dataset": True},
        max_length = 2048,
    ),
)

# --- 特訓の実行 ---
trainer_stats = trainer.train()

# --- 学習成果の保存(ローカルに保存） ---
model.save_pretrained("ministral_lora")  # Local saving
tokenizer.save_pretrained("ministral_lora")

from unsloth import FastVisionModel
from IPython.display import display, HTML
import html
import re

indices_to_test = [0]
results_compare = []
FastVisionModel.for_inference(model)

# --- テスト用データセットをロード ---
test_dataset = load_dataset("unsloth/LaTeX_OCR", split="test")

# ---【特訓前】元モデルでの推論 ---
print("元モデル（LoRAオフ）で推論中...")
model.disable_adapter_layers() # LoRAアダプターを一時的にオフにして使用しないように設定

for i in indices_to_test:
    image = test_dataset [i]["image"]
    instruction = "Write the LaTeX representation for this image."
    
    # 推論時にも学習時と同じシステムプロンプトを与える
    messages = [
        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": instruction}]}
    ]
    input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
    inputs = tokenizer(image, input_text, add_special_tokens=False, return_tensors="pt").to("cuda")
    
    outputs = model.generate(**inputs, max_new_tokens=256, use_cache=True, temperature=1.5, min_p=0.1)
    results_compare.append({
        "index": i, "true_latex": test_dataset [i]["text"],
        "before_raw": tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    })

# 元モデルでの推論の続き
# ---【特訓後】学習後モデルでの推論 ---
print("学習後モデル（LoRAオン）で推論中...")
model.enable_adapter_layers() # LoRAアダプターをオンにする

for i, res in enumerate(results_compare):
    idx = res["index"]
    image = test_dataset [idx]["image"]
    instruction = "Write the LaTeX representation for this image."
    
    messages = [
        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": instruction}]}
    ]
    input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
    inputs = tokenizer(image, input_text, add_special_tokens=False, return_tensors="pt").to("cuda")
    
    outputs = model.generate(**inputs, max_new_tokens=128, use_cache=True, temperature=1.5, min_p=0.1)
    results_compare[i]["after_raw"] = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)