260 lines
8.2 KiB
Python
260 lines
8.2 KiB
Python
|
"""
|
|||
|
|
gen_content_digest.py - 批量生成 22 期节目内容摘要卡
|
|||
|
|
用法:
|
|||
|
|
cd E:\tps-dashboard\ai-labeling
|
|||
|
|
python scripts/gen_content_digest.py
|
|||
|
|
|
|||
|
|
功能:
|
|||
|
|
- 读取 prompt4_content_digest.md 作为 system prompt
|
|||
|
|
- 遍历 doco/deliverables/ 下所有融合A稿 .docx 文件
|
|||
|
|
- 用 python-docx 提取文稿文本，调用 MiMo API 生成结构化摘要卡
|
|||
|
|
- 支持断点续跑（已存在的 digest 文件自动跳过）
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import sys
|
|||
|
|
sys.stdout.reconfigure(encoding='utf-8')
|
|||
|
|
sys.stderr.reconfigure(encoding='utf-8')
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import re
|
|||
|
|
import json
|
|||
|
|
import time
|
|||
|
|
from pathlib import Path
|
|||
|
|
from datetime import datetime
|
|||
|
|
from openai import OpenAI
|
|||
|
|
from dotenv import load_dotenv
|
|||
|
|
from docx import Document
|
|||
|
|
|
|||
|
|
# 加载 .env（优先加载 ai-labeling 目录下的 .env）
|
|||
|
|
SCRIPT_DIR = Path(__file__).parent
|
|||
|
|
BASE_DIR = SCRIPT_DIR.parent # ai-labeling/
|
|||
|
|
load_dotenv(BASE_DIR / ".env")
|
|||
|
|
load_dotenv() # 也尝试加载项目根目录的 .env
|
|||
|
|
|
|||
|
|
# 目录配置
|
|||
|
|
DELIVERABLES_DIR = BASE_DIR.parent / "doco" / "deliverables"
|
|||
|
|
PROMPTS_DIR = BASE_DIR / "prompts"
|
|||
|
|
EXPERIMENTS_DIR = BASE_DIR / "experiments" / "content_digests"
|
|||
|
|
|
|||
|
|
# MiMo API 配置（与 run_labeling.py 一致）
|
|||
|
|
MIMO_CONFIG = {
|
|||
|
|
"base_url": "https://api.xiaomimimo.com/v1",
|
|||
|
|
"model_name": "mimo-v2.5-pro",
|
|||
|
|
"api_key_env": "MIMO_API_KEY",
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_system_prompt() -> str:
|
|||
|
|
"""加载 prompt4_content_digest.md 作为 system prompt。"""
|
|||
|
|
prompt_file = PROMPTS_DIR / "prompt4_content_digest.md"
|
|||
|
|
if not prompt_file.exists():
|
|||
|
|
raise FileNotFoundError(f"找不到 prompt 文件: {prompt_file}")
|
|||
|
|
return prompt_file.read_text(encoding="utf-8")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def parse_filename(filename: str) -> dict:
|
|||
|
|
"""
|
|||
|
|
从文件名解析元信息。
|
|||
|
|
文件名格式: 第02期_20260113_武器进化论：海战颠覆者_付天雨_融合A稿.docx
|
|||
|
|
返回: {"ep": 2, "date": "2026-01-13", "title": "...", "editor": "..."}
|
|||
|
|
"""
|
|||
|
|
name = filename.replace(".docx", "")
|
|||
|
|
parts = name.split("_")
|
|||
|
|
if len(parts) < 4:
|
|||
|
|
return None
|
|||
|
|
# 解析期号
|
|||
|
|
ep_match = re.search(r'第(\d+)期', parts[0])
|
|||
|
|
if not ep_match:
|
|||
|
|
return None
|
|||
|
|
ep = int(ep_match.group(1))
|
|||
|
|
# 解析日期（YYYYMMDD -> YYYY-MM-DD）
|
|||
|
|
raw_date = parts[1]
|
|||
|
|
if len(raw_date) == 8 and raw_date.isdigit():
|
|||
|
|
date = f"{raw_date[:4]}-{raw_date[4:6]}-{raw_date[6:8]}"
|
|||
|
|
else:
|
|||
|
|
date = raw_date
|
|||
|
|
title = parts[2]
|
|||
|
|
editor = parts[3]
|
|||
|
|
return {"ep": ep, "date": date, "title": title, "editor": editor}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def extract_docx_text(filepath: Path) -> str:
|
|||
|
|
"""用 python-docx 提取 .docx 文件的全部文本，逐段拼接。"""
|
|||
|
|
doc = Document(str(filepath))
|
|||
|
|
paragraphs = []
|
|||
|
|
for para in doc.paragraphs:
|
|||
|
|
text = para.text.strip()
|
|||
|
|
if text:
|
|||
|
|
paragraphs.append(text)
|
|||
|
|
return "\n".join(paragraphs)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_user_message(meta: dict, body: str) -> str:
|
|||
|
|
"""构造 user message：元信息 + 文稿全文。"""
|
|||
|
|
return (
|
|||
|
|
f"期号：第{meta['ep']:02d}期\n"
|
|||
|
|
f"播出日期：{meta['date']}\n"
|
|||
|
|
f"节目名：{meta['title']}\n"
|
|||
|
|
f"编导：{meta['editor']}\n"
|
|||
|
|
f"\n以下是节目文稿全文：\n\n{body}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def extract_json_from_response(raw: str) -> dict:
|
|||
|
|
"""从模型响应中提取 JSON，兼容推理模型的<think>...输出。"""
|
|||
|
|
# 先去掉<think>...标签及其内容
|
|||
|
|
text = re.sub(r'<think>.*?</think>', '', raw, flags=re.DOTALL)
|
|||
|
|
text = text.strip()
|
|||
|
|
# 去掉 markdown 代码块
|
|||
|
|
text = re.sub(r'^```(?:json)?\s*', '', text)
|
|||
|
|
text = re.sub(r'\s*```$', '', text)
|
|||
|
|
text = text.strip()
|
|||
|
|
# 从第一个 { 开始，到最后一个 } 结束
|
|||
|
|
first_brace = text.find('{')
|
|||
|
|
last_brace = text.rfind('}')
|
|||
|
|
if first_brace != -1 and last_brace != -1 and last_brace >= first_brace:
|
|||
|
|
json_str = text[first_brace:last_brace + 1]
|
|||
|
|
return json.loads(json_str)
|
|||
|
|
# 兜底：直接尝试解析
|
|||
|
|
return json.loads(text)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def call_mimo(system_prompt: str, user_prompt: str) -> dict:
|
|||
|
|
"""调用 MiMo API 生成摘要卡，返回解析后的 JSON dict。"""
|
|||
|
|
api_key = os.environ.get(MIMO_CONFIG["api_key_env"])
|
|||
|
|
if not api_key:
|
|||
|
|
raise EnvironmentError(
|
|||
|
|
f"环境变量 {MIMO_CONFIG['api_key_env']} 未设置，请检查 ai-labeling/.env 或根目录 .env 文件"
|
|||
|
|
)
|
|||
|
|
client = OpenAI(
|
|||
|
|
api_key=api_key,
|
|||
|
|
base_url=MIMO_CONFIG["base_url"],
|
|||
|
|
)
|
|||
|
|
response = client.chat.completions.create(
|
|||
|
|
model=MIMO_CONFIG["model_name"],
|
|||
|
|
messages=[
|
|||
|
|
{"role": "system", "content": system_prompt},
|
|||
|
|
{"role": "user", "content": user_prompt},
|
|||
|
|
],
|
|||
|
|
temperature=0.0,
|
|||
|
|
# 关闭 thinking（与 run_labeling.py 一致）
|
|||
|
|
extra_body={"thinking": {"type": "disabled"}},
|
|||
|
|
)
|
|||
|
|
raw = response.choices[0].message.content
|
|||
|
|
return extract_json_from_response(raw)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def collect_docx_files() -> list[dict]:
|
|||
|
|
"""收集所有 .docx 文件并按期号排序。"""
|
|||
|
|
files = []
|
|||
|
|
for f in sorted(DELIVERABLES_DIR.iterdir()):
|
|||
|
|
if not f.name.endswith(".docx"):
|
|||
|
|
continue
|
|||
|
|
meta = parse_filename(f.name)
|
|||
|
|
if meta is None:
|
|||
|
|
print(f"⚠️ 跳过无法解析的文件: {f.name}")
|
|||
|
|
continue
|
|||
|
|
meta["filepath"] = f
|
|||
|
|
files.append(meta)
|
|||
|
|
# 按期号排序
|
|||
|
|
files.sort(key=lambda x: x["ep"])
|
|||
|
|
return files
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
# 确保输出目录存在
|
|||
|
|
EXPERIMENTS_DIR.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 加载 system prompt
|
|||
|
|
system_prompt = load_system_prompt()
|
|||
|
|
print(f"✅ 已加载 system prompt: prompt4_content_digest.md")
|
|||
|
|
|
|||
|
|
# 收集 docx 文件
|
|||
|
|
docx_files = collect_docx_files()
|
|||
|
|
print(f"✅ 找到 {len(docx_files)} 个融合A稿 docx 文件\n")
|
|||
|
|
|
|||
|
|
all_digests = []
|
|||
|
|
skipped = 0
|
|||
|
|
success = 0
|
|||
|
|
failed = 0
|
|||
|
|
|
|||
|
|
for i, meta in enumerate(docx_files):
|
|||
|
|
ep = meta["ep"]
|
|||
|
|
out_file = EXPERIMENTS_DIR / f"ep{ep:02d}_digest.json"
|
|||
|
|
|
|||
|
|
# 断点续跳：已存在则跳过
|
|||
|
|
if out_file.exists():
|
|||
|
|
print(f"[{i+1}/{len(docx_files)}] ep{ep:02d} 已存在，跳过")
|
|||
|
|
try:
|
|||
|
|
existing = json.loads(out_file.read_text(encoding="utf-8"))
|
|||
|
|
all_digests.append(existing)
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
skipped += 1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
print(f"[{i+1}/{len(docx_files)}] 正在处理 ep{ep:02d} - {meta['title']} ...")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 提取 docx 文本
|
|||
|
|
body = extract_docx_text(meta["filepath"])
|
|||
|
|
if not body.strip():
|
|||
|
|
print(f" ⚠️ ep{ep:02d} 文稿内容为空，跳过")
|
|||
|
|
failed += 1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 构造 user message
|
|||
|
|
user_msg = build_user_message(meta, body)
|
|||
|
|
|
|||
|
|
# 调用 MiMo
|
|||
|
|
result = call_mimo(system_prompt, user_msg)
|
|||
|
|
|
|||
|
|
# 构造输出
|
|||
|
|
output = {
|
|||
|
|
"ep": ep,
|
|||
|
|
"date": meta["date"],
|
|||
|
|
"title": meta["title"],
|
|||
|
|
"editor": meta["editor"],
|
|||
|
|
"filename": meta["filepath"].name,
|
|||
|
|
"digest": result,
|
|||
|
|
"generated_at": datetime.now().isoformat(),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 写入单期文件
|
|||
|
|
out_file.write_text(
|
|||
|
|
json.dumps(output, ensure_ascii=False, indent=2),
|
|||
|
|
encoding="utf-8",
|
|||
|
|
)
|
|||
|
|
all_digests.append(output)
|
|||
|
|
print(f" ✅ ep{ep:02d} 摘要卡已生成 -> {out_file.name}")
|
|||
|
|
success += 1
|
|||
|
|
|
|||
|
|
except json.JSONDecodeError as e:
|
|||
|
|
print(f" ⚠️ ep{ep:02d} LLM 返回的不是合法 JSON: {e}")
|
|||
|
|
failed += 1
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f" ❌ ep{ep:02d} 处理失败: {e}")
|
|||
|
|
failed += 1
|
|||
|
|
|
|||
|
|
# 限流保护：每期之间 sleep 1 秒
|
|||
|
|
if i < len(docx_files) - 1:
|
|||
|
|
time.sleep(1)
|
|||
|
|
|
|||
|
|
# 写入汇总文件
|
|||
|
|
summary_file = EXPERIMENTS_DIR / "_all_digests.json"
|
|||
|
|
summary_file.write_text(
|
|||
|
|
json.dumps(all_digests, ensure_ascii=False, indent=2),
|
|||
|
|
encoding="utf-8",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
print("\n" + "=" * 60)
|
|||
|
|
print(f"📊 处理完成:")
|
|||
|
|
print(f" 成功: {success}")
|
|||
|
|
print(f" 跳过(已存在): {skipped}")
|
|||
|
|
print(f" 失败: {failed}")
|
|||
|
|
print(f" 汇总文件: {summary_file}")
|
|||
|
|
print("=" * 60)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|