Late Chunking für RAG: Umsetzung mit Jina AI

Naives Chunking und seine Grenzen in RAG

In einer RAG-Pipeline werden die Dokumente in kleinere Teile zerlegt, bevor sie eingebettet und in einer Vektor-Datenbank. Jeder Chunk wird unabhängig verarbeitet und bei Abfragen zum Abruf verwendet. Bei diesem "naiven Chunking"-Ansatz gehen jedoch oft wichtige Zusammenhänge aus der Ferne verloren.

Das Problem entsteht, weil das traditionelle Chunking Dokumente aufteilt, ohne zu berücksichtigen, wie die Informationen zusammenhängen. In einem Dokument über Paris kann zum Beispiel der Ausdruck "die Stadt" an einer anderen Stelle stehen als "Paris". Ohne den vollständigen Kontext hat das Retrieval-Modell möglicherweise Schwierigkeiten, diese Verweise zu verknüpfen, was zu weniger genauen Ergebnissen führt. Noch schlimmer ist dieses Problem bei langen Dokumenten, bei denen der wichtige Kontext über mehrere Abschnitte verteilt ist.

Late Chunking: Erhaltung des Kontexts bei der Aufteilung von Dokumenten

Late Chunking löst das Problem, indem es den Zeitpunkt ändert, zu dem du das Dokument aufteilst. Anstatt das Dokument zuerst in Chunks zu zerlegen, wird beim Late Chunking das gesamte Dokument mit einem Long-Context-Modell eingebettet. Erst danach wird das Dokument in kleinere Stücke aufgeteilt.

Das sind die wichtigsten Vorteile des Late Chunking:

• Behält den Kontext bei: Late Chunking stellt sicher, dass jeder Chunk den Gesamtkontext behält, indem das gesamte Dokument zuerst eingebettet wird. Auf diese Weise bleiben die Verweise und Verbindungen im Text in den Chunk-Embeddings erhalten.
• Besserer Abruf: Die durch Late Chunking erstellten Chunk-Embeddings sind reichhaltiger und genauer und verbessern die Retrieval-Ergebnisse in RAG-Systemen, weil das Modell das Dokument besser versteht.
• Bearbeitet lange Texte: Es eignet sich hervorragend für sehr lange Dokumente, die herkömmliche Modelle aufgrund von Token-Limits nicht in einem Durchgang bearbeiten können.

Mit Long-Context-Modellen wie Jina's jinaai/jina-embeddings-v2-base-en, das bis zu 8192 Token unterstützt, ermöglicht Late Chunking die effektive Einbettung großer Textabschnitte, bevor sie in Chunks aufgeteilt werden.

Implementierung von Late Chunking

Hier ist eine Schritt-für-Schritt-Anleitung, die dir hilft, Late Chunking mit Jina's Long-Context Embedding Model umzusetzen. Du kannst Jinas API-Schlüssel kostenlos erhalten hierund wir werden den folgenden Eingabetext als Demo verwenden:

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."""

Schritt 1: Chunks und Span-Anmerkungen erhalten

Verwende zunächst deinen Jina-API-Schlüssel und die unten stehende Hilfsfunktion, um deinen Eingabetext in Teile zu zerlegen. Diese Chunks werden mit Span-Anmerkungen versehen, die später bei der Aufteilung der Dokumenteinbettung helfen. Jina's API nutzt natürliche Grenzen wie Absatz- oder Satzumbrüche, um sicherzustellen, dass die Chunks einen Sinn ergeben und ihre Bedeutung behalten.

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)

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

Schritt 2: Tokenisierung des Textes und Erzeugung von Dokumenteneinbettungen auf Token-Ebene

Verwende zunächst einen Tokenizer, der mit Long-Context-Modellen kompatibel ist, wie z.B. Jina's embeddings-v2-base-en, um das gesamte Dokument in Token zu zerlegen. Als Nächstes verwendest du ein Long-Context-Transformer-Modell, um Einbettungen für jedes Token zu erstellen. Das bedeutet, dass jedes Wort oder Token in deinem Dokument eine einzigartige Einbettung erhält, die seine Bedeutung erfasst.

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].shape

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

Schritt 3: Spätes Chunking

Sobald du die Token-Einbettungen für das gesamte Dokument hast, bist du bereit für das Late Chunking. Verwende die Span-Annotationen aus Schritt eins, um diese Token-Einbettungen in kleinere Teile zu zerlegen. Dann wendest du das Mean Pooling an, um den Durchschnitt der Einbettungen innerhalb jedes Chunks zu ermitteln und eine einzige Einbettung für jeden Chunk zu erstellen. Wir haben jetzt Chunk Embeddings mit starken Kontextinformationen über das gesamte Dokument.

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 outputs

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

3 # matches number of chunks in Step 1

Schritt 4: Ergebnisse von Late Chunking vs. traditionellem Chunking

Um die Vorteile des Late Chunking zu verstehen, vergleichen wir es mit dem traditionellen Chunking:

embeddings_traditional_chunking = model.encode(chunks)

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

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
------------------------------------------------------------------

Wie du im zweiten und dritten Chunk sehen kannst, zeigt das traditionelle Chunking Ähnlichkeitswerte von 70-75% im Vergleich zu dem Wort "Berlin". Mit Late Chunking, bei dem der Kontext des gesamten Dokuments erhalten bleibt, steigen diese Werte jedoch auf 82-84%. Dies zeigt, dass beim Late Chunking der Kontext besser erhalten bleibt und aussagekräftigere Einbettungen entstehen, was zu genaueren Suchergebnissen führt.

Fazit

Late Chunking ist eine wichtige Verbesserung für Dokumenten-Retrieval-Systeme, insbesondere in RAG-Pipelines. Wenn du mit der Aufteilung des Dokuments wartest, bis es vollständig eingebettet ist, bleibt beim Late Chunking der gesamte Kontext in jedem Chunk erhalten. Dies führt zu einer genaueren und aussagekräftigeren Einbettung.