134 lines
5.1 KiB
Python
134 lines
5.1 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
测试ContentGenerator组件
|
||
|
||
该脚本演示如何独立使用ContentGenerator组件生成海报文本配置。
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import json
|
||
import time
|
||
import logging
|
||
from pathlib import Path
|
||
|
||
# 将项目根目录添加到PATH
|
||
project_root = str(Path(__file__).parent.parent.absolute())
|
||
if project_root not in sys.path:
|
||
sys.path.append(project_root)
|
||
|
||
from core import contentGen
|
||
|
||
def load_config(config_path="poster_gen_config.json"):
|
||
"""加载配置文件"""
|
||
if not os.path.exists(config_path):
|
||
config_path = os.path.join(project_root, config_path)
|
||
if not os.path.exists(config_path):
|
||
raise FileNotFoundError(f"配置文件未找到: {config_path}")
|
||
|
||
with open(config_path, 'r', encoding='utf-8') as f:
|
||
config = json.load(f)
|
||
return config
|
||
|
||
def get_test_content():
|
||
"""返回一个用于测试的文章内容"""
|
||
return {
|
||
"title": "🌿清明遛娃天花板！悬空古寺+非遗探秘",
|
||
"content": """清明假期带娃哪里玩？泰宁甘露寺藏着明代建筑奇迹！一柱擎天的悬空阁楼+状元祈福传说，让孩子边玩边涨知识✨
|
||
|
||
🎒行程亮点：
|
||
✅ 安全科普第一站：讲解"一柱插地"千年不倒的秘密，用乐高积木模型让孩子理解力学原理
|
||
✅ 文化沉浸体验：穿汉服听"叶状元还愿建寺"故事，触摸3.38米粗的"状元柱"许愿
|
||
✅ 自然探索路线：连接金湖栈道徒步，观察丹霞地貌与古建筑的巧妙融合
|
||
|
||
📌实用攻略：
|
||
📍位置：福建省三明市泰宁县金湖西路（导航搜"甘露岩寺"）
|
||
🕒最佳时段：上午10点前抵达避开人流，下午可衔接参观明清园（80元/人）
|
||
⚠️注意事项：悬空栈道设置儿童安全绳租赁点，建议穿防滑鞋"""
|
||
}
|
||
|
||
def main():
|
||
"""主函数"""
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(levelname)s - %(message)s',
|
||
datefmt='%Y-%m-%d %H:%M:%S'
|
||
)
|
||
|
||
try:
|
||
# 加载配置
|
||
config = load_config()
|
||
print("配置加载成功")
|
||
|
||
# 准备参数
|
||
description_file = input("请输入描述文件路径 (直接回车使用默认示例): ").strip()
|
||
if not description_file:
|
||
# 尝试从配置中找一个描述文件作为示例
|
||
for dir_info in config.get("resource_dir", []):
|
||
if dir_info.get("type") == "Description" and dir_info.get("file_path"):
|
||
description_file = dir_info["file_path"][0]
|
||
break
|
||
|
||
if not description_file:
|
||
print("未找到描述文件，将创建一个空文件用于测试")
|
||
description_file = "temp_description.txt"
|
||
with open(description_file, "w", encoding="utf-8") as f:
|
||
f.write("测试描述文件")
|
||
|
||
# 确保文件存在
|
||
if not os.path.exists(description_file):
|
||
print(f"警告: 描述文件 {description_file} 不存在，将创建空文件")
|
||
with open(description_file, "w", encoding="utf-8") as f:
|
||
f.write("测试描述文件")
|
||
|
||
info_directory = [description_file]
|
||
poster_num = int(input("请输入要生成的海报数量 (默认2): ") or "2")
|
||
|
||
# 准备内容数据
|
||
content_data = get_test_content()
|
||
print(f"使用测试内容: {content_data['title']}")
|
||
|
||
# 创建ContentGenerator实例
|
||
content_gen = contentGen.ContentGenerator(
|
||
model_name=config.get("model", "qwenQWQ"),
|
||
api_base_url=config.get("api_url", "http://localhost:8000/v1"),
|
||
api_key=config.get("api_key", "EMPTY")
|
||
)
|
||
|
||
# 设置模型参数
|
||
content_gen.set_model_para(
|
||
temperature=config.get("content_temperature", 0.7),
|
||
top_p=config.get("content_top_p", 0.8),
|
||
presence_penalty=config.get("content_presence_penalty", 1.2)
|
||
)
|
||
|
||
print(f"开始生成海报文本配置，数量: {poster_num}")
|
||
start_time = time.time()
|
||
|
||
# 运行生成过程
|
||
poster_configs = content_gen.run(info_directory, poster_num, content_data)
|
||
|
||
elapsed = time.time() - start_time
|
||
print(f"生成完成，耗时: {elapsed:.2f}秒")
|
||
|
||
if poster_configs:
|
||
print(f"\n成功生成 {len(poster_configs) if isinstance(poster_configs, list) else 1} 个海报配置:")
|
||
print(json.dumps(poster_configs, ensure_ascii=False, indent=2))
|
||
|
||
# 保存到文件
|
||
output_file = f"poster_configs_{int(time.time())}.json"
|
||
with open(output_file, "w", encoding="utf-8") as f:
|
||
json.dump(poster_configs, f, ensure_ascii=False, indent=2)
|
||
print(f"配置已保存到: {output_file}")
|
||
else:
|
||
print("生成海报配置失败")
|
||
|
||
except Exception as e:
|
||
print(f"程序出错: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
if __name__ == "__main__":
|
||
main() |