# 安装依赖：pip install spacy sentence-transformers chromadb
# 下载spacy模型：python -m spacy download en_core_web_sm

import spacy
import re
from typing import List

# 加载spacy模型（检测句子和语义边界）
nlp = spacy.load("en_core_web_sm")

def semantic_chunking_with_sliding_window(
    text: str, 
    chunk_size: int = 300, 
    overlap: int = 50
) -> List[str]:
    """
    语义分块 + 滑动窗口切分文本
    参数：
        text: 输入文本
        chunk_size: 目标分块大小（字符数）
        overlap: 滑动窗口重叠长度（字符数）
    返回：
        chunks: 分块后的文本列表
    """
    # Step 1: 预处理（移除多余空格、换行符）
    text = re.sub(r'\s+', ' ', text).strip()
    
    # Step 2: 使用spacy分句（语义分块）
    doc = nlp(text)
    sentences = [sent.text.strip() for sent in doc.sents]
    
    # Step 3: 合并句子为初始块（基于chunk_size）
    chunks = []
    current_chunk = []
    current_length = 0
    
    for sent in sentences:
        sent_length = len(sent)
        if current_length + sent_length <= chunk_size:
            current_chunk.append(sent)
            current_length += sent_length
        else:
            chunks.append(' '.join(current_chunk))
            current_chunk = [sent]
            current_length = sent_length
    
    if current_chunk:
        chunks.append(' '.join(current_chunk))
    
    # Step 4: 应用滑动窗口（添加重叠区域）
    final_chunks = []
    for i in range(len(chunks)):
        current = chunks[i]
        if i > 0:
            # 取前一块的末尾overlap部分与当前块合并
            prev_chunk_tail = chunks[i-1][-overlap:]
            current = prev_chunk_tail + ' ' + current
        final_chunks.append(current[:chunk_size])  # 确保不超过chunk_size
    
    return final_chunks

# 测试示例
if __name__ == "__main__":
    sample_text = """
    Large language models (LLMs) have revolutionized AI applications. 
    However, their knowledge is static after training. Retrieval-Augmented 
    Generation (RAG) addresses this by dynamically retrieving relevant 
    documents during inference. For example, in medical diagnosis, 
    RAG systems can pull the latest research papers to provide up-to-date 
    recommendations. The key steps include: document chunking, vector 
    retrieval, and context-aware generation. Effective chunking strategies 
    like semantic segmentation with sliding windows improve retrieval accuracy.
    """
    
    chunks = semantic_chunking_with_sliding_window(sample_text, chunk_size=150, overlap=30)
    for i, chunk in enumerate(chunks):
        print(f"Chunk {i+1} ({len(chunk)} chars):\n{chunk}\n{'-'*50}")