RAGにおけるチャンキング戦略の比較：代替案、トレードオフ、および例

from __future__ import annotations

import hashlib
from dataclasses import dataclass
from typing import Any, Iterable, Optional

from transformers import AutoTokenizer  # pip install transformers

@dataclass(frozen=True)
class Chunk:
    text: str
    meta: dict[str, Any]

def sha1_id(*parts: str) -> str:
    h = hashlib.sha1()
    for p in parts:
        h.update(p.encode("utf-8"))
        h.update(b"\x1e")
    return h.hexdigest()

# 使用するトークナイザは、LLM/埋め込みのトークナイズにほぼ一致するものにしてください。
TOKENIZER = AutoTokenizer.from_pretrained("bert-base-uncased")

def token_len(text: str) -> int:
    return len(TOKENIZER.encode(text, add_special_tokens=False))

def chunk_fixed_tokens(
    text: str,
    *,
    chunk_size: int = 512,
) -> list[Chunk]:
    token_ids = TOKENIZER.encode(text, add_special_tokens=False)
    out: list[Chunk] = []

    for i in range(0, len(token_ids), chunk_size):
        window = token_ids[i : i + chunk_size]
        chunk_text = TOKENIZER.decode(window)
        out.append(
            Chunk(
                text=chunk_text,
                meta={"strategy": "fixed_tokens", "start_token": i, "end_token": i + len(window)},
            )
        )
    return out

def chunk_sliding_window(
    text: str,
    *,
    window_tokens: int = 512,
    stride_tokens: int = 384,  # ストライドが小さいほどオーバーラップが多い
) -> list[Chunk]:
    assert 1 <= stride_tokens <= window_tokens, "ストライドは(0, window]の範囲内でなければなりません"
    token_ids = TOKENIZER.encode(text, add_special_tokens=False)
    out: list[Chunk] = []

    start = 0
    while start < len(token_ids):
        end = min(start + window_tokens, len(token_ids))
        window = token_ids[start:end]
        out.append(
            Chunk(
                text=TOKENIZER.decode(window),
                meta={"strategy": "sliding_window", "start_token": start, "end_token": end},
            )
        )
        if end == len(token_ids):
            break
        start += stride_tokens

    return out

# pip install nltk
import nltk
from nltk.tokenize import sent_tokenize

nltk.download("punkt", quiet=True)

def chunk_by_sentences_nltk(
    text: str,
    *,
    max_tokens: int = 512,
    overlap_sentences: int = 1,
) -> list[Chunk]:
    sents = [s.strip() for s in sent_tokenize(text) if s.strip()]
    out: list[Chunk] = []

    buf: list[str] = []
    buf_tokens = 0

    def flush():
        nonlocal buf, buf_tokens
        if not buf:
            return
        chunk_text = " ".join(buf).strip()
        out.append(
            Chunk(
                text=chunk_text,
                meta={"strategy": "sentences_nltk", "sent_count": len(buf)},
            )
        )
        # オーバーラップするために最後のN文を保持
        if overlap_sentences > 0:
            buf = buf[-overlap_sentences:]
            buf_tokens = token_len(" ".join(buf))
        else:
            buf, buf_tokens = [], 0

    for s in sents:
        s_tokens = token_len(s)
        # 単一の文が予算を超える場合、そのスパンに対して固定サイズのチャンキングにフォールバック
        if s_tokens > max_tokens:
            flush()
            out.extend(chunk_fixed_tokens(s, chunk_size=max_tokens))
            continue

        if buf_tokens + s_tokens > max_tokens and buf:
            flush()

        buf.append(s)
        buf_tokens += s_tokens

    flush()
    return out

# pip install spacy
# python -m spacy download en_core_web_sm
import spacy

def chunk_by_sentences_spacy(
    text: str,
    *,
    max_tokens: int = 512,
) -> list[Chunk]:
    # 軽量なルールベースの分割（依存関係解析なし）を使用する場合、sentencizerを使用します。
    nlp = spacy.blank("en")
    nlp.add_pipe("sentencizer")  # ルールベースの文の境界検出
    doc = nlp(text)

    sents = [sent.text.strip() for sent in doc.sents if sent.text.strip()]
    return chunk_by_sentences_nltk(" ".join(sents), max_tokens=max_tokens, overlap_sentences=1)

# pip install langchain-text-splitters
from langchain_text_splitters import RecursiveCharacterTextSplitter

def chunk_recursive_langchain(
    text: str,
    *,
    chunk_size: int = 1200,
    chunk_overlap: int = 150,
) -> list[Chunk]:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=token_len,  # トークンを意識した予算配分
        separators=["\n\n", "\n", ". ", " ", ""],  # コンテンツに合わせてカスタマイズ（例：コード）
    )
    pieces = splitter.split_text(text)

    return [
        Chunk(text=p, meta={"strategy": "recursive_langchain", "chunk_size": chunk_size, "overlap": chunk_overlap})
        for p in pieces
    ]

# pip install openai numpy
import os
import numpy as np
from openai import OpenAI

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY", "YOUR_OPENAI_API_KEY"))

def embed_texts_openai(texts: list[str], *, model: str = "text-embedding-3-small") -> np.ndarray:
    # バッチが大きい場合、リクエストトークン制限とバッチサイズ制限を尊重してください。
    resp = client.embeddings.create(model=model, input=texts)
    embs = np.array([d.embedding for d in resp.data], dtype=np.float32)
    return embs

def cosine_sim(a: np.ndarray, b: np.ndarray) -> float:
    denom = (np.linalg.norm(a) * np.linalg.norm(b)) + 1e-8
    return float(np.dot(a, b) / denom)

def chunk_semantic(
    text: str,
    *,
    max_tokens: int = 800,
    breakpoint_threshold: float = 0.70,
) -> list[Chunk]:
    # 1) 文の候補から開始
    sents = [s.strip() for s in sent_tokenize(text) if s.strip()]
    if len(sents) <= 1:
        return [Chunk(text=text, meta={"strategy": "semantic", "note": "no_split"})]

    # 2) 各文の埋め込み
    embs = embed_texts_openai(sents)

    # 3) 類似度の減少を計算
    sims = [cosine_sim(embs[i], embs[i + 1]) for i in range(len(sents) - 1)]

    # 4) ブレイクポイントでセグメントを作成
    out: list[Chunk] = []
    buf: list[str] = []
    buf_tokens = 0

    for i, s in enumerate(sents):
        # 文を追加
        s_tok = token_len(s)
        if buf and buf_tokens + s_tok > max_tokens:
            out.append(Chunk(text=" ".join(buf), meta={"strategy": "semantic", "reason": "max_tokens"}))
            buf, buf_tokens = [], 0
        buf.append(s)
        buf_tokens += s_tok

        # 文iの後のブレイクポイントを決定（次の文との類似度に基づく）
        if i < len(sims) and sims[i] < breakpoint_threshold:
            out.append(
                Chunk(
                    text=" ".join(buf),
                    meta={"strategy": "semantic", "reason": "sim_drop", "sim_to_next": sims[i]},
                )
            )
            buf, buf_tokens = [], 0

    if buf:
        out.append(Chunk(text=" ".join(buf), meta={"strategy": "semantic", "reason": "final"}))

    return out

# pip install llama-index llama-index-llms-openai
from llama_index.core import Document, StorageContext, VectorStoreIndex
from llama_index.core.node_parser import HierarchicalNodeParser, get_leaf_nodes
from llama_index.core.retrievers import AutoMergingRetriever
from llama_index.core.storage.docstore import SimpleDocumentStore

def build_hierarchical_index(text: str):
    docs = [Document(text=text)]

    node_parser = HierarchicalNodeParser.from_defaults()  # デフォルトは粗から細のサイズ
    nodes = node_parser.get_nodes_from_documents(docs)

    docstore = SimpleDocumentStore()
    docstore.add_documents(nodes)
    storage_context = StorageContext.from_defaults(docstore=docstore)

    leaf_nodes = get_leaf_nodes(nodes)
    base_index = VectorStoreIndex(leaf_nodes, storage_context=storage_context)

    base_retriever = base_index.as_retriever(similarity_top_k=6)
    retriever = AutoMergingRetriever(base_retriever, storage_context, verbose=True)
    return retriever

# pip install docling
# NOTE: PDFのパース品質は環境（フォント、OCRなど）に依存します。
from docling.document_converter import DocumentConverter
from docling.transforms.chunker.hierarchical_chunker import HierarchicalChunker

def chunk_pdf_docling(pdf_path: str) -> list[Chunk]:
    converter = DocumentConverter()
    doc = converter.convert(pdf_path).document  # DoclingDocument
    chunker = HierarchicalChunker()
    doc_chunks = list(chunker.chunk(doc))

    out: list[Chunk] = []
    for c in doc_chunks:
        # c.textにはシリアル化されたチャンク内容が含まれています；c.metaには構造情報が含まれます
        out.append(Chunk(text=c.text, meta={"strategy": "docling_hierarchical", **dict(c.meta)}))
    return out

import ast

def chunk_python_by_ast(code: str, *, filepath: str = "<memory>") -> list[Chunk]:
    tree = ast.parse(code)
    out: list[Chunk] = []

    for node in tree.body:
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            start = getattr(node, "lineno", None)
            end = getattr(node, "end_lineno", None)
            if start is None or end is None:
                continue

            lines = code.splitlines()
            snippet = "\n".join(lines[start - 1 : end])

            kind = "class" if isinstance(node, ast.ClassDef) else "function"
            name = getattr(node, "name", "<anon>")
            out.append(
                Chunk(
                    text=snippet,
                    meta={
                        "strategy": "python_ast",
                        "kind": kind,
                        "name": name,
                        "filepath": filepath,
                        "start_line": start,
                        "end_line": end,
                    },
                )
            )

    return out

# pip install faiss-cpu numpy
import numpy as np
import faiss

def build_faiss_index(vectors: np.ndarray) -> faiss.Index:
    # vectors: shape (N, d), dtype float32
    d = vectors.shape[1]
    index = faiss.IndexFlatIP(d)  # 内積；ベクトルが正規化されている場合、コサイン類似度
    faiss.normalize_L2(vectors)
    index.add(vectors)
    return index

def faiss_search(index: faiss.Index, query_vec: np.ndarray, k: int = 5):
    q = query_vec.astype(np.float32).reshape(1, -1)
    faiss.normalize_L2(q)
    scores, ids = index.search(q, k)
    return ids[0].tolist(), scores[0].tolist()

# pip install chromadb
import chromadb

def build_chroma_collection(chunks: list[Chunk], embeddings: np.ndarray, *, path: str = "./chroma_store"):
    client = chromadb.PersistentClient(path=path)
    col = client.get_or_create_collection(name="docs")

    ids = [sha1_id(c.meta.get("strategy", "chunk"), str(i), c.text[:50]) for i, c in enumerate(chunks)]
    col.upsert(
        ids=ids,
        documents=[c.text for c in chunks],
        metadatas=[c.meta for c in chunks],
        embeddings=embeddings.tolist(),
    )
    return col

# pip install weaviate-client
import weaviate
from weaviate.classes.config import Configure

def weaviate_upsert_self_provided(
    chunks: list[Chunk],
    embeddings: np.ndarray,
):
    client = weaviate.connect_to_local()  # またはconnect_to_weaviate_cloud(...)
    try:
        collection = client.collections.create(
            name="Chunk",
            vector_config=Configure.Vectors.self_provided(),  # あなたがベクトルを提供
        )

        with collection.batch.dynamic() as batch:
            for c, v in zip(chunks, embeddings):
                batch.add_object(
                    properties={"text": c.text, **{f"m_{k}": str(vv) for k, vv in c.meta.items()}},
                    vector=v.tolist(),
                )

        # near-vectorでクエリ
        q = embeddings[0]
        res = collection.query.near_vector(near_vector=q.tolist(), limit=5)
        for obj in res.objects:
            print(obj.properties.get("text")[:200])
    finally:
        client.close()