239 lines
6.9 KiB
Python
239 lines
6.9 KiB
Python
"""批量翻译 + 事件抽取管道。
|
||
|
||
输入: data/deduped/{YYYYMMDD}/uniques/{url_hash}.json(M3 去重后唯一条目)
|
||
输出: data/events/{YYYYMMDD}/{url_hash}.json
|
||
"""
|
||
|
||
import json
|
||
import logging
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
import yaml
|
||
|
||
from crawler.utils import get_news_day
|
||
from extractor.models import ProcessedArticle
|
||
from llm.client import LLMConfig, load_llm_config, make_sync_client
|
||
from llm.extractor import PromptTemplate, translate_and_extract
|
||
from llm.models import EnTranslatedArticle, LLMCallError
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 默认并发数
|
||
_DEFAULT_CONCURRENCY = 3
|
||
|
||
|
||
def _load_concurrency() -> int:
|
||
"""从 system.yaml 读取 LLM 并发配置。"""
|
||
config_path = Path("configs/system.yaml")
|
||
if config_path.exists():
|
||
try:
|
||
with open(config_path, encoding="utf-8") as f:
|
||
raw = yaml.safe_load(f)
|
||
return int(raw.get("llm", {}).get("concurrency", _DEFAULT_CONCURRENCY))
|
||
except Exception:
|
||
pass
|
||
return _DEFAULT_CONCURRENCY
|
||
|
||
|
||
def _load_deduped_articles(
|
||
date_str: str,
|
||
) -> list[ProcessedArticle]:
|
||
"""加载指定日期的去重后唯一条目。
|
||
|
||
Args:
|
||
date_str: 日期 YYYYMMDD
|
||
|
||
Returns:
|
||
ProcessedArticle 列表
|
||
"""
|
||
base_dir = Path(f"data/deduped/{date_str}/uniques")
|
||
if not base_dir.exists():
|
||
return []
|
||
|
||
articles: list[ProcessedArticle] = []
|
||
for json_file in sorted(base_dir.glob("*.json")):
|
||
try:
|
||
data = json.loads(json_file.read_text(encoding="utf-8"))
|
||
articles.append(ProcessedArticle(**data))
|
||
except (json.JSONDecodeError, Exception) as e:
|
||
logger.warning("解析去重文章失败 %s: %s", json_file, e)
|
||
|
||
return articles
|
||
|
||
|
||
def _process_one(
|
||
article: ProcessedArticle,
|
||
client,
|
||
config: LLMConfig,
|
||
template: PromptTemplate,
|
||
) -> EnTranslatedArticle | None:
|
||
"""处理单篇文章的翻译 + 事件抽取。返回 None 表示失败。"""
|
||
try:
|
||
return translate_and_extract(
|
||
client=client,
|
||
config=config,
|
||
article=article,
|
||
template=template,
|
||
)
|
||
except LLMCallError as e:
|
||
logger.error(
|
||
"翻译抽取最终失败 url=%s: %s", article.url, e.reason
|
||
)
|
||
return None
|
||
except Exception as e:
|
||
logger.exception("翻译抽取未预期异常 url=%s: %s", article.url, e)
|
||
return None
|
||
|
||
|
||
def translate_all_deduped(
|
||
date_str: str | None = None,
|
||
*,
|
||
provider: str | None = None,
|
||
model: str | None = None,
|
||
concurrency: int | None = None,
|
||
) -> dict:
|
||
"""对去重后的所有唯一条目执行翻译 + 事件抽取。
|
||
|
||
Args:
|
||
date_str: 日期 YYYYMMDD,默认当前新闻日
|
||
provider: LLM provider,默认从 system.yaml 读取
|
||
model: LLM model,默认从 system.yaml 读取
|
||
concurrency: 并发数,默认从 system.yaml 读取
|
||
|
||
Returns:
|
||
统计 dict
|
||
"""
|
||
if date_str is None:
|
||
date_str = get_news_day()
|
||
|
||
concurrency = concurrency or _load_concurrency()
|
||
|
||
logger.info("══════ 开始翻译+事件抽取,日期: %s,并发: %d ══════",
|
||
date_str, concurrency)
|
||
|
||
# 加载去重后的文章
|
||
articles = _load_deduped_articles(date_str)
|
||
if not articles:
|
||
logger.warning("去重目录无文章: data/deduped/%s/uniques/", date_str)
|
||
return {"date": date_str, "total": 0, "success": 0, "failed": 0, "elapsed_sec": 0}
|
||
|
||
# 初始化 LLM 客户端
|
||
config = load_llm_config(provider=provider, model=model)
|
||
client = make_sync_client(config)
|
||
template = PromptTemplate()
|
||
|
||
# 输出目录
|
||
out_dir = Path(f"data/events/{date_str}")
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
start_time = datetime.now()
|
||
success = 0
|
||
failed = 0
|
||
|
||
# 增量:跳过已翻译的文章
|
||
new_articles = []
|
||
skipped = 0
|
||
for a in articles:
|
||
if (out_dir / f"{a.url_hash}.json").exists():
|
||
skipped += 1
|
||
else:
|
||
new_articles.append(a)
|
||
if skipped > 0:
|
||
logger.info("增量跳过 %d 篇已翻译,剩余 %d 篇待处理", skipped, len(new_articles))
|
||
articles = new_articles
|
||
|
||
# 并发处理
|
||
with ThreadPoolExecutor(max_workers=concurrency) as executor:
|
||
futures = {
|
||
executor.submit(_process_one, article, client, config, template): article
|
||
for article in articles
|
||
}
|
||
|
||
for future in as_completed(futures):
|
||
article = futures[future]
|
||
try:
|
||
result = future.result()
|
||
except Exception as e:
|
||
logger.error("并发任务异常 url=%s: %s", article.url, e)
|
||
failed += 1
|
||
continue
|
||
|
||
if result is None:
|
||
failed += 1
|
||
continue
|
||
|
||
# 写入输出文件
|
||
out_file = out_dir / f"{result.url_hash}.json"
|
||
out_file.write_text(
|
||
result.model_dump_json(indent=2, ensure_ascii=False),
|
||
encoding="utf-8",
|
||
)
|
||
success += 1
|
||
logger.info(
|
||
"[%s] ✅ %s → %s (%d events, %d zh chars)",
|
||
result.source_id,
|
||
result.title[:40],
|
||
result.title_zh[:30],
|
||
len(result.events),
|
||
result.word_count_zh,
|
||
)
|
||
|
||
elapsed = (datetime.now() - start_time).total_seconds()
|
||
|
||
# 写入事件索引
|
||
_write_event_index(date_str, success, failed, elapsed, config)
|
||
|
||
logger.info(
|
||
"══════ 翻译+事件抽取完成: 成功 %d / 失败 %d / 总计 %d,耗时 %.1f 秒 ══════",
|
||
success, failed, len(articles), elapsed,
|
||
)
|
||
|
||
return {
|
||
"date": date_str,
|
||
"total": len(articles),
|
||
"success": success,
|
||
"failed": failed,
|
||
"elapsed_sec": elapsed,
|
||
"provider": config.provider,
|
||
"model": config.model,
|
||
}
|
||
|
||
|
||
def _write_event_index(
|
||
date_str: str,
|
||
success: int,
|
||
failed: int,
|
||
elapsed_sec: float,
|
||
config: LLMConfig,
|
||
) -> None:
|
||
"""写入事件索引文件。
|
||
|
||
Args:
|
||
date_str: 日期
|
||
success: 成功数
|
||
failed: 失败数
|
||
elapsed_sec: 耗时(秒)
|
||
config: LLM 配置
|
||
"""
|
||
out_dir = Path(f"data/events/{date_str}")
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
index_data = {
|
||
"date": date_str,
|
||
"success": success,
|
||
"failed": failed,
|
||
"elapsed_sec": round(elapsed_sec, 1),
|
||
"provider": config.provider,
|
||
"model": config.model,
|
||
"generated_at": datetime.now().isoformat(),
|
||
}
|
||
|
||
index_path = out_dir / "index.json"
|
||
index_path.write_text(
|
||
json.dumps(index_data, indent=2, ensure_ascii=False),
|
||
encoding="utf-8",
|
||
)
|
||
logger.info("事件索引已写入: %s", index_path)
|