Ricerca con elenchi incorporati

Questa pagina spiega come impostare un sistema di ricerca testuale ColBERT e un sistema di ricerca testuale ColPali utilizzando l'array di strutture di Milvus, che consente di memorizzare un documento insieme ai suoi chunk vettoriali in liste di incorporamento.

Panoramica

Per costruire un sistema di recupero del testo, potrebbe essere necessario dividere i documenti in pezzi e memorizzare ogni pezzo insieme ai suoi embedding come entità in un database vettoriale per garantire precisione e accuratezza, soprattutto per i documenti lunghi in cui gli embedding a testo pieno potrebbero diluire la specificità semantica o superare i limiti di input del modello.

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

ColBERT (arXiv: 2004.12832) è un sistema di recupero di testi che offre una ricerca efficiente ed efficace dei passaggi attraverso interazioni tardive contestualizzate su BERT. Consente la codifica token-wise indipendente di query e documenti e ne calcola la similarità.

Codifica per token

Durante l'ingestione dei dati in ColBERT, ogni documento viene suddiviso in token, che vengono poi vettorializzati e memorizzati come elenco di incorporazioni, come in $\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 arriva una query, viene anche tokenizzata, vettorializzata e memorizzata come elenco di incorporazioni, come in $\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}$.

Nelle formule precedenti,

• $dd$ d: un documento

• $qq$ q: l'interrogazione

• $EdE_d$ Ed: la lista di incorporamento che rappresenta il documento.

• $EqE_q$ Eq: l'elenco di incorporazioni che rappresenta la query.

• $[e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$d1, ed2,$...$, edn]$\in$ R${}^{n\times d}$: il numero di incorporazioni vettoriali nell'elenco di incorporazioni che rappresenta il documento è compreso nell'intervallo $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}$: il numero di incorporazioni vettoriali nell'elenco di incorporazioni che rappresentano la query è compreso nell'intervallo $Rm×d\R^{m×d}$ R ${}^{m\times d}$.

Interazione tardiva

Una volta completata la vettorizzazione, l'elenco di incorporazioni della query viene confrontato con ogni elenco di incorporazioni del documento, token per token, per determinare il punteggio di similarità finale.

Interazione tardiva

Come mostrato nel diagramma precedente, la query contiene due token, ovvero machine e learning, mentre il documento nella finestra ha quattro token: neural, network, python, e tutorial. Una volta vettorializzati questi token, le incorporazioni vettoriali di ciascun token della query vengono confrontate con quelle del documento per ottenere un elenco di punteggi di somiglianza. I punteggi più alti di ciascun elenco vengono poi sommati per ottenere il punteggio finale. Il processo per determinare il punteggio finale di un documento è noto come somiglianza massima(MAX_SIM). Per maggiori dettagli sulla massima somiglianza, consultare la sezione Massima somiglianza.

Estensione ColPali

Basato su ColBERT, ColPali (arXiv: 2407.01449) propone un nuovo approccio al reperimento di documenti visivamente ricchi che sfrutta i Vision-Language Models (VLM). Durante l'ingestione dei dati, ogni pagina di documento viene renderizzata in un'immagine ad alta risoluzione e poi suddivisa in patch, piuttosto che tokenizzata. Ad esempio, un'immagine di 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, come il layout del documento, le immagini e le strutture delle tabelle, che vanno perse quando si utilizzano sistemi di recupero di solo testo.

Estensione Copali

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

Durante l'ingestione dei dati, una pagina di documento, rappresentata come immagine grezza, viene suddivisa in più patch visive, ognuna delle quali viene incorporata per generare un elenco di incorporazioni vettoriali. Quindi vengono proiettati nello spazio vettoriale del modello linguistico per ottenere l'elenco finale di incorporazioni, come in $\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 arriva una query, questa viene tokenizzata e ogni token viene incorporato per generare un elenco di incorporazioni vettoriali, come in $\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}$. Poi è stato applicato MAX_SIM per confrontare le due liste di incorporazioni e ottenere il punteggio finale tra la query e la pagina del documento.

Sistema di recupero del testo ColBERT

In questa sezione, verrà creato un sistema di reperimento di testi ColBERT utilizzando l'Array of Structs di Milvus. Prima di tutto, è necessario configurare un'istanza Milvus v2.6.x, un cluster Milvus Cloud compatibile con Milvus v2.6.x e ottenere un token di accesso Cohere.

Passo 1: installare le dipendenze

Eseguire il seguente comando per installare le dipendenze.

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

Passo 2: caricare il set di dati Cohere

In questo esempio, utilizzeremo il dataset di 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 di cui sopra scaricherà il dataset se non è disponibile localmente. Ogni record del dataset è un paragrafo di una pagina di Wikipedia. La tabella seguente mostra la struttura di questo set di dati.

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

Passo 3: raggruppare i paragrafi per titolo

Per cercare i documenti piuttosto che i paragrafi, dobbiamo 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.

Passo 4: Creare una raccolta per il set di dati Cohere

Una volta che i dati sono pronti, creeremo una collezione. Nella collezione, c'è un campo chiamato paragraphs, che è un array di strutture.

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
)

Passo 5: Inserire il set di dati Cohere nella raccolta

Ora possiamo inserire i dati preparati nella raccolta creata in precedenza.

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

Passo 6: Ricerca all'interno del set di dati Cohere

Secondo il progetto di ColBERT, il testo della query deve essere tokenizzato e poi incorporato in una EmbeddingList. In questo passo, utilizzeremo lo stesso modello che Cohere ha usato per generare embeddings per i paragrafi del 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 sono organizzati in token in query_inputs e incorporati in un elenco di vettori float. Poi si può usare EmbeddingList di Milvus per condurre una ricerca di somiglianza 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]",
    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}")

L'output del codice precedente è 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

Il punteggio di somiglianza del coseno va da -1 a 1, e i punteggi di somiglianza nell'output precedente mostrano chiaramente la somma di più punteggi di somiglianza a livello di token.

Sistema di recupero del testo ColPali

In questa sezione, verrà creato un sistema di recupero del testo basato su ColPali, utilizzando l'Array of Structs di Milvus. Prima di tutto, è necessario configurare un'istanza di Milvus v2.6.x e un cluster Milvus Cloud compatibile con Milvus v2.6.x.

Passo 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 chiamato vidore_v2_finance_it. Questo dataset è un corpus di rapporti annuali del settore bancario, destinato a compiti di comprensione di documenti lunghi. È uno dei 10 corpora che compongono il ViDoRe v3 Benchmark. I dettagli su questo set di dati sono disponibili a 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 di cui sopra scaricherà il dataset se non è disponibile localmente. Ogni record del dataset è una pagina di un rapporto finanziario. La tabella seguente mostra la struttura del dataset.

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 di pagina della pagina corrente del documento.

Fase 3: Generare le incorporazioni 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 di testo. In questa fase utilizzeremo l'ultimo modello ColPali vidore/colpali-v1.3. I dettagli su questo modello sono disponibili a 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, si può provare a generare patch per un'immagine specifica come segue.

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]

Nel codice qui sopra, il modello ColPali ridimensiona l'immagine a 448 x 448 pixel, quindi la divide in patch, ciascuna delle quali misura 14 x 14 pixel. Infine, queste patch vengono incorporate in 1.031 embeddings, ciascuno con 128 dimensioni.

È possibile generare embeddings per tutte le immagini utilizzando un ciclo come 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']
  })

Fase 4: Creare una raccolta per il set di dati dei rapporti finanziari

Una volta che i dati sono pronti, creeremo una raccolta. Nella raccolta, un campo chiamato patches è un array di strutture.

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
)

Passo 5: inserire i rapporti finanziari nella collezione

Ora possiamo inserire i rapporti finanziari preparati nella raccolta.

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

Dall'output, si può notare che tutte le pagine del dataset Vidore sono state inserite.

Fase 6: Ricerca all'interno dei rapporti finanziari

Una volta che i dati sono pronti, è possibile effettuare ricerche sui dati della 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].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"]
)