EmbeddingLists를 활용한 검색: ColBERT 및 ColPali

이 튜토리얼에서는 Milvus의 StructArray 벡터 하위 필드에 대해 EmbeddingList 검색을 사용하여 ColBERT 방식 및 ColPali 방식의 검색 시스템을 구축하는 방법을 설명합니다. 쿼리와 저장된 데이터가 모두 벡터 목록으로 표현되어 있으며, MAX_SIM* 메트릭을 활용한 엔티티 수준의 후기 상호작용 검색을 원하는 경우에 이 방법을 사용하십시오.

이 튜토리얼의 기반이 되는 StructArray 기본 개념에 대해서는 ‘StructArray 필드 생성’, ‘StructArray 필드 인덱싱’ 및 ‘StructArray를 사용한 기본 벡터 검색’을 참조하십시오. 이 튜토리얼은 일반적인 StructArray 구문보다는 ColBERT 및 ColPali 워크플로우에 중점을 둡니다.

개요

텍스트 검색 시스템을 구축하려면, 특히 전체 텍스트 임베딩이 의미적 특이성을 희석시키거나 모델 입력 제한을 초과할 수 있는 긴 문서의 경우, 정밀도와 정확도를 보장하기 위해 문서를 청크로 분할하고 각 청크를 임베딩과 함께 벡터 데이터베이스에 엔티티로 저장해야 할 수 있습니다.

그러나 데이터를 청크 단위로 저장하면 청크별 검색 결과가 나오게 되는데, 이는 검색 시 처음에 일관된 문서 전체가 아닌 관련 세그먼트만 식별된다는 것을 의미합니다. 이를 해결하려면 검색 후 추가 처리를 수행해야 합니다.

ColBERT(arXiv: 2004.12832)는 BERT를 기반으로 한 문맥적 후기 상호작용을 통해 효율적이고 효과적인 구절 검색을 제공하는 텍스트-텍스트 검색 시스템입니다. 이 시스템은 쿼리와 문서를 토큰 단위로 독립적으로 인코딩하고, 이들의 유사도를 계산합니다.

토큰 단위 인코딩

ColBERT에서 데이터를 수집하는 동안 각 문서는 토큰으로 분할되며, 이 토큰들은 벡터화되어 $∈Rn×dd \rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ E d = [ $e$ d1 $,{}^{}e$ d2 $,\dots ,e$ dn $]$ $\in$ R ${}^{n\times d}$. 쿼리가 도착하면, $∈ R^{m×d}$와 같이 토큰화되고 벡터화된 후 임베딩 목록으로 저장됩니다 $. \rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ $\to$ E q $=$ [e q1 $,e$ q2 $,\dots ,e$ qm $]\in$ R ${}^{m\times d}$.

위의 식에서,

• $dd$: 문서

• $qq$: 쿼리

• $EdE_d$: 문서를 나타내는 임베딩 목록.

• $EqE_q$: 쿼리를 나타내는 임베딩 목록.

• $} [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$$,e_{}$,…,및dn$]\in$ R${}^{n\times d}$: 문서를 나타내는 임베딩 목록 내 벡터 임베딩의 개수는 $\R^{n×d}$ 범위 내에 있다.

• $[e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$q1,eq2,…,eqm$]\in$ R${}^{m\times d}$: 쿼리를 나타내는 임베딩 목록 내 벡터 임베딩의 개수는 $\R^{m×d}$ 범위 내에 있습니다.

후기 상호작용

벡터화가 완료되면, 쿼리 임베딩 목록을 각 문서 임베딩 목록과 토큰 단위로 비교하여 최종 유사도 점수를 결정합니다.

위 다이어그램에서 볼 수 있듯이, 쿼리에는 machine 와 learning 라는 두 개의 토큰이 포함되어 있으며, 창 내의 문서에는 neural, network, python, tutorial 라는 네 개의 토큰이 있습니다. 이러한 토큰들이 벡터화되면, 각 쿼리 토큰의 벡터 임베딩이 문서 내의 벡터 임베딩과 비교되어 유사도 점수 목록이 생성됩니다. 그런 다음 각 점수 목록에서 가장 높은 점수들을 합산하여 최종 점수를 산출합니다. 문서의 최종 점수를 결정하는 이 과정을 최대 유사도(MAX_SIM)라고 합니다. 최대 유사도에 대한 자세한 내용은 ‘최대 유사도’를 참조하십시오.

ColPali 확장

ColBERT를 기반으로 하는 ColPali(arXiv: 2407.01449)는 비전-언어 모델(VLM)을 활용하는 시각적으로 풍부한 문서 검색을 위한 새로운 접근 방식을 제안합니다. 데이터 수집 과정에서 각 문서 페이지는 토큰화되지 않고 고해상도 이미지로 렌더링된 후 패치로 분할됩니다. 예를 들어, 448 x 448 픽셀 크기의 문서 페이지 이미지는 각각 14 x 14 픽셀 크기의 1,024개 패치를 생성할 수 있습니다.

이 방법은 텍스트 전용 검색 시스템을 사용할 때 손실되는 문서 레이아웃, 이미지, 표 구조와 같은 비텍스트 정보를 보존합니다.

ColPali에서 사용되는 VLM은 PaliGemma(arXiv: 2407.07726)라고 불리며, 위 다이어그램에서 볼 수 있듯이 이미지 인코더(SigLIP-400M), 디코더 전용 언어 모델(Gemma2-2B), 그리고 위 다이어그램에 표시된 대로 이미지 인코더의 출력을 언어 모델의 벡터 공간으로 투영하는 선형 레이어로 구성됩니다.

데이터 수집 과정에서 원본 이미지로 표현된 문서 페이지는 여러 시각적 패치로 분할되며, 각 패치는 임베딩을 거쳐 벡터 임베딩 목록을 생성합니다. 그런 다음, $∈Rn×dd \rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ E d = [ $e$ d1 $,e$ d2 $,\dots ,e$ dn $]$ $\in$ R ${}^{n\times d}$. 쿼리가 도착하면 토큰화되며, 각 토큰은 임베딩되어 벡터 임베딩 목록을 생성합니다. 예를 들어, $∈Rm×dq \rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ $\to$ E q $=$ [e q1 $,e$ q2 $,\dots ,e$ qm $]\in$ R ${}^{m\times d}$. 그런 다음, MAX_SIM을 적용하여 두 임베딩 목록을 비교하고 쿼리와 문서 페이지 간의 최종 점수를 산출했습니다.

ColBERT 텍스트 검색 시스템

이 섹션에서는 StructArray를 사용하여 ColBERT 텍스트 검색 시스템을 구축할 것입니다. 그 전에 Milvus v2.6.x 인스턴스를 설정하고 Cohere 액세스 토큰을 획득해야 합니다.

1단계: 종속성 설치

다음 명령어를 실행하여 종속성을 설치합니다.

pip install --upgrade huggingface-hub transformers datasets pymilvus cohere

2단계: Cohere 데이터셋 불러오기

이 예제에서는 Cohere의 위키백과 데이터셋을 사용하여 처음 10,000개의 레코드를 검색할 것입니다. 이 데이터셋에 대한 정보는 이 페이지에서 확인할 수 있습니다.

from datasets import load_dataset

lang = "simple"
docs = load_dataset(
    "Cohere/wikipedia-2023-11-embed-multilingual-v3",
    lang,
    split="train[:10000]"
)

위 스크립트를 실행하면 데이터셋이 로컬에 없는 경우 자동으로 다운로드됩니다. 데이터셋의 각 레코드는 위키백과 페이지의 한 단락입니다. 다음 표는 이 데이터셋의 구조를 보여줍니다.

열 이름	설명
_id	레코드 ID
url	현재 레코드의 URL입니다.
title	원본 문서의 제목입니다.
text	원본 문서의 한 단락.
emb	원본 문서의 텍스트에 대한 임베딩.

3단계: 제목별로 단락 그룹화하기

단락이 아닌 문서를 검색하려면 단락을 제목별로 그룹화해야 합니다.

df = docs.to_pandas()
groups = df.groupby('title')

data = []

for title, group in groups:
  data.append({
      "title": title,
      "paragraphs": [{
          "text": row['text'],
          'emb': row['emb']
      } for _, row in group.iterrows()]
  })

이 코드에서는 그룹화된 단락을 문서로 저장하고, 이를 ` data ` 목록에 포함시킵니다. 각 문서에는 단락 목록인 ` paragraphs ` 키가 있으며, 각 단락 객체에는 ` text ` 및 ` emb ` 키가 포함되어 있습니다.

4단계: Cohere 데이터셋용 컬렉션 생성

데이터 준비가 완료되면 컬렉션을 생성합니다. 이 컬렉션에서 ` paragraphs `는 StructArray 필드입니다. StructArray 스키마에 대한 일반적인 설명은 ‘StructArray 필드 생성’을 참조하십시오.

from pymilvus import MilvusClient, DataType

client = MilvusClient(
    uri="http://localhost:19530",
    token="root:Milvus"
)

# Create collection schema
schema = client.create_schema()

schema.add_field('id', DataType.INT64, is_primary=True, auto_id=True)
schema.add_field('title', DataType.VARCHAR, max_length=512)

# Create struct schema
struct_schema = client.create_struct_field_schema()
struct_schema.add_field('text', DataType.VARCHAR, max_length=65535)
struct_schema.add_field('emb', DataType.FLOAT_VECTOR, dim=512)

schema.add_field('paragraphs', DataType.ARRAY,
                 element_type=DataType.STRUCT,
                 struct_schema=struct_schema, max_capacity=200)

# Create index parameters
index_params = client.prepare_index_params()
index_params.add_index(
    field_name="paragraphs[emb]",
    index_type="AUTOINDEX",
    metric_type="MAX_SIM_COSINE"
)

# Create a collection
client.create_collection(
    collection_name='wiki_documents',
    schema=schema,
    index_params=index_params
)

5단계: Cohere 데이터셋을 컬렉션에 삽입

이제 준비된 데이터를 위에서 생성한 컬렉션에 삽입할 수 있습니다.

client.insert(
    collection_name='wiki_documents',
    data=data
)

6단계: Cohere 데이터셋 내에서 검색하기

ColBERT의 설계에 따르면, 쿼리 텍스트는 토큰화 된 후 EmbeddingList에 임베딩되어야 합니다. 이 단계에서는 Cohere가 위키백과 데이터셋의 단락에 대한 임베딩을 생성하는 데 사용한 것과 동일한 모델을 사용할 것입니다.

import cohere

co = cohere.ClientV2("COHERE_API_KEY")

query_inputs = [
    {
        'content': [
            {'type': 'text', 'text': 'Adobe'},
        ]
    },
    {
        'content': [
            {'type': 'text', 'text': 'software'}
        ]
    }
]

embeddings = co.embed(
    inputs=query_inputs,
    model='embed-multilingual-v3.0',
    input_type="classification",
    embedding_types=["float"],
)

코드에서 쿼리 텍스트는 ` query_inputs `에서 토큰으로 분할된 후 부동 소수점 벡터 목록에 임베딩됩니다. 그런 다음 Milvus의 `EmbeddingList`를 사용하여 다음과 같이 유사도 검색을 수행할 수 있습니다.

from pymilvus.client.embedding_list import EmbeddingList

query_emb_list = EmbeddingList()

if (embeddings.embeddings.float):
  query_emb_list.add_batch(embeddings.embeddings.float)

results = client.search(
    collection_name="wiki_documents",
    data=[query_emb_list],
    anns_field="paragraphs[emb]",
    limit=10,
    output_fields=["title"]
)

for hit in results[0]:
  print(f"Document {hit['entity']['title']}: {hit['distance']:.4f}")

위 코드의 출력 결과는 다음과 유사합니다:

# Document Software: 2.3035
# Document Application: 2.1875
# Document Adobe Illustrator: 2.1167
# Document Open source: 2.0542
# Document Computer: 1.9811
# Document Microsoft: 1.9784
# Document Web browser: 1.9655
# Document Program: 1.9627
# Document Website: 1.9594
# Document Computer science: 1.9460

각 쌍별 코사인 유사도 점수는 -1 에서 1 사이의 값을 가집니다. 최종 MAX_SIM_COSINE 점수는 여러 토큰 수준의 최대 유사도 점수를 집계하기 때문에 1 보다 클 수 있습니다.

ColPali 문서 검색 시스템

이 섹션에서는 StructArray를 사용하여 ColPali 기반의 문서 검색 시스템을 구축할 것입니다. 그 전에 Milvus v2.6.x 인스턴스를 설정하십시오.

1단계: 종속성 설치

pip install --upgrade huggingface-hub transformers datasets pymilvus 'colpali-engine>=0.3.0,<0.4.0'

2단계: Vidore 데이터셋 불러오기

이 섹션에서는 vidore_v2_finance_en이라는 Vidore 데이터셋을 사용할 것입니다. 이 데이터셋은 은행 부문의 연차 보고서 코퍼스로, 장문 문서 이해 작업을 위해 설계되었습니다. 이는 ViDoRe v3 벤치마크를 구성하는 10개의 코퍼스 중 하나입니다. 이 데이터셋에 대한 자세한 내용은 이 페이지에서 확인할 수 있습니다.

from datasets import load_dataset

ds = load_dataset("vidore/vidore_v3_finance_en", "corpus")
df = ds['test'].to_pandas()

위 스크립트를 실행하면 로컬에 데이터셋이 없는 경우 데이터셋이 다운로드됩니다. 데이터셋의 각 레코드는 재무 보고서의 한 페이지입니다. 다음 표는 이 데이터셋의 구조를 보여줍니다.

열 이름	설명
corpus_id	코퍼스 내의 레코드
image	바이트 단위의 페이지 이미지.
doc_id	설명 문서 ID.
page_number_in_doc	문서 내 현재 페이지의 페이지 번호.

3단계: 페이지 이미지에 대한 임베딩 생성

‘개요’ 섹션에서 설명한 바와 같이, ColPali 모델은 이미지를 텍스트 모델의 벡터 공간으로 투영하는 VLM입니다. 이 단계에서는 최신 ColPali 모델인 vidore/colpali-v1.3을 사용할 것입니다. 이 모델에 대한 자세한 내용은 이 페이지에서 확인할 수 있습니다.

import torch
from typing import cast
from colpali_engine.models import ColPali, ColPaliProcessor

model_name = "vidore/colpali-v1.3"

model = ColPali.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="cuda:0",  # or "mps" if on Apple Silicon
).eval()

processor = ColPaliProcessor.from_pretrained(model_name)

모델 준비가 완료되면 다음과 같이 특정 이미지에 대한 패치를 생성해 볼 수 있습니다.

from PIL import Image
from io import BytesIO

# Use the iterrows() generator to get the first row.
row = next(df.iterrows())[1]

# Decode the image bytes and generate patch embeddings.
images = [Image.open(BytesIO(row["image"]["bytes"]))]
batch_images = processor.process_images(images).to(model.device)

with torch.no_grad():
    patches_embeddings = model(**batch_images)[0]

# Check the shape of the embeddings generated for the patches.
print(patches_embeddings.shape)

# [1031, 128]

위의 코드에서 ColPali 모델은 이미지의 크기를 448 x 448 픽셀로 조정한 다음, 각각 14 x 14 픽셀 크기의 패치로 분할합니다. 마지막으로, 이러한 패치들은 각각 128차원의 1,031개 임베딩으로 변환됩니다.

다음과 같은 루프를 사용하여 모든 이미지에 대한 임베딩을 생성할 수 있습니다:

data = []

for _, row in df.iterrows():
    corpus_id = row["corpus_id"]
    images = [Image.open(BytesIO(row["image"]["bytes"]))]
    batch_images = processor.process_images(images).to(model.device)

    with torch.no_grad():
        patches = model(**batch_images)[0]

    doc_id = row["doc_id"]
    page_number_in_doc = row["page_number_in_doc"]

    data.append({
        "corpus_id": corpus_id,
        "patches": [
            {"emb": emb.float().cpu().tolist()}
            for emb in patches
        ],
        "doc_id": doc_id,
        "page_number_in_doc": page_number_in_doc,
    })

4단계: 재무 보고서 데이터셋용 컬렉션 생성

데이터 준비가 완료되면 컬렉션을 생성합니다. 이 컬렉션에서 ` patches `는 `StructArray` 필드입니다. 각 `Struct` 요소는 하나의 패치 임베딩을 저장합니다. `StructArray` 벡터 하위 필드의 인덱싱 요구 사항에 대해서는 ‘StructArray 필드 인덱싱’을 참조하십시오.

from pymilvus import MilvusClient, DataType

client = MilvusClient(
    uri=YOUR_CLUSTER_ENDPOINT,
    token=YOUR_API_KEY
)

schema = client.create_schema()

schema.add_field(
    field_name="corpus_id",
    datatype=DataType.INT64,
    is_primary=True
)

patch_schema = client.create_struct_field_schema()

patch_schema.add_field(
    field_name="emb",
    datatype=DataType.FLOAT_VECTOR,
    dim=128
)

schema.add_field(
    field_name="patches",
    datatype=DataType.ARRAY,
    element_type=DataType.STRUCT,
    struct_schema=patch_schema,
    max_capacity=1031
)

schema.add_field(
    field_name="doc_id",
    datatype=DataType.VARCHAR,
    max_length=512
)

schema.add_field(
    field_name="page_number_in_doc",
    datatype=DataType.INT64
)

index_params = client.prepare_index_params()

index_params.add_index(
    field_name="patches[emb]",
    index_type="AUTOINDEX",
    metric_type="MAX_SIM_COSINE"
)

client.create_collection(
    collection_name="financial_reports",
    schema=schema,
    index_params=index_params
)

5단계: 재무 보고서를 컬렉션에 삽입

이제 준비된 재무 보고서를 컬렉션에 삽입할 수 있습니다.

client.insert(
    collection_name="financial_reports",
    data=data
)

출력 결과를 보면 Vidore 데이터셋의 모든 페이지가 삽입된 것을 확인할 수 있습니다.

6단계: 재무 보고서 내 검색

데이터 준비가 완료되면 다음과 같이 컬렉션 내의 데이터를 대상으로 검색을 수행할 수 있습니다:

from pymilvus.client.embedding_list import EmbeddingList

queries = [
    "quarterly revenue growth chart"
]

batch_queries = processor.process_queries(queries).to(model.device)

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

query_emb_list = EmbeddingList()
query_emb_list.add_batch(query_embeddings[0].float().cpu().tolist())

results = client.search(
    collection_name="financial_reports",
    data=[query_emb_list],
    anns_field="patches[emb]",
    limit=10,
    output_fields=["doc_id", "page_number_in_doc"]
)