Update app.py

app.py

@@ -1,60 +1,62 @@
 import os
+import re
 import gradio as gr
 from transformers import pipeline
 import torch
 from openai import OpenAI
-from PIL import Image, ImageEnhance
+from PIL import Image
 from pypdf import PdfReader
 
+
 # =========================
-# 1. 设备 & Vision 模型
+# 1. Device & local vision model
 # =========================
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# local image->text (fast rough caption)
 vision_pipe = pipeline(
     "image-to-text",
     model="nlpconnect/vit-gpt2-image-captioning",
-    device=0 if device == "cuda" else -1,
+    device=0 if device == "cuda" else -1
 )
 
 # =========================
-# 2. 云雾 / DeepSeek API 设置
+# 2. Yunwu / OpenAI-compatible client
 # =========================
 
 YUNWU_API_KEY = os.environ.get("YUNWU_API_KEY")
-
 if not YUNWU_API_KEY:
     raise RuntimeError(
-        "环境变量 YUNWU_API_KEY 未设置。请在 Hugging Face Space 的 "
-        "Settings → Variables and secrets 中添加：Name=YUNWU_API_KEY, Type=Secret。"
+        "YUNWU_API_KEY not set. Add it in HF Space Settings → Variables and secrets."
     )
 
-# 注意：这里用的是云雾文档里的 openai 兼容端点
 client = OpenAI(
     api_key=YUNWU_API_KEY,
-    base_url="https://yunwu.ai/v1",
+    base_url="https://yunwu.ai/v1"
 )
 
-# 你可以根据需要调不同模型，例如：
-MODEL_EXPLANATION = "deepseek-chat"   # Step 1：解释图
-MODEL_ANNOTATION = "deepseek-chat"    # Step 2：标注建议
-MODEL_PARAPHRASE = "deepseek-chat"    # Step 3：改写
+# =========================
+# 3. Per-step model routing (edit here)
+# =========================
+MODEL_STEP1 = "gpt-5.1"       # strongest reasoning for paper+figure meaning
+MODEL_STEP2 = "deepseek-chat" # cheap/fast for annotation suggestions
+MODEL_STEP3B = "gpt-4o"       # good fluency for plain-language explanation
 
 
 def call_llm(prompt: str,
              model: str,
-             temperature: float = 0.25,
-             max_tokens: int = 600) -> str:
-    """统一封装一次聊天调用。"""
+             temperature: float = 0.2,
+             max_tokens: int = 512) -> str:
     resp = client.chat.completions.create(
         model=model,
         messages=[
             {
                 "role": "system",
                 "content": (
-                    "You are an academic writing assistant for scientific figures. "
-                    "Always reply in clear, professional English suitable for university-level work."
+                    "You are a helpful academic assistant. "
+                    "Be accurate, concrete, and avoid hallucinating numbers. "
+                    "Do not use markdown unless asked."
                 ),
             },
             {"role": "user", "content": prompt},
@@ -66,17 +68,15 @@ def call_llm(prompt: str,
 
 
 # =========================
-# 3. 工具函数：PDF 文本、图像增强
+# 4. Utilities: PDF context + output cleaning
 # =========================
 
 def extract_pdf_snippet(pdf_path: str | None, max_chars: int = 2500) -> str:
-    """读取 PDF，抽取前几页文本作为上下文。"""
     if not pdf_path:
         return ""
-
     try:
         reader = PdfReader(pdf_path)
-        texts: list[str] = []
+        texts = []
         for page in reader.pages:
             txt = page.extract_text() or ""
             texts.append(txt)
@@ -88,227 +88,174 @@ def extract_pdf_snippet(pdf_path: str | None, max_chars: int = 2500) -> str:
         return ""
 
 
-def enhance_image_simple(image: Image.Image | None) -> Image.Image | None:
-    """简单增强亮度&对比度（现在只是备选，不再作为单独一步展示）。"""
-    if image is None:
-        return None
-    try:
-        img = image.convert("RGB")
-        img = ImageEnhance.Brightness(img).enhance(1.05)
-        img = ImageEnhance.Contrast(img).enhance(1.15)
-        return img
-    except Exception:
-        return image
+def clean_text(text: str) -> str:
+    """
+    Remove markdown-ish / AI-ish formatting and obvious repetition.
+    """
+    if not text:
+        return text
+
+    # remove bold/italic markers
+    text = text.replace("**", "").replace("__", "")
+
+    # remove leading markdown bullets like "- ", "* ", "• "
+    text = re.sub(r"(?m)^\s*[-•*]\s+", "", text)
+
+    # collapse repeated whitespace
+    text = re.sub(r"\n{3,}", "\n\n", text)
+    text = re.sub(r"[ \t]{2,}", " ", text)
+
+    return text.strip()
 
 
 # =========================
-# 4. 主逻辑：3 大输出
+# 5. Main workflow
 # =========================
 
 def run_workflow(image, keywords, style, pdf_path):
-    """
-    返回 4 个东西：
-    1. 原始图像
-    2. Step 1: 图表含义解释（结合论文）
-    3. Step 2: 如何 annotate 的建议
-    4. Step 3: 论文风格解释 + 根据 style 改写的版本
-    """
     if image is None:
-        return None, "Please upload a figure first.", "", "", ""
+        return "", "", ""
 
-    # -------- Vision caption ----------
-    try:
-        vision_raw = vision_pipe(image)[0]["generated_text"]
-    except Exception as e:
-        vision_raw = f"(Vision model failed: {e})"
+    # Step 0: rough vision caption
+    vision_raw = vision_pipe(image)[0]["generated_text"]
 
+    # contexts
     kw_text = keywords.strip() if keywords else ""
     pdf_context = extract_pdf_snippet(pdf_path)
 
-    # -------- Step 1: 解释图表含义 ----------
+    # style control for Step1 only
+    if style == "formal":
+        style_instruction = (
+            "Write in formal academic English suitable for a Results/Figure explanation."
+        )
+    elif style == "fluency":
+        style_instruction = (
+            "Write in smooth, natural academic English, clear and readable."
+        )
+    else:  # simple
+        style_instruction = (
+            "Write clearly with simpler wording, but still accurate."
+        )
+
+    # -------- Step 1: Explanation of what the figure shows (paper+figure) --------
     step1_prompt = f"""
-You are given a scientific figure from a research paper.
+You are given a scientific figure and some related paper text.
 
-Rough description from a vision model:
+Rough vision caption of the image:
 \"\"\"{vision_raw}\"\"\"
 
-User-provided figure keywords / variables (comma-separated):
+Figure keywords:
 \"\"\"{kw_text}\"\"\"
 
-Context from the paper (may be noisy or partial):
+Relevant paper context (may be partial/noisy):
 \"\"\"{pdf_context}\"\"\"
 
 Task:
-In about 5–8 sentences of clear academic English:
-- Explain what the figure is about.
-- State what the x-axis represents (in generic terms, e.g. temperature treatments, species, experimental groups).
-- State what the y-axis represents (e.g. relative bleaching, response intensity).
-- Describe the main patterns and comparisons (which treatments are higher / lower, more or less sensitive).
-- Connect the pattern to the scientific meaning (e.g. higher temperature + high CO2 leads to more bleaching).
-
-Do NOT invent exact numbers or p-values.
-Do NOT describe details of the full experimental protocol (no sample size, location, etc.).
-"""
-    explanation_en = call_llm(
-        prompt=step1_prompt,
-        model=MODEL_EXPLANATION,
-        max_tokens=550,
-    )
+- Explain what this figure means in 6–8 sentences.
+- Say what is compared on the x-axis (or panels) and what the y-axis measures.
+- Describe the main trends and the scientific takeaway.
+- If the paper context implies a specific mechanism, mention it briefly.
+- Do NOT invent exact numbers, statistics, or p-values.
+- Do NOT describe unrelated parts of the paper.
 
-    # -------- Step 2: 标注建议 ----------
-    step2_prompt = f"""
-Here is an explanation of a scientific figure:
-
-\"\"\"{explanation_en}\"\"\"
-
-The figure is a chart/diagram used in a research paper.  
-Suggest how a student should annotate this figure to make it clearer for presentations or homework.
-
-In 4–7 bullet points of concise English:
-- Propose specific labels, arrows, or text boxes to add.
-- Indicate WHAT to label (e.g. treatment groups, axes, key contrasts).
-- Indicate WHERE to place the annotations (e.g. near bars with highest bleaching, next to control group, above x-axis groups).
-- Mention any colour or symbol conventions that would help (e.g. “use red arrows for stress conditions”).
-
-Write only the bullet-point suggestions.
+{style_instruction}
+Return plain text only.
 """
-    annotation_suggestions = call_llm(
-        prompt=step2_prompt,
-        model=MODEL_ANNOTATION,
-        max_tokens=420,
-    )
-
-    # -------- Step 3A: 论文风格解释 ----------
-    step3_paper_prompt = f"""
-Using the same figure, write a short text that could appear in the Results section of a scientific paper.
+    step1_text = call_llm(step1_prompt, model=MODEL_STEP1, max_tokens=520)
+    step1_text = clean_text(step1_text)
 
-Base it on the following explanation and context:
-
-Current explanation:
-\"\"\"{explanation_en}\"\"\"
+    # -------- Step 2: Suggestions for annotating the figure (less AI-ish) --------
+    step2_prompt = f"""
+You are helping someone improve a scientific figure for a presentation.
 
-Paper context:
-\"\"\"{pdf_context}\"\"\"
+Figure meaning (for context):
+\"\"\"{step1_text}\"\"\"
 
 Task:
-- Write a concise paragraph (4–6 sentences) in formal academic English.
-- Describe what is on the x-axis and y-axis in generic terms.
-- Describe the main trends and key contrasts shown in the figure.
-- End with one sentence that interprets what these patterns imply for the biological / scientific question.
+Give 5–7 practical, specific suggestions for how to annotate this figure directly on the image
+(e.g., labels, arrows, callouts, highlights), so the key message is immediately obvious.
 
-Do NOT include exact numeric values or statistics.
-"""
-    paper_style_explanation = call_llm(
-        prompt=step3_paper_prompt,
-        model=MODEL_PARAPHRASE,
-        max_tokens=450,
-    )
+Rules:
+- Write like a human TA giving advice.
+- Avoid fancy formatting, no markdown, no bold, no bullet symbols like "-" or "•".
+- Each suggestion should be one short sentence.
+- Do not invent values; only suggest how to visually emphasize real trends.
 
-    # -------- Step 3B: 根据用户选择风格改写 ----------
-    if style == "formal":
-        style_instruction = (
-            "Make it slightly more concise and polished, but keep a formal academic tone "
-            "suitable for a written assignment or report."
-        )
-    elif style == "fluency":
-        style_instruction = (
-            "Make it smoother and more speech-like, as if the student is explaining the figure "
-            "in an oral presentation, while still sounding professional."
-        )
-    else:  # "simple"
-        style_instruction = (
-            "Rewrite it in simpler English for a non-expert audience, such as classmates, "
-            "while keeping the scientific meaning accurate."
-        )
+Return plain text, one suggestion per line.
+"""
+    step2_text = call_llm(step2_prompt, model=MODEL_STEP2, temperature=0.3, max_tokens=260)
+    step2_text = clean_text(step2_text)
 
-    step3_para_prompt = f"""
-Here is a paper-style explanation of a figure:
+    # -------- Step 3B: Plain-language explanation (replaces old paraphrase) --------
+    step3b_prompt = f"""
+Explain this same figure to a smart high-school or first-year university student.
 
-\"\"\"{paper_style_explanation}\"\"\"
+Figure meaning:
+\"\"\"{step1_text}\"\"\"
 
-Rewrite this explanation according to the following requirement:
-{style_instruction}
+Rules:
+- 4–6 sentences.
+- Use simple, conversational English.
+- Keep the science correct but avoid jargon unless necessary.
+- No numbers or stats unless they were explicitly in the meaning above.
 
-Keep the main scientific meaning and relationships between conditions.
-Keep it in English. Output a single coherent paragraph.
+Return plain text only.
 """
-    paraphrased_explanation = call_llm(
-        prompt=step3_para_prompt,
-        model=MODEL_PARAPHRASE,
-        max_tokens=350,
-    )
+    step3b_text = call_llm(step3b_prompt, model=MODEL_STEP3B, temperature=0.4, max_tokens=240)
+    step3b_text = clean_text(step3b_text)
 
-    # 返回顺序一定要和 outputs 一致！
-    return (
-        image,                    # 原图预览
-        explanation_en,           # Step 1
-        annotation_suggestions,   # Step 2
-        paper_style_explanation,  # Step 3A
-        paraphrased_explanation,  # Step 3B
-    )
+    return step1_text, step2_text, step3b_text
 
 
 # =========================
-# 5. Gradio UI
+# 6. Gradio UI
 # =========================
 
 with gr.Blocks() as demo:
     gr.Markdown("## ChartSmith – AI Figure Explainer (multi-model workflow)")
 
     with gr.Row():
-        # 左边：输入
         with gr.Column():
-            img_in = gr.Image(
-                type="pil",
-                label="Upload your scientific figure (screenshot is fine)"
-            )
+            img_in = gr.Image(type="pil", label="Upload your scientific figure (screenshot is fine)")
 
             keywords = gr.Textbox(
-                label="Figure keywords / variables (English, comma-separated, e.g. bleaching, coral, temperature, CO2)",
-                placeholder="bleaching, coral, temperature, CO2",
+                label="Figure keywords / variables (English, comma-separated)",
+                placeholder="bleaching, coral, temperature, CO2"
             )
 
             style = gr.Radio(
                 ["formal", "fluency", "simple"],
                 value="formal",
-                label="Step 3: Explanation style",
+                label="Explanation style (for Step 1)"
             )
 
             pdf_in = gr.File(
                 label="Upload the paper PDF (for context, optional)",
-                type="filepath",
+                type="filepath"
             )
 
             run_btn = gr.Button("Run workflow", variant="primary")
 
-        # 右边：输出
         with gr.Column():
-            orig_img = gr.Image(label="Figure preview")
-
             step1_box = gr.Textbox(
-                label="Step 1: Explanation of what the figure shows (English)",
-                lines=9,
+                label="Step 1: Explanation of what the figure shows (paper-aware)",
+                lines=10
             )
 
             step2_box = gr.Textbox(
                 label="Step 2: Suggestions for annotating the figure",
-                lines=8,
-            )
-
-            step3_paper_box = gr.Textbox(
-                label="Step 3A: Paper-style explanation of the figure (Results-style)",
-                lines=8,
+                lines=8
             )
 
-            step3_para_box = gr.Textbox(
-                label="Step 3B: Paraphrased explanation for presentations / assignments",
-                lines=8,
+            step3b_box = gr.Textbox(
+                label="Step 3: Plain-language explanation (for class/presentation)",
+                lines=6
             )
 
     run_btn.click(
         run_workflow,
         inputs=[img_in, keywords, style, pdf_in],
-        outputs=[orig_img, step1_box, step2_box, step3_paper_box, step3_para_box],
+        outputs=[step1_box, step2_box, step3b_box],
     )