RAG के लिए लेट चंकिंग: Jina AI के साथ इम्प्लीमेंटेशन

साधारण (Naive) चंकिंग और RAG में उसकी सीमाएँ

RAG पाइपलाइन में, दस्तावेज़ों को छोटे-छोटे चंक्स में तोड़ा जाता है, फिर उन्हें एम्बेड कर वेक्टर डेटाबेस में रखा जाता है। हर chunk को स्वतंत्र रूप से प्रोसेस किया जाता है और क्वेरी पर रिट्रीवल के लिए इस्तेमाल होता है। लेकिन यह "साधारण चंकिंग" तरीका अक्सर लंबी दूरी वाले संदर्भ को खो देता है।

समस्या इसलिए आती है क्योंकि पारंपरिक चंकिंग यह नहीं देखती कि जानकारी आपस में कैसे जुड़ी हुई है। उदाहरण के लिए, पेरिस पर एक दस्तावेज़ में "the city" जैसा वाक्यांश उस chunk में जा सकता है जो "Paris" के उल्लेख वाले chunk से अलग हो। पूरे संदर्भ के बिना, रिट्रीवल मॉडल इन संदर्भों को जोड़ने में संघर्ष कर सकता है, जिससे कम सटीक परिणाम मिलते हैं। यह समस्या लंबे दस्तावेज़ों में और गंभीर हो जाती है, जहाँ मुख्य संदर्भ कई अनुभागों में फैला होता है।

लेट चंकिंग: दस्तावेज़ विभाजन में संदर्भ सुरक्षित रखना

लेट चंकिंग समस्या को इस तरह हल करती है कि आप दस्तावेज़ को कब विभाजित करते हैं, यह बदल देती है। दस्तावेज़ को पहले chunks में तोड़ने के बजाय, लेट चंकिंग लंबे-संदर्भ मॉडल से पूरे दस्तावेज़ को एम्बेड करती है। इसके बाद ही इसे छोटे चंक्स में बांटा जाता है।

लेट चंकिंग के मुख्य लाभ:

• संदर्भ बना रहता है: पहले पूरे दस्तावेज़ को एम्बेड करने से हर chunk में समग्र संदर्भ सुरक्षित रहता है। इससे पाठ में फैले संदर्भ और कनेक्शन chunk embeddings में अक्षुण्ण रहते हैं।
• बेहतर रिट्रीवल: लेट चंकिंग से बने chunk embeddings अधिक समृद्ध और सटीक होते हैं, जिससे RAG सिस्टम में रिट्रीवल के परिणाम सुधरते हैं क्योंकि मॉडल दस्तावेज़ को बेहतर समझता है।
• लंबे पाठ को संभालता है: यह बहुत लंबे दस्तावेज़ों के लिए उपयोगी है जिन्हें पारंपरिक मॉडल टोकन सीमाओं के कारण एक बार में संभाल नहीं पाते।

Jina के jinaai/jina-embeddings-v2-base-en जैसे लंबे-संदर्भ मॉडल, जो 8192 टोकन तक सपोर्ट करता है, का उपयोग करके लेट चंकिंग बड़े पाठ खंडों को पहले प्रभावी ढंग से एम्बेड करने और फिर उन्हें चंक्स में बाँटने की सुविधा देती है।

लेट चंकिंग को लागू करना

यहाँ Jina के लंबे-संदर्भ embedding मॉडल का उपयोग करके लेट चंकिंग लागू करने की चरण-दर-चरण गाइड दी गई है। आप Jina की API कुंजी निःशुल्क यहाँ से प्राप्त कर सकते हैं, और हम निम्न इनपुट पाठ को डेमो के रूप में उपयोग करेंगे:

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

चरण 1: चंक्स और स्पैन एनोटेशन प्राप्त करें

सबसे पहले, अपनी Jina API कुंजी और नीचे दिए गए हेल्पर फ़ंक्शन का उपयोग करके इनपुट टेक्स्ट को चंक्स में तोड़ें। इन चंक्स के साथ स्पैन एनोटेशन मिलते हैं, जो आगे चलकर दस्तावेज़ embedding को विभाजित करने में मदद करेंगे। Jina की API पैराग्राफ़ या वाक्य जैसी प्राकृतिक सीमाओं का उपयोग करती है ताकि चंक्स सार्थक रहें और अपना अर्थ बनाए रखें।

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

चरण 2: टेक्स्ट को टोकनाइज़ करें और टोकन-स्तरीय दस्तावेज़ embeddings बनाएँ

सबसे पहले, Jina के embeddings-v2-base-en जैसे लंबे-संदर्भ मॉडल के अनुकूल tokenizer का उपयोग करके पूरे दस्तावेज़ को टोकन्स में तोड़ें। फिर, लंबे-संदर्भ ट्रांसफॉर्मर मॉडल से हर टोकन के लिए embeddings बनाएँ। इसका मतलब है कि आपके दस्तावेज़ के हर शब्द या टोकन को उसका एक यूनिक embedding मिलेगा जो उसके अर्थ को कैप्चर करता है।

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

चरण 3: लेट चंकिंग

जब पूरे दस्तावेज़ के टोकन embeddings तैयार हो जाएँ, तब आप लेट चंकिंग के लिए तैयार हैं। चरण एक से मिले स्पैन एनोटेशन का उपयोग करके इन टोकन embeddings को छोटे चंक्स में विभाजित करें। फिर, प्रत्येक chunk के भीतर embeddings का औसत लेने के लिए मीन पूलिंग लागू करें, ताकि हर chunk का एकल embedding बने। अब हमारे पास ऐसे chunk embeddings हैं जिनमें पूरे दस्तावेज़ का प्रबल संदर्भगत जानकारी समाहित है।

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

चरण 4: लेट चंकिंग बनाम पारंपरिक चंकिंग के परिणाम

लेट चंकिंग के फायदों को समझने के लिए इसे पारंपरिक चंकिंग से तुलना करते हैं:

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

जैसा कि आप दूसरे और तीसरे चंक में देख सकते हैं, पारंपरिक चंकिंग में "Berlin" शब्द की तुलना में समानता स्कोर 70–75% के बीच हैं। वहीं, पूरे दस्तावेज़ के संदर्भ को बनाए रखने वाली लेट चंकिंग के साथ ये स्कोर 82–84% तक बढ़ जाते हैं। यह दर्शाता है कि लेट चंकिंग संदर्भ को बेहतर तरीके से सुरक्षित रखती है और अधिक सार्थक embeddings बनाती है, जिसके परिणामस्वरूप खोज परिणाम अधिक सटीक होते हैं।

निष्कर्ष

लेट चंकिंग दस्तावेज़ रिट्रीवल सिस्टम, विशेषकर RAG पाइपलाइनों, के लिए एक बड़ा सुधार है। दस्तावेज़ को पूरी तरह एम्बेड करने के बाद ही उसे विभाजित करने से, लेट चंकिंग हर chunk में पूर्ण संदर्भ बनाए रखती है। इसका नतीजा है अधिक सटीक और सार्थक embeddings।