import re
import nltk
from nltk.tokenize import sent_tokenize, word_tokenize
from transformers import AutoTokenizer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# 確保已經(jīng)下載了nltk的punkt分詞器
try:
    nltk.data.find('tokenizers/punkt')
except nltk.downloader.DownloadError:
    nltk.download('punkt')

# 模擬一些文本數(shù)據(jù)
sample_text_long = """
RAG（Retrieval-Augmented Generation）是一種結(jié)合了檢索和生成能力的AI技術(shù)。它的核心思想是，當(dāng)大語言模型（LLM）需要回答問題時，它不再是憑空生成答案，而是首先從一個龐大的知識庫中檢索出最相關(guān)的上下文信息。

這些檢索到的信息隨后會被作為輸入的一部分，提供給LLM，LLM再基于這些“額外”的知識來生成最終的回答。這種方式能夠顯著減少模型“胡說八道”（幻覺）的現(xiàn)象，并使得模型能夠回答特定領(lǐng)域的問題，因為它的知識不再僅僅局限于訓(xùn)練時的數(shù)據(jù)。

在實際應(yīng)用中，RAG系統(tǒng)涉及到多個關(guān)鍵組件。首先是**數(shù)據(jù)攝取和預(yù)處理**，這通常包括將原始文檔（如PDF、網(wǎng)頁、數(shù)據(jù)庫記錄等）清洗、解析，并轉(zhuǎn)換為適合存儲和檢索的格式。接著是**切塊（Chunking）**，這是RAG流程中至關(guān)重要的一步。切塊的目的是將長文本分割成較小的、易于管理和檢索的片段。

切塊策略的選擇直接影響檢索的質(zhì)量。如果切塊過大，可能包含太多無關(guān)信息，增加LLM處理的負(fù)擔(dān)；如果切塊過小，則可能丟失上下文，導(dǎo)致LLM無法理解完整語義。

然后是**嵌入（Embedding）**，將文本切塊轉(zhuǎn)換為高維向量，這些向量能夠捕捉文本的語義信息。接著是**向量數(shù)據(jù)庫（Vector Database）**，用于高效存儲和檢索這些嵌入向量。當(dāng)用戶提出問題時，**檢索器（Retriever）**會根據(jù)用戶問題的嵌入向量，在向量數(shù)據(jù)庫中查找最相似的文本切塊。

最后，**生成器（Generator）**，也就是LLM，會結(jié)合用戶的原始問題和檢索到的相關(guān)上下文信息，生成最終的答案。整個RAG流程的優(yōu)化是一個迭代的過程，需要不斷地調(diào)整各個環(huán)節(jié)，才能達(dá)到最佳效果。
"""

sample_text_structured = """
# 第一章 引言
## 1.1 RAG的魅力
RAG技術(shù)有效解決了大模型幻覺問題。

## 1.2 本文目的
本文將深入探討RAG中的切塊策略。

# 第二章 切塊基礎(chǔ)
## 2.1 樸素切塊
這是一種簡單的切塊方法。

## 2.2 固定大小切塊
我們將會看到固定大小切塊的實現(xiàn)。

---
**Note:** 本文檔旨在提供RAG切塊的全面指南。
"""

sample_text_with_tables = """
這是一段關(guān)于公司業(yè)績的文本。

| Month | Sales | Profit |
|---|---|---|
| Jan | 100 | 20 |
| Feb | 120 | 25 |
| Mar | 110 | 22 |

以上是第一季度的財務(wù)數(shù)據(jù)。

下面是團隊成員信息：

| Name | Role |
|---|---|
| Alice | Engineer |
| Bob | Designer |
"""

sample_text_mixed_format = """
這是一個段落。它包含一些重要的信息，需要被完整保留。

1.  這是一個列表項。
2.  這是另一個列表項。
    * 子項1
    * 子項2

## 重要提示
請注意以下表格數(shù)據(jù)：

| Item | Quantity | Price |
|---|---|---|
| Apple | 10 | 1.0 |
| Banana | 5 | 0.5 |
"""

def naive_chunking(text):
    """
    按行分割文本。
    """
    return text.split('\n')

print("--- 1. 樸素切塊 ---")
chunks_naive = naive_chunking(sample_text_long)
for i, chunk in enumerate(chunks_naive[:5]): # 只打印前5個示例
    print(f"Chunk {i+1}: '{chunk.strip()}'")
print("...")

def fixed_size_chunking(text, chunk_size, overlap=0):
    """
    按固定大小切塊，可選擇重疊。
    chunk_size: 每個切塊的字符數(shù)
    overlap: 重疊的字符數(shù)
    """
    chunks = []
    text_len = len(text)
    # 注意：這里 i 的步長是 chunk_size - overlap，確保重疊
    for i in range(0, text_len, chunk_size - overlap):
        chunk = text[i:i + chunk_size]
        chunks.append(chunk)
        if i + chunk_size >= text_len: # 防止最后一個切塊不完整，但又跳過重疊部分
            break
    return chunks

print("\n--- 2. 固定大小切塊 ---")
chunks_fixed = fixed_size_chunking(sample_text_long, chunk_size=200, overlap=0)
for i, chunk in enumerate(chunks_fixed[:3]):
    print(f"Chunk {i+1}: '{chunk.strip()}'")
print("...")

def sliding_window_chunking(text, chunk_size, overlap):
    """
    滑動窗口切塊，每個切塊與前一個重疊。
    chunk_size: 每個切塊的字符數(shù)
    overlap: 重疊的字符數(shù)
    """
    chunks = []
    text_len = len(text)
    start = 0
    while start < text_len:
        end = min(start + chunk_size, text_len)
        chunk = text[start:end]
        chunks.append(chunk)
        if end == text_len:
            break
        start += (chunk_size - overlap)
    return chunks

print("\n--- 3. 滑動窗口切塊 ---")
chunks_sliding = sliding_window_chunking(sample_text_long, chunk_size=200, overlap=50)
for i, chunk in enumerate(chunks_sliding[:3]):
    print(f"Chunk {i+1}: '{chunk.strip()}'")
print("...")

def sentence_based_chunking(text):
    """
    使用nltk進行句子級別的切塊。
    """
    return sent_tokenize(text)

print("\n--- 4. 基于句子切塊 ---")
chunks_sentence = sentence_based_chunking(sample_text_long)
for i, chunk in enumerate(chunks_sentence[:5]):
    print(f"Chunk {i+1}: '{chunk.strip()}'")
print("...")

def paragraph_based_chunking(text):
    """
    按雙換行符分割文本（通常表示段落）。
    """
    # 使用正則表達(dá)式匹配一個或多個換行符，并移除空字符串
    paragraphs = [p.strip() for p in re.split(r'\n{2,}', text) if p.strip()]
    return paragraphs

print("\n--- 5. 基于段落切塊 ---")
chunks_paragraph = paragraph_based_chunking(sample_text_long)
for i, chunk in enumerate(chunks_paragraph[:3]):
    print(f"Chunk {i+1}: '{chunk.strip()}'")
print("...")

def page_based_chunking(text, page_delimiter="---PAGE_BREAK---"):
    """
    模擬頁面切塊，假設(shè)文本中存在頁面分隔符。
    實際應(yīng)用中需要從PDF等文件讀取。
    """
    # 模擬一個多頁文本
    multi_page_text = (
        "這是第一頁的內(nèi)容。\n一些重要的信息在這里。\n" +
        page_delimiter + "\n" +
        "這是第二頁的內(nèi)容。\n繼續(xù)重要的討論。\n" +
        page_delimiter + "\n" +
        "第三頁是總結(jié)。\n全文到此結(jié)束。"
    )
    return [p.strip() for p in multi_page_text.split(page_delimiter) if p.strip()]

print("\n--- 6. 基于頁面切塊 ---")
chunks_page = page_based_chunking(sample_text_long) # 使用模擬文本
for i, chunk in enumerate(chunks_page):
    print(f"Page {i+1}:\n'{chunk}'")

def structured_chunking(text):
    """
    根據(jù)Markdown標(biāo)題結(jié)構(gòu)進行切塊。
    """
    chunks = []
    # 匹配Markdown標(biāo)題，同時捕獲標(biāo)題和其后的內(nèi)容
    # 注意：這里會把每個標(biāo)題下的內(nèi)容切成一個塊
    sections = re.split(r'^(#+ .*)$', text, flags=re.MULTILINE)
    
    current_heading = ""
    current_content = []
    
    # 第一個元素通常是空字符串或標(biāo)題之前的內(nèi)容
    if sections[0].strip():
        chunks.append(sections[0].strip())

    for i in range(1, len(sections)):
        part = sections[i].strip()
        if part.startswith('#'): # 這是一個標(biāo)題
            if current_content: # 如果有之前收集的內(nèi)容，先作為一個塊
                chunks.append("\n".join(current_content).strip())
                current_content = [] # 重置
            current_heading = part
            current_content.append(current_heading) # 將標(biāo)題也包含在塊內(nèi)
        else: # 這是一個標(biāo)題下的內(nèi)容
            current_content.append(part)
    
    if current_content: # 添加最后一個塊
        chunks.append("\n".join(current_content).strip())

    return [chunk for chunk in chunks if chunk] # 過濾空塊

print("\n--- 7. 結(jié)構(gòu)化切塊 ---")
chunks_structured = structured_chunking(sample_text_structured)
for i, chunk in enumerate(chunks_structured):
    print(f"Chunk {i+1}:\n'{chunk}'")

def document_based_chunking(text):
    """
    基于文檔的自然結(jié)構(gòu)（如Markdown的章節(jié)和子章節(jié)）進行切塊。
    這里我們將捕獲頂級標(biāo)題下的所有內(nèi)容作為一個塊，直到下一個同級或更高級的標(biāo)題。
    """
    chunks = []
    lines = text.split('\n')
    current_chunk_lines = []
    
    for line in lines:
        if line.startswith('#'): # 匹配任何級別的標(biāo)題
            if current_chunk_lines: # 如果當(dāng)前塊有內(nèi)容，就結(jié)束并添加
                chunks.append("\n".join(current_chunk_lines).strip())
                current_chunk_lines = []
            current_chunk_lines.append(line) # 將標(biāo)題作為新塊的開始
        else:
            current_chunk_lines.append(line)
            
    if current_chunk_lines: # 添加最后一個塊
        chunks.append("\n".join(current_chunk_lines).strip())
        
    return [chunk for chunk in chunks if chunk]

print("\n--- 8. 基于文檔結(jié)構(gòu)切塊 ---")
chunks_doc_struct = document_based_chunking(sample_text_structured)
for i, chunk in enumerate(chunks_doc_struct):
    print(f"Chunk {i+1}:\n'{chunk}'")

def keyword_based_chunking(text, keywords):
    """
    在指定關(guān)鍵詞處進行切塊。
    """
    chunks = []
    # 構(gòu)建正則表達(dá)式，匹配所有關(guān)鍵詞并保留關(guān)鍵詞本身
    # 使用非捕獲組 (?:...) 結(jié)合 | 運算符
    pattern = '|'.join(re.escape(k) for k in keywords)
    parts = re.split(f'({pattern})', text) # 使用捕獲組保留分隔符

    current_chunk = ""
    for part in parts:
        if part.strip() in keywords: # 如果當(dāng)前部分是關(guān)鍵詞
            if current_chunk.strip(): # 將之前的累積作為新塊
                chunks.append(current_chunk.strip())
            current_chunk = part # 關(guān)鍵詞作為新塊的開始
        else:
            current_chunk += part
    
    if current_chunk.strip(): # 添加最后一個塊
        chunks.append(current_chunk.strip())
        
    return [chunk for chunk in chunks if chunk] # 過濾空塊

print("\n--- 9. 基于關(guān)鍵詞切塊 ---")
keywords_for_chunking = ["Note:", "首先是", "接著是"]
chunks_keyword = keyword_based_chunking(sample_text_long + sample_text_structured, keywords_for_chunking)
for i, chunk in enumerate(chunks_keyword):
    print(f"Chunk {i+1}:\n'{chunk}'")

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

try:
    nlp = spacy.load("en_core_web_sm")
except OSError:
    print("Downloading spacy model 'en_core_web_sm'...")
    spacy.cli.download("en_core_web_sm")
    nlp = spacy.load("en_core_web_sm")

def entity_based_chunking(text):
    """
    使用NER模型識別實體，并圍繞實體聚合文本。
    這里為了簡化，我們找到實體所在句子，并以句子為單位聚合。
    """
    doc = nlp(text)
    entities = {} # {entity_text: [sentences containing this entity]}

    for sent in doc.sents:
        found_entities_in_sent = False
        for ent in sent.ents:
            if ent.label_ in ["PERSON", "ORG", "GPE", "PRODUCT"]: # 關(guān)注人、組織、地理、產(chǎn)品等實體
                if ent.text notin entities:
                    entities[ent.text] = []
                entities[ent.text].append(sent.text.strip())
                found_entities_in_sent = True
        ifnot found_entities_in_sent: # 如果句子沒有實體，作為獨立塊或添加到“無實體”塊
            if"NO_ENTITY"notin entities:
                entities["NO_ENTITY"] = []
            entities["NO_ENTITY"].append(sent.text.strip())
            
    # 將字典轉(zhuǎn)換為列表，每個實體或無實體組一個塊
    chunks = []
    for entity, sents in entities.items():
        chunk_content = f"Related to {entity}:\n" + "\n".join(list(set(sents))) # 使用set去重
        chunks.append(chunk_content)
    
    return chunks

print("\n--- 10. 基于實體切塊 ---")
sample_ner_text = "Apple公司發(fā)布了新的iPhone 15。Tim Cook在發(fā)布會上強調(diào)了其強大的A17芯片。用戶可以在紐約的Apple Store購買。"
chunks_entity = entity_based_chunking(sample_ner_text)
for i, chunk in enumerate(chunks_entity):
    print(f"Chunk {i+1}:\n'{chunk}'")

# pip install transformers
def token_based_chunking(text, tokenizer_name="bert-base-uncased", max_tokens=128):
    """
    使用分詞器按token數(shù)量進行切塊。
    """
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
    tokens = tokenizer.encode(text, add_special_tokens=False) # 不添加特殊token
    
    chunks = []
    for i in range(0, len(tokens), max_tokens):
        chunk_tokens = tokens[i:i + max_tokens]
        chunk_text = tokenizer.decode(chunk_tokens)
        chunks.append(chunk_text)
    return chunks

print("\n--- 11. 基于Token切塊 ---")
chunks_token = token_based_chunking(sample_text_long, max_tokens=60)
for i, chunk in enumerate(chunks_token[:3]):
    print(f"Chunk {i+1}:\n'{chunk}'")
print("...")

# pip install scikit-learn
def topic_based_chunking(text, min_sentences_per_topic=3):
    """
    通過句子相似度模擬主題切塊。
    將相似的句子聚類成一個主題塊。
    """
    sentences = sent_tokenize(text)
    if len(sentences) < min_sentences_per_topic:
        return [text] # 句子太少，無法有效分主題

    # 使用TF-IDF向量化句子
    vectorizer = TfidfVectorizer().fit(sentences)
    sentence_vectors = vectorizer.transform(sentences)

    chunks = []
    current_topic_sentences = [sentences[0]]
    
    for i in range(1, len(sentences)):
        # 計算當(dāng)前句子與當(dāng)前主題塊中所有句子的平均相似度
        current_sentence_vector = sentence_vectors[i]
        
        # 將當(dāng)前主題塊的句子向量合并
        current_topic_vectors = vectorizer.transform(current_topic_sentences)
        avg_similarity = np.mean(cosine_similarity(current_sentence_vector, current_topic_vectors))

        # 如果相似度低于某個閾值，或者當(dāng)前主題塊句子太多，就認(rèn)為主題切換
        # 這里閾值和數(shù)量都是啟發(fā)式的，實際應(yīng)用中需調(diào)優(yōu)
        if avg_similarity < 0.3or len(current_topic_sentences) >= 5: # 假設(shè)相似度低于0.3或句子多于5句視為新主題
            chunks.append(" ".join(current_topic_sentences))
            current_topic_sentences = [sentences[i]]
        else:
            current_topic_sentences.append(sentences[i])
            
    if current_topic_sentences:
        chunks.append(" ".join(current_topic_sentences))
        
    return chunks

print("\n--- 12. 基于主題切塊 ---")
chunks_topic = topic_based_chunking(sample_text_long)
for i, chunk in enumerate(chunks_topic):
    print(f"Chunk {i+1}:\n'{chunk}'")

def table_aware_chunking(text):
    """
    識別并單獨切塊表格。將表格內(nèi)容轉(zhuǎn)換為Markdown格式。
    """
    chunks = []
    # 匹配Markdown表格的正則表達(dá)式
    # 捕獲表格內(nèi)容，包括表頭、分隔線和行
    table_pattern = re.compile(r'(\|.*\|\n\|[-: ]+\|\n(?:\|.*\|\n?)+)', re.MULTILINE)
    
    last_end = 0
    for match in table_pattern.finditer(text):
        # 添加表格前的內(nèi)容
        pre_table_text = text[last_end:match.start()].strip()
        if pre_table_text:
            chunks.append(pre_table_text)
        
        # 添加表格本身
        chunks.append(match.group(0).strip())
        last_end = match.end()
        
    # 添加最后一個表格后的內(nèi)容
    post_table_text = text[last_end:].strip()
    if post_table_text:
        chunks.append(post_table_text)
        
    return [chunk for chunk in chunks if chunk]

print("\n--- 13. 表格感知切塊 ---")
chunks_table_aware = table_aware_chunking(sample_text_with_tables)
for i, chunk in enumerate(chunks_table_aware):
    print(f"Chunk {i+1}:\n'{chunk}'")

def content_aware_chunking(text):
    """
    根據(jù)內(nèi)容類型（段落、列表、表格、標(biāo)題等）應(yīng)用不同的切塊規(guī)則。
    這是一個結(jié)合了多種策略的示例。
    """
    chunks = []
    lines = text.split('\n')
    current_chunk_lines = []
    in_table = False
    
    for line in lines:
        stripped_line = line.strip()

        # 檢查是否是表格行（簡單的啟發(fā)式判斷）
        if stripped_line.startswith('|') and'|'in stripped_line[1:]:
            ifnot in_table: # 如果剛進入表格，先結(jié)束前一個非表格塊
                if current_chunk_lines:
                    chunks.append("\n".join(current_chunk_lines).strip())
                    current_chunk_lines = []
                in_table = True
            current_chunk_lines.append(line)
        elif in_table: # 如果在表格中，但當(dāng)前行不是表格行，則表格結(jié)束
            if current_chunk_lines: # 添加完整的表格塊
                chunks.append("\n".join(current_chunk_lines).strip())
                current_chunk_lines = []
            in_table = False
            current_chunk_lines.append(line) # 當(dāng)前行作為新塊的開始
        elif stripped_line.startswith('#'): # 匹配標(biāo)題
            if current_chunk_lines: # 如果有內(nèi)容，結(jié)束前一個塊
                chunks.append("\n".join(current_chunk_lines).strip())
                current_chunk_lines = []
            current_chunk_lines.append(line) # 標(biāo)題作為新塊的開始
        elifnot stripped_line and current_chunk_lines: # 空行作為段落分隔符
            if current_chunk_lines[-1].strip() != "": # 避免連續(xù)空行導(dǎo)致空塊
                chunks.append("\n".join(current_chunk_lines).strip())
                current_chunk_lines = []
        else: # 普通文本行
            current_chunk_lines.append(line)
            
    if current_chunk_lines: # 添加最后一個塊
        chunks.append("\n".join(current_chunk_lines).strip())
        
    return [chunk for chunk in chunks if chunk['content']]

print("\n--- 14. 內(nèi)容感知切塊 ---")
chunks_content_aware = content_aware_chunking(sample_text_mixed_format)
for i, chunk in enumerate(chunks_content_aware):
    print(f"Chunk {i+1}:\n'{chunk}'")

def mock_llm_add_context(chunk_text, knowledge_base_overview):
    """
    模擬LLM為每個切塊添加相關(guān)上下文。
    在實際中，這需要調(diào)用一個真正的LLM。
    """
    if"RAG"in chunk_text and"檢索"in chunk_text:
        returnf"Context: This chunk details the core retrieval mechanism of RAG and its purpose related to knowledge bases. ---\n{chunk_text}"
    elif"切塊"in chunk_text and"影響"in chunk_text:
        returnf"Context: This chunk elaborates on the criticality of chunking strategies and their impact on LLM performance and context. ---\n{chunk_text}"
    else:
        # 模擬一個通用上下文
        returnf"Context: This text fragment discusses general AI concepts or system components. ---\n{chunk_text}"

def contextual_chunking(text, base_chunking_strategy=paragraph_based_chunking, knowledge_base_overview="Overview of AI and RAG systems."):
    """
    先進行基礎(chǔ)切塊，然后用LLM為每個切塊添加上下文。
    """
    base_chunks = base_chunking_strategy(text, **{}) # 確?？梢詡魅肟兆值?    contextualized_chunks = [mock_llm_add_context(chunk, knowledge_base_overview) for chunk in base_chunks]
    return contextualized_chunks

print("\n--- 15. 上下文切塊 ---")
chunks_contextual = contextual_chunking(sample_text_long)
for i, chunk in enumerate(chunks_contextual[:3]):
    print(f"Chunk {i+1}:\n'{chunk}'")
print("...")

# pip install sentence-transformers # 實際應(yīng)用會用這個
from sentence_transformers import SentenceTransformer

# 加載一個預(yù)訓(xùn)練的句子嵌入模型 (首次運行可能需要下載)
try:
    embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
except Exception:
    print("Failed to load sentence-transformers model. Please ensure you have internet or download it manually.")
    # 提供一個備用/跳過策略
    embedding_model = None

def semantic_chunking(text, model=embedding_model, similarity_threshold=0.7):
    """
    先嵌入所有句子，然后根據(jù)相似度聚合。
    """
    if model isNone:
        print("Embedding model not loaded, skipping embedding chunking demo.")
        return [text] # 返回原始文本或進行其他默認(rèn)切塊

    sentences = sent_tokenize(text)
    if len(sentences) <= 1:
        return sentences

    sentence_embeddings = model.encode(sentences)
    
    chunks = []
    current_chunk_sentences = [sentences[0]]

    for i in range(1, len(sentences)):
        # 計算當(dāng)前句子與前一個句子嵌入的余弦相似度
        similarity = cosine_similarity([sentence_embeddings[i]], [sentence_embeddings[i-1]])[0][0]
        
        if similarity < similarity_threshold:
            # 如果相似度低，則認(rèn)為語義不連續(xù)，結(jié)束當(dāng)前塊
            chunks.append(" ".join(current_chunk_sentences))
            current_chunk_sentences = [sentences[i]]
        else:
            current_chunk_sentences.append(sentences[i])
            
    if current_chunk_sentences: # 添加最后一個塊
        chunks.append(" ".join(current_chunk_sentences))
        
    return chunks

print("\n--- 16. 語義切塊 ---")
if embedding_model:
    chunks_semantic = semantic_chunking(sample_text_long, similarity_threshold=0.5) # 調(diào)整閾值以觀察不同效果
    for i, chunk in enumerate(chunks_semantic[:3]):
        print(f"Chunk {i+1}:\n'{chunk}'")
    print("...")
else:
    print("跳過嵌入切塊演示，因為SentenceTransformer模型未加載。")

def recursive_chunking(text, separators, max_chunk_size_char=500):
    """
    遞歸切塊，嘗試不同的分隔符，直到塊大小符合要求。
    separators: 分隔符列表，從大到小排列 (如 ['\n\n', '\n', '. ', ' '])
    max_chunk_size_char: 最大切塊字符數(shù)
    """
    chunks = []
    ifnot text:
        return []

    # 如果文本已經(jīng)小于最大塊大小，直接返回
    if len(text) <= max_chunk_size_char:
        return [text]

    # 嘗試當(dāng)前最大的分隔符
    if separators:
        current_separator = separators[0]
        remaining_separators = separators[1:]
        
        parts = text.split(current_separator)
        
        for part in parts:
            part_stripped = part.strip()
            if part_stripped: # 確保不是空字符串
                if len(part_stripped) > max_chunk_size_char:
                    # 如果部分仍然太大，遞歸調(diào)用更小的分隔符
                    chunks.extend(recursive_chunking(part_stripped, remaining_separators, max_chunk_size_char))
                else:
                    chunks.append(part_stripped)
    else: # 沒有更多分隔符可用，直接按字符切分（作為兜底）
        for i in range(0, len(text), max_chunk_size_char):
            chunks.append(text[i:i + max_chunk_size_char])
            
    return [chunk for chunk in chunks if chunk] # 過濾空塊

print("\n--- 17. 遞歸切塊 ---")
# 模擬一個非常長的段落，需要遞歸切分
long_paragraph = "這是一個非常非常長的段落，它包含了多句話，并且可能在語義上可以被分割。我們希望這個段落能夠被智能地切分成更小的部分，以便于RAG系統(tǒng)處理。如果直接固定大小切塊，可能會切斷句子的上下文，導(dǎo)致信息丟失。所以，我們需要一個更靈活的策略來處理這種長文本。RAG的成功很大程度上取決于切塊的質(zhì)量。我們在這里模擬一個非常長的輸入，以測試遞歸切塊的能力。請注意，這個段落的長度遠(yuǎn)遠(yuǎn)超過了我們設(shè)定的最大塊大小，所以它將被進一步切分。切塊的藝術(shù)在于平衡信息的完整性和粒度。適當(dāng)?shù)那袎K能夠幫助大模型更好地理解檢索到的信息，從而生成更準(zhǔn)確、更相關(guān)的回答。這是一項技術(shù)挑戰(zhàn)，也是RAG優(yōu)化的關(guān)鍵一步。通過不同的分隔符進行遞歸切分，我們可以確保每個塊都不會過大，同時盡量保持語義的完整性。當(dāng)遇到一個超長的段落時，首先嘗試用段落符切分，如果還超長，就用句號切分，再超長就用逗號，直到達(dá)到預(yù)設(shè)的最大長度。"

separators = ['\n\n', '. ', '，', ' '] # 嘗試從大到小的分隔符
chunks_recursive = recursive_chunking(long_paragraph, separators, max_chunk_size_char=100)
for i, chunk in enumerate(chunks_recursive):
    print(f"Chunk {i+1} (len={len(chunk)}):\n'{chunk}'")

# pip install sentence-transformers
from sentence_transformers import SentenceTransformer

# 加載一個預(yù)訓(xùn)練的句子嵌入模型 (首次運行可能需要下載)
try:
    embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
except Exception:
    print("Failed to load sentence-transformers model. Please ensure you have internet or download it manually.")
    # 提供一個備用/跳過策略
    embedding_model = None

def embedding_chunking(text, model=embedding_model, similarity_threshold=0.7):
    """
    先嵌入所有句子，然后根據(jù)相似度聚合。
    """
    if model isNone:
        print("Embedding model not loaded, skipping embedding chunking demo.")
        return [text] # 返回原始文本或進行其他默認(rèn)切塊

    sentences = sent_tokenize(text)
    if len(sentences) <= 1:
        return sentences

    sentence_embeddings = model.encode(sentences)
    
    chunks = []
    current_chunk_sentences = [sentences[0]]

    for i in range(1, len(sentences)):
        # 計算當(dāng)前句子與前一個句子嵌入的余弦相似度
        similarity = cosine_similarity([sentence_embeddings[i]], [sentence_embeddings[i-1]])[0][0]
        
        if similarity < similarity_threshold:
            # 如果相似度低，則認(rèn)為語義不連續(xù)，結(jié)束當(dāng)前塊
            chunks.append(" ".join(current_chunk_sentences))
            current_chunk_sentences = [sentences[i]]
        else:
            current_chunk_sentences.append(sentences[i])
            
    if current_chunk_sentences: # 添加最后一個塊
        chunks.append(" ".join(current_chunk_sentences))
        
    return chunks

print("\n--- 18. 嵌入切塊 ---")
if embedding_model:
    chunks_embedding = embedding_chunking(sample_text_long, similarity_threshold=0.5) # 調(diào)整閾值以觀察不同效果
    for i, chunk in enumerate(chunks_embedding[:3]):
        print(f"Chunk {i+1}:\n'{chunk}'")
    print("...")
else:
    print("跳過嵌入切塊演示，因為SentenceTransformer模型未加載。")

def mock_llm_chunking_decision(text_segment):
    """
    模擬LLM決定如何切塊。
    在實際中，需要給LLM提供文本段落和切塊規(guī)則，讓它返回分割點或直接返回切好的塊。
    例如，可以給LLM一個Prompt：
    "給定以下文本，請將其分割成語義連貫的、不超過200字的獨立片段，并以'---CHUNK---'作為分隔符返回："
    """
    # 模擬LLM智能地將文本分割成幾個邏輯塊
    if"RAG"in text_segment and"幻覺"in text_segment:
        return ["RAG技術(shù)有效解決了大模型幻覺問題。", "其核心在于結(jié)合檢索和生成能力。"]
    elif"切塊"in text_segment and"影響"in text_segment:
        return ["切塊策略直接影響檢索的質(zhì)量。", "如果切塊過大，會增加LLM處理負(fù)擔(dān)；如果過小，則可能丟失上下文。"]
    else:
        # 如果LLM無法智能切分，就回退到句子切塊
        return sent_tokenize(text_segment)

def agentic_llm_based_chunking(text, max_segment_for_llm=500):
    """
    使用LLM來決定切塊邊界。
    由于LLM調(diào)用成本，通常我們會先將大文本切分成適合LLM處理的段落，
    然后讓LLM對這些段落進行細(xì)粒度切塊。
    """
    # 先進行一個粗粒度的切塊（例如，按段落或固定大小），確保每個段落大小適合LLM處理
    coarse_chunks = fixed_size_chunking(text, max_segment_for_llm, overlap=0)
    
    final_chunks = []
    for chunk in coarse_chunks:
        # 模擬LLM對每個粗粒度塊進行智能切分
        llm_decided_sub_chunks = mock_llm_chunking_decision(chunk)
        final_chunks.extend(llm_decided_sub_chunks)
        
    return [c.strip() for c in final_chunks if c.strip()]

print("\n--- 19. Agentic / 基于LLM切塊 ---")
chunks_llm_based = agentic_llm_based_chunking(sample_text_long, max_segment_for_llm=300)
for i, chunk in enumerate(chunks_llm_based[:5]):
    print(f"Chunk {i+1}:\n'{chunk}'")
print("...")

class HierarchicalChunk:
    def __init__(self, content, level, title=None, children=None):
        self.content = content
        self.level = level
        self.title = title
        self.children = children if children isnotNoneelse []

    def __repr__(self):
        returnf"Level {self.level} '{self.title or self.content[:30]}...'"

def parse_markdown_hierarchy(text):
    """
    解析Markdown文本，構(gòu)建分層結(jié)構(gòu)。
    返回一個包含頂級HierarchicalChunk對象的列表。
    """
    lines = text.split('\n')
    
    # 存儲當(dāng)前的層級路徑，方便構(gòu)建嵌套結(jié)構(gòu)
    # Stack stores (level, parent_chunk)
    root_chunks = []
    current_path = [(0, None)] # (level, parent_chunk)

    for line in lines:
        stripped_line = line.strip()
        ifnot stripped_line:
            continue

        match = re.match(r'^(#+)\s*(.*)$', stripped_line)
        if match:
            level = len(match.group(1)) # 標(biāo)題級別
            title = match.group(2).strip()
            
            new_chunk = HierarchicalChunk(cnotallow="", level=level, title=title)
            
            # 回溯到正確的父級
            while current_path and current_path[-1][0] >= level:
                current_path.pop()

            if current_path and current_path[-1][1]: # 有父級
                current_path[-1][1].children.append(new_chunk)
            else: # 頂級標(biāo)題
                root_chunks.append(new_chunk)
            
            current_path.append((level, new_chunk))
        else: # 普通內(nèi)容，添加到當(dāng)前最低層級塊的內(nèi)容
            if current_path and current_path[-1][1]:
                # 如果是第一個內(nèi)容行，直接賦值，否則追加
                if current_path[-1][1].content:
                    current_path[-1][1].content += "\n" + line
                else:
                    current_path[-1][1].content = line
            else: # 沒有標(biāo)題的開頭內(nèi)容，作為頂級塊
                ifnot root_chunks or root_chunks[-1].level != 0or root_chunks[-1].title: # 如果沒有頂級塊或者上一個是標(biāo)題，就創(chuàng)建一個新的
                     root_chunks.append(HierarchicalChunk(cnotallow=line, level=0))
                else: # 追加到第一個無標(biāo)題頂級塊
                     root_chunks[-1].content += "\n" + line

    # 遞歸清理并整合內(nèi)容
    def consolidate_chunks(chunk_list):
        final_chunks = []
        for chunk in chunk_list:
            # 將標(biāo)題本身和內(nèi)容整合到content中
            full_content = ""
            if chunk.title:
                full_content += "#" * chunk.level + " " + chunk.title + "\n"
            full_content += chunk.content.strip()
            
            if full_content: # 確保內(nèi)容不為空
                 final_chunks.append(HierarchicalChunk(full_content, chunk.level, chunk.title))
            
            if chunk.children:
                final_chunks.extend(consolidate_chunks(chunk.children))
        return final_chunks

    return consolidate_chunks(root_chunks)


print("\n--- 20. 分層切塊 ---")
# 使用一個更適合分層切塊的結(jié)構(gòu)化文本
hierarchical_text = """
# 第一章 RAG概述
RAG是一種強大的AI技術(shù)。

## 1.1 RAG的原理
結(jié)合檢索和生成。
### 1.1.1 檢索部分
從知識庫中獲取信息。

## 1.2 RAG的優(yōu)勢
減少幻覺，提升準(zhǔn)確性。

# 第二章 切塊策略
切塊是RAG的關(guān)鍵一步。
"""
hierarchical_chunks = parse_markdown_hierarchy(hierarchical_text)
for i, chunk_obj in enumerate(hierarchical_chunks):
    print(f"Chunk {i+1} (Level {chunk_obj.level}, Title: '{chunk_obj.title}'):\n'{chunk_obj.content}'")

def modality_aware_chunking(text):
    """
    分離不同模態(tài)的內(nèi)容（文本、表格）。
    這里只處理文本和Markdown表格。
    """
    chunks = []
    lines = text.split('\n')
    current_chunk_lines = []
    
    in_table_block = False

    for line in lines:
        stripped_line = line.strip()

        # 檢查是否是Markdown表格行
        is_table_line = stripped_line.startswith('|') and'|'in stripped_line[1:]

        if is_table_line:
            ifnot in_table_block:
                # 結(jié)束之前的文本塊
                if current_chunk_lines:
                    chunks.append({"type": "text", "content": "\n".join(current_chunk_lines).strip()})
                    current_chunk_lines = []
                in_table_block = True
            current_chunk_lines.append(line)
        else:
            if in_table_block:
                # 結(jié)束表格塊
                if current_chunk_lines:
                    chunks.append({"type": "table", "content": "\n".join(current_chunk_lines).strip()})
                    current_chunk_lines = []
                in_table_block = False
            current_chunk_lines.append(line)

    # 處理最后一個塊
    if current_chunk_lines:
        chunk_type = "table"if in_table_block else"text"
        chunks.append({"type": chunk_type, "content": "\n".join(current_chunk_lines).strip()})
        
    return [chunk for chunk in chunks if chunk['content']]

print("\n--- 21. 模態(tài)感知切塊 ---")
chunks_modality = modality_aware_chunking(sample_text_mixed_format)
for i, chunk in enumerate(chunks_modality):
    print(f"Chunk {i+1} (Type: {chunk['type']}):\n'{chunk['content']}'")

def hybrid_chunking(text, primary_strategy, secondary_strategy, primary_args={}, secondary_args={}):
    """
    混合切塊策略示例：先用一種策略粗切，再用另一種策略細(xì)切。
    primary_strategy: 第一階段切塊函數(shù) (如 paragraph_based_chunking)
    secondary_strategy: 第二階段切塊函數(shù) (如 recursive_chunking)
    """
    # 步驟1：用主要策略進行粗粒度切塊
    coarse_chunks = primary_strategy(text, **primary_args)
    
    final_chunks = []
    # 步驟2：對每個粗粒度塊，再用次要策略進行細(xì)粒度切塊
    for chunk in coarse_chunks:
        # 如果粗粒度塊仍然太大或需要進一步細(xì)分
        if len(chunk) > 500: # 假設(shè)一個啟發(fā)式條件，可以根據(jù)token數(shù)或語義復(fù)雜度來定
            # 注意：這里需要確保secondary_strategy能夠處理傳入的參數(shù)
            fine_grained_chunks = secondary_strategy(chunk, **secondary_args)
            final_chunks.extend(fine_grained_chunks)
        else:
            final_chunks.append(chunk)
            
    return [chunk for chunk in final_chunks if chunk.strip()]

print("\n--- BONUS: 混合切塊 ---")
# 示例：先按段落切塊，然后對超過一定長度的段落進行遞歸切塊
chunks_hybrid = hybrid_chunking(
    sample_text_long + sample_text_structured,
    primary_strategy=paragraph_based_chunking,
    primary_args={}, # 段落切塊不需要額外參數(shù)
    secondary_strategy=recursive_chunking,
    secondary_args={'separators': ['. ', '，'], 'max_chunk_size_char': 200}
)

for i, chunk in enumerate(chunks_hybrid[:5]):
    print(f"Chunk {i+1} (len={len(chunk)}):\n'{chunk}'")
print("...")