시각 문서 검색

시각 문서 검색 (Visual document retrieval)

문서에 텍스트 외에도 차트, 표, 시각 자료가 포함되면 다중 모달(multimodal) 데이터를 담고 있다고 볼 수 있어요. 텍스트 검색 모델만으로는 시각 데이터를 다룰 수 없고, 이미지 검색 모델은 세밀함과 문서 처리 능력이 부족하기 때문에 이런 문서들에서 정보를 검색하는 것은 까다로워요.

출처: 문서

본문

시각 문서 검색(visual document retrieval)은 다중 모달 검색 증강 생성(RAG, retrieval augmented generation)을 포함해 모든 유형의 문서에서 정보를 검색하는 데 도움을 줄 수 있어요. 이러한 모델은 문서(이미지로)와 텍스트를 받아들여 그 사이의 유사도(similarity) 점수를 계산해요.

이 가이드는 ColPali로 문서를 인덱싱하고 검색하는 방법을 보여줘요.

[!TIP] 대규모 사용 사례라면 벡터 데이터베이스로 문서를 인덱싱하고 검색하고 싶을 수도 있어요.

Transformers와 Datasets이 설치되어 있는지 확인해요.

pip install -q datasets transformers

UFO 목격과 관련된 문서 데이터셋을 인덱싱할 거예요. 관심 있는 열이 없는 예시는 필터링해요. 이 데이터셋에는 여러 열이 있는데, 우리가 관심 있는 것은 문서의 짧은 요약이 담긴 specific_detail_query 열과 문서가 담긴 image 열이에요.

from datasets import load_dataset

dataset = load_dataset("davanstrien/ufo-ColPali")
dataset = dataset["train"]
dataset = dataset.filter(lambda example: example["specific_detail_query"] is not None)
print(dataset)

Dataset({
    features: ['image', 'raw_queries', 'broad_topical_query', 'broad_topical_explanation', 'specific_detail_query', 'specific_detail_explanation', 'visual_element_query', 'visual_element_explanation', 'parsed_into_json'],
    num_rows: 2172
})

모델과 토크나이저를 로드해요.

import torch
from transformers import ColPaliForRetrieval, ColPaliProcessor
from accelerate import Accelerator

device = Accelerator().device

model_name = "vidore/colpali-v1.2-hf"

processor = ColPaliProcessor.from_pretrained(model_name)

model = ColPaliForRetrieval.from_pretrained(
    model_name,
    dtype=torch.bfloat16,
    device_map="auto",
).eval()

텍스트 쿼리를 프로세서에 전달하고 모델에서 인덱싱된 텍스트 임베딩을 반환해요. 이미지-텍스트 검색의 경우 ColPaliProcessor의 text 파라미터를 images 파라미터로 바꿔서 이미지를 전달하면 돼요.

inputs = processor(text="a document about Mars expedition").to(model.device)
with torch.no_grad():
  text_embeds = model(**inputs, return_tensors="pt").embeddings

이미지를 오프라인으로 인덱싱하고, 추론 중에는 쿼리 텍스트 임베딩을 반환해서 가장 가까운 이미지 임베딩을 알아내요.

아래처럼 map으로 이미지와 이미지 임베딩을 데이터셋에 작성해서 저장해요. 인덱싱된 임베딩이 담긴 embeddings 열을 추가해요. ColPali 임베딩은 저장 공간을 많이 차지하므로, 가속기에서 제거하고 CPU에 NumPy 벡터로 저장해요.

ds_with_embeddings = dataset.map(lambda example: {'embeddings': model(**processor(images=example["image"]).to(device), return_tensors="pt").embeddings.to(torch.float32).detach().cpu().numpy()})

온라인 추론을 위해, 배치 단위로 이미지 임베딩을 검색해서 가장 관련성 높은 k개의 이미지를 가져오는 함수를 만들어요. 아래 함수는 주어진 인덱싱된 데이터셋, 텍스트 임베딩, 상위 결과 개수, 배치 크기에 대해 데이터셋의 인덱스와 그 점수를 반환해요.

def find_top_k_indices_batched(dataset, text_embedding, processor, k=10, batch_size=4):
    scores_and_indices = []

    for start_idx in range(0, len(dataset), batch_size):

        end_idx = min(start_idx + batch_size, len(dataset))
        batch = dataset[start_idx:end_idx]        
        batch_embeddings = [torch.tensor(emb[0], dtype=torch.float32) for emb in batch["embeddings"]]
        scores = processor.score_retrieval(text_embedding.to("cpu").to(torch.float32), batch_embeddings)

        if hasattr(scores, "tolist"):
            scores = scores.tolist()[0]

        for i, score in enumerate(scores):
            scores_and_indices.append((score, start_idx + i))

    sorted_results = sorted(scores_and_indices, key=lambda x: -x[0])

    topk = sorted_results[:k]
    indices = [idx for _, idx in topk]
    scores = [score for score, _ in topk]

    return indices, scores

텍스트 임베딩을 생성하고 위 함수에 전달하면 데이터셋 인덱스와 점수를 반환해요.

with torch.no_grad():
  text_embeds = model(**processor(text="a document about Mars expedition").to(model.device), return_tensors="pt").embeddings
indices, scores = find_top_k_indices_batched(ds_with_embeddings, text_embeds, processor, k=3, batch_size=4)
print(indices, scores)

([440, 442, 443],
 [14.370786666870117,
  13.675487518310547,
  12.9899320602417])

이미지를 표시해서 화성 관련 문서를 확인해요.

for i in indices:
  display(dataset[i]["image"])

<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/doc_1.png"
     alt="Document 1"
     style="height: 200px; object-fit: contain; margin-right: 10px;">
<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/doc_2.png"
     alt="Document 2"
     style="height: 200px; object-fit: contain;">
<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/doc_3.png"
     alt="Document 3"
     style="height: 200px; object-fit: contain;">