Late Chunking untuk RAG: Implementasi dengan Jina AI

Chunking Naif dan Keterbatasannya dalam RAG

Dalam pipeline RAG, dokumen dipecah menjadi potongan yang lebih kecil sebelum dibuat embedding dan disimpan dalam database vektor. Setiap potongan diproses secara independen dan digunakan untuk penelusuran saat kueri diajukan. Namun, pendekatan "chunking naif" ini sering kali kehilangan konteks jarak jauh yang penting.

Masalah muncul karena chunking tradisional membagi dokumen tanpa mempertimbangkan bagaimana informasi saling terhubung. Misalnya, dalam sebuah dokumen tentang Paris, frasa "kota tersebut" bisa berada di potongan yang berbeda dari tempat "Paris" disebutkan. Tanpa konteks penuh, model penelusuran dapat kesulitan mengaitkan rujukan ini, sehingga menghasilkan hasil yang kurang akurat. Masalah ini semakin parah pada dokumen panjang di mana konteks kunci tersebar di beberapa bagian.

Late Chunking: Menjaga Konteks dalam Pemecahan Dokumen

Late chunking menyelesaikan masalah dengan mengubah waktu Anda membagi dokumen. Alih-alih memecah dokumen menjadi potongan terlebih dahulu, late chunking membuat embedding seluruh dokumen menggunakan model ber-konteks panjang. Setelah itu barulah dokumen dipecah menjadi potongan yang lebih kecil.

Berikut manfaat utama late chunking:

• Menjaga konteks: Late chunking memastikan setiap potongan mempertahankan konteks keseluruhan dengan terlebih dahulu membuat embedding seluruh dokumen. Dengan begitu, rujukan dan keterkaitan sepanjang teks tetap utuh dalam embedding potongan.
• Penelusuran lebih baik: Embedding potongan yang dihasilkan melalui late chunking lebih kaya dan akurat, sehingga meningkatkan hasil penelusuran dalam sistem RAG karena model lebih memahami dokumen.
• Menangani teks panjang: Sangat cocok untuk dokumen yang sangat panjang yang tidak dapat ditangani model tradisional sekaligus karena batas token.

Dengan menggunakan model ber-konteks panjang seperti jinaai/jina-embeddings-v2-base-en dari Jina, yang mendukung hingga 8192 token, late chunking memungkinkan pembuatan embedding bagian teks besar secara efektif sebelum membaginya menjadi potongan.

Mengimplementasikan Late Chunking

Berikut panduan langkah demi langkah untuk membantu Anda mengimplementasikan late chunking menggunakan model embedding ber-konteks panjang dari Jina. Anda dapat memperoleh API key Jina secara gratis di sini, dan kita akan menggunakan teks masukan berikut sebagai 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."""

Langkah 1: Dapatkan potongan dan anotasi rentang

Pertama, gunakan API key Jina Anda dan fungsi pembantu di bawah ini untuk memecah teks masukan menjadi potongan. Potongan ini disertai anotasi rentang yang membantu membagi embedding dokumen nanti. API Jina menggunakan batas alami seperti jeda paragraf atau kalimat untuk memastikan potongan masuk akal dan mempertahankan maknanya.

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

Langkah 2: Tokenisasi teks dan buat embedding dokumen pada level token

Pertama, gunakan tokenizer yang kompatibel dengan model ber-konteks panjang, seperti embeddings-v2-base-en dari Jina, untuk memecah seluruh dokumen menjadi token. Selanjutnya, gunakan model transformer ber-konteks panjang untuk membuat embedding bagi setiap token. Artinya, setiap kata atau token dalam dokumen Anda memperoleh embedding unik yang menangkap maknanya.

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

Langkah 3: Late chunking

Setelah Anda memiliki embedding token untuk seluruh dokumen, Anda siap melakukan late chunking. Gunakan anotasi rentang dari langkah pertama untuk membagi embedding token ini menjadi potongan yang lebih kecil. Kemudian, terapkan mean pooling untuk merata-ratakan embedding dalam setiap potongan, sehingga menghasilkan satu embedding untuk tiap potongan. Kini kita memiliki embedding potongan dengan informasi kontekstual yang kuat tentang seluruh dokumen.

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

Langkah 4: Hasil late chunking vs. chunking tradisional

Untuk memahami manfaat late chunking, mari bandingkan dengan chunking tradisional:

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

Seperti terlihat pada potongan kedua dan ketiga, chunking tradisional menunjukkan skor kemiripan 70–75% saat dibandingkan dengan kata "Berlin." Namun, dengan late chunking, yang mempertahankan konteks seluruh dokumen, skor ini naik menjadi 82–84%. Ini menunjukkan bahwa late chunking lebih baik dalam menjaga konteks dan menghasilkan embedding yang lebih bermakna, sehingga menghasilkan hasil pencarian yang lebih akurat.

Kesimpulan

Late chunking merupakan peningkatan besar untuk sistem penelusuran dokumen, khususnya dalam pipeline RAG. Dengan menunda pemecahan dokumen hingga setelah seluruhnya dibuat embedding, late chunking menjaga konteks penuh di setiap potongan. Hasilnya adalah embedding yang lebih akurat dan bermakna.