Update app.py

app.py

@@ -6,38 +6,25 @@ from openai import OpenAI
 from pypdf import PdfReader
 
 
-# =========================
-# 1. Device & local vision captioner
-# =========================
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Lightweight image->text captioner for rough visual grounding
 vision_pipe = pipeline(
     "image-to-text",
     model="nlpconnect/vit-gpt2-image-captioning",
     device=0 if device == "cuda" else -1
 )
 
-# =========================
-# 2. Yunwu / OpenAI-compatible API client
-# =========================
-YUNWU_API_KEY = os.environ.get("YUNWU_API_KEY")
-if not YUNWU_API_KEY:
-    raise RuntimeError(
-        "YUNWU_API_KEY not set. Add it in HF Space: Settings → Variables and secrets."
-    )
+api_key = os.environ.get("YUNWU_API_KEY")
+if not api_key:
+    raise RuntimeError("YUNWU_API_KEY not set in Space secrets.")
 
 client = OpenAI(
-    api_key=YUNWU_API_KEY,
-    base_url="https://yunwu.ai/v1"  # Yunwu OpenAI-compatible endpoint
+    api_key=api_key,
+    base_url="https://yunwu.ai/v1"
 )
 
-def call_llm(
-    prompt: str,
-    model: str = "deepseek-chat",
-    temperature: float = 0.2,
-    max_tokens: int = 512
-) -> str:
+
+def call_llm(prompt, model="deepseek-chat", temperature=0.2, max_tokens=512):
     resp = client.chat.completions.create(
         model=model,
         messages=[
@@ -56,15 +43,12 @@ def call_llm(
     return resp.choices[0].message.content.strip()
 
 
-# =========================
-# 3. Utils: PDF snippet
-# =========================
-def extract_pdf_snippet(pdf_path: str | None, max_chars: int = 2500) -> str:
+def extract_pdf_snippet(pdf_path, max_chars=2500):
     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)
@@ -76,31 +60,13 @@ def extract_pdf_snippet(pdf_path: str | None, max_chars: int = 2500) -> str:
         return ""
 
 
-# =========================
-# 4. Main workflow
-# =========================
 def analyze_figure(image, style, pdf_path):
-    """
-    Outputs:
-    - Step 1: Paper-style explanation of what the figure shows (English)
-    - Step 2: Human-like suggestions for annotating the figure (English)
-    - Step 3: Plain-language explanation for presentations/assignments,
-              adapted to style (formal / fluency / simple)
-    """
-
     if image is None:
         return None, "Please upload a figure first.", "", ""
 
-    # Step 0: rough visual caption
     vision_raw = vision_pipe(image)[0]["generated_text"]
-
-    # Paper context
     pdf_context = extract_pdf_snippet(pdf_path)
 
-    # -------------------------
-    # Step 1: Figure meaning (paper-style)
-    # Use a stronger model if you like, e.g. "gpt-5.1" on Yunwu.
-    # -------------------------
     step1_prompt = f"""
 You are looking at a scientific figure from a paper.
 
@@ -121,10 +87,6 @@ Write in formal academic English, but keep it readable.
 """
     step1_text = call_llm(step1_prompt, model="deepseek-chat", max_tokens=420)
 
-    # -------------------------
-    # Step 2: Annotation suggestions (natural, not “AI-ish”)
-    # Keep it short, practical, no bold/markdown.
-    # -------------------------
     step2_prompt = f"""
 You are helping a student annotate this scientific figure for a presentation.
 
@@ -142,13 +104,8 @@ Constraints:
 - Focus on labels, arrows, callouts, grouping, legend clarity, and highlighting key contrasts.
 """
     step2_text = call_llm(step2_prompt, model="deepseek-chat", temperature=0.4, max_tokens=260)
-
-    # Light cleanup to remove accidental markdown/bold
     step2_text = step2_text.replace("**", "").replace("*", "").strip()
 
-    # -------------------------
-    # Step 3: Plain-language explanation in chosen style
-    # -------------------------
     style_map = {
         "formal": "formal but still plain language, suitable for a report",
         "fluency": "smooth, narrative, easy to speak aloud in a presentation",
@@ -176,9 +133,6 @@ Now paraphrase/explain the figure in {style_instruction}.
     return image, step1_text, step2_text, step3_text
 
 
-# =========================
-# 5. Gradio UI
-# =========================
 with gr.Blocks() as demo:
     gr.Markdown("## ChartSmith – AI Figure Explainer (multi-model workflow)")