Pesquisa com listas de incorporação

Esta página explica como configurar um sistema de recuperação de texto ColBERT e um sistema de recuperação de texto ColPali utilizando a matriz de structs em Milvus, que permite armazenar um documento juntamente com os seus pedaços vectorizados em listas de incorporação.

Síntese

Para construir um sistema de recuperação de texto, pode ser necessário dividir os documentos em partes e armazenar cada parte juntamente com as suas incorporações como uma entidade numa base de dados vetorial para garantir a precisão e exatidão, especialmente para documentos longos em que as incorporações de texto completo podem diluir a especificidade semântica ou exceder os limites de entrada do modelo.

No entanto, o armazenamento de dados em partes leva a resultados de pesquisa por partes, o que significa que a recuperação identifica inicialmente segmentos relevantes em vez de documentos coesos. Para resolver este problema, é necessário efetuar um processamento adicional após a pesquisa.

O ColBERT (arXiv: 2004.12832) é um sistema de recuperação de texto-texto que oferece uma pesquisa de passagens eficiente e eficaz através de interações tardias contextualizadas sobre o BERT. Permite a codificação independente de consultas e documentos através de tokens e calcula a sua similaridade.

Codificação por token

Durante a ingestão de dados no ColBERT, cada documento é dividido em tokens, que são depois vectorizados e armazenados como uma lista de incorporação, como em $\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}$. Quando uma consulta chega, ela também é tokenizada, vetorizada e armazenada como uma lista de embedding, como em $\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}$.

Nas fórmulas acima,

• $dd$ d: um documento

• $qq$ q: a consulta

• $EdE_d$ Ed: a lista de incorporação que representa o documento.

• $EqE_q$ Eq: a lista de incorporação que representa a consulta.

• $[e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$d1, ed2,$...$, edn]$\in$ R${}^{n\times d}$: o número de vetores de incorporação na lista de incorporação que representa o documento está dentro do intervalo de $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}$: o número de incorporações vectoriais na lista de incorporações que representa a consulta está dentro do intervalo $Rm×d\R^{m×d}$ R ${}^{m\times d}$.

Interação tardia

Quando a vectorização estiver concluída, a lista de incorporação da consulta é comparada com a lista de incorporação de cada documento, token a token, para determinar a pontuação final de semelhança.

Interação tardia

Como mostra o diagrama acima, a consulta contém dois tokens, nomeadamente machine e learning, e o documento na janela tem quatro tokens: neural, network, python, e tutorial. Depois de estes tokens serem vectorizados, os vectores de incorporação de cada token da consulta são comparados com os do documento para obter uma lista de pontuações de semelhança. Em seguida, as pontuações mais altas de cada lista de pontuações são somadas para produzir a pontuação final. O processo para determinar a pontuação final de um documento é conhecido como similaridade máxima(MAX_SIM). Para obter detalhes sobre a similaridade máxima, consulte Similaridade máxima.

Extensão ColPali

Baseado no ColBERT, o ColPali (arXiv: 2407.01449) propõe uma nova abordagem para a recuperação de documentos visualmente ricos que utiliza modelos de visão-linguagem (VLMs). Durante a ingestão de dados, cada página de documento é renderizada numa imagem de alta resolução e depois dividida em patches, em vez de ser tokenizada. Por exemplo, uma imagem de página de documento com 448 x 448 pixéis pode produzir 1024 fragmentos, cada um com 14 x 14 pixéis.

Este método preserva informações não textuais, como a disposição do documento, imagens e estruturas de tabelas, que se perdem quando se utilizam sistemas de recuperação apenas de texto.

Extensão Copali

O VLM utilizado no ColPali chama-se PaliGemma (arXiv: 2407.07726), que inclui um codificador de imagem(SigLIP-400M), um modelo de linguagem apenas descodificador(Gemma2-2B) e uma camada linear que projecta a saída do codificador de imagem no espaço vetorial do modelo de linguagem, como se mostra no diagrama acima.

Durante a ingestão de dados, uma página de um documento, representada como uma imagem em bruto, é dividida em vários fragmentos visuais, cada um dos quais é incorporado para gerar uma lista de incorporações vectoriais. Em seguida, são projectados no espaço vetorial do modelo de linguagem para obter a lista final de incorporação, como em $\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}$. Quando uma consulta chega, ela é tokenizada, e cada token é embutido para gerar uma lista de embeddings vetoriais, como em $\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}$. Em seguida, o MAX_SIM foi aplicado para comparar as duas listas de incorporação e obter a pontuação final entre a consulta e a página do documento.

Sistema de recuperação de texto ColBERT

Nesta secção, vamos criar um sistema de recuperação de texto ColBERT utilizando a matriz de estruturas de Milvus. Antes disso, configure um cluster Milvus v2.6.x compatível com Milvus v2.6.x e obtenha um token de acesso Cohere.

Passo 1: Instalar as dependências

Execute o seguinte comando para instalar as dependências.

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

Etapa 2: carregar o conjunto de dados Cohere

Neste exemplo, vamos utilizar o conjunto de dados da Wikipédia do Cohere e recuperar os primeiros 10.000 registos. Você pode encontrar informações sobre esse conjunto de dados nesta página.

from datasets import load_dataset

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

A execução dos scripts acima fará o download do conjunto de dados se ele não estiver disponível localmente. Cada registo no conjunto de dados é um parágrafo de uma página da Wikipédia. A tabela seguinte mostra a estrutura deste conjunto de dados.

Coluna Nome	Descrição
_id	ID de um registo
url	O URL do registo atual.
title	O título do documento de origem.
text	Um parágrafo do documento de origem.
emb	Embeddings do texto do documento de origem.

Passo 3: Agrupar parágrafos por título

Para pesquisar documentos em vez de parágrafos, devemos agrupar os parágrafos por título.

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()]
  })

Neste código, armazenamos os parágrafos agrupados como documentos e incluímo-los na lista data. Cada documento tem uma chave paragraphs, que é uma lista de parágrafos; cada objeto de parágrafo contém chaves text e emb.

Passo 4: Criar uma coleção para o conjunto de dados Cohere

Quando os dados estiverem prontos, criaremos uma coleção. Na coleção, existe um campo denominado paragraphs, que é uma matriz de estruturas.

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
)

Etapa 5: Inserir o conjunto de dados Cohere na coleção

Agora podemos inserir os dados preparados na coleção que criámos acima.

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

Etapa 6: Pesquisar no conjunto de dados Cohere

De acordo com a conceção do ColBERT, o texto da consulta deve ser tokenizado e, em seguida, incorporado numa EmbeddingList. Neste passo, utilizaremos o mesmo modelo que o Cohere utilizou para gerar embeddings para os parágrafos do conjunto de dados da Wikipédia.

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"],
)

No código, os textos de consulta são organizados em tokens em query_inputs e incorporados numa lista de vectores de flutuação. Em seguida, pode utilizar a EmbeddingList do Milvus para efetuar uma pesquisa de semelhança da seguinte forma.

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}")

O resultado do código acima é semelhante ao seguinte:

# 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

A pontuação de semelhança de cosseno varia entre -1 e 1, e as pontuações de semelhança na saída acima demonstram claramente a soma de várias pontuações de semelhança a nível de token.

Sistema de recuperação de texto ColPali

Nesta secção, vamos configurar um sistema de recuperação de texto baseado no ColPali utilizando a matriz de estruturas do Milvus. Antes disso, configure um cluster do Milvus v2.6.x instanceZilliz Cloud compatível com o Milvus v2.6.x.

Passo 1: Instalar as dependências

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

Passo 2: Carregar o conjunto de dados Vidore

Nesta secção, utilizaremos um conjunto de dados Vidore denominado vidore_v2_finance_en. Este conjunto de dados é um corpus de relatórios anuais do sector bancário, destinado a tarefas de compreensão de documentos longos. É um dos 10 corpora que compõem o ViDoRe v3 Benchmark. Pode encontrar detalhes sobre este conjunto de dados nesta página.

from datasets import load_dataset

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

Ao executar os scripts acima, o conjunto de dados será descarregado se não estiver disponível localmente. Cada registo no conjunto de dados é uma página de um relatório financeiro. A tabela a seguir mostra a estrutura desse conjunto de dados.

Nome da coluna	Descrição
corpus_id	Um registo no corpus
image	A imagem da página em bytes.
doc_id	O ID descritivo do documento.
page_number_in_doc	O número de página da página atual no documento.

Passo 3: Gerar embeddings para as imagens das páginas

Tal como ilustrado na secção Visão geral, o modelo ColPali é um VLM que projecta imagens no espaço vetorial de um modelo de texto. Nesta etapa, usaremos o modelo ColPali mais recente vidore/colpali-v1.3. Pode encontrar detalhes sobre este modelo nesta página.

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)

Quando o modelo estiver pronto, pode tentar gerar correcções para uma imagem específica da seguinte forma.

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]

No código acima, o modelo ColPali redimensiona a imagem para 448 x 448 pixels, depois divide-a em patches, cada um medindo 14 x 14 pixels. Por fim, esses patches são incorporados em 1.031 embeddings, cada um com 128 dimensões.

Pode gerar embeddings para todas as imagens utilizando um ciclo como se segue:

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']
  })

Passo 4: Criar uma coleção para o conjunto de dados dos relatórios financeiros

Quando os dados estiverem prontos, criaremos uma coleção. Na coleção, um campo denominado patches é uma matriz de estruturas.

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
)

Etapa 5: Inserir os relatórios financeiros na coleção

Agora, podemos inserir os relatórios financeiros preparados na coleção.

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

A partir do resultado, pode verificar que todas as páginas do conjunto de dados Vidore foram inseridas.

Passo 6: Pesquisar nos relatórios financeiros

Quando os dados estiverem prontos, podemos efetuar pesquisas nos dados da coleção da seguinte forma:

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"]
)