임베딩 목록으로 검색하기

이 페이지에서는 문서를 벡터화된 청크와 함께 임베딩 목록에 저장할 수 있는 Milvus의 구조 배열을 사용하여 ColBERT 텍스트 검색 시스템과 ColPali 텍스트 검색 시스템을 설정하는 방법에 대해 설명합니다.

개요

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

그러나 데이터를 청크 단위로 저장하면 청크 단위의 검색 결과가 나오므로 검색 시 일관된 문서가 아닌 관련 세그먼트가 먼저 식별됩니다. 이 문제를 해결하려면 추가적인 검색 후 처리를 수행해야 합니다.

ColBERT(arXiv: 2004.12832)는 BERT를 통해 문맥화된 후기 상호작용을 통해 효율적이고 효과적인 구절 검색을 제공하는 텍스트-텍스트 검색 시스템입니다. 쿼리와 문서를 독립적인 토큰 단위로 인코딩하고 유사성을 계산할 수 있습니다.

토큰 단위 인코딩

콜버트에서 데이터를 수집하는 동안 각 문서는 토큰으로 분할되며, 각 토큰은 벡터화되어 임베딩 목록으로 저장됩니다(예: $d\to Ed=[{\mathrm{ed1}}_{},{\mathrm{ed2}}_{},...,{\mathrm{edn}}_{}]{\mathrm{\in Rn\times dd}}^{}$ $\rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ d $\to$ E d = $[e$ d1, e d2, $...$, e dn] $\in$ R ${}^{n\times d입니다}$. 쿼리가 도착하면 토큰화되고 벡터화되어 $q\to Eq=[{\mathrm{eq1}}_{},{\mathrm{eq2}}_{},...,{\mathrm{eqm}}_{}]{\mathrm{\in Rm\times dq}}^{}$ $\rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ q $\to$ E q = $[e$ q1, e q2, $...$, e qm] $\in$ R ${}^{m\times d입니다}$.

위의 공식에서,

• $dd$ d: 문서

• $qq$ q: 쿼리

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

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

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

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

후기 상호 작용

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

후기 상호 작용

위의 다이어그램에서 볼 수 있듯이 쿼리에는 machine 과 learning 이라는 두 개의 토큰이 포함되어 있고, 창에 있는 문서에는 4개의 토큰이 있습니다: 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), 이미지 인코더의 출력을 언어 모델의 벡터 공간으로 투영하는 선형 계층으로 구성되어 있습니다.

데이터 수집 중에 원시 이미지로 표시되는 문서 페이지가 여러 개의 시각적 패치로 나뉘며, 각 패치는 벡터 임베딩 목록을 생성하기 위해 임베딩됩니다. 그런 다음 언어 모델의 벡터 공간에 투영하여 최종 임베딩 목록을 얻습니다(예: $d\to Ed=[{\mathrm{ed1}}_{},{\mathrm{ed2}}_{},...,{\mathrm{edn}}_{}]{\mathrm{\in Rn\times dd}}^{}$ $\rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ d $\to$ E d = $[e$ d1, e d2, $...$, e dn] $\in$ R ${}^{n\times d입니다}$. 쿼리가 도착하면 토큰화되고 각 토큰은 벡터 임베딩의 목록을 생성하기 위해 임베딩된다(예: $q\to Eq=[{\mathrm{eq1}}_{},{\mathrm{eq2}}_{},...,{\mathrm{eqm}}_{}]{\mathrm{\in Rm\times dq}}^{}$ $\rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ q $\to$ E q = $[e$ q1, e q2, $...$, e qm] $\in$ R ${}^{m\times d입니다}$. 그런 다음 MAX_SIM을 적용하여 두 임베딩 목록을 비교하고 쿼리와 문서 페이지 간의 최종 점수를 얻습니다.

콜버트 텍스트 검색 시스템

이 섹션에서는 Milvus의 구조 배열을 사용하여 ColBERT 텍스트 검색 시스템을 설정해 보겠습니다. 그 전에 Milvus v2.6.x와 호환되는 Milvus v2.6.x 인스턴스Zilliz Cloud 클러스터를 설정하고 Cohere 액세스 토큰을 받습니다.

1단계: 종속성 설치

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

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

2단계: Cohere 데이터 세트 로드하기

이 예에서는 Cohere의 Wikipedia 데이터 세트를 사용하여 처음 10,000개의 레코드를 검색하겠습니다. 이 데이터 세트에 대한 정보는 이 페이지에서 찾을 수 있습니다.

from datasets import load_dataset

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

위의 스크립트를 실행하면 로컬에서 사용할 수 없는 경우 데이터 세트가 다운로드됩니다. 데이터 집합의 각 레코드는 Wikipedia 페이지의 단락입니다. 다음 표는 이 데이터 집합의 구조를 보여줍니다.

열 이름	설명
_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 이라는 필드가 있습니다.

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 데이터 세트 내에서 검색하기

콜베르트의 설계에 따르면 쿼리 텍스트는 토큰화된 다음 EmbeddingList에 임베드되어야 합니다. 이 단계에서는 Cohere가 Wikipedia 데이터 세트의 단락에 대한 임베딩을 생성하는 데 사용한 것과 동일한 모델을 사용합니다.

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]",
    search_params={
        "metric_type": "MAX_SIM_COSINE"
    },
    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 이며, 위 출력의 유사도 점수는 여러 토큰 수준의 유사도 점수의 합을 명확하게 보여줍니다.

ColPali 텍스트 검색 시스템

이 섹션에서는 Milvus의 구조 배열을 사용하여 ColPali-기반 텍스트 검색 시스템을 설정하겠습니다. 그 전에 Milvus v2.6.x와 호환되는 Milvus v2.6.x 인스턴스Zilliz Cloud 클러스터를 설정합니다.

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 iterrow() generator to get the first row
row = next(df.iterrows())[1]

# Include the image in the above row in a list
images = [ Image.open(row['image']['bytes'] ]
patches = processor.process_images(images).to(model.device)
patches_embeddings = model(**patches_in_pixels)[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 index, row in df.iterrows():
  row = next(df.iterrows())[1]
  corpus_id = row['corpus_id']
  
  images = [Image.open(BytesIO(row['image']['bytes']))]
  batch_images = processor.process_images(images).to(model.device)
  patches = model(**batch_images)[0]

  doc_id = row['doc_id']
  markdown = row['markdown']
  page_number_in_doc = row['page_number_in_doc']

  data.append({
      "corpus_id": corpus_id,
      "patches": [ {"emb": emb} for emb in patches ],
      "doc_id": markdown,
      "page_number_in_doc": row['page_number_in_doc']
  })

4단계: 재무 보고서 데이터 집합을 위한 컬렉션 만들기

데이터가 준비되면 컬렉션을 만듭니다. 컬렉션에서 patches 이라는 필드는 구조 배열입니다.

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].cpu())

results = client.search(
    collection_name="financial_reports",
    data=[query_emb_list],
    anns_field="patches[emb]",
    search_params={
        "metric_type": "MAX_SIM_COSINE"
    },
    limit=10,
    output_fields=["doc_id", "page_number_in_doc"]
)