Late Chunking per RAG: implementazione con Jina AI

Chunking ingenuo e i suoi limiti in RAG

In una pipeline RAG, i documenti vengono suddivisi in chunk più piccoli prima di essere embedded e archiviati in un database vettoriale. Ogni chunk è elaborato in modo indipendente e usato per il recupero quando vengono effettuate query. Tuttavia, questo approccio di "chunking ingenuo" spesso perde un importante contesto a lungo raggio.

Il problema nasce perché il chunking tradizionale divide i documenti senza considerare come le informazioni siano collegate. Ad esempio, in un documento su Parigi, l’espressione "la città" potrebbe finire in un chunk diverso rispetto a dove viene menzionata "Parigi". Senza il contesto completo, il modello di retrieval può faticare a collegare questi riferimenti, producendo risultati meno accurati. Questo problema è ancora più grave nei documenti lunghi in cui il contesto chiave è distribuito su più sezioni.

Late chunking: preservare il contesto nella suddivisione dei documenti

Il late chunking risolve il problema cambiando il momento in cui suddividi il documento. Invece di spezzare subito il documento in chunk, il late chunking esegue l’embedding dell’intero documento usando un modello a lungo contesto. Solo dopo lo divide in chunk più piccoli.

Questi sono i principali vantaggi del late chunking:

• Mantiene il contesto: il late chunking fa sì che ogni chunk conservi il contesto generale eseguendo prima l’embedding dell’intero documento. In questo modo, riferimenti e connessioni lungo il testo restano intatti negli embedding dei chunk.
• Recupero migliore: gli embedding dei chunk creati con il late chunking sono più ricchi e accurati, migliorando i risultati di retrieval nei sistemi RAG perché il modello comprende meglio il documento.
• Gestisce testi lunghi: è ideale per documenti molto lunghi che i modelli tradizionali non riescono a gestire in una sola passata a causa dei limiti di token.

Usando modelli a lungo contesto come il jinaai/jina-embeddings-v2-base-en di Jina, che supporta fino a 8192 token, il late chunking consente di eseguire efficacemente l’embedding di ampie sezioni di testo prima di suddividerle in chunk.

Implementare il late chunking

Ecco una guida passo passo per implementare il late chunking usando il modello di embedding a lungo contesto di Jina. Puoi ottenere gratuitamente la chiave API di Jina qui e useremo il seguente testo di input come demo:

input_text = """Berlin is the capital and largest city of Germany, both by area and by population.
Its more than 3.85 million inhabitants make it the European Union's most populous city, as measured by population within city limits.
The city is also one of the states of Germany, and is the third smallest state in the country in terms of area."""Offerto da

Passaggio 1: ottieni i chunk e le annotazioni di intervallo

Per prima cosa, usa la tua chiave API di Jina e la funzione di supporto qui sotto per suddividere il testo di input in chunk. Questi chunk includono annotazioni di intervallo che aiuteranno a suddividere l’embedding del documento più avanti. L’API di Jina usa confini naturali come interruzioni di paragrafo o di frase per assicurare che i chunk abbiano senso e mantengano il loro significato.

import json
import requests

def custom_tokenize_jina_api(input_text: str):
    url = '<https://segment.jina.ai/>'
    headers = {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer ENTER_YOUR_JINA_API_KEY'
    }
    data = {
        "content": input_text,
        "tokenizer": "o200k_base",
        "return_tokens": "true",
        "return_chunks": "true",
        "max_chunk_length": "1000"
    }
    # Make the API request
    response = requests.post(url, headers=headers, json=data)
    response_data = response.json()
    chunks = response_data.get("chunks", [])
    i = 1
    j = 1
    span_annotations = []
    for x in response_data['tokens']:
        if j == 1:
            j = len(x)
        else:
            j = len(x) + i
        span_annotations.append((i, j))
        i = j
    return chunks, span_annotations
chunks, span_annotations = custom_tokenize_jina_api(input_text)

print(chunks)
print(span_annotations)Offerto da

['Berlin is the capital and largest city of Germany, both by area and by population.\\n\\n', "Its more than 3.85 million inhabitants make it the European Union's most populous city, as measured by population within city limits.\\n\\n", 'The city is also one of the states of Germany, and is the third smallest state in the country in terms of area.']
[(1, 17), (17, 44), (44, 69)]Offerto da

Passaggio 2: tokenizza il testo e genera embedding a livello di token

Per prima cosa, usa un tokenizer compatibile con modelli a lungo contesto, come l’embeddings-v2-base-en di Jina, per suddividere l’intero documento in token. Poi usa un modello transformer a lungo contesto per creare gli embedding per ciascun token. In pratica, ogni parola o token del documento ottiene un embedding specifico che ne cattura il significato.

from transformers import AutoModel
from transformers import AutoTokenizer

# load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)
inputs = tokenizer(input_text, return_tensors='pt')
model_output = model(**inputs)
model_output[0].shapeOfferto da

torch.Size([1, 71, 768]) # 71 represents number of tokens in the entire documentOfferto da

Passaggio 3: late chunking

Una volta ottenuti gli embedding dei token per l’intero documento, puoi eseguire il late chunking. Usa le annotazioni di intervallo del primo passaggio per suddividere questi embedding di token in chunk più piccoli. Quindi applica il mean pooling per fare la media degli embedding all’interno di ciascun chunk, creando un embedding unico per ogni chunk. Ora abbiamo embedding di chunk con una forte informazione contestuale sull’intero documento.

def late_chunking(
    model_output: 'BatchEncoding', span_annotation: list, max_length=None
):
    token_embeddings = model_output[0]
    outputs = []
    for embeddings, annotations in zip(token_embeddings, span_annotation):
        if (
            max_length is not None
        ):  # remove annotations which go bejond the max-length of the model
            annotations = [
                (start, min(end, max_length - 1))
                for (start, end) in annotations
                if start < (max_length - 1)
            ]
        pooled_embeddings = [
            embeddings[start:end].sum(dim=0) / (end - start)
            for start, end in annotations
            if (end - start) >= 1
        ]
        pooled_embeddings = [
            embedding.detach().cpu().numpy() for embedding in pooled_embeddings
        ]
        outputs.append(pooled_embeddings)
    return outputsOfferto da

embeddings = late_chunking(model_output, [span_annotations])[0]
len(embeddings)Offerto da

3 # matches number of chunks in Step 1Offerto da

Passaggio 4: risultati del late chunking vs. chunking tradizionale

Per capire i vantaggi del late chunking, confrontiamolo con il chunking tradizionale:

embeddings_traditional_chunking = model.encode(chunks)Offerto da

import numpy as np

cos_sim = lambda x, y: np.dot(x, y) / (np.linalg.norm(x) * np.linalg.norm(y))
q = "Berlin"
berlin_embedding = model.encode(q)

print(q)
print('\n')
for chunk, new_embedding, trad_embeddings in zip(chunks, embeddings, embeddings_traditional_chunking):
  print(chunk.strip())
  print(f'Late chunking:', cos_sim(berlin_embedding, new_embedding))
  print(f'Traditional chunking:', cos_sim(berlin_embedding, trad_embeddings))
  print("------------------------------------------------------------------")Offerto da

Berlin
Berlin is the capital and largest city of Germany, both by area and by population.
Late chunking: 0.84954596
Traditional chunking: 0.84862185
------------------------------------------------------------------
Its more than 3.85 million inhabitants make it the European Union's most populous city, as measured by population within city limits.
Late chunking: 0.82489026
Traditional chunking: 0.70843375
------------------------------------------------------------------
The city is also one of the states of Germany, and is the third smallest state in the country in terms of area.
Late chunking: 0.84980094
Traditional chunking: 0.7534553
------------------------------------------------------------------Offerto da

Come puoi vedere nel secondo e terzo chunk, il chunking tradizionale mostra punteggi di similarità del 70-75% rispetto alla parola "Berlin". Con il late chunking, che mantiene il contesto dell’intero documento, questi punteggi salgono all’82-84%. Ciò dimostra che il late chunking preserva meglio il contesto e crea embedding più significativi, con risultati di ricerca più accurati.

Conclusione

Il late chunking è un grande passo avanti per i sistemi di document retrieval, soprattutto nelle pipeline RAG. Aspettando a suddividere il documento finché non è stato completamente embedded, il late chunking mantiene il contesto completo in ogni chunk. Il risultato sono embedding più accurati e significativi.