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

## はじめに

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

https://learnopencv.com/multimodal-rag-with-colpali/

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

https://arxiv.org/pdf/2407.01449

## 環境

- Python 3.10
- ColPali
- Gemini Flash
- PyTorch
- Transformers

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

### 1. RetrievalEvaluatorの初期化エラー

#### エラーメッセージ
```python
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`を作成し、直接実装しました。

```python
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を生成するように修正しました。

```python
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の次元を下記のように処理し、形状が変化しないようにしました。

```python
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つのエラーを解決した実装例です。

```python
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. テンソルのスカラー変換エラー

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