Update app.py

app.py

@@ -3,22 +3,19 @@ import tempfile
 from pathlib import Path
 import json
 import torch
-import re
-from PIL import Image
 
 from detection_ocr import DetectionOCR
 from repair_head import RepairHead
 from semantic_head import ExplanationLLM, FormulaEvaluator
 
-# ============================================================
+# =========================
 # Device
-# ============================================================
+# =========================
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"[System] Using device: {DEVICE}")
 
-# ============================================================
+# =========================
 # Lazy singletons
-# ============================================================
+# =========================
 _ocr = None
 _repair = None
 _semantic = None
@@ -40,15 +37,15 @@ def get_semantic():
     if _semantic is None:
         explainer = ExplanationLLM(
             base_model="Qwen/Qwen2.5-Math-1.5B-Instruct",
-            lora_dir=None,
+            lora_dir="src/repair/repair_training/Repair_improved_1.5B_v1",
             device="cpu",
         )
         _semantic = FormulaEvaluator(explainer)
     return _semantic
 
-# ============================================================
+# =========================
 # Helpers
-# ============================================================
+# =========================
 def is_reasonable_math(expr: str) -> bool:
     if not expr:
         return False
@@ -58,31 +55,17 @@ def is_reasonable_math(expr: str) -> bool:
         return False
     return True
 
-def wrap_equation(expr: str) -> str:
-    return f"\\begin{{equation}}\n{expr}\n\\end{{equation}}"
-
-def merge_bbox_images_to_pdf(output_dir: Path) -> Path:
-    images = sorted(output_dir.glob("page_*_bbox.png"))
-    if not images:
-        return None
-
-    pil_images = [Image.open(p).convert("RGB") for p in images]
-    pdf_path = output_dir / "merged_bbox.pdf"
-    pil_images[0].save(
-        pdf_path,
-        save_all=True,
-        append_images=pil_images[1:],
-    )
-    return pdf_path
-
-# ============================================================
-# Core pipeline
-# ============================================================
+# =========================
+# Main pipeline (simple & reliable)
+# =========================
 def run_pipeline(pdf_file):
-    logs = []
+    if pdf_file is None:
+        return "No PDF uploaded.", None, None, None
 
+    logs = []
     def log(msg):
         logs.append(msg)
+        return "\n".join(logs)
 
     tmpdir = Path(tempfile.mkdtemp(prefix="pdf_ocr_"))
     outdir = tmpdir / "output"
@@ -90,129 +73,123 @@ def run_pipeline(pdf_file):
 
     log(f"📁 Output dir: {outdir}")
 
-    # ---------------- OCR ----------------
+    # -------- OCR --------
     log("\n🚀 Stage 1: OCR")
     ocr = get_ocr()
     ocr.process_pdf(pdf_file, output_dir=outdir)
+    log("✅ OCR done")
 
-    page_jsons = sorted(outdir.glob("page_*_results.json"))
-    log(f"✅ OCR done. Found {len(page_jsons)} pages")
-
-    # ---------------- Merge raw OCR ----------------
+    # -------- Merge OCR results --------
+    log("\n🔄 Stage 2: Merge OCR results")
     merged = {}
-    for p in page_jsons:
+    page_files = sorted(outdir.glob("page_*_results.json"))
+
+    for p in page_files:
         page = p.stem.replace("_results", "")
         data = json.loads(p.read_text(encoding="utf-8"))
-        merged[page] = {
-            "formulas": data.get("formulas", []),
-            "total": len(data.get("formulas", [])),
-        }
+        merged[page] = data["formulas"]
+        log(f"  {page}: {len(data['formulas'])} formulas")
 
-    overall_raw = outdir / "overall_ocr_results.json"
-    overall_raw.write_text(json.dumps(merged, indent=2, ensure_ascii=False), "utf-8")
-    log("💾 Saved overall_ocr_results.json")
+    overall_ocr = outdir / "overall_ocr_results.json"
+    overall_ocr.write_text(
+        json.dumps(merged, indent=2, ensure_ascii=False),
+        encoding="utf-8",
+    )
 
-    # ---------------- Repair ----------------
-    log("\n🔧 Stage 2: Repairing formulas")
+    # -------- Repair --------
+    log("\n🔧 Stage 3: Repair formulas")
     repair = get_repair()
 
     repaired_all = {}
-    merged_equations = []
-
-    for page, content in merged.items():
+    for page, formulas in merged.items():
         repaired_all[page] = []
-        for fm in content["formulas"]:
+        log(f"\n==== {page} ====")
+        for i, fm in enumerate(formulas, 1):
             raw = fm.get("raw_latex", "")
+            log(f"[{i}/{len(formulas)}] RAW: {raw[:80]}")
             repaired = repair.repair(raw)
-
             if is_reasonable_math(repaired):
+                log("   ✔ repaired")
                 repaired_all[page].append({
                     "raw_latex": raw,
                     "repaired": repaired
                 })
-                merged_equations.append(repaired)
             else:
+                log("   ✖ noise")
                 repaired_all[page].append({
                     "raw_latex": raw,
                     "repaired": "noise"
                 })
 
     repaired_json = outdir / "overall_ocr_repaired.json"
-    repaired_json.write_text(json.dumps(repaired_all, indent=2, ensure_ascii=False), "utf-8")
-    log("💾 Saved overall_ocr_repaired.json")
+    repaired_json.write_text(
+        json.dumps(repaired_all, indent=2, ensure_ascii=False),
+        encoding="utf-8",
+    )
 
-    # ---------------- Semantic ----------------
-    log("\n🧠 Stage 3: Semantic evaluation")
-    evaluator = get_semantic()
+    # -------- Semantic --------
+    log("\n🧠 Stage 4: Semantic evaluation")
+    semantic = get_semantic()
 
     semantic_results = []
     for page, items in repaired_all.items():
-        for item in items:
-            if item["repaired"] == "noise":
+        for r in items:
+            if r["repaired"] == "noise":
                 continue
-            res = evaluator.evaluate(item["repaired"])
-            semantic_results.append({
-                "page": page,
-                **res
-            })
+            res = semantic.evaluate(r["repaired"])
+            res["page"] = page
+            semantic_results.append(res)
+            log(f"   [{page}] {res['semantic_verdict']} (conf={res['confidence']})")
 
     semantic_json = outdir / "semantic_results.json"
-    semantic_json.write_text(json.dumps(semantic_results, indent=2, ensure_ascii=False), "utf-8")
-    log("💾 Saved semantic_results.json")
-
-    # ---------------- merge.tex ----------------
-    merge_tex = outdir / "merge.tex"
-    with merge_tex.open("w", encoding="utf-8") as f:
-        f.write("\\documentclass{article}\n")
-        f.write("\\usepackage{amsmath,amssymb}\n")
-        f.write("\\begin{document}\n\n")
-        for eq in merged_equations:
-            f.write(wrap_equation(eq) + "\n\n")
-        f.write("\\end{document}\n")
-    log("💾 Saved merge.tex")
-
-    # ---------------- bbox pdf ----------------
-    bbox_pdf = merge_bbox_images_to_pdf(outdir)
-    if bbox_pdf:
-        log("💾 Saved merged_bbox.pdf")
+    semantic_json.write_text(
+        json.dumps(semantic_results, indent=2, ensure_ascii=False),
+        encoding="utf-8",
+    )
 
     log("\n🎉 DONE")
 
     return (
         "\n".join(logs),
+        str(overall_ocr),
         str(repaired_json),
         str(semantic_json),
-        str(merge_tex),
-        str(bbox_pdf) if bbox_pdf else None,
     )
 
-# ============================================================
-# Gradio UI
-# ============================================================
+# =========================
+# Gradio UI (simple, stable)
+# =========================
 with gr.Blocks() as demo:
-    gr.Markdown("# 📘 PDF Formula OCR → Repair → Semantic")
+    gr.Markdown("# 📘 PDF → OCR → Repair → Semantic (Simple)")
 
     pdf = gr.File(label="Upload PDF", file_types=[".pdf"])
     run = gr.Button("Run")
 
-    log_box = gr.Textbox(label="Log", lines=25)
-
-    with gr.Row():
-        repaired_file = gr.File(label="overall_ocr_repaired.json")
-        semantic_file = gr.File(label="semantic_results.json")
+    live_log = gr.Textbox(
+        label="Live Log",
+        lines=22,
+        interactive=False
+    )
 
     with gr.Row():
-        tex_file = gr.File(label="merge.tex")
-        bbox_file = gr.File(label="merged_bbox.pdf")
+        ocr_json = gr.File(label="overall_ocr_results.json")
+        repaired_json = gr.File(label="overall_ocr_repaired.json")
+        semantic_json = gr.File(label="semantic_results.json")
 
     run.click(
         fn=run_pipeline,
         inputs=pdf,
-        outputs=[log_box, repaired_file, semantic_file, tex_file, bbox_file],
+        outputs=[
+            live_log,
+            ocr_json,
+            repaired_json,
+            semantic_json,
+        ],
     )
 
-# ============================================================
-# Main
-# ============================================================
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+    )