# 这个脚本讲义的代码架构图没有体现，需要进行补充
import os
from langchain_community.document_loaders import TextLoader
from langchain_community.document_loaders.markdown import UnstructuredMarkdownLoader # NLTK
from langchain_text_splitters import MarkdownTextSplitter
from datetime import datetime
import sys
# 获取当前文件所在目录的绝对路径
current_dir = os.path.dirname(os.path.abspath(__file__))
# print(f'current_dir--》{current_dir}')
# 获取core文件所在的目录的绝对路径
rag_qa_path = os.path.dirname(current_dir)
# print(f'rag_qa_path--》{rag_qa_path}')
sys.path.insert(0, rag_qa_path)
# 获取根目录文件所在的绝对位置
project_root = os.path.dirname(rag_qa_path)
sys.path.insert(0, project_root)
from rag_qa.edu_document_loaders import OCRPDFLoader, OCRDOCLoader, OCRPPTLoader, OCRIMGLoader
from rag_qa.edu_text_spliter import ChineseRecursiveTextSplitter
from base import logger, Config
conf = Config()
# 定义支持的文件类型及其对应的加载器字典
document_loaders = {
# 文本文件使用 TextLoader
".txt": TextLoader,
# PDF 文件使用 OCRPDFLoader
".pdf": OCRPDFLoader,
# Word 文件使用 OCRDOCLoader
".docx": OCRDOCLoader,
# PPT 文件使用 OCRPPTLoader
".ppt": OCRPPTLoader,
# PPTX 文件使用 OCRPPTLoader
".pptx": OCRPPTLoader,
# JPG 文件使用 OCRIMGLoader
".jpg": OCRIMGLoader,
# PNG 文件使用 OCRIMGLoader
".png": OCRIMGLoader,
# Markdown 文件使用 UnstructuredMarkdownLoader
".md": UnstructuredMarkdownLoader
}
# 定义函数，从指定文件夹加载多种类型文件并添加元数据
def load_documents_from_directory(directory_path):
# 初始化空列表，用于存储加载的文档
documents = []
# 获取支持的文件扩展名集合
supported_extensions = document_loaders.keys()
# print(f'supported_extensions--》{supported_extensions}')
# 从目录名提取学科类别（如 "ai_data" -> "ai"）
# print(f'1---》{os.path.basename(directory_path)}')
source = os.path.basename(directory_path).replace("_data", "")
# print(f'source-->{source}')
# 遍历指定目录及其子目录
for root, _, files in os.walk(directory_path):
# print(f'root---》{root}')
# print(f'files---》{files}')
# 遍历当前目录下的所有文件
for file in files:
# 构造文件的完整路径
file_path = os.path.join(root, file)
# print(f'file_path--》{file_path}')
# print(os.path.splitext(file_path))
# 获取文件扩展名并转换为小写
file_extension = os.path.splitext(file_path)[1].lower()
# print(f'file_extension--》{file_extension}')
# 检查文件类型是否在支持的扩展名列表中
if file_extension in supported_extensions:
# 使用 try-except 捕获加载过程中的异常
try:
# 根据文件扩展名获取对应的加载器类
loader_class = document_loaders[file_extension]
# 实例化加载器对象，传入文件路径
if file_extension == ".txt":
loader = loader_class(file_path, encoding="utf-8")
else:
loader = loader_class(file_path)
# 调用加载器加载文档内容，返回文档列表
loaded_docs = loader.load()
# print(f'loaded_docs--》{loaded_docs}')
# print(f'loaded_docs--》{len(loaded_docs)}')
for doc in loaded_docs:
# 为文档添加学科类别元数据
doc.metadata["source"] = source
# 为文档添加文件路径元数据
doc.metadata["file_path"] = file_path
# 为文档添加当前时间戳元数据
doc.metadata["timestamp"] = datetime.now().isoformat()
# print(f'loaded_docs111--》{loaded_docs}')
documents.extend(loaded_docs)
# 记录成功加载文件的日志
logger.info(f"成功加载文件: {file_path}")
except Exception as e:
logger.error(f"加载文件 {file_path} 失败: {str(e)}")
# 如果文件类型不在支持列表中
else:
# 记录警告日志，提示不支持的文件类型
logger.warning(f"不支持的文件类型: {file_path}")
# 返回加载的所有文档列表
return documents
# 定义函数，处理文档并进行分层切分，返回子块结果
def process_documents(directory_path,
parent_chunk_size=conf.PARENT_CHUNK_SIZE,
child_chunk_size=conf.CHILD_CHUNK_SIZE,
chunk_overlap=conf.CHUNK_OVERLAP):
# 从指定目录加载所有文档
documents = load_documents_from_directory(directory_path)
# 记录加载的文档总数日志
logger.info(f"加载的文档数量: {len(documents)}")
# 初始化父块和子块分词器（通用）
parent_splitter = ChineseRecursiveTextSplitter(chunk_size=parent_chunk_size, chunk_overlap=chunk_overlap)
child_splitter = ChineseRecursiveTextSplitter(chunk_size=child_chunk_size, chunk_overlap=chunk_overlap)
# 初始化 Markdown 专用分词器
markdown_parent_splitter = MarkdownTextSplitter(chunk_size=parent_chunk_size, chunk_overlap=chunk_overlap)
markdown_child_splitter = MarkdownTextSplitter(chunk_size=child_chunk_size, chunk_overlap=chunk_overlap)
# 初始化空列表，用于存储所有子块
child_chunks = []
# 遍历每个原始文档，带上索引 i
for i, doc in enumerate(documents):
# print(f'doc--》{doc}')
# print(doc.metadata.get("file_path", ''))
# # 获取文件的扩展名
# print(os.path.splitext(doc.metadata.get("file_path", '')))
file_extension = os.path.splitext(doc.metadata.get("file_path", ''))[1].lower()
# print(f'file_extension--》{file_extension}')
# 选择分词器
is_markdown = (file_extension == '.md')
# print(f'is_markdown--》{is_markdown}')
parent_splitter_to_use = markdown_parent_splitter if is_markdown else parent_splitter
# print(f'parent_splitter_to_use-->{parent_splitter_to_use}')
child_splitter_to_use = markdown_child_splitter if is_markdown else child_splitter
logger.info(
f"处理文档: {doc.metadata['file_path']}, 使用切分器: {'Markdown' if is_markdown else 'ChineseRecursive'}")
# 使用父块切分器将文档切分为父块
parent_docs = parent_splitter_to_use.split_documents([doc])
# print(f'parent_docs--》{parent_docs}')
# print(f'parent_docs--》{len(parent_docs)}')
# 遍历每个父块，带上索引 j
for j, parent_doc in enumerate(parent_docs):
# 为父块生成唯一 ID，格式为 "doc_i_parent_j"
parent_id = f"doc_{i}_parent_{j}"
# # 将父块 ID 添加到元数据
parent_doc.metadata["parent_id"] = parent_id
# # 将父块内容存储到元数据
parent_doc.metadata["parent_content"] = parent_doc.page_content
# 使用子块分词器将父块切分为子块
sub_chunks = child_splitter_to_use.split_documents([parent_doc])
# print(f'sub_chunks--》{sub_chunks}')
# print(f'sub_chunks--》{len(sub_chunks)}')
# 遍历每个子块，为子块主要添加对应的父块文档
for k, sub_chunk in enumerate(sub_chunks):
# print(f'原始的sub_chunk--》{sub_chunk}')
# 为子块添加父块的ID
sub_chunk.metadata["parent_id"] = parent_id
# 为子块添加对应的父块文档（元数据）
sub_chunk.metadata["parent_content"] = parent_doc.page_content
# 为子块生成一个唯一的ID,格式为 "parent_id_child_k"
sub_chunk.metadata["id"] = f"{parent_id}_child_{k}"
# print(f'修改后的sub_chunk--》{sub_chunk}')
# 将子块添加到子块列表中
child_chunks.append(sub_chunk)
# 记录子块总数日志
logger.info(f"子块数量: {len(child_chunks)}")
# 返回所有子块列表
return child_chunks
if __name__ == '__main__':
directory_path = r'D:/code/workspace7/integrated_qa_system/rag_qa/data/ai_data'
# documents = load_documents_from_directory(directory_path)
# print('documents-->', documents)
child_chunks = process_documents(directory_path)
print("child_chunks--", child_chunks)