Late Chunking cho RAG: Triển khai với Jina AI

Chunking ngây thơ và những hạn chế trong RAG

Trong một pipeline RAG, tài liệu được chia thành các đoạn nhỏ trước khi được nhúng và lưu trong cơ sở dữ liệu vector. Mỗi đoạn được xử lý độc lập và dùng để truy xuất khi có truy vấn. Tuy nhiên, cách làm "chunking ngây thơ" này thường làm mất đi ngữ cảnh tầm xa quan trọng.

Vấn đề nảy sinh vì chunking truyền thống tách tài liệu mà không xét đến sự liên kết thông tin. Ví dụ, trong một tài liệu về Paris, cụm "thành phố" có thể nằm ở một đoạn khác với nơi "Paris" được nhắc đến. Thiếu ngữ cảnh đầy đủ, mô hình truy xuất có thể khó liên kết các tham chiếu này, dẫn tới kết quả kém chính xác. Điều này còn tệ hơn với các tài liệu dài, nơi ngữ cảnh quan trọng trải rộng trên nhiều phần.

Late Chunking: Giữ nguyên ngữ cảnh khi tách tài liệu

Late chunking giải quyết vấn đề bằng cách thay đổi thời điểm bạn tách tài liệu. Thay vì tách tài liệu thành đoạn trước, late chunking sẽ nhúng toàn bộ tài liệu bằng một mô hình ngữ cảnh dài. Chỉ sau đó mới tách tài liệu thành các đoạn nhỏ hơn.

Đây là những lợi ích chính của late chunking:

• Giữ ngữ cảnh: Late chunking đảm bảo mỗi đoạn vẫn giữ được ngữ cảnh tổng thể bằng cách nhúng toàn bộ tài liệu trước. Nhờ vậy, các tham chiếu và liên kết xuyên suốt văn bản vẫn được bảo toàn trong embedding của từng đoạn.
• Truy xuất tốt hơn: Embedding của các đoạn được tạo ra qua late chunking phong phú và chính xác hơn, cải thiện kết quả truy xuất trong hệ thống RAG vì mô hình hiểu tài liệu tốt hơn.
• Xử lý văn bản dài: Rất phù hợp cho các tài liệu rất dài mà mô hình truyền thống không thể xử lý một lần do giới hạn token.

Sử dụng các mô hình ngữ cảnh dài như jinaai/jina-embeddings-v2-base-en của Jina, hỗ trợ tới 8192 token, late chunking cho phép nhúng hiệu quả các đoạn văn bản lớn trước khi chia chúng thành các đoạn nhỏ.

Triển khai Late Chunking

Dưới đây là hướng dẫn từng bước giúp bạn triển khai late chunking bằng mô hình embedding ngữ cảnh dài của Jina. Bạn có thể lấy API key của Jina miễn phí tại đây, và chúng ta sẽ dùng đoạn văn bản đầu vào sau để 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."""

Bước 1: Lấy các đoạn và chú giải span

Trước tiên, dùng Jina API key của bạn và hàm trợ giúp dưới đây để tách văn bản đầu vào thành các đoạn. Những đoạn này kèm theo chú giải span để hỗ trợ việc tách embedding tài liệu ở bước sau. API của Jina sử dụng các ranh giới tự nhiên như ngắt đoạn hoặc ngắt câu để đảm bảo các đoạn có ý nghĩa và giữ được ngữ nghĩa.

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

Bước 2: Tokenize văn bản và tạo embedding cấp độ token cho tài liệu

Trước tiên, sử dụng một tokenizer tương thích với các mô hình ngữ cảnh dài, như embeddings-v2-base-en của Jina, để phân tách toàn bộ tài liệu thành các token. Tiếp theo, dùng một mô hình transformer ngữ cảnh dài để tạo embedding cho từng token. Điều này có nghĩa là mỗi từ hoặc token trong tài liệu của bạn sẽ có embedding riêng, nắm bắt ý nghĩa của nó.

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

Bước 3: Late chunking

Khi đã có embedding cấp token cho toàn bộ tài liệu, bạn sẵn sàng thực hiện late chunking. Dùng các chú giải span từ bước một để tách các embedding token này thành các đoạn nhỏ hơn. Sau đó, áp dụng mean pooling để trung bình các embedding trong mỗi đoạn, tạo ra một embedding duy nhất cho từng đoạn. Giờ đây, chúng ta có embedding của các đoạn với thông tin ngữ cảnh mạnh mẽ về toàn bộ tài liệu.

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

Bước 4: So sánh late chunking và chunking truyền thống

Để hiểu lợi ích của late chunking, hãy so sánh với chunking truyền thống:

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

Như bạn thấy ở đoạn thứ hai và thứ ba, chunking truyền thống cho điểm tương đồng khoảng 70–75% khi so với từ "Berlin". Tuy nhiên, với late chunking, vốn duy trì ngữ cảnh của toàn bộ tài liệu, các điểm số này tăng lên 82–84%. Điều này cho thấy late chunking làm tốt hơn trong việc bảo toàn ngữ cảnh và tạo embedding có ý nghĩa hơn, dẫn tới kết quả tìm kiếm chính xác hơn.

Kết luận

Late chunking là một cải tiến lớn cho các hệ thống truy xuất tài liệu, đặc biệt trong pipeline RAG. Bằng cách trì hoãn việc tách tài liệu cho đến sau khi đã nhúng toàn bộ, late chunking giữ được ngữ cảnh đầy đủ trong mỗi đoạn. Kết quả là các embedding chính xác và có ý nghĩa hơn.