Búsqueda con EmbeddingLists: ColBERT y ColPali

Este tutorial muestra cómo crear sistemas de recuperación al estilo ColBERT y ColPali mediante la búsqueda con EmbeddingLists en subcampos vectoriales de StructArray en Milvus. Úsalo cuando tanto tu consulta como los datos almacenados se representen como listas de vectores y desees una recuperación con interacción tardía a nivel de entidad con métricas de « MAX_SIM* ».

Para conocer los conceptos básicos de StructArray en los que se basa este tutorial, consulta «Crear un campo StructArray», «Indexar campos StructArray» y «Búsqueda vectorial básica con StructArray». Este tutorial se centra en los flujos de trabajo de ColBERT y ColPali, más que en la sintaxis general de StructArray.

Descripción general

Para crear un sistema de recuperación de texto, es posible que tengas que dividir los documentos en fragmentos y almacenar cada fragmento, junto con sus incrustaciones, como una entidad en una base de datos vectorial para garantizar la precisión y la exactitud, especialmente en el caso de documentos largos en los que las incrustaciones de texto completo podrían diluir la especificidad semántica o superar los límites de entrada del modelo.

Sin embargo, el almacenamiento de datos en fragmentos da lugar a resultados de búsqueda por fragmentos, lo que significa que la recuperación identifica inicialmente segmentos relevantes en lugar de documentos cohesionados. Para solucionar esto, debes realizar un procesamiento adicional tras la búsqueda.

ColBERT (arXiv: 2004.12832) es un sistema de recuperación de texto-texto que ofrece una búsqueda eficiente y eficaz de pasajes mediante interacciones contextuales tardías sobre BERT. Permite la codificación independiente, token por token, de consultas y documentos, y calcula su similitud.

Codificación por tokens

Durante la ingesta de datos en ColBERT, cada documento se divide en tokens, que a continuación se vectorizan y se almacenan como una lista de incrustaciones, tal y como se muestra en $∈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}$. Cuando llega una consulta, también se tokeniza, se vectoriza y se almacena como una lista de incrustaciones, tal y como en $∈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}$.

En las fórmulas anteriores,

• $dd$: un documento

• $qq$: la consulta

• $EdE_d$: la lista de incrustaciones que representa el documento.

• $EqE_q$: la lista de incrustaciones que representa la consulta.

• $[e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$$,e_{}$,…,ydn$]\in$ R${}^{n\times d}$: el número de incrustaciones vectoriales en la lista de incrustaciones que representa el documento se encuentra dentro del rango de $\R^{n×d}$.

• $[e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$q1,eq2,…,eqm$]\in$ R${}^{m\times d}$: el número de incrustaciones vectoriales en la lista de incrustaciones que representa la consulta se encuentra dentro del rango de $\R^{m×d}$.

Interacción tardía

Una vez completada la vectorización, la lista de incrustaciones de la consulta se compara con la lista de incrustaciones de cada documento, token por token, para determinar la puntuación final de similitud.

Como se muestra en el diagrama anterior, la consulta contiene dos tokens, a saber, machine y learning, y el documento de la ventana tiene cuatro tokens: neural, network, python y tutorial. Una vez vectorizados estos tokens, las representaciones vectoriales de cada token de la consulta se comparan con las del documento para obtener una lista de puntuaciones de similitud. A continuación, se suman las puntuaciones más altas de cada lista de puntuaciones para obtener la puntuación final. El proceso para determinar la puntuación final de un documento se conoce como «similitud máxima» (MAX_SIM). Para obtener más información sobre la similitud máxima, consulta «Similitud máxima».

Extensión ColPali

Basándose en ColBERT, ColPali (arXiv: 2407.01449) propone un enfoque novedoso para la recuperación de documentos con gran riqueza visual que aprovecha los modelos de visión y lenguaje (VLM). Durante la ingesta de datos, cada página del documento se convierte en una imagen de alta resolución y, a continuación, se divide en fragmentos, en lugar de tokenizarse. Por ejemplo, la imagen de una página de documento de 448 x 448 píxeles puede generar 1.024 fragmentos, cada uno de 14 x 14 píxeles.

Este método conserva la información no textual, como el diseño del documento, las imágenes y las estructuras de las tablas, que se pierden cuando se utilizan sistemas de recuperación basados únicamente en texto.

El VLM utilizado en ColPali se denomina PaliGemma (arXiv: 2407.07726), y consta de un codificador de imágenes (SigLIP-400M), un modelo de lenguaje de solo decodificador (Gemma2-2B), y una capa lineal que proyecta la salida del codificador de imágenes en el espacio vectorial del modelo de lenguaje, tal y como se muestra en el diagrama anterior.

Durante la ingesta de datos, una página de documento, representada como una imagen sin procesar, se divide en múltiples fragmentos visuales, cada uno de los cuales se incrusta para generar una lista de incrustaciones vectoriales. A continuación, se proyectan en el espacio vectorial del modelo de lenguaje para obtener la lista de incrustaciones final, como en $∈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}$. Cuando llega una consulta, se tokeniza y cada token se incrusta para generar una lista de incrustaciones vectoriales, como en $∈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}$. A continuación, se ha aplicado MAX_SIM para comparar las dos listas de representaciones y obtener la puntuación final entre la consulta y la página del documento.

Sistema de recuperación de texto ColBERT

En esta sección, vamos a configurar un sistema de recuperación de texto ColBERT utilizando StructArray. Antes de eso, configura una instancia de Milvus v2.6.x y obtén un token de acceso de Cohere.

Paso 1: Instalar las dependencias

Ejecuta el siguiente comando para instalar las dependencias.

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

Paso 2: Cargar el conjunto de datos de Cohere

En este ejemplo, vamos a utilizar el conjunto de datos de Wikipedia de Cohere y a recuperar los primeros 10 000 registros. Puedes encontrar información sobre este conjunto de datos en esta página.

from datasets import load_dataset

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

Al ejecutar los scripts anteriores, se descargará el conjunto de datos si no está disponible localmente. Cada registro del conjunto de datos es un párrafo de una página de Wikipedia. La siguiente tabla muestra la estructura de este conjunto de datos.

Nombre de la columna	Descripción
_id	ID del registro
url	La URL del registro actual.
title	El título del documento de origen.
text	Un párrafo del documento original.
emb	Incrustaciones del texto del documento original.

Paso 3: Agrupar párrafos por título

Para buscar documentos en lugar de párrafos, debemos agrupar los párrafos 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()]
  })

En este código, almacenamos los párrafos agrupados como documentos y los incluimos en la lista « data ». Cada documento tiene una clave « paragraphs », que es una lista de párrafos; cada objeto de párrafo contiene las claves « text » y « emb ».

Paso 4: Crear una colección para el conjunto de datos de Cohere

Una vez que los datos estén listos, crearemos una colección. En la colección, « paragraphs » es un campo StructArray. Para obtener una explicación general de los esquemas StructArray, consulta «Crear un campo 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
)

Paso 5: Insertar el conjunto de datos de Cohere en la colección

Ahora podemos insertar los datos preparados en la colección que hemos creado anteriormente.

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

Paso 6: Realizar una búsqueda en el conjunto de datos de Cohere

Según el diseño de ColBERT, el texto de la consulta debe tokenizarse y, a continuación, integrarse en una EmbeddingList. En este paso, utilizaremos el mismo modelo que utilizó Cohere para generar representaciones para los párrafos del conjunto de datos de 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"],
)

En el código, los textos de consulta se organizan en tokens en un query_inputs o y se incrustan en una lista de vectores de tipo float. A continuación, puedes utilizar la EmbeddingList de Milvus para realizar una búsqueda por similitud de la siguiente manera.

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

El resultado del código anterior es similar al siguiente:

# 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

Cada puntuación de similitud coseno por pares oscila entre -1 y 1. La puntuación final de « MAX_SIM_COSINE » puede ser superior a 1, ya que agrega múltiples puntuaciones de similitud máxima a nivel de token.

Sistema de recuperación de documentos ColPali

En esta sección, configuraremos un sistema de recuperación de documentos basado en ColPali utilizando StructArray. Antes de eso, configura una instancia de Milvus v2.6.x.

Paso 1: Instalar las dependencias

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

Paso 2: Cargar el conjunto de datos de Vidore

En esta sección, utilizaremos un conjunto de datos de Vidore denominado vidore_v2_finance_en. Este conjunto de datos es un corpus de informes anuales del sector bancario, destinado a tareas de comprensión de documentos largos. Es uno de los 10 corpus que componen el ViDoRe v3 Benchmark. Puedes encontrar más detalles sobre este conjunto de datos en esta página.

from datasets import load_dataset

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

Al ejecutar los scripts anteriores, se descargará el conjunto de datos si no está disponible localmente. Cada registro del conjunto de datos corresponde a una página de un informe financiero. La siguiente tabla muestra la estructura de este conjunto de datos.

Nombre de la columna	Descripción
corpus_id	Un registro del corpus
image	La imagen de la página en bytes.
doc_id	El identificador descriptivo del documento.
page_number_in_doc	El número de página de la página actual en el documento.

Paso 3: Generar representaciones vectoriales de las imágenes de las páginas

Tal y como se ilustra en la sección «Descripción general», el modelo ColPali es un VLM que proyecta imágenes en el espacio vectorial de un modelo de texto. En este paso, utilizaremos la última versión del modelo ColPali : vidore/colpali-v1.3. Puedes encontrar más detalles sobre este modelo en esta 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)

Una vez que el modelo esté listo, puedes intentar generar fragmentos para una imagen específica de la siguiente manera.

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]

En el código anterior, el modelo ColPali redimensiona la imagen a 448 x 448 píxeles y, a continuación, la divide en fragmentos de 14 x 14 píxeles cada uno. Por último, estos fragmentos se incrustan en 1.031 representaciones, cada una con 128 dimensiones.

Puedes generar representaciones para todas las imágenes utilizando un bucle de la siguiente manera:

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

Paso 4: Crear una colección para el conjunto de datos de informes financieros

Una vez que los datos estén listos, crearemos una colección. En la colección, « patches » es un campo de tipo StructArray. Cada elemento de Struct almacena una incrustación de fragmento. Para conocer los requisitos de indexación de los subcampos vectoriales de StructArray, consulta «Indexar campos de 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
)

Paso 5: Insertar los informes financieros en la colección

Ahora podemos insertar los informes financieros preparados en la colección.

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

En la salida, podrá comprobar que se han insertado todas las páginas del conjunto de datos de Vidore.

Paso 6: Realizar búsquedas en los informes financieros

Una vez que los datos estén listos, podemos realizar búsquedas en los datos de la colección de la siguiente manera:

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