269 lines
9.4 KiB
Python
269 lines
9.4 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
分批处理管理器 - 管理菜谱处理的分批保存和断点续传功能
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import sys
|
|
import argparse
|
|
from datetime import datetime
|
|
from recipe_ai_agent import KimiRecipeAgent, RecipeKnowledgeGraphBuilder
|
|
|
|
def load_config():
|
|
"""加载配置文件"""
|
|
config_file = "config.json"
|
|
if os.path.exists(config_file):
|
|
with open(config_file, 'r', encoding='utf-8') as f:
|
|
return json.load(f)
|
|
else:
|
|
print("❌ 未找到config.json配置文件")
|
|
sys.exit(1)
|
|
|
|
def show_progress_status(output_dir: str):
|
|
"""显示处理进度状态"""
|
|
progress_file = os.path.join(output_dir, "progress.json")
|
|
|
|
if not os.path.exists(progress_file):
|
|
print("📋 未找到进度文件，没有正在进行的任务")
|
|
return
|
|
|
|
try:
|
|
with open(progress_file, 'r', encoding='utf-8') as f:
|
|
progress_data = json.load(f)
|
|
|
|
print("处理进度状态:")
|
|
print(f"总文件: {progress_data.get('total_files', 'N/A')}, 已处理: {progress_data.get('processed_count', 0)}")
|
|
|
|
status = progress_data.get('current_file', 'N/A')
|
|
if status == 'COMPLETED':
|
|
print("状态: 已完成")
|
|
elif status == 'INTERRUPTED':
|
|
print("状态: 被中断")
|
|
else:
|
|
print("状态: 进行中")
|
|
|
|
# 检查批次文件
|
|
batch_dirs = [d for d in os.listdir(output_dir)
|
|
if d.startswith("batch_") and os.path.isdir(os.path.join(output_dir, d))]
|
|
|
|
if batch_dirs:
|
|
print(f"已保存批次: {len(batch_dirs)} 个")
|
|
|
|
except Exception as e:
|
|
print(f"❌ 读取进度文件失败: {str(e)}")
|
|
|
|
def clean_progress(output_dir: str):
|
|
"""清理进度文件，重新开始处理"""
|
|
progress_file = os.path.join(output_dir, "progress.json")
|
|
|
|
if os.path.exists(progress_file):
|
|
confirm = input("⚠️ 确认要清理进度文件吗？这将删除所有处理进度 (y/N): ").strip().lower()
|
|
if confirm == 'y':
|
|
os.remove(progress_file)
|
|
print("✅ 进度文件已清理")
|
|
else:
|
|
print("取消操作")
|
|
else:
|
|
print("📋 未找到进度文件")
|
|
|
|
def clean_batches(output_dir: str):
|
|
"""清理所有批次数据"""
|
|
batch_dirs = [d for d in os.listdir(output_dir)
|
|
if d.startswith("batch_") and os.path.isdir(os.path.join(output_dir, d))]
|
|
|
|
if not batch_dirs:
|
|
print("📁 未找到批次数据")
|
|
return
|
|
|
|
print(f"找到 {len(batch_dirs)} 个批次目录:")
|
|
for batch_dir in sorted(batch_dirs):
|
|
print(f" - {batch_dir}")
|
|
|
|
confirm = input("\n⚠️ 确认要删除所有批次数据吗？ (y/N): ").strip().lower()
|
|
if confirm == 'y':
|
|
import shutil
|
|
for batch_dir in batch_dirs:
|
|
batch_path = os.path.join(output_dir, batch_dir)
|
|
shutil.rmtree(batch_path)
|
|
print(f"🗑️ 已删除: {batch_dir}")
|
|
print("✅ 所有批次数据已清理")
|
|
else:
|
|
print("取消操作")
|
|
|
|
def merge_batches(output_dir: str):
|
|
"""手动合并所有批次数据"""
|
|
config = load_config()
|
|
api_key = config["kimi"].get("api_key")
|
|
|
|
if not api_key:
|
|
print("❌ 未找到API密钥配置")
|
|
return
|
|
|
|
try:
|
|
ai_agent = KimiRecipeAgent(api_key)
|
|
builder = RecipeKnowledgeGraphBuilder(ai_agent, output_dir)
|
|
|
|
print("合并批次数据...")
|
|
total_concepts, total_relationships = builder.merge_all_batches()
|
|
|
|
if total_concepts > 0 or total_relationships > 0:
|
|
print(f"合并完成: {total_concepts} 概念, {total_relationships} 关系")
|
|
|
|
# 生成Neo4j格式
|
|
format_type = config["output"].get("format", "neo4j")
|
|
if format_type == "neo4j":
|
|
builder.export_to_neo4j_csv(output_dir, merge_batches=False)
|
|
else:
|
|
print("未找到有效的批次数据")
|
|
|
|
except Exception as e:
|
|
print(f"❌ 合并失败: {str(e)}")
|
|
|
|
def continue_processing(recipe_dir: str, output_dir: str):
|
|
"""继续处理中断的任务"""
|
|
config = load_config()
|
|
api_key = config["kimi"].get("api_key")
|
|
|
|
if not api_key:
|
|
print("❌ 未找到API密钥配置")
|
|
return
|
|
|
|
try:
|
|
ai_agent = KimiRecipeAgent(api_key)
|
|
batch_size = config.get("processing", {}).get("batch_size", 20)
|
|
builder = RecipeKnowledgeGraphBuilder(ai_agent, output_dir, batch_size)
|
|
|
|
print("继续处理中断的任务...")
|
|
processed, failed = builder.batch_process_recipes(recipe_dir, resume=True)
|
|
|
|
print(f"处理完成: 总数 {processed}, 失败 {failed}")
|
|
|
|
# 自动合并数据
|
|
print("自动合并批次数据...")
|
|
merge_batches(output_dir)
|
|
|
|
except Exception as e:
|
|
print(f"❌ 处理失败: {str(e)}")
|
|
|
|
def show_batch_details(output_dir: str, batch_num: int = None):
|
|
"""显示批次详细信息"""
|
|
batch_dirs = [d for d in os.listdir(output_dir)
|
|
if d.startswith("batch_") and os.path.isdir(os.path.join(output_dir, d))]
|
|
|
|
if not batch_dirs:
|
|
print("📁 未找到批次数据")
|
|
return
|
|
|
|
batch_dirs.sort()
|
|
|
|
if batch_num is not None:
|
|
target_batch = f"batch_{batch_num:03d}"
|
|
if target_batch not in batch_dirs:
|
|
print(f"❌ 未找到批次 {batch_num}")
|
|
return
|
|
batch_dirs = [target_batch]
|
|
|
|
import pandas as pd
|
|
|
|
for batch_dir in batch_dirs:
|
|
print(f"\n📁 {batch_dir}:")
|
|
batch_path = os.path.join(output_dir, batch_dir)
|
|
|
|
# 概念文件
|
|
concepts_file = os.path.join(batch_path, "concepts.csv")
|
|
if os.path.exists(concepts_file):
|
|
df = pd.read_csv(concepts_file)
|
|
print(f" 概念数量: {len(df)}")
|
|
|
|
# 按类型统计
|
|
if 'concept_type' in df.columns:
|
|
type_counts = df['concept_type'].value_counts()
|
|
for concept_type, count in type_counts.items():
|
|
print(f" - {concept_type}: {count}")
|
|
|
|
# 关系文件
|
|
relationships_file = os.path.join(batch_path, "relationships.csv")
|
|
if os.path.exists(relationships_file):
|
|
df = pd.read_csv(relationships_file)
|
|
print(f" 关系数量: {len(df)}")
|
|
|
|
# 按类型统计
|
|
if 'relationship_type' in df.columns:
|
|
type_counts = df['relationship_type'].value_counts()
|
|
for rel_type, count in type_counts.items():
|
|
print(f" - {rel_type}: {count}")
|
|
|
|
def main():
|
|
"""主函数"""
|
|
parser = argparse.ArgumentParser(
|
|
description='分批处理管理器 - 管理菜谱处理的分批保存和断点续传',
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
使用示例:
|
|
python batch_manager.py status # 查看处理状态
|
|
python batch_manager.py continue ./HowToCook-master # 继续中断的处理
|
|
python batch_manager.py merge # 合并批次数据
|
|
python batch_manager.py clean-progress # 清理进度文件
|
|
python batch_manager.py clean-batches # 清理批次数据
|
|
python batch_manager.py details # 显示所有批次详情
|
|
python batch_manager.py details -b 1 # 显示指定批次详情
|
|
""")
|
|
|
|
subparsers = parser.add_subparsers(dest='command', help='可用命令')
|
|
|
|
# status命令
|
|
subparsers.add_parser('status', help='查看处理进度状态')
|
|
|
|
# continue命令
|
|
continue_parser = subparsers.add_parser('continue', help='继续中断的处理')
|
|
continue_parser.add_argument('recipe_dir', help='菜谱目录路径')
|
|
|
|
# merge命令
|
|
subparsers.add_parser('merge', help='合并所有批次数据')
|
|
|
|
# clean命令
|
|
subparsers.add_parser('clean-progress', help='清理进度文件')
|
|
subparsers.add_parser('clean-batches', help='清理所有批次数据')
|
|
|
|
# details命令
|
|
details_parser = subparsers.add_parser('details', help='显示批次详细信息')
|
|
details_parser.add_argument('-b', '--batch', type=int, help='指定批次编号')
|
|
|
|
# 全局参数
|
|
parser.add_argument('-o', '--output', default='./ai_output', help='输出目录路径')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if not args.command:
|
|
parser.print_help()
|
|
return
|
|
|
|
# 确保输出目录存在
|
|
os.makedirs(args.output, exist_ok=True)
|
|
|
|
print("🛠️ 分批处理管理器")
|
|
print("=" * 50)
|
|
|
|
try:
|
|
if args.command == 'status':
|
|
show_progress_status(args.output)
|
|
elif args.command == 'continue':
|
|
continue_processing(args.recipe_dir, args.output)
|
|
elif args.command == 'merge':
|
|
merge_batches(args.output)
|
|
elif args.command == 'clean-progress':
|
|
clean_progress(args.output)
|
|
elif args.command == 'clean-batches':
|
|
clean_batches(args.output)
|
|
elif args.command == 'details':
|
|
show_batch_details(args.output, args.batch)
|
|
|
|
except KeyboardInterrupt:
|
|
print(f"\n\n⏹️ 操作被用户中断")
|
|
except Exception as e:
|
|
print(f"\n❌ 操作失败: {str(e)}")
|
|
|
|
if __name__ == "__main__":
|
|
main() |