Exploring Semantic Chunking (실습)

해당 포스트는 Medium 'Florian June'이 작성한 Advanced RAG 포스트 시리즈 그 다섯번째 내용을 실제로 실습하는 과정을 정리하며, 의미 기반 청킹 여러 가지 방법을 사용해본다. 이론적인 부분은 다음을 참고하자.

Exploring Semantic Chunking

---

Embedding-based methods

pip install llama-index-core
pip install llama-index-readers-file
pip install llama-index-embeddings-openai

from llama_index.core.node_parser import SentenceSplitter, SemanticSplitterNodeParser,
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import SimpleDirectoryReader
import os
os.environ["OPENAI_API_KEY"] = "YOUR_OPEN_AI_KEY"
 
# load documents
dir_path = "YOUR_DIR_PATH"
documents = SimpleDirectoryReader(dir_path).load_data()
 
embed_model = OpenAIEmbedding()
splitter = SemanticSplitterNodeParser(
    buffer_size=1, breakpoint_percentile_threshold=95, embed_model=embed_model
)
 
nodes = splitter.get_nodes_from_documents(documents)
for node in nodes:
    print('-' * 100)
    print(node.get_content())

임베딩 기반의 의미 청킹은 본질적으로 슬라이딩 윈도우(combined_sentence)를 기반으로 유사성을 계산하는 것을 포함하며, 인접하고 임계값을 충족하는 문장들이 하나의 청크로 분류된다. 청킹 결과를 살펴보면 청크의 세분성이 비교적 거친 것으로 나타난다. 해당 방법은 페이지 기반이며, 여러 페이지에 걸쳐 있는 청크 문제를 직접적으로 해결하지 않는다.

Model based method

Model based 

Model based 방법론(Naive BERT, Cross Segment Attention, SeqModel)은 이론 부분에서 살펴봤을 때, 성능이 크게 좋은 모델이 없었기에 실습 파트에서는 다루지 않는다.

LLM based method

논문 'Dense X Retrieval: What Retrieval Granularity Should We Use?'에서는 새로운 검색 단위인 '명제'를 소개한다. 명제는 텍스트 내에서 독립적인 사실을 담고 있는 원자적 표현으로 정의되며, 간결하고 독립적인 자연어 형식으로 제시된다. 

Env Setting

!pip install llama-index-core
!pip install llama-index-readers-file
!pip install llama-index-embeddings-openai

from llama_index.core.node_parser import SentenceSplitter, SemanticSplitterNodeParser
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import SimpleDirectoryReader
from llama_index.core.async_utils import run_jobs
from llama_index.core import ServiceContext, VectorStoreIndex
from llama_index.core.retrievers import RecursiveRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.llama_pack import download_llama_pack
from llama_index.llms.openai import OpenAI
from llama_index.core.text_splitter import SentenceSplitter
from llama_index.core import Document
from llama_index.core.schema import TextNode
import os
import nest_asyncio 

nest_asyncio.apply()

Class DenseXRetrieval

class DenseXRetrievalPack(BaseLlamaPack):
    def __init__(
        self,
        documents, #: List[Document],
        proposition_llm, # : Optional[LLM] = None,
        query_llm, #: Optional[LLM] = None,
        embed_model, #: Optional[BaseEmbedding] = None,
        text_splitter, #: TextSplitter = SentenceSplitter(),
        similarity_top_k: int = 4,
    ) -> None:
        """Init params."""
        self._proposition_llm = proposition_llm or OpenAI(
            model="gpt-3.5-turbo",
            temperature=0.1,
            max_tokens=750,
        )
        embed_model = embed_model or OpenAIEmbedding(embed_batch_size=128)
        nodes = text_splitter.get_nodes_from_documents(documents)   # 노드 추출
        sub_nodes = self._gen_propositions(nodes)   # 명제 획득
        self.sub_nodes = sub_nodes

        all_nodes = nodes + sub_nodes
        all_nodes_dict = {n.node_id: n for n in all_nodes}
        service_context = ServiceContext.from_defaults(    # 임베딩 서비스  
            llm=query_llm or OpenAI(),
            embed_model=embed_model,
            num_output=self._proposition_llm.metadata.num_output,
        )
        self.vector_index = VectorStoreIndex(
            all_nodes, service_context=service_context, show_progress=True
        )
        self.retriever = RecursiveRetriever(   # 추출기 선언 
            "vector",
            retriever_dict={
                "vector": self.vector_index.as_retriever(
                    similarity_top_k=similarity_top_k
                )
            },
            node_dict=all_nodes_dict,
        )
        self.query_engine = RetrieverQueryEngine.from_args(
            self.retriever, service_context=service_context
        )
        
    async def _aget_proposition(self, node: TextNode) -> List[TextNode]:
        """Get proposition."""
        inital_output = await self._proposition_llm.apredict(
            PROPOSITIONS_PROMPT, node_text=node.text
        )
        outputs = inital_output.split("\n")

        all_propositions = []   # 명제 저장
        for output in outputs:
            if not output.strip():
                continue
            if not output.strip().endswith("]"):
                if not output.strip().endswith('"') and not output.strip().endswith(
                    ","
                ):
                    output = output + '"'
                output = output + " ]"
            if not output.strip().startswith("["):
                if not output.strip().startswith('"'):
                    output = '"' + output
                output = "[ " + output
            try:
                propositions = json.loads(output)
            except Exception:
                # fallback to yaml
                try:
                    propositions = yaml.safe_load(output)
                except Exception:
                    # fallback to next output
                    continue
            if not isinstance(propositions, list):
                continue

            all_propositions.extend(propositions)
        assert isinstance(all_propositions, list)
        nodes = [TextNode(text=prop) for prop in all_propositions if prop]
        return [IndexNode.from_text_node(n, node.node_id) for n in nodes]

    def _gen_propositions(self, nodes: List[TextNode]) -> List[TextNode]:
        """ 명제 추출 """
        sub_nodes = asyncio.run(
            run_jobs(
                [self._aget_proposition(node) for node in nodes],
                show_progress=True,
                workers=8,
            )
        )
        print(f'sub_nodes: {sub_nodes}')
        # Flatten list
        return [node for sub_node in sub_nodes for node in sub_node]

DenseXretrieval을 활용해 쿼리해보자

dense_pack = DenseXRetrievalPack(
    documents, 
    proposition_llm=OpenAI(model="gpt-3.5-turbo"),
    query_llm=OpenAI(model="gpt-3.5-turbo"),
    embed_model=None,
    text_splitter=SentenceSplitter(chunk_size=1024)
)
dense_query_engine = dense_pack.query_engine
response = dense_query_engine.query("are you depressed ?")
response

왜인지는 모르겠지만, sub_nodes (명제)는 계속해서 추출되지 않는다. 이 부분은 추후에 더 살펴봐야 할 거 같다.

다음과 같이 보다 간단하게 사용하는 방법도 존재한다.

DenseXRetrievalPack = download_llama_pack(
  "DenseXRetrievalPack", "./dense_pack"
)

documents = SimpleDirectoryReader("../data").load_data(['applsci_aud.pdf'])
# uses the LLM to extract propositions from every document/node!
# dense_pack = DenseXRetrievalPack(documents)   
dense_pack = DenseXRetrievalPack(
  documents, 
  proposition_llm=OpenAI(model="gpt-3.5-turbo"),
  query_llm=OpenAI(model="gpt-3.5-turbo"),
  text_splitter=SentenceSplitter(chunk_size=1024)
)
dense_query_engine = dense_pack.query_engine
response = dense_query_engine.query("How can we measure depression intensity ?")

사실 아직 DenseXRetrieval에 대해 잘 모르겠다. 좀 더 공부하고 내용 보완할 예정이다.