RobotDaily/arxiv-digest/scripts/enrich_papers.py

#!/usr/bin/env python3
"""Translate abstracts, generate tags, and produce short explanations."""

from __future__ import annotations

import argparse
import json
from typing import Any, Dict, List

from fetch_arxiv import DOMAIN_CONFIGS
from utils import log, normalize_space, ollama_generate_json, read_json, truncate, write_json

FALLBACK_TAGS = {
    "embodied": ["具身智能", "机器人", "真实部署", "操控", "导航"],
    "representation": ["表征学习", "潜在空间", "世界模型", "预训练", "对象中心"],
    "reinforcement": ["强化学习", "策略优化", "奖励设计", "离线 RL", "模仿学习"],
}


def build_prompt(paper: Dict[str, Any]) -> str:
    domain_label = DOMAIN_CONFIGS[paper["domain"]]["label_zh"]
    return f"""
你是 RobotDaily 的论文晨报编辑。请根据给定的英文标题与英文摘要，输出严格 JSON。

只输出一个 JSON 对象，结构如下：
{{
  "translated_abstract_zh": "...",
  "brief_explanation_zh": "...",
  "tags": ["标签 1", "标签 2", "标签 3", "标签 4", "标签 5"]
}}

要求：
1. translated_abstract_zh：忠实翻译原摘要，不要增加原文没有的实验结果；控制在 180-400 个中文字符，必须完整覆盖原文摘要的所有要点。
2. brief_explanation_zh：40-90 个中文字符，说明为什么值得读，尽量偏应用价值和创新点。
3. tags：给 4-6 个适合直接贴在移动端卡片上的简短标签；尽量用中文，必要时保留通用英文术语，如 World Model、Offline RL。
4. 语气务实、技术导向，不要夸张。
5. 不要输出 Markdown，不要输出代码块。

领域：{domain_label}
标题：{paper['title']}
英文摘要：{paper['summary']}
""".strip()


def fallback_enrichment(paper: Dict[str, Any]) -> Dict[str, Any]:
    tags = FALLBACK_TAGS.get(paper["domain"], ["AI 论文", "机器学习", "应用研究"])
    title = paper.get("title", "")
    summary = paper.get("summary", "")
    
    # 从标题提取核心方法
    brief = title.split(':')[0].strip()
    if '.' in brief:
        brief = brief.split('.')[0].strip()
    
    # 判断方法类型
    summary_lower = summary.lower()
    if "diffusion" in summary_lower:
        method = "扩散模型"
    elif "reinforcement learning" in summary_lower:
        method = "强化学习"
    elif "imitation learning" in summary_lower:
        method = "模仿学习"
    elif "contrastive" in summary_lower:
        method = "对比学习"
    elif "transformer" in summary_lower:
        method = "Transformer"
    elif "self-supervised" in summary_lower:
        method = "自监督学习"
    elif "representation" in summary_lower:
        method = "表征学习"
    else:
        method = "多种技术"
    
    # 判断应用领域
    if "robot" in summary_lower or "manipulation" in summary_lower:
        field = "机器人操作"
    elif "navigation" in summary_lower or "driving" in summary_lower:
        field = "导航控制"
    elif "translation" in summary_lower or "generation" in summary_lower:
        field = "生成任务"
    else:
        field = "相关任务"
    
    # 判断结果
    if "real-world" in summary_lower or "deployment" in summary_lower:
        result = "真实部署"
    elif "zero-shot" in summary_lower:
        result = "零样本泛化"
    elif "first" in paper.get("matched_innovation_terms", []) or "novel" in paper.get("matched_innovation_terms", []):
        result = "首次提出"
    elif "improve" in summary_lower or "better" in summary_lower:
        result = "性能提升"
    else:
        result = "性能优化"
    
    brief = f"{brief}，采用{method}解决{field}，实现{result}"
    
    return {
        "translated_abstract_zh": f"【LLM 暂不可用，先保留英文摘要要点】{truncate(summary, 220)}",
        "brief_explanation_zh": truncate(brief, 86),
        "tags": tags[:5],
    }


def enrich_paper(paper: Dict[str, Any], model_names: List[str]) -> Dict[str, Any]:
    prompt = build_prompt(paper)
    result = None
    used_model = ""
    for model in model_names:
        model = normalize_space(model)
        if not model:
            continue
        log(f"Enriching {paper['arxiv_id']} with {model}")
        result = ollama_generate_json(prompt, model=model, timeout=150)
        if result:
            used_model = model
            break

    enriched = dict(paper)
    payload = result or fallback_enrichment(paper)
    tags = [normalize_space(tag).lstrip("#") for tag in payload.get("tags", []) if normalize_space(tag)]
    if not tags:
        tags = FALLBACK_TAGS.get(paper["domain"], [])[:5]

    enriched["translated_abstract_zh"] = normalize_space(payload.get("translated_abstract_zh", "")) or fallback_enrichment(paper)["translated_abstract_zh"]
    enriched["brief_explanation_zh"] = normalize_space(payload.get("brief_explanation_zh", "")) or fallback_enrichment(paper)["brief_explanation_zh"]
    enriched["tags"] = tags[:6]
    enriched["enrichment_model"] = used_model or "fallback"
    return enriched


def enrich_selection(selection_payload: Dict[str, Any], model_names: List[str]) -> Dict[str, Any]:
    papers = selection_payload.get("papers", [])
    enriched_papers = [enrich_paper(paper, model_names=model_names) for paper in papers]

    by_domain: Dict[str, List[Dict[str, Any]]] = {domain: [] for domain in selection_payload.get("selected_by_domain", {})}
    for paper in enriched_papers:
        by_domain.setdefault(paper["domain"], []).append(paper)

    output = dict(selection_payload)
    output["papers"] = enriched_papers
    output["selected_by_domain"] = by_domain
    output["configured_models"] = model_names
    output["effective_models_used"] = list(
        dict.fromkeys(
            paper.get("enrichment_model", "")
            for paper in enriched_papers
            if paper.get("enrichment_model")
        )
    )
    return output


def main() -> None:
    parser = argparse.ArgumentParser(description="Enrich RobotDaily papers with zh translation and tags")
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", default="")
    parser.add_argument("--models", default="qwen3.5:27b")
    args = parser.parse_args()

    payload = read_json(args.input, default={}) or {}
    models = [item.strip() for item in args.models.split(",") if item.strip()]
    enriched = enrich_selection(payload, model_names=models)

    if args.output:
        write_json(args.output, enriched)
    else:
        print(json.dumps(enriched, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()