src/chunking/splitter.py

"""Knowledge-base construction and chunking, as used by Notebooks 2-6.

`build_knowledge_base` dedupes every document referenced anywhere in a
RAGBench config into a single corpus, keyed by md5 of its text.
`chunk_documents` splits that corpus with a recursive, tokenizer-aware
splitter (the only chunking strategy validated so far).
"""

import hashlib

import pandas as pd
from langchain_text_splitters import RecursiveCharacterTextSplitter
from sentence_transformers import SentenceTransformer


def build_knowledge_base(df: pd.DataFrame) -> pd.DataFrame:
    docs = {}
    for documents in df["documents"]:
        for doc in documents:
            doc_id = hashlib.md5(doc.encode("utf-8")).hexdigest()
            docs.setdefault(doc_id, doc)
    return pd.DataFrame([{"doc_id": k, "text": v} for k, v in docs.items()])


def chunk_documents(
    docs_df: pd.DataFrame,
    embed_model: SentenceTransformer,
    chunk_size: int,
    chunk_overlap: int,
) -> pd.DataFrame:
    splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
        embed_model.tokenizer,
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
    )

    chunk_records = []
    for _, row in docs_df.iterrows():
        for i, chunk_text in enumerate(splitter.split_text(row["text"])):
            chunk_records.append({
                "doc_id": row["doc_id"],
                "chunk_id": f"{row['doc_id']}_{i}",
                "text": chunk_text,
            })

    return pd.DataFrame(chunk_records)