Ricerca con EmbeddingLists: ColBERT e ColPali

Questo tutorial illustra come realizzare sistemi di recupero in stile ColBERT e ColPali utilizzando la ricerca con EmbeddingList sui sottocampi vettoriali di StructArray in Milvus. Utilizzatelo quando sia la query che i dati memorizzati sono rappresentati come elenchi di vettori e desiderate un recupero a livello di entità con interazione tardiva e metriche di tipo " MAX_SIM* ".

Per le nozioni di base su StructArray alla base di questo tutorial, consultare Creare un campo StructArray, Indicizzare campi StructArray e Ricerca vettoriale di base con StructArray. Questo tutorial si concentra sui flussi di lavoro ColBERT e ColPali piuttosto che sulla sintassi generale di StructArray.

Panoramica

Per realizzare un sistema di recupero di testo, potrebbe essere necessario suddividere i documenti in blocchi e memorizzare ciascun blocco insieme alle relative incorporazioni come entità in un database vettoriale, al fine di garantire precisione e accuratezza, specialmente nel caso di documenti lunghi in cui le incorporazioni a testo intero potrebbero diluire la specificità semantica o superare i limiti di input del modello.

Tuttavia, la memorizzazione dei dati in blocchi porta a risultati di ricerca a livello di blocco, il che significa che il recupero identifica inizialmente segmenti rilevanti piuttosto che documenti coerenti. Per ovviare a questo problema, è necessario eseguire un'ulteriore elaborazione post-ricerca.

ColBERT (arXiv: 2004.12832) è un sistema di recupero testo-testo che offre una ricerca efficiente ed efficace di passaggi attraverso interazioni contestualizzate in fase avanzata su BERT. Consente la codifica indipendente a livello di token di query e documenti e ne calcola la somiglianza.

Codifica a livello di token

Durante l’acquisizione dei dati in ColBERT, ogni documento viene suddiviso in token, che vengono poi vettorializzati e memorizzati come lista di embedding, come in $∈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}$. Quando arriva una query, viene anch’essa tokenizzata, vettorializzata e memorizzata come lista di embedding, come in $∈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}$.

Nelle formule sopra riportate,

• $dd$: un documento

• $qq$: la query

• $EdE_d$: l’elenco di embedding che rappresenta il documento.

• $EqE_q$: la lista di embedding che rappresenta la query.

• $[e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$$,e_{}$,…,edn$]\in$ R${}^{n\times d}$: il numero di vettori di embedding nell’elenco di embedding che rappresenta il documento rientra nell’intervallo di $\R^{n×d}$.

• $[e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$q1,eq2,…,eqm$]\in$ R${}^{m\times d}$: il numero di vettori di embedding nell’elenco di embedding che rappresenta la query rientra nell’intervallo $\R^{m×d}$.

Interazione tardiva

Una volta completata la vettorizzazione, l’elenco di embedding della query viene confrontato con ciascun elenco di embedding dei documenti, token per token, per determinare il punteggio di similarità finale.

Come mostrato nel diagramma sopra, la query contiene due token, ovvero machine e learning, mentre il documento nella finestra presenta quattro token: neural, network, python e tutorial. Una volta vettorializzati questi token, le rappresentazioni vettoriali di ciascun token della query vengono confrontate con quelle del documento per ottenere un elenco di punteggi di similarità. Successivamente, i punteggi più alti di ciascun elenco vengono sommati per ottenere il punteggio finale. Il processo di determinazione del punteggio finale di un documento è noto come «massima somiglianza» (MAX_SIM). Per ulteriori dettagli sulla massima somiglianza, consultare la sezione «Massima somiglianza».

Estensione ColPali

Basandosi su ColBERT, ColPali (arXiv: 2407.01449) propone un approccio innovativo al recupero di documenti visivamente ricchi che sfrutta i modelli visione-linguaggio (VLM). Durante l’acquisizione dei dati, ogni pagina del documento viene renderizzata in un’immagine ad alta risoluzione, quindi suddivisa in patch, anziché essere tokenizzata. Ad esempio, l’immagine di una pagina di documento di 448 x 448 pixel può produrre 1.024 patch, ciascuna delle quali misura 14 x 14 pixel.

Questo metodo preserva le informazioni non testuali, quali l’impaginazione del documento, le immagini e le strutture delle tabelle, che vanno perse quando si utilizzano sistemi di recupero basati esclusivamente sul testo.

Il VLM utilizzato in ColPali si chiama PaliGemma (arXiv: 2407.07726), che comprende un codificatore di immagini (SigLIP-400M), un modello linguistico solo decodificatore (Gemma2-2B), e uno strato lineare che proietta l’output del codificatore di immagini nello spazio vettoriale del modello linguistico, come mostrato nel diagramma sopra.

Durante l’acquisizione dei dati, una pagina di documento, rappresentata come immagine grezza, viene suddivisa in più patch visive, ciascuna delle quali viene incorporata per generare un elenco di incorporamenti vettoriali. Successivamente, queste vengono proiettate nello spazio vettoriale del modello linguistico per ottenere l’elenco finale di rappresentazioni, come in $∈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}$. Quando arriva una query, questa viene tokenizzata e ogni token viene incorporato per generare un elenco di vettori di incorporamento, come in $∈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}$. Successivamente, è stato applicato MAX_SIM per confrontare i due elenchi di embedding e ottenere il punteggio finale tra la query e la pagina del documento.

Sistema di recupero testuale ColBERT

In questa sezione, configureremo un sistema di recupero testuale ColBERT utilizzando StructArray. Prima di ciò, è necessario configurare un'istanza di Milvus v2.6.x e ottenere un token di accesso Cohere.

Passaggio 1: Installare le dipendenze

Eseguire il seguente comando per installare le dipendenze.

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

Passaggio 2: Caricare il set di dati Cohere

In questo esempio utilizzeremo il set di dati Wikipedia di Cohere e recupereremo i primi 10.000 record. È possibile trovare informazioni su questo set di dati in questa pagina.

from datasets import load_dataset

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

L’esecuzione degli script sopra riportati scaricherà il set di dati se non è disponibile localmente. Ogni record del set di dati corrisponde a un paragrafo tratto da una pagina di Wikipedia. La tabella seguente illustra la struttura di questo set di dati.

Nome colonna	Descrizione
_id	ID del record
url	L'URL del record corrente.
title	Il titolo del documento di origine.
text	Un paragrafo tratto dal documento di origine.
emb	Incorporamenti del testo dal documento di origine.

Fase 3: raggruppare i paragrafi per titolo

Per cercare documenti anziché paragrafi, è necessario raggruppare i paragrafi per titolo.

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

In questo codice, memorizziamo i paragrafi raggruppati come documenti e li includiamo nell'elenco data. Ogni documento ha una chiave paragraphs, che è un elenco di paragrafi; ogni oggetto paragrafo contiene le chiavi text e emb.

Fase 4: Creazione di una collezione per il set di dati Cohere

Una volta che i dati sono pronti, creeremo una collezione. Nella collezione, ` paragraphs ` è un campo `StructArray`. Per una spiegazione generale degli schemi `StructArray`, consultare Creare 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
)

Passaggio 5: Inserire il set di dati Cohere nella raccolta

Ora possiamo inserire i dati preparati nella raccolta che abbiamo creato in precedenza.

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

Passaggio 6: Effettuare una ricerca all’interno del set di dati Cohere

In base alla progettazione di ColBERT, il testo della query deve essere tokenizzato e quindi incorporato in un EmbeddingList. In questo passaggio utilizzeremo lo stesso modello impiegato da Cohere per generare gli embedding dei paragrafi nel dataset di 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"],
)

Nel codice, i testi delle query vengono suddivisi in token in un query_inputs e e incorporati in un elenco di vettori float. È quindi possibile utilizzare l’EmbeddingList di Milvus per eseguire una ricerca per similarità come segue.

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

L’output del codice sopra riportato è simile al seguente:

# 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

Ogni punteggio di similarità coseno a coppie varia da -1 a 1. Il punteggio finale di " MAX_SIM_COSINE " può essere superiore a 1 poiché aggrega più punteggi di similarità massima a livello di token.

Sistema di recupero dei documenti ColPali

In questa sezione, configureremo un sistema di recupero dei documenti basato su ColPali utilizzando StructArray. Prima di ciò, configurare un’istanza di Milvus v2.6.x.

Passaggio 1: Installare le dipendenze

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

Passo 2: Caricare il dataset Vidore

In questa sezione utilizzeremo un dataset Vidore denominato vidore_v2_finance_en. Questo dataset è un corpus di relazioni annuali del settore bancario, destinato a compiti di comprensione di documenti lunghi. È uno dei 10 corpora che compongono il benchmark ViDoRe v3. È possibile trovare dettagli su questo dataset in questa pagina.

from datasets import load_dataset

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

L'esecuzione degli script sopra riportati scaricherà il set di dati se non è disponibile localmente. Ogni record del set di dati corrisponde a una pagina di una relazione finanziaria. La tabella seguente mostra la struttura di questo set di dati.

Nome colonna	Descrizione
corpus_id	Un record del corpus
image	L'immagine della pagina in byte.
doc_id	L'ID descrittivo del documento.
page_number_in_doc	Il numero della pagina corrente nel documento.

Fase 3: Generazione degli embedding per le immagini delle pagine

Come illustrato nella sezione Panoramica, il modello ColPali è un VLM che proietta le immagini nello spazio vettoriale di un modello testuale. In questo passaggio utilizzeremo l’ultima versione del modello ColPali, vidore/colpali-v1.3. È possibile trovare ulteriori dettagli su questo modello in questa pagina.

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 volta che il modello è pronto, è possibile provare a generare patch per un'immagine specifica come segue.

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]

Nel codice sopra riportato, il modello ColPali ridimensiona l’immagine a 448 x 448 pixel, quindi la suddivide in patch, ciascuna delle dimensioni di 14 x 14 pixel. Infine, queste patch vengono incorporate in 1.031 embedding, ciascuno con 128 dimensioni.

È possibile generare le rappresentazioni per tutte le immagini utilizzando un ciclo come segue:

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

Fase 4: Creazione di una raccolta per il set di dati dei rendiconti finanziari

Una volta che i dati sono pronti, creeremo una collezione. All’interno della collezione, ` patches ` è un campo `StructArray`. Ogni elemento di `Struct` memorizza un embedding di patch. Per i requisiti di indicizzazione dei sottocampi vettoriali di `StructArray`, consultare Indicizzazione dei campi `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
)

Passaggio 5: Inserimento dei report finanziari nella collezione

Ora possiamo inserire i rendiconti finanziari preparati nella collezione.

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

Dall’output è possibile verificare che tutte le pagine del set di dati Vidore siano state inserite.

Passaggio 6: Ricerca all’interno dei rapporti finanziari

Una volta che i dati sono pronti, è possibile effettuare ricerche sui dati presenti nella raccolta come segue:

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