TravelContentCreator/examples/test_content_generator.py

134 lines
5.1 KiB
Python
Raw Normal View History

2025-04-22 21:57:06 +08:00
#!/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()