ColPaliとGemini Flashを使用してRAG実装してたら遭遇したエラーたち

はじめに

ColPaliというVision Language Modelをいじっていて、エラーにかなり悩まされたので記事にしていきます！下記サイトをベースに、Gemini Flashと組み合わせてRAGを構築していたのですが、実装時にいくつか厄介なエラーに遭遇してしまったので、内容とその解決方法をご紹介します。

ColPaliの論文も下記に掲載します。

環境

• Python 3.10
• ColPali
• Gemini Flash
• PyTorch
• Transformers

発生したエラーと解決方法

1. RetrievalEvaluatorの初期化エラー

エラーメッセージ

TypeError: RetrievalEvaluator.__init__() missing 1 required positional argument: 'retriever'

原因

CustomEvaluatorクラスの初期化時に、MTEBのRetrievalEvaluatorに必要なretriever引数が不足していました。

colpali_engine/trainer/retrieval_evaluator.pyの中でCustomEvaluatorというクラスがあり、その__init__でRetrievalEvaluatorが使われているのですが、retriever引数を与えていないからエラーになってしまいました。

解決方法

上記の原因から、最初はRetrievalEvaluatorを直接編集しようと思ったのですが、ライブラリの更新時に修正が上書きされてしまったり、他の内部実装に影響を与えてしまうことが考えられたので、新しくCustomeEvaluetorを作成し、直接実装しました。

class CustomEvaluator:
    def __init__(self, is_multi_vector=False):
        self.is_multi_vector = is_multi_vector
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
    def evaluate(self, query_embeddings, document_embeddings):
        # クエリと文書のembeddingsを使用して類似度スコアを計算
        query_embeddings = torch.stack(query_embeddings)
        document_embeddings = torch.stack(document_embeddings)
        
        # コサイン類似度の計算
        similarity_scores = torch.nn.functional.cosine_similarity(
            query_embeddings.unsqueeze(1),
            document_embeddings.unsqueeze(0),
            dim=2
        )
        
        return similarity_scores

2. Embeddingサイズの不一致エラー

エラーメッセージ

RuntimeError: The size of tensor a (1031) must match the size of tensor b (1030) at non-singleton dimension 2

原因

クエリと文書のembeddingサイズが異なっていました（1031次元 vs 1030次元）。

解決方法

processorの設定を調整し、一貫したサイズのembeddingを生成するように修正しました。

def process_queries(processor, queries, placeholder_image):
    queries = [f"<image> {query}" for query in queries]
    
    inputs = processor(
        text=queries,
        images=[placeholder_image] * len(queries),
        return_tensors="pt",
        padding='max_length',
        max_length=1030,
        truncation=True
    )
    
    return inputs

3. テンソルのスカラー変換エラー

エラーメッセージ

RuntimeError: a Tensor with 128 elements cannot be converted to Scalar

原因

embeddingの形状が異なっており（2054x128 vs 1030x128）、テンソルからスカラー値への変換に問題が発生していました。

解決方法

embeddingの次元を下記のように処理し、形状が変化しないようにしました。

def evaluate(self, query_embeddings, document_embeddings):
    # embeddingの形状を2D形状に変換
    query_embeddings = torch.stack([emb.mean(dim=0) for emb in query_embeddings])
    document_embeddings = torch.stack([emb.mean(dim=0) for emb in document_embeddings])
    
    # コサイン類似度の計算
    similarity_scores = torch.nn.functional.cosine_similarity(
        query_embeddings.unsqueeze(1),
        document_embeddings.unsqueeze(0),
        dim=2
    )
    
    return similarity_scores

最終的にはこんな実装になりました

以下が、上記で紹介した3つのエラーを解決した実装例です。

class CustomEvaluator:
    def __init__(self, is_multi_vector=False):
        self.is_multi_vector = is_multi_vector
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
    def evaluate(self, query_embeddings, document_embeddings):
        # 各embeddingを2D形状に変換
        query_embeddings = torch.stack([emb.mean(dim=0) for emb in query_embeddings])
        document_embeddings = torch.stack([emb.mean(dim=0) for emb in document_embeddings])
        
        # コサイン類似度の計算
        similarity_scores = torch.nn.functional.cosine_similarity(
            query_embeddings.unsqueeze(1),
            document_embeddings.unsqueeze(0),
            dim=2
        )
        
        return similarity_scores

def retrieve_top_document(
    query: str,
    document_embeddings: List[torch.Tensor],
    document_images: List[Image.Image],
) -> Tuple[str, Image.Image]:
    query_embeddings = []
    placeholder_image = Image.new("RGB", (448, 448), (255, 255, 255))
 
    with torch.no_grad():
        query_batch = process_queries(paligemma_processor, [query], placeholder_image)
        query_batch = {key: value.to(device) for key, value in query_batch.items()}
        query_embeddings_tensor = retrieval_model(**query_batch)
        query_embeddings = list(torch.unbind(query_embeddings_tensor.to("cpu")))
 
    evaluator = CustomEvaluator(is_multi_vector=True)
    similarity_scores = evaluator.evaluate(query_embeddings, document_embeddings)
    best_index = similarity_scores[0].argmax().item()
 
    return document_images[best_index], best_index

def process_queries(processor, queries, placeholder_image):
    # 各クエリに<image>トークンを追加
    queries = [f"<image> {query}" for query in queries]
    
    # プロセッサーに渡す
    inputs = processor(
        text=queries,
        images=[placeholder_image] * len(queries),
        return_tensors="pt",
        padding=True
    )
    
    return inputs

まとめ

ColPaliとGemini Flashを使用したRAG実装時に遭遇したエラーは以下の3つです。

1. RetrievalEvaluatorの初期化エラー
2. Embeddingサイズの不一致エラー
3. テンソルのスカラー変換エラー

本記事がエラー解決の手助けになれば幸いです！