How to Choose the Right Chunking Strategy for Your LLM Application

1
! pip install -qU langchain langchain-openai langchain-mongodb langchain-experimental ragas pymongo tqdm

1
import getpass
2
MONGODB_URI = getpass.getpass("Enter your MongoDB connection string:")

1
import os
2
from openai import OpenAI
3
os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter your OpenAI API Key:")
4
openai_client = OpenAI()

1
from langchain_community.document_loaders import WebBaseLoader
2
3
web_loader = WebBaseLoader(
4
    [
5
        "https://peps.python.org/pep-0483/",
6
        "https://peps.python.org/pep-0008/",
7
        "https://peps.python.org/pep-0257/",
8
    ]
9
)
10
11
pages = web_loader.load()

1
Document(
2
    page_content="\nThe goal of this PEP is to propose such a systematic way of defining types\nfor type annotations of variables and functions using PEP 3107 syntax.",
3
    metadata= {'source': 'https://peps.python.org/pep-0483/', 'title': 'PEP 483 – The Theory of Type Hints | peps.python.org', 'description': 'Python Enhancement Proposals (PEPs)', 'language': 'en'}
4
)

1
from langchain.text_splitter import TokenTextSplitter
2
3
def fixed_token_split(
4
    docs: List[Document], chunk_size: int, chunk_overlap: int
5
) -> List[Document]:
6
    """
7
    Fixed token chunking
8
9
    Args:
10
        docs (List[Document]): List of documents to chunk
11
        chunk_size (int): Chunk size (number of tokens)
12
        chunk_overlap (int): Token overlap between chunks
13
14
    Returns:
15
        List[Document]: List of chunked documents
16
    """
17
    splitter = TokenTextSplitter(
18
        encoding_name="cl100k_base", chunk_size=chunk_size, chunk_overlap=chunk_overlap
19
    )
20
    return splitter.split_documents(docs)

1
from langchain.text_splitter import (
2
    Language,
3
    RecursiveCharacterTextSplitter,
4
)
5
6
def recursive_split(
7
    docs: List[Document],
8
    chunk_size: int,
9
    chunk_overlap: int,
10
    language: Optional[Language] = None,
11
) -> List[Document]:
12
    """
13
    Recursive chunking
14
15
    Args:
16
        docs (List[Document]): List of documents to chunk
17
        chunk_size (int): Chunk size (number of tokens)
18
        chunk_overlap (int): Token overlap between chunks
19
        language (Optional[Language], optional): Programming language enum. Defaults to None.
20
21
    Returns:
22
        List[Document]: List of chunked documents
23
    """
24
    separators = ["\n\n", "\n", " ", ""]
25
26
    if language is not None:
27
        try:
28
            separators = RecursiveCharacterTextSplitter.get_separators_for_language(
29
                language
30
            )
31
        except (NameError, ValueError) as e:
32
            print(f"No separators found for language {language}. Using defaults.")
33
34
    splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
35
        encoding_name="cl100k_base",
36
        chunk_size=chunk_size,
37
        chunk_overlap=chunk_overlap,
38
        separators=separators,
39
    )
40
    return splitter.split_documents(docs)

1
from langchain_experimental.text_splitter import SemanticChunker
2
from langchain_openai.embeddings import OpenAIEmbeddings
3
4
def semantic_split(docs: List[Document]) -> List[Document]:
5
    """
6
    Semantic chunking
7
8
    Args:
9
        docs (List[Document]): List of documents to chunk
10
11
    Returns:
12
        List[Document]: List of chunked documents
13
    """
14
    splitter = SemanticChunker(
15
        OpenAIEmbeddings(), breakpoint_threshold_type="percentile"
16
    )
17
    return splitter.split_documents(docs)

1
from ragas import RunConfig
2
from ragas.testset.generator import TestsetGenerator
3
from ragas.testset.evolutions import simple, reasoning, multi_context
4
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
5
6
RUN_CONFIG = RunConfig(max_workers=4, max_wait=180)
7
8
# generator with openai models
9
generator_llm = ChatOpenAI(model="gpt-3.5-turbo-16k")
10
critic_llm = ChatOpenAI(model="gpt-4")
11
embeddings = OpenAIEmbeddings()
12
13
generator = TestsetGenerator.from_langchain(generator_llm, critic_llm, embeddings)
14
15
# Change resulting question type distribution
16
distributions = {simple: 0.5, multi_context: 0.4, reasoning: 0.1}
17
18
testset = generator.generate_with_langchain_docs(
19
    pages, 10, distributions, run_config=RUN_CONFIG
20
)

1
{
2
  "fields": [
3
    {
4
      "numDimensions": 1536,
5
      "path": "embedding",
6
      "similarity": "cosine",
7
      "type": "vector"
8
    }
9
  ]
10
}

1
from langchain_mongodb import MongoDBAtlasVectorSearch
2
from pymongo import MongoClient
3
4
client = MongoClient(MONGODB_URI)
5
DB_NAME = "evals"
6
COLLECTION_NAME = "chunking"
7
ATLAS_VECTOR_SEARCH_INDEX_NAME = "vector_index"
8
MONGODB_COLLECTION = client[DB_NAME][COLLECTION_NAME]
9
10
def create_vector_store(docs: List[Document]) -> MongoDBAtlasVectorSearch:
11
    """
12
    Create MongoDB Atlas vector store
13
14
    Args:
15
        docs (List[Document]): List of documents to create the vector store
16
17
    Returns:
18
        MongoDBAtlasVectorSearch: MongoDB Atlas vector store
19
    """
20
    vector_store = MongoDBAtlasVectorSearch.from_documents(
21
        documents=docs,
22
        embedding=OpenAIEmbeddings(model="text-embedding-3-small"),
23
        collection=MONGODB_COLLECTION,
24
        index_name=ATLAS_VECTOR_SEARCH_INDEX_NAME,
25
    )
26
27
    return vector_store

1
from tqdm import tqdm
2
from datasets import Dataset
3
from ragas import evaluate
4
from ragas.metrics import context_precision, context_recall
5
import nest_asyncio
6
7
# Allow nested use of asyncio (used by Ragas)
8
nest_asyncio.apply()
9
10
# Disable tqdm locks
11
tqdm.get_lock().locks = []
12
13
QUESTIONS = testset.question.to_list()
14
GROUND_TRUTH = testset.ground_truth.to_list()
15
16
def perform_eval(docs: List[Document]) -> Dict[str, float]:
17
    """
18
    Perform Ragas evaluation
19
20
    Args:
21
        docs (List[Document]): List of documents to create the vector store
22
23
    Returns:
24
        Dict[str, float]: Dictionary of evaluation metrics
25
    """
26
    eval_data = {
27
        "question": QUESTIONS,
28
        "ground_truth": GROUND_TRUTH,
29
        "contexts": [],
30
    }
31
32
    print(f"Deleting existing documents in the collection {DB_NAME}.{COLLECTION_NAME}")
33
    MONGODB_COLLECTION.delete_many({})
34
    print(f"Deletion complete")
35
    vector_store = create_vector_store(docs)
36
37
    # Getting relevant documents for the evaluation dataset
38
    print(f"Getting contexts for evaluation set")
39
    for question in tqdm(QUESTIONS):
40
        eval_data["contexts"].append(
41
            [doc.page_content for doc in retriever.similarity_search(question, k=3)]
42
        )
43
    # RAGAS expects a Dataset object
44
    dataset = Dataset.from_dict(eval_data)
45
46
    print(f"Running evals")
47
    result = evaluate(
48
        dataset=dataset,
49
        metrics=[context_precision, context_recall],
50
        run_config=RUN_CONFIG,
51
        raise_exceptions=False,
52
    )
53
    return result

1
for chunk_size in [100, 200, 500, 1000]:
2
    chunk_overlap = int(0.15 * chunk_size)
3
    print(f"CHUNK SIZE: {chunk_size}")
4
    print("------ Fixed token without overlap ------")
5
    print(f"Result: {perform_eval(fixed_token_split(pages, chunk_size, 0))}")
6
    print("------ Fixed token with overlap ------")
7
    print(
8
        f"Result: {perform_eval(fixed_token_split(pages, chunk_size, chunk_overlap))}"
9
    )
10
    print("------ Recursive with overlap ------")
11
    print(f"Result: {perform_eval(recursive_split(pages, chunk_size, chunk_overlap))}")
12
    print("------ Recursive Python splitter with overlap ------")
13
    print(
14
        f"Result: {perform_eval(recursive_split(pages, chunk_size, chunk_overlap, Language.PYTHON))}"
15
    )
16
print("------ Semantic chunking ------")
17
print(f"Result: {perform_eval(semantic_split(pages))}")