Use ColPali para recuperação multimodal com Milvus
Os modelos de recuperação modernos utilizam normalmente um único embedding para representar texto ou imagens. O ColBERT, no entanto, é um modelo neural que utiliza uma lista de embeddings para cada instância de dados e emprega uma operação "MaxSim" para calcular a semelhança entre dois textos. Para além dos dados textuais, as figuras, tabelas e diagramas também contêm informações ricas, que são frequentemente ignoradas na recuperação de informações baseada em texto.
A função MaxSim compara uma consulta com um documento (o que está a ser pesquisado) através da análise dos seus token embeddings. Para cada palavra na consulta, seleciona a palavra mais semelhante do documento (utilizando a semelhança de cosseno ou a distância L2 ao quadrado) e soma estas semelhanças máximas entre todas as palavras na consulta
O ColPali é um método que combina a representação multi-vetorial do ColBERT com o PaliGemma (um modelo de linguagem multimodal de grande dimensão) para tirar partido das suas fortes capacidades de compreensão. Esta abordagem permite que uma página com texto e imagens seja representada utilizando uma incorporação multi-vetorial unificada. As incorporações dentro desta representação multi-vetorial podem captar informações detalhadas, melhorando o desempenho da geração aumentada de recuperação (RAG) para dados multimodais.
Neste caderno, referimo-nos a este tipo de representação multi-vetorial como "ColBERT embeddings" por uma questão de generalidade. No entanto, o modelo real que está a ser utilizado é o modelo ColPali. Vamos demonstrar como utilizar o Milvus para a recuperação multi-vetorial. Com base nisso, apresentaremos como utilizar o ColPali para recuperar páginas com base numa determinada consulta.
Preparação
$ pip install pdf2image
$ pip pymilvus
$ pip install colpali_engine
$ pip install tqdm
$ pip instal pillow
Preparar os dados
Utilizaremos o PDF RAG como exemplo. Pode descarregar o documento ColBERT e colocá-lo em ./pdf
. O ColPali não processa diretamente o texto; em vez disso, toda a página é rasterizada numa imagem. O modelo ColPali é excelente para compreender a informação textual contida nestas imagens. Portanto, converteremos cada página do PDF em uma imagem para processamento.
from pdf2image import convert_from_path
pdf_path = "pdfs/2004.12832v2.pdf"
images = convert_from_path(pdf_path)
for i, image in enumerate(images):
image.save(f"pages/page_{i + 1}.png", "PNG")
Em seguida, inicializaremos um banco de dados usando o Milvus Lite. Pode mudar facilmente para uma instância completa do Milvus definindo o uri para o endereço apropriado onde o seu serviço Milvus está alojado.
from pymilvus import MilvusClient, DataType
import numpy as np
import concurrent.futures
client = MilvusClient(uri="milvus.db")
- Se apenas necessitar de uma base de dados vetorial local para dados de pequena escala ou prototipagem, definir o uri como um ficheiro local, por exemplo
./milvus.db
, é o método mais conveniente, uma vez que utiliza automaticamente o Milvus Lite para armazenar todos os dados neste ficheiro. - Se tiver uma grande escala de dados, digamos mais de um milhão de vectores, pode configurar um servidor Milvus mais eficiente em Docker ou Kubernetes. Nesta configuração, use o endereço e a porta do servidor como seu uri, por exemplo,
http://localhost:19530
. Se ativar a funcionalidade de autenticação no Milvus, utilize "<your_username>:<your_password>" como token, caso contrário não defina o token. - Se utilizar o Zilliz Cloud, o serviço de nuvem totalmente gerido para o Milvus, ajuste os campos
uri
etoken
, que correspondem ao Public Endpoint e à chave API no Zilliz Cloud.
Vamos definir uma classe MilvusColbertRetriever para envolver o cliente Milvus na recuperação de dados multi-vectoriais. A implementação aplaina os embeddings ColBERT e insere-os numa coleção, em que cada linha representa um embedding individual da lista de embeddings ColBERT. Também regista o doc_id e o seq_id para rastrear a origem de cada embedding.
Ao pesquisar com uma lista de inserções ColBERT, serão efectuadas várias pesquisas - uma para cada inserção ColBERT. Os doc_ids recuperados serão então deduplicados. Será efectuado um processo de nova classificação, em que as incorporações completas para cada doc_id são obtidas e a pontuação MaxSim é calculada para produzir os resultados finais classificados.
class MilvusColbertRetriever:
def __init__(self, milvus_client, collection_name, dim=128):
# Initialize the retriever with a Milvus client, collection name, and dimensionality of the vector embeddings.
# If the collection exists, load it.
self.collection_name = collection_name
self.client = milvus_client
if self.client.has_collection(collection_name=self.collection_name):
self.client.load_collection(collection_name)
self.dim = dim
def create_collection(self):
# Create a new collection in Milvus for storing embeddings.
# Drop the existing collection if it already exists and define the schema for the collection.
if self.client.has_collection(collection_name=self.collection_name):
self.client.drop_collection(collection_name=self.collection_name)
schema = self.client.create_schema(
auto_id=True,
enable_dynamic_fields=True,
)
schema.add_field(field_name="pk", datatype=DataType.INT64, is_primary=True)
schema.add_field(
field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=self.dim
)
schema.add_field(field_name="seq_id", datatype=DataType.INT16)
schema.add_field(field_name="doc_id", datatype=DataType.INT64)
schema.add_field(field_name="doc", datatype=DataType.VARCHAR, max_length=65535)
self.client.create_collection(
collection_name=self.collection_name, schema=schema
)
def create_index(self):
# Create an index on the vector field to enable fast similarity search.
# Releases and drops any existing index before creating a new one with specified parameters.
self.client.release_collection(collection_name=self.collection_name)
self.client.drop_index(
collection_name=self.collection_name, index_name="vector"
)
index_params = self.client.prepare_index_params()
index_params.add_index(
field_name="vector",
index_name="vector_index",
index_type="HNSW", # or any other index type you want
metric_type="IP", # or the appropriate metric type
params={
"M": 16,
"efConstruction": 500,
}, # adjust these parameters as needed
)
self.client.create_index(
collection_name=self.collection_name, index_params=index_params, sync=True
)
def create_scalar_index(self):
# Create a scalar index for the "doc_id" field to enable fast lookups by document ID.
self.client.release_collection(collection_name=self.collection_name)
index_params = self.client.prepare_index_params()
index_params.add_index(
field_name="doc_id",
index_name="int32_index",
index_type="INVERTED", # or any other index type you want
)
self.client.create_index(
collection_name=self.collection_name, index_params=index_params, sync=True
)
def search(self, data, topk):
# Perform a vector search on the collection to find the top-k most similar documents.
search_params = {"metric_type": "IP", "params": {}}
results = self.client.search(
self.collection_name,
data,
limit=int(50),
output_fields=["vector", "seq_id", "doc_id"],
search_params=search_params,
)
doc_ids = set()
for r_id in range(len(results)):
for r in range(len(results[r_id])):
doc_ids.add(results[r_id][r]["entity"]["doc_id"])
scores = []
def rerank_single_doc(doc_id, data, client, collection_name):
# Rerank a single document by retrieving its embeddings and calculating the similarity with the query.
doc_colbert_vecs = client.query(
collection_name=collection_name,
filter=f"doc_id in [{doc_id}]",
output_fields=["seq_id", "vector", "doc"],
limit=1000,
)
doc_vecs = np.vstack(
[doc_colbert_vecs[i]["vector"] for i in range(len(doc_colbert_vecs))]
)
score = np.dot(data, doc_vecs.T).max(1).sum()
return (score, doc_id)
with concurrent.futures.ThreadPoolExecutor(max_workers=300) as executor:
futures = {
executor.submit(
rerank_single_doc, doc_id, data, client, self.collection_name
): doc_id
for doc_id in doc_ids
}
for future in concurrent.futures.as_completed(futures):
score, doc_id = future.result()
scores.append((score, doc_id))
scores.sort(key=lambda x: x[0], reverse=True)
if len(scores) >= topk:
return scores[:topk]
else:
return scores
def insert(self, data):
# Insert ColBERT embeddings and metadata for a document into the collection.
colbert_vecs = [vec for vec in data["colbert_vecs"]]
seq_length = len(colbert_vecs)
doc_ids = [data["doc_id"] for i in range(seq_length)]
seq_ids = list(range(seq_length))
docs = [""] * seq_length
docs[0] = data["filepath"]
# Insert the data as multiple vectors (one for each sequence) along with the corresponding metadata.
self.client.insert(
self.collection_name,
[
{
"vector": colbert_vecs[i],
"seq_id": seq_ids[i],
"doc_id": doc_ids[i],
"doc": docs[i],
}
for i in range(seq_length)
],
)
Utilizaremos o motor colpali_engine para extrair listas de incorporação para duas consultas e obter as informações relevantes das páginas PDF.
from colpali_engine.models import ColPali
from colpali_engine.models.paligemma.colpali.processing_colpali import ColPaliProcessor
from colpali_engine.utils.processing_utils import BaseVisualRetrieverProcessor
from colpali_engine.utils.torch_utils import ListDataset, get_torch_device
from torch.utils.data import DataLoader
import torch
from typing import List, cast
device = get_torch_device("cpu")
model_name = "vidore/colpali-v1.2"
model = ColPali.from_pretrained(
model_name,
torch_dtype=torch.bfloat16,
device_map=device,
).eval()
queries = [
"How to end-to-end retrieval with ColBert?",
"Where is ColBERT performance table?",
]
processor = cast(ColPaliProcessor, ColPaliProcessor.from_pretrained(model_name))
dataloader = DataLoader(
dataset=ListDataset[str](queries),
batch_size=1,
shuffle=False,
collate_fn=lambda x: processor.process_queries(x),
)
qs: List[torch.Tensor] = []
for batch_query in dataloader:
with torch.no_grad():
batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
embeddings_query = model(**batch_query)
qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
Além disso, precisamos de extrair a lista de incorporação para cada página e esta mostra que existem 1030 incorporações de 128 dimensões para cada página.
from tqdm import tqdm
from PIL import Image
import os
images = [Image.open("./pages/" + name) for name in os.listdir("./pages")]
dataloader = DataLoader(
dataset=ListDataset[str](images),
batch_size=1,
shuffle=False,
collate_fn=lambda x: processor.process_images(x),
)
ds: List[torch.Tensor] = []
for batch_doc in tqdm(dataloader):
with torch.no_grad():
batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
embeddings_doc = model(**batch_doc)
ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
print(ds[0].shape)
0%| | 0/10 [00:00<?, ?it/s]
100%|██████████| 10/10 [01:22<00:00, 8.24s/it]
torch.Size([1030, 128])
Vamos criar uma coleção chamada "colpali" utilizando o MilvusColbertRetriever.
retriever = MilvusColbertRetriever(collection_name="colpali", milvus_client=client)
retriever.create_collection()
retriever.create_index()
Vamos inserir as listas de ligações na base de dados Milvus.
filepaths = ["./pages/" + name for name in os.listdir("./pages")]
for i in range(len(filepaths)):
data = {
"colbert_vecs": ds[i].float().numpy(),
"doc_id": i,
"filepath": filepaths[i],
}
retriever.insert(data)
Agora, podemos procurar a página mais relevante utilizando a lista de incorporação da consulta.
for query in qs:
query = query.float().numpy()
result = retriever.search(query, topk=1)
print(filepaths[result[0][1]])
./pages/page_5.png
./pages/page_7.png
Por fim, recuperamos o nome da página original. Com o ColPali, podemos recuperar documentos multimodais sem a necessidade de técnicas de processamento complexas para extrair texto e imagens dos documentos. Ao tirar partido de modelos de visão de grande dimensão, é possível analisar mais informações - como tabelas e figuras - sem perdas significativas de informação.