Update app.py

app.py

@@ -35,13 +35,13 @@ if not YUNWU_API_KEY:
         "请在 Hugging Face Space 的 Settings → Variables and secrets 中添加 YUNWU_API_KEY。"
     )
 
-# 这里用云雾的 OpenAI 兼容地址
+# 云雾的 OpenAI 兼容地址
 client = OpenAI(
     api_key=YUNWU_API_KEY,
     base_url="https://yunwu.ai/v1",
 )
 
-# 便宜模型：用于检索、抽取、翻译、风格调整
+# 便宜模型：用于检索、翻译、粗 paraphrase
 MODEL_CHEAP = os.environ.get("MODEL_CHEAP", "deepseek-chat")
 
 # 强模型：用于核心学术解释、annotation 设计、严肃 paraphrase
@@ -107,20 +107,21 @@ def extract_pdf_snippet(pdf_path: Optional[str], max_chars: int = 4000) -> str:
 # 4. 主逻辑函数：多模型协作的图表解释工作流
 # =========================
 
-def chart_assistant_workflow(image, keywords, style, pdf_path):
+def chart_assistant_workflow(image, keywords, style, language, pdf_path):
     """
     主函数：
-    1. 结合论文说明图表的含义（英文 + 中文）
-    2. 给出如何 annotate 图表的建议（英文 + 中文）
-    3. 按 style（formal / fluency / simple）改写论文解释（英文 + 中文）
+    1. 结合论文说明图表含义（根据 language 输出 English 或 中文）
+    2. 给出如何 annotate 图表的建议（同样根据 language）
+    3. 根据 style（formal / fluency / simple）paraphrase 论文解释（同样根据 language）
     """
 
     if image is None:
         return (
             None,
-            "", "",   # bilingual explanation
-            "", "",   # annotation suggestions
-            "", "", ""  # paraphrase: base EN, styled EN, styled ZH
+            "",  # explanation_out
+            "",  # annotation_out
+            "",  # paper_explanation_out
+            "",  # styled_explanation_out
         )
 
     # -------- Step 0：Vision 粗描述 + PDF 上下文 ----------
@@ -128,7 +129,7 @@ def chart_assistant_workflow(image, keywords, style, pdf_path):
     pdf_context = extract_pdf_snippet(pdf_path)
     kw_text = (keywords or "").strip()
 
-    # ========== Step 1：图表含义解释（英文 + 中文） ==========
+    # ========== Step 1：图表含义解释（英文为母版，按 language 决定最终输出） ==========
 
     # 1.1 强模型生成英文解释（基于图 + 论文）
     explain_en_prompt = f"""
@@ -160,8 +161,11 @@ Write in clear, formal academic English, suitable for a Results-style explanatio
         max_tokens=350,
     )
 
-    # 1.2 便宜模型翻译为中文（适合直接拿去讲解）
-    explain_zh_prompt = f"""
+    # 1.2 按语言选择输出
+    if language == "English":
+        explanation_out = explanation_en
+    else:
+        explain_zh_prompt = f"""
 你是一名科学写作助手。请将下面这段英文图表解释翻译成自然、清晰的中文，适合大学生在课堂讲解或书面作业中使用。
 
 要求：
@@ -173,16 +177,16 @@ Write in clear, formal academic English, suitable for a Results-style explanatio
 
 请给出中文译文：
 """
-    explanation_zh = call_llm(
-        explain_zh_prompt,
-        model=MODEL_CHEAP,
-        max_tokens=400,
-    )
+        explanation_out = call_llm(
+            explain_zh_prompt,
+            model=MODEL_CHEAP,
+            max_tokens=400,
+        )
 
-    # ========== Step 2：给出 annotate 建议（英文 + 中文） ==========
+    # ========== Step 2：给出 annotate 建议（根据 language 输出） ==========
 
     annotate_prompt = f"""
-You are designing annotation suggestions for a scientific figure that compares treatments in a coral reef or similar experiment.
+You are designing annotation suggestions for a scientific figure that compares treatments in an experiment.
 
 Here is an English explanation of what the figure shows:
 \"\"\"{explanation_en}\"\"\"
@@ -196,49 +200,44 @@ User keywords:
 Task:
 Propose how a student should annotate this figure to make it easier to understand in a presentation or assignment.
 
-1) In ENGLISH, give 3–5 bullet-point suggestions. Each bullet should follow this pattern:
-   - what to highlight (e.g. "the bar with the highest bleaching under high temperature"),
-   - how to annotate it (arrow / circle / label),
-   - why (e.g. "this is the strongest response and supports the main conclusion").
-
-2) Then, in CHINESE, give对应的 3–5 条标注建议，语言简洁、适合直接告诉同学「你应该如何在图上标注」。
+Give 3–5 bullet-point suggestions in ENGLISH. Each bullet should follow this pattern:
+- what to highlight (e.g. "the bar with the highest bleaching under high temperature"),
+- how to annotate it (arrow / circle / label),
+- why (e.g. "this is the strongest response and supports the main conclusion").
 
-Return your answer in EXACTLY this format:
-
-ANNOTATION_EN:
-- ...
-
-ANNOTATION_ZH:
-- ...
+Return only the bullet list in English.
 """
-    annotation_raw = call_llm(
+    annotation_en = call_llm(
         annotate_prompt,
         model=MODEL_STRONG,
-        max_tokens=500,
+        max_tokens=450,
     )
 
-    # 解析英文 / 中文标注建议
-    annotation_en = ""
-    annotation_zh = ""
-    text_ann = annotation_raw
-
-    if "ANNOTATION_EN:" in text_ann:
-        after_en = text_ann.split("ANNOTATION_EN:", 1)[1]
+    if language == "English":
+        annotation_out = annotation_en
     else:
-        after_en = text_ann
+        annotate_zh_prompt = f"""
+你是一名科学可视化助手。
 
-    if "ANNOTATION_ZH:" in after_en:
-        annotation_en, annotation_zh = after_en.split("ANNOTATION_ZH:", 1)
-    else:
-        annotation_en, annotation_zh = after_en, ""
+下面是一组关于“如何在图表上添加标注”的英文建议（每条都是 bullet）：
+\"\"\"{annotation_en}\"\"\"
 
-    annotation_en = annotation_en.strip()
-    annotation_zh = annotation_zh.strip()
+请将这些建议转写为中文 bullet list，要求：
+- 每条都说明「标注什么、如何标注、为什么要这样标注」
+- 语言简洁、适合给学生作为操作说明
+
+请直接给出中文 bullet 列表：
+"""
+        annotation_out = call_llm(
+            annotate_zh_prompt,
+            model=MODEL_CHEAP,
+            max_tokens=450,
+        )
 
     # ========== Step 3：根据风格 paraphrase 论文解释 ==========
-    # 目标：给用户一个“论文风格解释”的基础版本 + 按 style 改写的版本（英 + 中）
+    # 先得到“论文风格英文解释”，再按 style 改写，最后按 language 选择输出
 
-    # 3.1 便宜模型从论文中抽一段「原论文解释风格」的英文说明（基于 pdf_context）
+    # 3.1 便宜模型从论文中抽一段「论文风格」的英文说明
     paper_exp_prompt = f"""
 You are given part of a scientific paper and a description of a figure.
 
@@ -299,15 +298,37 @@ Give ONLY the rewritten paragraph in English.
         max_tokens=320,
     )
 
-    # 3.3 便宜模型将风格化后的英文解释翻译/改写成中文
+    # 3.3 按语言输出 paper-style 和 styled 版本
+
+    if language == "English":
+        paper_explanation_out = paper_explanation_en
+        styled_explanation_out = styled_explanation_en
+    else:
+        style_label_zh = {
+            "formal": "正式学术风格",
+            "fluency": "口语讲解风格",
+            "simple": "简明易懂风格",
+        }.get(style, "简明易懂风格")
+
+        paper_zh_prompt = f"""
+你是一名科学写作助手。
+
+下面是一段英文解释，是「论文风格」的图表说明：
+\"\"\"{paper_explanation_en}\"\"\"
+
+请将这段英文翻译为中文：
+- 保持学术论文 Results 段落的大致风格
+- 保留科学含义，不添加论文里没有的信息
 
-    style_label_zh = {
-        "formal": "正式学术风格",
-        "fluency": "口语讲解风格",
-        "simple": "简明易懂风格",
-    }.get(style, "简明易懂风格")
+请直接给出中文版本：
+"""
+        paper_explanation_out = call_llm(
+            paper_zh_prompt,
+            model=MODEL_CHEAP,
+            max_tokens=320,
+        )
 
-    style_prompt_zh = f"""
+        styled_zh_prompt = f"""
 你是一名科学写作助手。
 
 下面是一段英文解释，是针对某个图表、已经按「{style_label_zh}」改写过的版本：
@@ -320,31 +341,25 @@ Give ONLY the rewritten paragraph in English.
 
 请直接给出中文版本：
 """
-    styled_explanation_zh = call_llm(
-        style_prompt_zh,
-        model=MODEL_CHEAP,
-        max_tokens=350,
-    )
+        styled_explanation_out = call_llm(
+            styled_zh_prompt,
+            model=MODEL_CHEAP,
+            max_tokens=350,
+        )
 
     # ========== 最终返回给 Gradio 的输出 ==========
     # 1. 原图
-    # 2. Step 1 英文解释
-    # 3. Step 1 中文解释
-    # 4. Step 2 英文 annotation 建议
-    # 5. Step 2 中文 annotation 建议
-    # 6. Step 3 原论文风格英文解释
-    # 7. Step 3 按风格改写英文解释
-    # 8. Step 3 按风格改写的中文解释
+    # 2. 图表含义解释（当前语言）
+    # 3. Annotation 建议（当前语言）
+    # 4. 论文风格解释（当前语言）
+    # 5. 按 style 改写后的解释（当前语言）
 
     return (
         image,
-        explanation_en,
-        explanation_zh,
-        annotation_en,
-        annotation_zh,
-        paper_explanation_en,
-        styled_explanation_en,
-        styled_explanation_zh,
+        explanation_out,
+        annotation_out,
+        paper_explanation_out,
+        styled_explanation_out,
     )
 
 
@@ -353,81 +368,73 @@ Give ONLY the rewritten paragraph in English.
 # =========================
 
 with gr.Blocks() as demo:
-    gr.Markdown("## ChartSmith v3 – 多模型协作的论文图表理解助手")
+    gr.Markdown("## ChartSmith v3 – Multi-model Chart Explanation Assistant")
 
     with gr.Row():
         # 左侧：输入
         with gr.Column():
             img_in = gr.Image(
                 type="pil",
-                label="上传论文中的图表截图（Figure）"
+                label="Upload a figure from the paper"
             )
 
             keywords = gr.Textbox(
-                label="图表关键词 / 变量（英文，逗号分隔，例如 bleaching, coral, temperature, CO2）",
+                label="Figure keywords / variables (English, comma-separated, e.g. bleaching, coral, temperature, CO2)",
                 placeholder="bleaching, coral, temperature, CO2",
             )
 
             style = gr.Radio(
                 choices=["formal", "fluency", "simple"],
                 value="formal",
-                label="Step 3: 解释风格（formal / fluency / simple）"
+                label="Step 3: Explanation style"
+            )
+
+            language = gr.Radio(
+                choices=["English", "中文"],
+                value="English",
+                label="Output language"
             )
 
             pdf_in = gr.File(
-                label="上传论文 PDF（用于结合上下文解释图表）",
+                label="Upload the paper PDF (for context)",
                 type="filepath"
             )
 
-            run_btn = gr.Button("运行多模型工作流", variant="primary")
+            run_btn = gr.Button("Run workflow", variant="primary")
 
         # 右侧：输出
         with gr.Column():
-            orig_img = gr.Image(label="原始图表预览")
+            orig_img = gr.Image(label="Figure preview")
 
-            explain_en_box = gr.Textbox(
-                label="Step 1A: 图表含义解释（English）",
-                lines=6
-            )
-            explain_zh_box = gr.Textbox(
-                label="Step 1B: 图表含义解释（中文）",
+            explain_box = gr.Textbox(
+                label="Step 1: Explanation of what the figure shows",
                 lines=6
             )
 
-            ann_en_box = gr.Textbox(
-                label="Step 2A: Annotation 建议（English）",
-                lines=6
-            )
-            ann_zh_box = gr.Textbox(
-                label="Step 2B: 图表标注建议（中文）",
+            ann_box = gr.Textbox(
+                label="Step 2: Suggestions for annotating the figure",
                 lines=6
             )
 
             paper_exp_box = gr.Textbox(
-                label="Step 3A: 论文风格英文解释（近似原文风格）",
-                lines=6
-            )
-            styled_en_box = gr.Textbox(
-                label="Step 3B: 按所选风格改写后的英文解释",
+                label="Step 3A: Paper-style explanation of the figure",
                 lines=6
             )
-            styled_zh_box = gr.Textbox(
-                label="Step 3C: 按所选风格改写后的中文解释",
+
+            styled_exp_box = gr.Textbox(
+                label="Step 3B: Style-adjusted explanation (for report / talk / simple notes)",
                 lines=6
             )
 
     run_btn.click(
         chart_assistant_workflow,
-        inputs=[img_in, keywords, style, pdf_in],
+        inputs=[img_in, keywords, style, language, pdf_in],
         outputs=[
             orig_img,
-            explain_en_box,
-            explain_zh_box,
-            ann_en_box,
-            ann_zh_box,
+            explain_box,
+            ann_box,
             paper_exp_box,
-            styled_en_box,
-            styled_zh_box,
+            styled_exp_box,
         ],
     )