LoRAアダプターのホットスワッピングを使ったColQwen2単体でのRAGを試す

!pip install -q -U "colpali-engine>=0.3.1,<0.4.0" requests

from io import BytesIO
from typing import Any, List, cast

import requests
import torch
from colpali_engine.models import ColQwen2, ColQwen2Processor
from colpali_engine.utils.torch_utils import get_torch_device
from IPython.display import display
from peft import LoraConfig
from PIL import Image
from transformers.models.qwen2_vl import Qwen2VLForConditionalGeneration, Qwen2VLProcessor

def load_image_from_url(url: str) -> Image.Image:
    """
    有効なURLからPILイメージをロードする
    """
    response = requests.get(url)
    return Image.open(BytesIO(response.content))


def scale_image(image: Image.Image, new_height: int = 1024) -> Image.Image:
    """
    アスペクト比を維持しながら画像を新しい高さにスケーリングする
    """
    width, height = image.size
    aspect_ratio = width / height
    new_width = int(new_height * aspect_ratio)

    scaled_image = image.resize((new_width, new_height))

    return scaled_image

class ColQwen2ForRAG(ColQwen2):
    """
    ColQwen2モデルは、検索と生成の両方で使用できる実装です。
    検索モードと生成モードの切り替えが可能です。
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._is_retrieval_enabled = True

    def forward(self, *args, **kwargs) -> Any:
        """
        現在のモードに基づいて、Qwen2VLForConditionalGeneration.forward（生成用）
        またはColQwen2.forward（検索用）を呼び出すフォワードパス。
        """
        if self.is_retrieval_enabled:
            return ColQwen2.forward(self, *args, **kwargs)
        else:
            return Qwen2VLForConditionalGeneration.forward(self, *args, **kwargs)

    def generate(self, *args, **kwargs):
        """
        Qwen2VLForConditionalGeneration.generateを使用してテキストを生成
        """
        if not self.is_generation_enabled:
            raise ValueError(
                "generate()を呼び出す前に、enable_generation()を呼び出してモデルを生成モードに設定してください。"
            )
        return super().generate(*args, **kwargs)

    @property
    def is_retrieval_enabled(self) -> bool:
        return self._is_retrieval_enabled

    @property
    def is_generation_enabled(self) -> bool:
        return not self.is_retrieval_enabled

    def enable_retrieval(self) -> None:
        """
        検索モードに切り替え
        """
        self.enable_adapters()
        self._is_retrieval_enabled = True

    def enable_generation(self) -> None:
        """
        生成モードに切り替え
        """
        self.disable_adapters()
        self._is_retrieval_enabled = False

model_name = "vidore/colqwen2-v1.0"
device = get_torch_device("auto")

print(f"Using device: {device}")

# 事前学習済みの検索モデルからLoRA設定を取得
lora_config = LoraConfig.from_pretrained(model_name)

# プロセッサをロード
processor_retrieval = cast(ColQwen2Processor, ColQwen2Processor.from_pretrained(model_name))
processor_generation = cast(Qwen2VLProcessor, Qwen2VLProcessor.from_pretrained(lora_config.base_model_name_or_path))

# 事前学習済みのアダプターを検索用にロードしてモデルを読み込み
model = cast(
    ColQwen2ForRAG,
    ColQwen2ForRAG.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16,
        device_map=device,
    ),
)

query = "カザフスタンの石油生産量のうち、海底油田の占める割合は？"

images: List[Image.Image] = [
    load_image_from_url(
        "https://raw.githubusercontent.com/kun432/colpali-cookbooks-jp-files/refs/heads/main/data/energy_electricity_generation_ja.jpg"
    ),
    load_image_from_url(
        "https://raw.githubusercontent.com/kun432/colpali-cookbooks-jp-files/refs/heads/main/data/shift_kazakhstan_ja.jpg"
    ),
]

images = [scale_image(image, new_height=512) for image in images]

for image in images:
    display(scale_image(image, new_height=256))

# 入力を処理
batch_images = processor_retrieval.process_images(images).to(model.device)
batch_queries = processor_retrieval.process_queries([query]).to(model.device)

# フォワードパス
model.enable_retrieval()

with torch.no_grad():
    image_embeddings = model.forward(**batch_images)
    query_embeddings = model.forward(**batch_queries)

# 類似度スコアを計算
scores = processor_retrieval.score_multi_vector(query_embeddings, image_embeddings)

# 1位のページ画像を取得
retrieved_image_index = scores.argmax().item()
retrieved_image = images[retrieved_image_index]

print(f"次のクエリに対して取得された画像: `{query}`")
display(scale_image(retrieved_image, new_height=256))

# 入力を処理
conversation = [
    {
        "role": "user",
        "content": [
            {
                "type": "image",
            },
            {
                "type": "text",
                "text": f"入力画像を使って次の質問に答えてください: {query}",
            },
        ],
    }
]
text_prompt = processor_generation.apply_chat_template(conversation, add_generation_prompt=True)
inputs_generation = processor_generation(
    text=[text_prompt],
    images=[image],
    padding=True,
    return_tensors="pt",
).to(device)

# RAGのレスポンスを生成
model.enable_generation()
output_ids = model.generate(**inputs_generation, max_new_tokens=128)

# 生成された`output_ids`から新たに生成されたトークンIDのみが保持されていることを確認
generated_ids = [output_ids[len(input_ids) :] for input_ids, output_ids in zip(inputs_generation.input_ids, output_ids)]

# RAGのレスポンスをデコード
output_text = processor_generation.batch_decode(
    generated_ids,
    skip_special_tokens=True,
    clean_up_tokenization_spaces=True,
)

print(output_text)

['カザフスタンの石油生産量のうち、海底油田の占める割合は約84%を占めています。']

query = "2019年に最も多くの電力が生成された時間帯はどの時間ですか？"

['2019年に最も多くの電力が生成された時間帯は、午前8時から午前10時です。']