364 lines
13 KiB
Python
364 lines
13 KiB
Python
"""
|
|
数据准备模块
|
|
"""
|
|
|
|
import logging
|
|
import hashlib
|
|
from pathlib import Path
|
|
from typing import List, Dict, Any
|
|
|
|
from langchain_text_splitters import MarkdownHeaderTextSplitter
|
|
from langchain_core.documents import Document
|
|
from pathlib import Path
|
|
import uuid
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class DataPreparationModule:
|
|
"""数据准备模块 - 负责数据加载、清洗和预处理"""
|
|
# 统一维护的分类与难度配置，供外部复用，避免关键词重复定义
|
|
CATEGORY_MAPPING = {
|
|
'meat_dish': '荤菜',
|
|
'vegetable_dish': '素菜',
|
|
'soup': '汤品',
|
|
'dessert': '甜品',
|
|
'breakfast': '早餐',
|
|
'staple': '主食',
|
|
'aquatic': '水产',
|
|
'condiment': '调料',
|
|
'drink': '饮品'
|
|
}
|
|
CATEGORY_LABELS = list(set(CATEGORY_MAPPING.values()))
|
|
DIFFICULTY_LABELS = ['非常简单', '简单', '中等', '困难', '非常困难']
|
|
|
|
def __init__(self, data_path: str):
|
|
"""
|
|
初始化数据准备模块
|
|
|
|
Args:
|
|
data_path: 数据文件夹路径
|
|
"""
|
|
self.data_path = data_path
|
|
self.documents: List[Document] = [] # 父文档（完整食谱）
|
|
self.chunks: List[Document] = [] # 子文档（按标题分割的小块）
|
|
self.parent_child_map: Dict[str, str] = {} # 子块ID -> 父文档ID的映射
|
|
|
|
def load_documents(self) -> List[Document]:
|
|
"""
|
|
加载文档数据
|
|
|
|
Returns:
|
|
加载的文档列表
|
|
"""
|
|
logger.info(f"正在从 {self.data_path} 加载文档...")
|
|
|
|
# 直接读取Markdown文件以保持原始格式
|
|
documents = []
|
|
data_path_obj = Path(self.data_path)
|
|
|
|
for md_file in data_path_obj.rglob("*.md"):
|
|
try:
|
|
# 直接读取文件内容，保持Markdown格式
|
|
with open(md_file, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
|
|
# 为每个父文档分配确定性的唯一ID（基于数据根目录的相对路径）
|
|
try:
|
|
data_root = Path(self.data_path).resolve()
|
|
relative_path = Path(md_file).resolve().relative_to(data_root).as_posix()
|
|
except Exception:
|
|
relative_path = Path(md_file).as_posix()
|
|
parent_id = hashlib.md5(relative_path.encode("utf-8")).hexdigest()
|
|
|
|
# 创建Document对象
|
|
doc = Document(
|
|
page_content=content,
|
|
metadata={
|
|
"source": str(md_file),
|
|
"parent_id": parent_id,
|
|
"doc_type": "parent" # 标记为父文档
|
|
}
|
|
)
|
|
documents.append(doc)
|
|
|
|
except Exception as e:
|
|
logger.warning(f"读取文件 {md_file} 失败: {e}")
|
|
|
|
# 增强文档元数据
|
|
for doc in documents:
|
|
self._enhance_metadata(doc)
|
|
|
|
self.documents = documents
|
|
logger.info(f"成功加载 {len(documents)} 个文档")
|
|
return documents
|
|
|
|
def _enhance_metadata(self, doc: Document):
|
|
"""
|
|
增强文档元数据
|
|
|
|
Args:
|
|
doc: 需要增强元数据的文档
|
|
"""
|
|
file_path = Path(doc.metadata.get('source', ''))
|
|
path_parts = file_path.parts
|
|
|
|
# 提取菜品分类
|
|
doc.metadata['category'] = '其他'
|
|
for key, value in self.CATEGORY_MAPPING.items():
|
|
if key in path_parts:
|
|
doc.metadata['category'] = value
|
|
break
|
|
|
|
# 提取菜品名称
|
|
doc.metadata['dish_name'] = file_path.stem
|
|
|
|
# 分析难度等级
|
|
content = doc.page_content
|
|
if '★★★★★' in content:
|
|
doc.metadata['difficulty'] = '非常困难'
|
|
elif '★★★★' in content:
|
|
doc.metadata['difficulty'] = '困难'
|
|
elif '★★★' in content:
|
|
doc.metadata['difficulty'] = '中等'
|
|
elif '★★' in content:
|
|
doc.metadata['difficulty'] = '简单'
|
|
elif '★' in content:
|
|
doc.metadata['difficulty'] = '非常简单'
|
|
else:
|
|
doc.metadata['difficulty'] = '未知'
|
|
|
|
@classmethod
|
|
def get_supported_categories(cls) -> List[str]:
|
|
"""对外提供支持的分类标签列表"""
|
|
return cls.CATEGORY_LABELS
|
|
|
|
@classmethod
|
|
def get_supported_difficulties(cls) -> List[str]:
|
|
"""对外提供支持的难度标签列表"""
|
|
return cls.DIFFICULTY_LABELS
|
|
|
|
def chunk_documents(self) -> List[Document]:
|
|
"""
|
|
Markdown结构感知分块
|
|
|
|
Returns:
|
|
分块后的文档列表
|
|
"""
|
|
logger.info("正在进行Markdown结构感知分块...")
|
|
|
|
if not self.documents:
|
|
raise ValueError("请先加载文档")
|
|
|
|
# 使用Markdown标题分割器
|
|
chunks = self._markdown_header_split()
|
|
|
|
# 为每个chunk添加基础元数据
|
|
for i, chunk in enumerate(chunks):
|
|
if 'chunk_id' not in chunk.metadata:
|
|
# 如果没有chunk_id（比如分割失败的情况），则生成一个
|
|
chunk.metadata['chunk_id'] = str(uuid.uuid4())
|
|
chunk.metadata['batch_index'] = i # 在当前批次中的索引
|
|
chunk.metadata['chunk_size'] = len(chunk.page_content)
|
|
|
|
self.chunks = chunks
|
|
logger.info(f"Markdown分块完成，共生成 {len(chunks)} 个chunk")
|
|
return chunks
|
|
|
|
def _markdown_header_split(self) -> List[Document]:
|
|
"""
|
|
使用Markdown标题分割器进行结构化分割
|
|
|
|
Returns:
|
|
按标题结构分割的文档列表
|
|
"""
|
|
# 定义要分割的标题层级
|
|
headers_to_split_on = [
|
|
("#", "主标题"), # 菜品名称
|
|
("##", "二级标题"), # 必备原料、计算、操作等
|
|
("###", "三级标题") # 简易版本、复杂版本等
|
|
]
|
|
|
|
# 创建Markdown分割器
|
|
markdown_splitter = MarkdownHeaderTextSplitter(
|
|
headers_to_split_on=headers_to_split_on,
|
|
strip_headers=False # 保留标题，便于理解上下文
|
|
)
|
|
|
|
all_chunks = []
|
|
|
|
for doc in self.documents:
|
|
try:
|
|
# 检查文档内容是否包含Markdown标题
|
|
content_preview = doc.page_content[:200]
|
|
has_headers = any(line.strip().startswith('#') for line in content_preview.split('\n'))
|
|
|
|
if not has_headers:
|
|
logger.warning(f"文档 {doc.metadata.get('dish_name', '未知')} 内容中没有发现Markdown标题")
|
|
logger.debug(f"内容预览: {content_preview}")
|
|
|
|
# 对每个文档进行Markdown分割
|
|
md_chunks = markdown_splitter.split_text(doc.page_content)
|
|
|
|
logger.debug(f"文档 {doc.metadata.get('dish_name', '未知')} 分割成 {len(md_chunks)} 个chunk")
|
|
|
|
# 如果没有分割成功，说明文档可能没有标题结构
|
|
if len(md_chunks) <= 1:
|
|
logger.warning(f"文档 {doc.metadata.get('dish_name', '未知')} 未能按标题分割，可能缺少标题结构")
|
|
|
|
# 为每个子块建立与父文档的关系
|
|
parent_id = doc.metadata["parent_id"]
|
|
|
|
for i, chunk in enumerate(md_chunks):
|
|
# 为子块分配唯一ID
|
|
child_id = str(uuid.uuid4())
|
|
|
|
# 合并原文档元数据和新的标题元数据
|
|
chunk.metadata.update(doc.metadata)
|
|
chunk.metadata.update({
|
|
"chunk_id": child_id,
|
|
"parent_id": parent_id,
|
|
"doc_type": "child", # 标记为子文档
|
|
"chunk_index": i # 在父文档中的位置
|
|
})
|
|
|
|
# 建立父子映射关系
|
|
self.parent_child_map[child_id] = parent_id
|
|
|
|
all_chunks.extend(md_chunks)
|
|
|
|
except Exception as e:
|
|
logger.warning(f"文档 {doc.metadata.get('source', '未知')} Markdown分割失败: {e}")
|
|
# 如果Markdown分割失败，将整个文档作为一个chunk
|
|
all_chunks.append(doc)
|
|
|
|
logger.info(f"Markdown结构分割完成，生成 {len(all_chunks)} 个结构化块")
|
|
return all_chunks
|
|
|
|
def filter_documents_by_category(self, category: str) -> List[Document]:
|
|
"""
|
|
按分类过滤文档
|
|
|
|
Args:
|
|
category: 菜品分类
|
|
|
|
Returns:
|
|
过滤后的文档列表
|
|
"""
|
|
return [doc for doc in self.documents if doc.metadata.get('category') == category]
|
|
|
|
def filter_documents_by_difficulty(self, difficulty: str) -> List[Document]:
|
|
"""
|
|
按难度过滤文档
|
|
|
|
Args:
|
|
difficulty: 难度等级
|
|
|
|
Returns:
|
|
过滤后的文档列表
|
|
"""
|
|
return [doc for doc in self.documents if doc.metadata.get('difficulty') == difficulty]
|
|
|
|
def get_statistics(self) -> Dict[str, Any]:
|
|
"""
|
|
获取数据统计信息
|
|
|
|
Returns:
|
|
统计信息字典
|
|
"""
|
|
if not self.documents:
|
|
return {}
|
|
|
|
categories = {}
|
|
difficulties = {}
|
|
|
|
for doc in self.documents:
|
|
# 统计分类
|
|
category = doc.metadata.get('category', '未知')
|
|
categories[category] = categories.get(category, 0) + 1
|
|
|
|
# 统计难度
|
|
difficulty = doc.metadata.get('difficulty', '未知')
|
|
difficulties[difficulty] = difficulties.get(difficulty, 0) + 1
|
|
|
|
return {
|
|
'total_documents': len(self.documents),
|
|
'total_chunks': len(self.chunks),
|
|
'categories': categories,
|
|
'difficulties': difficulties,
|
|
'avg_chunk_size': sum(chunk.metadata.get('chunk_size', 0) for chunk in self.chunks) / len(self.chunks) if self.chunks else 0
|
|
}
|
|
|
|
def export_metadata(self, output_path: str):
|
|
"""
|
|
导出元数据到JSON文件
|
|
|
|
Args:
|
|
output_path: 输出文件路径
|
|
"""
|
|
import json
|
|
|
|
metadata_list = []
|
|
for doc in self.documents:
|
|
metadata_list.append({
|
|
'source': doc.metadata.get('source'),
|
|
'dish_name': doc.metadata.get('dish_name'),
|
|
'category': doc.metadata.get('category'),
|
|
'difficulty': doc.metadata.get('difficulty'),
|
|
'content_length': len(doc.page_content)
|
|
})
|
|
|
|
with open(output_path, 'w', encoding='utf-8') as f:
|
|
json.dump(metadata_list, f, ensure_ascii=False, indent=2)
|
|
|
|
logger.info(f"元数据已导出到: {output_path}")
|
|
|
|
def get_parent_documents(self, child_chunks: List[Document]) -> List[Document]:
|
|
"""
|
|
根据子块获取对应的父文档（智能去重）
|
|
|
|
Args:
|
|
child_chunks: 检索到的子块列表
|
|
|
|
Returns:
|
|
对应的父文档列表（去重，按相关性排序）
|
|
"""
|
|
# 统计每个父文档被匹配的次数（相关性指标）
|
|
parent_relevance = {}
|
|
parent_docs_map = {}
|
|
|
|
# 收集所有相关的父文档ID和相关性分数
|
|
for chunk in child_chunks:
|
|
parent_id = chunk.metadata.get("parent_id")
|
|
if parent_id:
|
|
# 增加相关性计数
|
|
parent_relevance[parent_id] = parent_relevance.get(parent_id, 0) + 1
|
|
|
|
# 缓存父文档（避免重复查找）
|
|
if parent_id not in parent_docs_map:
|
|
for doc in self.documents:
|
|
if doc.metadata.get("parent_id") == parent_id:
|
|
parent_docs_map[parent_id] = doc
|
|
break
|
|
|
|
# 按相关性排序（匹配次数多的排在前面）
|
|
sorted_parent_ids = sorted(parent_relevance.keys(),
|
|
key=lambda x: parent_relevance[x],
|
|
reverse=True)
|
|
|
|
# 构建去重后的父文档列表
|
|
parent_docs = []
|
|
for parent_id in sorted_parent_ids:
|
|
if parent_id in parent_docs_map:
|
|
parent_docs.append(parent_docs_map[parent_id])
|
|
|
|
# 收集父文档名称和相关性信息用于日志
|
|
parent_info = []
|
|
for doc in parent_docs:
|
|
dish_name = doc.metadata.get('dish_name', '未知菜品')
|
|
parent_id = doc.metadata.get('parent_id')
|
|
relevance_count = parent_relevance.get(parent_id, 0)
|
|
parent_info.append(f"{dish_name}({relevance_count}块)")
|
|
|
|
logger.info(f"从 {len(child_chunks)} 个子块中找到 {len(parent_docs)} 个去重父文档: {', '.join(parent_info)}")
|
|
return parent_docs
|