Qwen2-VLの仕組みと概念モデルの実装

# ==============================================================================
# 本記事の概念モデル実装コードは、Apache License 2.0 で公開されている Qwen2-VL の
# 公式モデリングコードおよび、Hugging Face 上の zhouyik/DenseLabelDev の実装を
# 一部参照・改変して作成されています。
# ==============================================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import AdamW
import torchvision.transforms as T
from datasets import load_dataset
from transformers import AutoTokenizer
import matplotlib.pyplot as plt
from PIL import Image
import math


class MiniQwen2VL(nn.Module):
    def __init__(self, vocab_size, vision_dim=256, llm_dim=512, nhead=4):
        super().__init__()
        # 画像パッチから特徴を抽出するモジュールを定義
        self.patch_embed = DynamicPatchEmbedding(embed_dim=vision_dim)
        # 視覚トークンをマージして圧縮し、LLMの入力次元に合わせるモジュールを定義
        self.merger = PatchMerger(in_dim=vision_dim, out_dim=llm_dim)
        # テキストトークンをベクトル埋め込みに変換する層
        self.text_embedding = nn.Embedding(vocab_size, llm_dim)

        # M-RoPEのセクション分割 (head_dim = 128 = 32 + 48 + 48) -> *2すると 64, 96, 96 (元の仕様に準拠)
        # ※head_dim // 2 に対して分割を定義します。ここでは head_dim=128 なので、半分は64。
        # 64を Temporal, Height, Width に割り当てます: [16, 24, 24]
        self.mrope_section = [16, 24, 24]
        # マルチモーダル回転位置埋め込み（M-RoPE）の計算用クラス
        self.rotary_emb = Qwen2VLRotaryEmbedding(dim=llm_dim // nhead)

        # トランスフォーマーデコーダー層を4層作成
        self.layers = nn.ModuleList([
            CustomTransformerLayer(d_model=llm_dim, nhead=nhead, mrope_section=self.mrope_section)
            for _ in range(4)
        ])
        # 最終層出力の正規化層
        self.norm = nn.LayerNorm(llm_dim)
        # 最終出力を語彙空間のロジットに変換する全結合層
        self.lm_head = nn.Linear(llm_dim, vocab_size)

    def get_3d_position_ids(self, img_grid_size, text_len, device, batch_size=1):
        """画像のH, Wグリッドとテキスト長から3D position IDsを生成"""
        H, W = img_grid_size
        img_len = H * W

        # 1. Vision用のPosition ID (T, H, W)
        t_pos = torch.zeros(img_len, dtype=torch.long) # 静止画なので時間軸Tは0で固定
        h_pos = torch.arange(H).view(-1, 1).expand(H, W).flatten() # 高さ（Y座標）方向のインデックスを生成
        w_pos = torch.arange(W).view(1, -1).expand(H, W).flatten() # 幅（X座標）方向のインデックスを生成
        vision_pos_ids = torch.stack([t_pos, h_pos, w_pos]) # 3つの軸の位置情報を結合して (3, img_len) にする

        # 2. Text用のPosition ID (T, H, W は同じ値を取る)
        # 画像トークンの次のインデックスからテキストのインデックスを開始する
        start_idx = vision_pos_ids.max() + 1
        text_pos_ids = torch.arange(start_idx, start_idx + text_len).expand(3, text_len) # 各軸共通で割り当て (3, text_len)

        # 3. 視覚とテキストの位置IDをシーケンス方向に結合: (3, seq_len)
        position_ids = torch.cat([vision_pos_ids, text_pos_ids], dim=1).to(device)

        # バッチ次元 (B) を追加して (3, B, seq_len) の形状にして返す
        return position_ids.unsqueeze(1).expand(3, batch_size, -1)

    def forward(self, image, text_seq):
        # 1. ビジョンエンコーダに画像を送り、動的解像度パッチ特徴とグリッドサイズ（H, W）を抽出
        img_feats, grid_size = self.patch_embed(image)
        # 2. 視覚特徴トークンを1/4に圧縮し、LLMの次元（llm_dim）にマッピング
        merged_img_feats, merged_grid_size = self.merger(img_feats, grid_size)
        # 3. 入力されたテキストトークンIDを埋め込みベクトルに変換
        text_feats = self.text_embedding(text_seq)

        # 4. 圧縮された画像特徴とテキスト特徴をシーケンス方向（dim=1）に直列結合
        multimodal_inputs = torch.cat([merged_img_feats, text_feats], dim=1)
        total_len = multimodal_inputs.size(1)
        B = image.size(0) # バッチサイズを取得

        # 5. 自己回帰的なデコードのための因果マスク（Causal Mask）を生成
        mask = nn.Transformer.generate_square_subsequent_mask(total_len).to(image.device)

        # 6. 位置IDを生成し、対応するRoPE（M-RoPE）のsin/cos周波数を計算
        position_ids = self.get_3d_position_ids(merged_grid_size, text_seq.size(1), image.device, batch_size=B)
        cos, sin = self.rotary_emb(position_ids) # (3, B, seq_len, head_dim)

        # 7. トランスフォーマー層を順番に適用
        hidden_states = multimodal_inputs
        for layer in self.layers:
            hidden_states = layer(hidden_states, cos, sin, mask)

        # 8. 正規化と出力層の適用
        hidden_states = self.norm(hidden_states)
        logits = self.lm_head(hidden_states) # (B, seq_len, vocab_size)

        img_len = merged_img_feats.size(1) # 画像由来のトークン数を返す（出力の切り出しに必要）
        return logits, img_len

class DynamicPatchEmbedding(nn.Module):
    def __init__(self, in_channels=3, embed_dim=256, patch_size=14):
        super().__init__()
        self.patch_size = patch_size
        # 重なりのない畳み込み層を使い、画像をパッチに分割すると同時に特徴を抽出
        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, images):
        # images形状: [Batch, Channels, Height, Width]
        x = self.proj(images) # 出力形状: [Batch, embed_dim, H_p, W_p]
        H_p, W_p = x.shape[2], x.shape[3] # 畳み込み後のグリッドサイズ（高さ・幅）を保持
        # [Batch, embed_dim, H_p, W_p] -> [Batch, embed_dim, H_p * W_p] -> [Batch, H_p * W_p, embed_dim]
        x = x.flatten(2).transpose(1, 2)
        return x, (H_p, W_p)

class PatchMerger(nn.Module):
    def __init__(self, in_dim=256, out_dim=512):
        super().__init__()
        # 結合前の入力を正規化するための層
        self.ln_q = nn.LayerNorm(in_dim)
        # 2x2（=4個）のトークンをチャネル結合した 4*in_dim 次元を、LLMの次元である out_dim に圧縮投影するMLP
        self.mlp = nn.Sequential(
            nn.Linear(in_dim * 4, in_dim * 4),
            nn.GELU(),
            nn.Linear(in_dim * 4, out_dim)
        )

    def forward(self, x, grid_size):
        B, L, C = x.shape
        H_p, W_p = grid_size
        # 1次元のシーケンス形状を元の2次元のグリッド形状 [Batch, H_p, W_p, C] に変形
        x = x.view(B, H_p, W_p, C)
        x = self.ln_q(x)
        
        # 2x2のローカルグリッドトークンをまとめるため、テンソルを細分化して変形
        # H_p // 2 と W_p // 2 のグリッドに分割し、それぞれ2ピクセルずつのセクションを作成
        x = x.view(B, H_p // 2, 2, W_p // 2, 2, C)
        # 各ブロック内の2x2（=4個）のピクセルが連続して結合されるように、並び替えを実施
        # 置換後の形状: [Batch, H_p//2, W_p//2, 2, 2, C]
        x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
        # 2x2=4個のトークン特徴（チャネル数 C）を平坦化して 4*C 次元のチャンネルにマージ
        # 形状は [Batch, (H_p // 2) * (W_p // 2), 4 * C] になり、トークン数が 1/4 に圧縮されます
        x = x.view(B, (H_p // 2) * (W_p // 2), 4 * C)
        
        # マージした 4*C 次元の特徴量を MLP を通してアラインメントと次元削減
        merged_features = self.mlp(x)
        # 圧縮後の特徴量と、1/2 になった新しいグリッドサイズを返す
        return merged_features, (H_p // 2, W_p // 2)

class MRopeAttention(nn.Module):
    def __init__(self, d_model, nhead, mrope_section):
        super().__init__()
        self.d_model = d_model
        self.nhead = nhead
        self.head_dim = d_model // nhead
        self.mrope_section = mrope_section

        # Query, Key, Value への射影線形層を定義
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        # アテンションの出力を結合し、元の次元に戻す線形層
        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x, cos, sin, mask=None):
        B, L, C = x.shape
        # Q, K, V を計算し、マルチヘッドの形状に変換
        # [Batch, Length, C] -> [Batch, Length, Heads, Head_dim] -> [Batch, Heads, Length, Head_dim]
        q = self.q_proj(x).view(B, L, self.nhead, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, L, self.nhead, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, L, self.nhead, self.head_dim).transpose(1, 2)

        # M-RoPEを適用し、Query と Key に対して時空間の3次元回転位置情報を注入
        q, k = apply_multimodal_rotary_pos_emb(q, k, cos, sin, self.mrope_section, unsqueeze_dim=1)

        # スケールド・ドットプロダクト・アテンションの重みを計算
        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        # 因果マスク等がある場合、アテンションの重みに加算（未注目位置は負の無限大方向へシフト）
        if mask is not None:
            attn_weights = attn_weights + mask

        # ソフトマックスにより確率分布へ変換
        attn_probs = F.softmax(attn_weights, dim=-1)
        # 重みに基づいて Value ベクトルを加算
        attn_output = torch.matmul(attn_probs, v)
        # マルチヘッドの出力を元の次元に再構成
        # [Batch, Heads, Length, Head_dim] -> [Batch, Length, Heads, Head_dim] -> [Batch, Length, C]
        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L, C)
        # 最終的な全結合層を通す
        return self.out_proj(attn_output)

class CustomTransformerLayer(nn.Module):
    def __init__(self, d_model, nhead, mrope_section):
        super().__init__()
        # カスタム M-RoPE アテンションブロック
        self.attn = MRopeAttention(d_model, nhead, mrope_section)
        # 前段の正規化（Pre-LN）
        self.norm1 = nn.LayerNorm(d_model)
        # 特徴量の表現力を拡大する2層のMLP
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Linear(d_model * 4, d_model)
        )
        # 後段の正規化（Pre-LN）
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x, cos, sin, mask=None):
        # 1. 前正規化 -> アテンション計算 -> 残差接続
        x = x + self.attn(self.norm1(x), cos, sin, mask)
        # 2. 前正規化 -> MLP計算 -> 残差接続
        x = x + self.mlp(self.norm2(x))
        return x

def rotate_half(x):
    """特徴次元の半分にマイナスを付与して入れ替えることで、複素数空間の90度回転を再現"""
    x1 = x[..., : x.shape[-1] // 2] # 前半の次元を取得
    x2 = x[..., x.shape[-1] // 2 :] # 後半の次元を取得
    return torch.cat((-x2, x1), dim=-1) # 後半に負号をつけ、前半と結合

def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1):
    """
    QueryとKeyにMultimodal RoPEを適用する関数
    cos, sin: (3, B, seq_len, head_dim) の形状を想定
    """
    # 割り当てられたセクション（時間, 高さ, 幅）の次元数を2倍にする（cos/sinの次元はすでにペア化され2倍拡張されているため）
    mrope_section = [s * 2 for s in mrope_section]

    # 時間(T), 高さ(H), 幅(W)の各軸ごとの周波数に分割
    cos_parts = cos.split(mrope_section, dim=-1)
    sin_parts = sin.split(mrope_section, dim=-1)

    # 各スプリットから時間、高さ、幅の対応する成分を選択してチャネル方向に再結合
    # これにより、3つの軸の位置情報が混ざり合った統合位置埋め込み（cos/sin）が作成されます
    cos = torch.cat([cos_parts[i][i % 3] for i in range(len(cos_parts))], dim=-1).unsqueeze(unsqueeze_dim)
    sin = torch.cat([sin_parts[i][i % 3] for i in range(len(sin_parts))], dim=-1).unsqueeze(unsqueeze_dim)

    # オイラーの公式に基づく回転公式を適用: (q * cos) + (rotate_half(q) * sin)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

class Qwen2VLRotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000.0):
        super().__init__()
        self.dim = dim
        self.base = base
        # 逆周波数（inv_freq）を算出（周波数の幾何級数的なステップを定義）
        inv_freq = 1.0 / (self.base ** (torch.arange(0, dim, 2).float() / dim))
        # 逆周波数をバッファとして登録（学習パラメータではないが保存・復元できるようにする）
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, position_ids):
        # position_idsの形状: (3, B, seq_len) - 時間、高さ、幅の3次元座標
        # self.inv_freqの形状: (dim // 2)

        # position_ids の末尾に次元を追加し、逆周波数 inv_freq と要素積を取って回転角を算出
        # freqs の形状: (3, B, seq_len, dim // 2)
        freqs = position_ids.unsqueeze(-1).float() * self.inv_freq

        # 1組の角度(freqs, freqs)をチャネル方向に結合して、(3, B, seq_len, dim) の形状にする
        emb = torch.cat((freqs, freqs), dim=-1)
        # それぞれ余弦（cos）と正弦（sin）を計算して返す
        return emb.cos(), emb.sin()

def combine_images_with_aspect_ratio(images, target_height=224, patch_multiple=28):
    """
    複数画像をアスペクト比を保持したまま結合する関数。
    各画像を target_height に縮小し、幅はパッチサイズの倍数に丸めます。
    """
    processed_images = []
    for img in images:
        w_orig, h_orig = img.size
        # アスペクト比を考慮して新しい幅を計算
        w_new = int(target_height * (w_orig / h_orig))
        # 幅が patch_multiple (28) の倍数になるように丸める (最小値は 28)
        w_rounded = max(patch_multiple, round(w_new / patch_multiple) * patch_multiple)
        # 指定されたサイズに高品質リサイズ
        resized_img = img.resize((w_rounded, target_height), Image.Resampling.LANCZOS)
        processed_images.append(resized_img)

    # 結合後の画像の合計幅を計算
    total_width = sum(img.size[0] for img in processed_images)
    # 1つの横長画像を新規作成 (RGB)
    combined_img = Image.new('RGB', (total_width, target_height))
    
    # 各画像を横方向に隙間なく貼り付け
    x_offset = 0
    for img in processed_images:
        combined_img.paste(img, (x_offset, 0))
        x_offset += img.size[0]
    return combined_img

# 利用可能な実行デバイス (GPU/CPU) の選択
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# トークナイザのロードと語彙サイズの取得
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
vocab_size = len(tokenizer)

# 【注意】本記事で使用する「JA-Multi-Image-VQA」データセットは、Sakana AIによって提供されています。
# 画像のライセンスは Unsplash License、テキスト等のデータは Apache License 2.0 に基づいています。
# 詳細は記事の末尾に記載しています。

# データセットのロード (JA-Multi-Image-VQA)
ds = load_dataset("SakanaAI/JA-Multi-Image-VQA", split="test")
sample_indices = [0, 6] # 過学習実験に用いるサンプルインデックス
prepared_data = []

# 画像テンソル変換と標準化設定
transform = T.Compose([
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# サンプルデータの前処理ループ
for i in sample_indices:
    sample = ds[i]
    # 各画像データを RGB 形式に変換
    raw_images = [img.convert("RGB") for img in sample['images']]
    # 複数画像を1枚にアスペクト比維持結合
    combined_image = combine_images_with_aspect_ratio(raw_images)

    question = sample['question']
    answer = sample['answer']

    print(f"\n================ Index {i} ================")
    print(f"Original Images count : {len(raw_images)}")
    print(f"Combined Image Size   : {combined_image.size} (W x H)")
    print(f"Question              : {question}")
    print(f"Ground Truth          : {answer}")

    # 画像を可視化して確認
    plt.figure(figsize=(4 * len(raw_images), 3))
    plt.imshow(combined_image)
    plt.title(f"Dataset Index {i} (Aspect Ratio Preserved)")
    plt.axis("off")
    plt.show()

    # チャットフォーマットに基づくプロンプト文と正解データの結合
    prompt_text = (
        f"<|im_start|>system\nあなたは優秀なAIアシスタントです。<|im_end|>\n"
        f"<|im_start|>user\n質問: {question}<|im_end|>\n"
        f"<|im_start|>assistant\n回答: "
    )
    train_text = prompt_text + f"{answer}<|im_end|>"

    # テキストをトークンIDに変換
    token_ids = tokenizer(train_text, return_tensors="pt").input_ids.to(device)
    prompt_ids = tokenizer(prompt_text, return_tensors="pt").input_ids.to(device)

    # 訓練用データを格納 (input_idsは最後のトークンを除外し、target_idsは最初のトークンを除外：次トークン予測用)
    prepared_data.append({
        "index": i,
        "image_tensor": transform(combined_image).unsqueeze(0).to(device),
        "input_ids": token_ids[:, :-1],
        "target_ids": token_ids[:, 1:],
        "prompt_ids": prompt_ids,
        "answer": answer
    })

# 概念モデルを初期化してデバイスへ転送
model = MiniQwen2VL(vocab_size=vocab_size).to(device)
# オプティマイザの定義
optimizer = AdamW(model.parameters(), lr=1e-3)
# 損失関数の定義 (多クラス分類のためクロスエントロピーを使用)
criterion = nn.CrossEntropyLoss()

print(f"\n--- Training Start (Overfitting on indices {sample_indices}) ---")
model.train() # モデルを訓練モードに設定
epochs = 40
for epoch in range(epochs):
    optimizer.zero_grad() # 勾配の初期化
    total_loss = 0

    for data in prepared_data:
        # モデルの順伝播を実行。戻り値はロジット全体と画像パッチトークン長
        logits, img_len = model(data["image_tensor"], data["input_ids"])
        # 画像トークンを除去し、テキスト部分のロジットのみを切り出す
        pred_logits = logits[:, img_len:, :]
        # クロスエントロピー損失の計算
        loss = criterion(pred_logits.reshape(-1, vocab_size), data["target_ids"].reshape(-1))
        loss.backward() # 勾配計算 (逆伝播)
        total_loss += loss.item()

    optimizer.step() # 重みの更新

    # 10エポックごとに進捗（平均損失）を出力
    if (epoch + 1) % 10 == 0:
        avg_loss = total_loss / len(sample_indices)
        print(f"Epoch [{epoch+1}/{epochs}], Average Loss: {avg_loss:.4f}")

print("\n--- Inference (Auto-regressive Generation) ---")
model.eval() # モデルを評価モードに設定
stop_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>") # 生成終了を判断するための特殊トークンID

for data in prepared_data:
    print(f"\n[Inferencing Index {data['index']}...]")
    # システムプロンプト＋ユーザー質問までのトークンIDを初期配列としてコピー
    generated_ids = data["prompt_ids"][0].tolist()

    with torch.no_grad(): # 推論中につき勾配計算を無効化
        # 最大 30 トークンの自己回帰生成ループ
        for _ in range(30):
            # 現在の生成トークンIDシーケンスをテンソル化して転送
            input_seq = torch.tensor([generated_ids], dtype=torch.long).to(device)
            logits, img_len = model(data["image_tensor"], input_seq)

            # シーケンスの最後尾（最新のトークン予測）における確率分布ロジットを取得
            next_token_logits = logits[0, -1, :]
            # 最も確率（スコア）が高いトークンIDを選択
            next_token_id = torch.argmax(next_token_logits).item()
            # 予測したトークンIDを履歴に追加
            generated_ids.append(next_token_id)

            # 終了トークン（<|im_end|>）が予測された場合は、ループを終了
            if next_token_id == stop_token_id:
                break

    prompt_length = data["prompt_ids"].shape[1] # プロンプト（質問）部分のトークン長を取得
    # 生成されたトークン全体からプロンプト部分を除外し、回答テキストのみをデコード
    generated_text = tokenizer.decode(generated_ids[prompt_length:], skip_special_tokens=True)

    print(f"Ground Truth : {data['answer']}")
    print(f"Model Output : {generated_text}")

Using device: cuda

================ Index 0 ================
Original Images count : 2
Combined Image Size   : (560, 224) (W x H)
Question              : <image> <image> すべての画像を合わせると何匹の犬が見えますか？
Ground Truth          : 3匹の犬が写っています。

================ Index 6 ================
Original Images count : 3
Combined Image Size   : (420, 224) (W x H)
Question              : <image> <image> <image> 何枚目の写真の猫が一番飼い猫の可能性が高いですか？理由も一緒に簡潔に教えてください。
Ground Truth          : 2枚目の写真の猫です。理由は人と一緒に寝ているからです。

--- Training Start (Overfitting on indices [0, 6]) ---
Epoch [10/40], Average Loss: 0.9393
Epoch [20/40], Average Loss: 0.0936
Epoch [30/40], Average Loss: 0.0101
Epoch [40/40], Average Loss: 0.0034

--- Inference (Auto-regressive Generation) ---

[Inferencing Index 0...]
Ground Truth : 3匹の犬が写っています。
Model Output : 3匹の犬が写っています。

[Inferencing Index 6...]
Ground Truth : 2枚目の写真の猫です。理由は人と一緒に寝ているからです。
Model Output : 2枚目の写真の猫です。理由は人と一緒に寝ているからです。

import torch
from datasets import load_dataset
from PIL import Image
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
import matplotlib.pyplot as plt


# 複数画像をアスペクト比を維持したまま、横一列に統合する関数
def combine_images_with_aspect_ratio(
    images, target_height=224, patch_multiple=28
):
    processed_images = []
    for img in images:
        w_orig, h_orig = img.size
        # アスペクト比を維持した新しい幅を算出
        w_new = int(target_height * (w_orig / h_orig))
        # 畳み込み処理のパッチサイズの倍数に適合するように幅を調整
        w_rounded = max(
            patch_multiple, round(w_new / patch_multiple) * patch_multiple
        )
        # 高品質リサンプリングを用いて画像をリサイズ
        resized_img = img.resize(
            (w_rounded, target_height), Image.Resampling.LANCZOS
        )
        processed_images.append(resized_img)

    # 結合用キャンバスの幅（全画像の合計幅）を計算
    total_width = sum(img.size[0] for img in processed_images)
    # 横長の白紙キャンバス画像を新規作成
    combined_img = Image.new("RGB", (total_width, target_height))
    x_offset = 0
    # 各画像を横にスライドさせながら貼り付け
    for img in processed_images:
        combined_img.paste(img, (x_offset, 0))
        x_offset += img.size[0]
    return combined_img


# 1. 2Bパラメータの事前学習済みQwen2-VL指示チューニングモデルとプロセッサをロード
model_id = "Qwen/Qwen2-VL-2B-Instruct"
model = Qwen2VLForConditionalGeneration.from_pretrained(
    model_id, 
    torch_dtype=torch.bfloat16,  # 計算メモリ節約と実行安定化のためにbfloat16を選択
    device_map="auto"           # 利用可能な最速デバイスに各レイヤーを自動で割り当て
)
processor = AutoProcessor.from_pretrained(model_id)

# 2. 推論を行うサンプルデータのロードと前処理
dataset = load_dataset("SakanaAI/JA-Multi-Image-VQA", split="test")
i = 1 # 1番目のサンプルを選択
sample = dataset[i]

images = sample["images"]
question = sample["question"]

print(f"質問: {question}")
print(f"元の画像枚数: {len(images)}枚")

# 複数画像を前処理して1枚のパノラマ画像にアラインメント
combined_image = combine_images_with_aspect_ratio(images)
print(f"統合後の画像サイズ: {combined_image.size}")

# 画像を画面に表示
plt.figure(figsize=(4 * len(images), 3))
plt.imshow(combined_image)
plt.title(f"Dataset Index {i} (Aspect Ratio Preserved)")
plt.axis("off")
plt.show()


# 3. チャット形式の構造テンプレートに画像枠と質問のプロンプトを挿入
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image"},  # 画像を1枚に統合したため、イメージスロットは1つで処理します
            {"type": "text", "text": question},
        ],
    }
]

# モデル特有 of 会話タグ (<|im_start|>, <|im_end|>) を適用してプレーンテキストに展開
text = processor.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
)

# プロセッサを利用してテキストトークン化と、ビジョン（統合画像）の処理を統合テンソルに変換
inputs = processor(
    text=[text], images=combined_image, padding=True, return_tensors="pt"
).to(model.device)

# 4. モデルによる自己回帰生成の実行
with torch.no_grad(): # 推論のため評価モード＆勾配計算オフ
    generated_ids = model.generate(**inputs, max_new_tokens=512)

# 出力されたトークンID履歴から、入力プロンプト部分を差し引いて回答文の部分のみにトリミング
generated_ids_trimmed = [
    out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
]
# トークンIDから文字列テキストへとデコード
output_text = processor.batch_decode(
    generated_ids_trimmed,
    skip_special_tokens=True, # 特殊制御用トークンをスキップ
    clean_up_tokenization_spaces=False,
)

print("\nモデルの出力:")
print(output_text[0])

print(f"\n（参考）データセットに用意されている正解: {sample['answer']}")

質問: <image> <image> どちらの画像にも写っている動物は何ですか？一言で答えてください。
元の画像枚数: 2枚
統合後の画像サイズ: (560, 224)
モデルの出力:
犬

（参考）データセットに用意されている正解: 犬