Búsqueda con listas incrustadas

Esta página explica cómo configurar un sistema de recuperación de texto ColBERT y un sistema de recuperación de texto ColPali utilizando la matriz de structs en Milvus, que le permite almacenar un documento junto con sus trozos vectorizados en listas incrustadas.

Visión general

Para construir un sistema de recuperación de texto, puede que necesite dividir los documentos en trozos y almacenar cada trozo junto con sus incrustaciones como una entidad en una base de datos vectorial para garantizar la precisión y la exactitud, especialmente para documentos largos en los que las incrustaciones de texto completo podrían diluir la especificidad semántica o exceder los límites de entrada del modelo.

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

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

Codificación por token

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

En las fórmulas anteriores

• $dd$ d: un documento

• $qq$ q: la consulta

• $EdE_d$ Ed: la lista de incrustación que representa el documento.

• $EqE_q$ Eq: la lista de incrustación que representa la consulta.

• $[e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$d1, ed2,$...$, edn]$\in$ R${}^{n\times d}$: el número de incrustaciones vectoriales en la lista de incrustaciones que representa el documento está dentro del 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}$: el número de incrustaciones vectoriales en la lista de incrustaciones que representa la consulta está dentro del intervalo de $Rm×d\R^{m×d}$ R ${}^{m\times d}$.

Interacción tardía

Una vez finalizada la vectorización, la lista de incrustación de la consulta se compara con cada lista de incrustación del documento, símbolo a símbolo, para determinar la puntuación final de similitud.

Interacción tardía

Como se muestra en el diagrama anterior, la consulta contiene dos tokens, machine y learning, y el documento de la ventana tiene cuatro tokens: neural network , python y tutorial. Una vez vectorizados estos tokens, los vectores de cada token de la consulta se comparan con los del documento para obtener una lista de puntuaciones de similitud. A continuación, se suman las puntuaciones más altas de cada lista 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 más detalles sobre la máxima similitud, consulte Máxima similitud.

Extensión ColPali

Basado en ColBERT, ColPali (arXiv: 2407.01449) propone un enfoque novedoso para la recuperación de documentos visualmente ricos 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 parches, en lugar de tokenizarse. Por ejemplo, una imagen de página de documento de 448 x 448 píxeles puede producir 1.024 parches, cada uno de 14 x 14 píxeles.

Este método preserva la información no textual, como la disposición del documento, las imágenes y las estructuras de tablas, que se pierden al utilizar sistemas de recuperación de sólo texto.

Extensión Copali

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 sólo decodificador(Gemma2-2B) y una capa lineal que proyecta la salida del codificador de imágenes en el espacio vectorial del modelo de lenguaje, como se muestra en el diagrama anterior.

Durante la ingesta de datos, la página de un documento, representada como una imagen en bruto, se divide en múltiples parches 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 lingüístico para obtener la lista final de incrustación, como en $\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}$. Cuando llega una consulta, se tokeniza, y cada token se incrusta para generar una lista de incrustaciones vectoriales, como en $\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}$. A continuación, se ha aplicado MAX_SIM para comparar las dos listas de incrustación y obtener la puntuación final entre la consulta y la página del documento.

Sistema de recuperación de textos ColBERT

En esta sección, vamos a configurar un sistema de recuperación de texto ColBERT utilizando la matriz de estructuras de Milvus. Antes de eso, configure una instancia Milvus v2.6.xZilliz Cloud cluster compatible con Milvus v2.6.x, obtenga un token de acceso Cohere.

Paso 1: Instalar las dependencias

Ejecute el siguiente comando para instalar las dependencias.

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

Paso 2: Cargar el conjunto de datos Cohere

En este ejemplo, vamos a utilizar el conjunto de datos de Wikipedia de Cohere y recuperar los primeros 10.000 registros. Puede 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]"
)

La ejecución de los scripts anteriores 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.

Columna Nombre	Descripción
_id	ID del registro
url	La URL del registro actual.
title	El título del documento fuente.
text	Un párrafo del documento fuente.
emb	Incrustaciones del texto del documento fuente.

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 párrafo contiene unas claves text y emb.

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

Una vez que los datos estén listos, crearemos una colección. En la colección, hay un campo llamado paragraphs, que es un Array of Structs.

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 Cohere en la colección

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

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

Paso 6: Buscar en el conjunto de datos Cohere

De acuerdo con el diseño de ColBERT, el texto de la consulta debe ser tokenizado y luego incrustado en un EmbeddingList. En este paso, utilizaremos el mismo modelo que Cohere utilizó para generar incrustaciones para los párrafos en el 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 query_inputs y se incrustan en una lista de vectores flotantes. A continuación, se puede utilizar EmbeddingList de Milvus para realizar una búsqueda de similitud como se indica a continuación.

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

La salida del código anterior es similar a la 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

La puntuación de similitud coseno va de -1 a 1, y las puntuaciones de similitud en la salida anterior muestran claramente la suma de múltiples puntuaciones de similitud a nivel de token.

Sistema de recuperación de texto ColPali

En esta sección, crearemos un sistema de recuperación de texto basado en ColPali utilizando la matriz de estructuras de Milvus. Antes de eso, configure una instancia de Milvus v2.6.xZilliz Cloud cluster compatible con 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 Vidore

En esta sección, utilizaremos un conjunto de datos Vidore llamado 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. Encontrará más información 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()

La ejecución de los scripts anteriores descargará el conjunto de datos si no está disponible localmente. Cada registro del conjunto de datos es una página de un informe financiero. La siguiente tabla muestra la estructura de este conjunto de datos.

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

Paso 3: Generar incrustaciones para las imágenes de página

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 el último modelo ColPali vidore/colpali-v1.3. Puede encontrar 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, puede intentar generar parches para una imagen específica de la siguiente manera.

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]

En el código anterior, el modelo ColPali redimensiona la imagen a 448 x 448 píxeles, luego la divide en parches, cada uno de los cuales mide 14 x 14 píxeles. Por último, estos parches se incrustan en 1.031 incrustaciones, cada una con 128 dimensiones.

Puede generar incrustaciones para todas las imágenes utilizando un bucle como el siguiente:

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

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, un campo llamado patches es un Array of Structs.

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, puede ver que se han insertado todas las páginas del conjunto de datos Vidore.

Paso 6: Buscar 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].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"]
)