Update app.py

app.py

@@ -3,8 +3,7 @@ from transformers import pipeline
 from PIL import Image, ImageEnhance
 from pypdf import PdfReader
 
-# ===================== 1. 加载模型（CPU + flan-t5-base） =====================
-
+# ========== 1. 小模型：方便在 CPU Basic 上跑 ==========
 vision_pipe = pipeline(
     "image-to-text",
     model="nlpconnect/vit-gpt2-image-captioning",
@@ -13,13 +12,12 @@ vision_pipe = pipeline(
 
 text_pipe = pipeline(
     "text2text-generation",
-    model="t5-small",
+    model="google/flan-t5-small",
     max_length=256,
     device=-1
 )
 
-# ===================== 2. 论文 PDF 上下文处理 =====================
-
+# ========== 2. 论文 PDF → 文本 ==========
 def extract_paper_context(paper_file):
     if paper_file is None:
         return ""
@@ -27,30 +25,27 @@ def extract_paper_context(paper_file):
         path = paper_file if isinstance(paper_file, str) else paper_file.name
         reader = PdfReader(path)
         text = ""
-        for i, page in enumerate(reader.pages[:3]):
+        for page in reader.pages[:3]:
             page_text = page.extract_text() or ""
             text += page_text + "\n"
-        return text.strip()[:4000]
-    except:
+        text = text.strip()
+        return text[:3000]   # 截短一点，避免太长
+    except Exception as e:
+        print("PDF error:", e)
         return ""
 
 def summarize_paper_context(raw_text):
     if not raw_text:
         return ""
-    prompt = f"""
-Summarize in 4-6 sentences the research background most relevant for understanding
-a result figure. Focus on variables, comparisons, and experimental conditions.
-
-Paper text:
-{raw_text}
-
-Summary:
-"""
+    prompt = (
+        "Here is an excerpt from a scientific paper. "
+        "In 4-5 sentences, summarize the experimental setup and the main "
+        "variables or comparisons that the results focus on.\n\n"
+        f"{raw_text}\n\nSummary:"
+    )
     return text_pipe(prompt)[0]["generated_text"]
 
-
-# ===================== 3. 核心工作流 =====================
-
+# ========== 3. 核心工作流 ==========
 def analyze_and_enhance(image, style, language, paper_file):
     if image is None:
         return None, "请先上传图表。", "", None, "", ""
@@ -59,108 +54,115 @@ def analyze_and_enhance(image, style, language, paper_file):
     paper_raw = extract_paper_context(paper_file)
     paper_context = summarize_paper_context(paper_raw) if paper_raw else ""
 
-    # 图像描述
+    # Step 2: 图像粗略描述
     raw_caption = vision_pipe(image)[0]["generated_text"]
     rc = raw_caption.lower()
     if not any(k in rc for k in ["bar", "chart", "graph", "plot", "curve", "line"]):
-        raw_caption = "a bar chart comparing multiple experimental groups under different conditions"
-
-    # Academic explanation
-    if paper_context:
-        context_part = f"Here is relevant context from the paper:\n{paper_context}\n"
-    else:
-        context_part = "No additional paper context available.\n"
-
-    prompt_academic = f"""
-You are writing the Results section of a scientific paper.
-
-{context_part}
-
-Write a structured academic explanation (3–4 sentences) of the figure:
-
-1. Identify the figure type.
-2. Describe generically what is on the horizontal axis.
-3. Describe generically what is on the vertical axis.
-4. Describe overall trends.
-5. Summarize the main comparison, consistent with the paper context.
-
-Do NOT repeat sentences or invent specific numbers.
-
-Figure description: {raw_caption}
-"""
+        # 避免“clock and a building”这类离谱 caption
+        raw_caption = "a bar chart showing experimental results described in the paper"
+
+    # Step 3: 学术化解释（简单指令，不再列一堆 1,2,3,4）
+    context_part = (
+        f"Paper context: {paper_context}\n\n"
+        if paper_context else
+        ""
+    )
+    prompt_academic = (
+        "You are writing the Results section of a scientific paper.\n\n"
+        f"{context_part}"
+        f"Figure description: {raw_caption}\n\n"
+        "In 3–4 sentences, explain what this figure shows and what the main "
+        "pattern or comparison is. Use formal academic English. "
+        "Answer directly with the explanation:"
+    )
     academic_desc = text_pipe(prompt_academic)[0]["generated_text"]
 
-    # Caption
-    caption_prompt = f"""
-Write a concise 1–2 sentence academic figure caption.
-Keep it consistent with the paper context.
-
-Description: {academic_desc}
-
-Caption:
-"""
+    # Step 5: Caption（基于刚才那段 explanation）
+    caption_prompt = (
+        "Write a 1–2 sentence figure caption for an academic paper based on "
+        "the explanation below. Mention the type of figure and the main "
+        "comparison or trend.\n\n"
+        f"{academic_desc}\n\nCaption:"
+    )
     caption_text = text_pipe(caption_prompt)[0]["generated_text"]
 
-    # Summary paragraph
-    summary_prompt = f"""
-Write a 3–4 sentence paragraph explaining the key message of this figure,
-consistent with the paper context. Do not fabricate numbers.
-
-Description: {academic_desc}
-
-Paragraph:
-"""
+    # Step 5: Summary / 讨论段落
+    summary_prompt = (
+        "Rewrite the following explanation as a short paragraph (3–4 sentences) "
+        "for the Results section of a paper, highlighting the key trend and conclusion.\n\n"
+        f"{academic_desc}\n\nParagraph:"
+    )
     summary_text = text_pipe(summary_prompt)[0]["generated_text"]
 
-    # Language translation
+    # 语言切换
     if language == "中文":
-        academic_desc = text_pipe(f"Translate into Chinese:\n{academic_desc}")[0]["generated_text"]
-        caption_text = text_pipe(f"Translate into Chinese:\n{caption_text}")[0]["generated_text"]
-        summary_text = text_pipe(f"Translate into Chinese:\n{summary_text}")[0]["generated_text"]
+        academic_desc_out = text_pipe(
+            f"Translate the following academic explanation into Chinese:\n{academic_desc}"
+        )[0]["generated_text"]
+        caption_text = text_pipe(
+            f"Translate the following figure caption into Chinese:\n{caption_text}"
+        )[0]["generated_text"]
+        summary_text = text_pipe(
+            f"Translate the following paragraph into Chinese:\n{summary_text}"
+        )[0]["generated_text"]
+    else:
+        academic_desc_out = academic_desc
 
-    # 图像增强
+    # Step 4: 图像增强
     img_rgb = image.convert("RGB")
     img_c = ImageEnhance.Contrast(img_rgb).enhance(1.4)
     enhanced_image = ImageEnhance.Brightness(img_c).enhance(1.1)
 
-    return image, raw_caption, academic_desc, enhanced_image, caption_text, summary_text
-
+    return image, raw_caption, academic_desc_out, enhanced_image, caption_text, summary_text
 
-# ===================== 4. Gradio UI（宽文本框 + 全宽布局） =====================
-
-with gr.Blocks() as demo:
+# ========== 4. Gradio UI ==========
+with gr.Blocks(css="""
+.wide_textbox textarea {
+    font-size: 14px;
+    line-height: 1.5;
+    min-height: 120px;
+}
+""") as demo:
     gr.Markdown("# ChartSmith – AI 论文图表生成助手")
 
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
             img_in = gr.Image(type="pil", label="上传你的学术图表（截图也可以）")
-
             style = gr.Dropdown(
                 ["Formal academic", "Infographic", "Magazine-style"],
                 value="Formal academic",
-                label="重绘风格"
+                label="重绘风格（当前版本仅用于说明）"
             )
-
-            language = gr.Radio(
-                ["English", "中文"],
-                value="English",
-                label="输出语言"
-            )
-
+            language = gr.Radio(["English", "中文"], value="English", label="输出语言")
             paper_file = gr.File(
-                label="上传相关论文（可选，用于更精准解释）",
-                file_types=[".txt", ".pdf", ".md"]
+                label="上传相关论文 PDF（可选，用于更精准解释）",
+                file_types=[".pdf"]
             )
+            btn = gr.Button("分析并美化图表", variant="primary")
 
-            btn = gr.Button("分析并美化图表")
-
-        with gr.Column():
+        with gr.Column(scale=1.2):
             orig_img = gr.Image(label="原始图表")
-            raw_caption_box = gr.Textbox(label="Step 2: 初步自动描述（Vision-LLM）", lines=4)
-            academic_box = gr.Textbox(label="Step 3: 学术化解释（结合论文上下文）", lines=6)
-            enhanced_img = gr.Image(label="Step 4: 强化后的图表（亮度/对比度提升）")
-            caption_box = gr.Textbox(label="Step 5: 自动生成图注（Caption）", lines=4)
-            summary_box = gr.Textbox(label="Step 5: 图表简短摘要 / 讨论", lines=6)
+            raw_caption_box = gr.Textbox(
+                label="Step 2: 初步自动描述（Vision-LLM）",
+                lines=3,
+                elem_classes=["wide_textbox"]
+            )
+            academic_box = gr.Textbox(
+                label="Step 3: 学术化解释（结合论文上下文）",
+                lines=5,
+                elem_classes=["wide_textbox"]
+            )
+            enhanced_img = gr.Image(label="Step 4: 增强后的图表（亮度/对比度提升）")
+            caption_box = gr.Textbox(
+                label="Step 5: 自动生成图注（Caption）",
+                lines=3,
+                elem_classes=["wide_textbox"]
+            )
+            summary_box = gr.Textbox(
+                label="Step 5: 图表相关简短摘要 / 讨论",
+                lines=5,
+                elem_classes=["wide_textbox"]
+            )
 
     btn.click(
         analyze_and_enhance,