Update app.py

app.py

@@ -1,20 +1,13 @@
-import gradio as gr
-from transformers import pipeline
-import fitz
-import tempfile
 import os
 import re
 import time
-import base64
+import fitz
+import gradio as gr
 
-# Load model 
-import time as _t
-_t0 = _t.time()
-ner = pipeline("ner", model="cahya/NusaBert-ner-v1.3", aggregation_strategy="simple")
-print(f"[STARTUP] Model loaded in {_t.time()-_t0:.2f}s — running warmup...")
-ner("PT Bank Indonesia")  # warmup: paksa model fully loaded ke memori
-print(f"[STARTUP] Warmup selesai. Total: {_t.time()-_t0:.2f}s")
+from transformers import pipeline
 
+# Load model 
+ner = pipeline("ner", model="cahya/NusaBert-ner-v1.3", aggregation_strategy="simple") 
 
 # Label mapping
 LABEL_MAP = {
@@ -26,19 +19,6 @@ LABEL_MAP = {
     "REG": "Agama", "TIM": "Waktu", "WOA": "Karya Seni", "LAN": "Bahasa",
 }
 
-RTF_COLOR_TABLE = [
-    (255,191,191),(174,221,255),(141,199,255),(184,255,184),(153,242,204),
-    (217,255,165),(255,229,140),(255,204,102),(204,255,218),(242,204,255),
-    (224,191,255),(255,242,178),(199,242,242),(255,186,186),(255,217,165),
-    (191,223,255),(230,218,255),(255,230,218),(204,255,242),
-]
-
-LABEL_RTF_COLOR_IDX = {
-    "PER":1,"ORG":2,"NOR":3,"LOC":4,"GPE":5,"FAC":6,"DAT":7,"TIM":8,
-    "MON":9,"CRD":10,"ORD":11,"PRC":12,"QTY":13,"LAW":14,"EVT":15,
-    "PRD":16,"REG":17,"WOA":18,"LAN":19,
-}
-
 LABEL_HEX = {
     "PER":"#FFBFBF","ORG":"#AEDDFF","NOR":"#8DC7FF","LOC":"#B8FFB8",
     "GPE":"#99F2CC","FAC":"#D9FFA5","DAT":"#FFE58C","TIM":"#FFCC66",
@@ -91,23 +71,7 @@ def get_label_id(raw_label: str) -> str:
     label_id = raw_label.replace("B-","").replace("I-","").replace("B_","").replace("I_","")
     return label_id.split("-")[-1].upper().strip()
 
-def escape_rtf(text: str) -> str:
-    out = []
-    for ch in text:
-        if ch == '\\': out.append('\\\\')
-        elif ch == '{': out.append('\\{')
-        elif ch == '}': out.append('\\}')
-        elif ord(ch) > 127: out.append(f'\\u{ord(ch)}?')
-        else: out.append(ch)
-    return ''.join(out)
-
-def build_rtf_color_table() -> str:
-    entries = "\\red0\\green0\\blue0;"
-    for r,g,b in RTF_COLOR_TABLE:
-        entries += f"\\red{r}\\green{g}\\blue{b};"
-    return "{\\colortbl ;" + entries + "}"
-
-def text_to_rtf_with_highlights(text: str, entity_map: dict) -> str:
+def highlight_html(text: str, entity_map: dict) -> str:
     sorted_entities = sorted(entity_map.items(), key=lambda x: len(x[0]), reverse=True)
     spans = []
     used = [False] * len(text)
@@ -123,140 +87,55 @@ def text_to_rtf_with_highlights(text: str, entity_map: dict) -> str:
             for i in range(s, e):
                 used[i] = True
     spans.sort(key=lambda x: x[0])
-    rtf_parts = []
+    parts = []
     cursor = 0
     for s, e, label_id in spans:
         if cursor < s:
-            rtf_parts.append(escape_rtf(text[cursor:s]))
-        color_idx = LABEL_RTF_COLOR_IDX.get(label_id, 1)
-        highlighted = escape_rtf(text[s:e])
-        rtf_parts.append(f"{{\\highlight{color_idx} {highlighted}}}")
+            parts.append(text[cursor:s].replace("\n", "<br>"))
+        hex_color = LABEL_HEX.get(label_id, "#e2e8f0")
+        label_idn = LABEL_MAP.get(label_id, label_id)
+        word = text[s:e]
+        parts.append(
+            f'<mark style="background:{hex_color};border-radius:4px;padding:1px 4px;'
+            f'font-weight:600;" title="{label_idn}">{word}</mark>'
+        )
         cursor = e
     if cursor < len(text):
-        rtf_parts.append(escape_rtf(text[cursor:]))
-    return "".join(rtf_parts)
-
-def wrap_rtf(content: str) -> str:
-    color_table = build_rtf_color_table()
-    content_rtf = content.replace("\r\n","\n").replace("\r","\n").replace("\n","\\par\n")
+        parts.append(text[cursor:].replace("\n", "<br>"))
     return (
-        "{\\rtf1\\ansi\\ansicpg1252\\deff0\n"
-        "{\\fonttbl{\\f0\\froman\\fcharset0 Times New Roman;}}\n"
-        f"{color_table}\n"
-        "\\widowctrl\\wpaper12240\\wpaperh15840\n"
-        "\\margl1800\\margr1800\\margt1440\\margb1440\n"
-        "\\f0\\fs24\\sl360\\slmult1\n"
-        f"{content_rtf}"
-        "}"
-    )
-
-def _build_debug_html(lines: list) -> str:
-    rows = "".join(
-        f'<div style="padding:6px 10px; border-bottom:1px solid rgba(255,255,255,0.07); '
-        f'font-size:13px; color:#cbd5e1; font-family:monospace;">{line}</div>'
-        for line in lines
+        '<div style="font-family:\'DM Sans\',sans-serif;font-size:14px;line-height:2;'
+        'color:#1e293b;background:#fff;border-radius:12px;padding:20px 24px;'
+        'border:1px solid #e2e8f0;white-space:pre-wrap;">'
+        + "".join(parts) + "</div>"
     )
-    return f"""
-    <div style="margin-top:12px; background:rgba(15,23,42,0.85); border:1px solid rgba(99,102,241,0.35);
-                border-radius:12px; overflow:hidden;">
-        <div style="background:linear-gradient(135deg,#1e40af,#6d28d9); padding:8px 14px;
-                    font-size:11px; font-weight:700; letter-spacing:0.1em; color:#fff;">
-            🔍 DEBUG — WAKTU PER LANGKAH
-        </div>
-        {rows}
-    </div>"""
-
+    
 # NER Teks 
-def run_ner(text: str):
-    if not text or not text.strip():
-        return "<p style='color:#94a3b8; font-style:italic; padding:16px;'>Masukkan teks terlebih dahulu.</p>"
-    results = ner(text.strip())
-    if not results:
-        return "<p style='color:#94a3b8; font-style:italic; padding:16px;'>Tidak ada entitas yang ditemukan.</p>"
-
-    rows_html = ""
-    row_num = 1
-    seen_words = set()
-    for ent in results:
-        raw_label = ent["entity_group"]
-        label_id  = get_label_id(raw_label)
-        label_idn = LABEL_MAP.get(label_id, raw_label)
-        word      = clean_word(ent["word"])
-        if not word: continue
-        word_key = word.lower()
-        if word_key in seen_words: continue
-        seen_words.add(word_key)
-        hex_color = LABEL_HEX.get(label_id, "#e2e8f0")
-        score     = f"{ent['score']:.2%}"
-        row_bg    = "#f8faff" if row_num % 2 == 0 else "#ffffff"
-        rows_html += f"""
-        <tr style="background:{row_bg};">
-            <td style="padding:9px 14px;border-bottom:1px solid #e8edf5;text-align:center;color:#64748b;font-size:12px;">{row_num}</td>
-            <td style="padding:9px 14px;border-bottom:1px solid #e8edf5;font-weight:600;color:#1e293b;">{word}</td>
-            <td style="padding:9px 14px;border-bottom:1px solid #e8edf5;">
-                <span style="background:{hex_color};padding:3px 10px;border-radius:20px;font-size:12px;font-weight:600;color:#1e293b;">{label_idn}</span>
-            </td>
-            <td style="padding:9px 14px;border-bottom:1px solid #e8edf5;font-size:12px;color:#64748b;text-align:center;">{score}</td>
-        </tr>"""
-        row_num += 1
-
-    if not rows_html:
-        return "<p style='color:#94a3b8; font-style:italic; padding:16px;'>Tidak ada entitas yang ditemukan.</p>"
-
-    return f"""
-    <div style="overflow-x:auto;margin-top:4px;border-radius:12px;border:1px solid #e2e8f0;box-shadow:0 2px 12px rgba(0,0,0,0.06);">
-        <table style="width:100%;border-collapse:collapse;font-size:14px;font-family:'Segoe UI',sans-serif;">
-            <thead>
-                <tr style="background:linear-gradient(135deg,#1e40af,#6d28d9);">
-                    <th style="padding:12px 14px;color:#fff;width:55px;font-weight:600;font-size:12px;letter-spacing:0.05em;">NO</th>
-                    <th style="padding:12px 14px;color:#fff;text-align:left;font-weight:600;font-size:12px;letter-spacing:0.05em;">KATA / FRASA</th>
-                    <th style="padding:12px 14px;color:#fff;text-align:left;font-weight:600;font-size:12px;letter-spacing:0.05em;">ENTITAS</th>
-                    <th style="padding:12px 14px;color:#fff;width:90px;font-weight:600;font-size:12px;letter-spacing:0.05em;">SKOR</th>
-                </tr>
-            </thead>
-            <tbody>{rows_html}</tbody>
-        </table>
-    </div>"""
-
-
-# NER PDF → RTF
-def run_ner_pdf(pdf_file):
-    if pdf_file is None:
-        return "", "", "<p style='color:#94a3b8;font-style:italic;padding:16px;'>Unggah file PDF terlebih dahulu.</p>", ""
-
-    pdf_path = pdf_file if isinstance(pdf_file, str) else pdf_file.name
-    doc = fitz.open(pdf_path)
-    page_count = len(doc)
-
-    if page_count > MAX_PDF_PAGES:
+def run_ner_file(upload_file):
+    if upload_file is None:
+        return "<p style='color:#94a3b8;font-style:italic;padding:16px;'>Unggah file terlebih dahulu.</p>", ""
+
+    file_path = upload_file if isinstance(upload_file, str) else upload_file.name
+    ext = os.path.splitext(file_path)[-1].lower()
+
+    if ext == ".txt":
+        with open(file_path, "r", encoding="utf-8", errors="replace") as f:
+            full_text = f.read()
+        page_count = 1
+    elif ext == ".pdf":
+        doc = fitz.open(file_path)
+        page_count = len(doc)
+        if page_count > MAX_PDF_PAGES:
+            doc.close()
+            return (
+                f"<p style='color:#be123c;padding:16px;'>PDF terlalu banyak halaman "
+                f"({page_count}). Maks {MAX_PDF_PAGES} halaman.</p>", ""
+            )
+        full_text = "\n\n".join(page.get_text() for page in doc)
         doc.close()
-        return "", "", f"""
-        <div style="background:#fff1f2;border:1.5px solid #fda4af;border-radius:12px;
-                    padding:16px 20px;display:flex;align-items:center;gap:12px;">
-            <span style="font-size:24px;">⚠️</span>
-            <div>
-                <p style="margin:0;font-weight:700;color:#be123c;font-size:15px;">PDF Terlalu Banyak Halaman</p>
-                <p style="margin:4px 0 0;color:#9f1239;font-size:13px;">
-                    File Anda memiliki <strong>{page_count} halaman</strong>.
-                    Maksimal <strong>{MAX_PDF_PAGES} halaman</strong>.
-                </p>
-            </div>
-        </div>""", ""
-
-    debug_lines = []
-    t_total = time.time()
-
-    # LANGKAH 1: Ekstrak teks
-    t0 = time.time()
-    pages_text = [page.get_text() for page in doc]
-    doc.close()
-    full_text = "\n\n".join(pages_text)
-    char_count = len(full_text)
-    t1 = time.time()
-    debug_lines.append(f"✅ <b>Langkah 1</b> — Ekstrak teks PDF: <b>{t1-t0:.3f}s</b> | {char_count:,} karakter | {page_count} halaman")
-
-    # ── LANGKAH 2: Chunking ──
-    t0 = time.time()
+    else:
+        return "<p style='color:#be123c;padding:16px;'>Format tidak didukung.</p>", ""
+
+    # Chunking
     chunks = []
     start = 0
     while start < len(full_text):
@@ -265,137 +144,42 @@ def run_ner_pdf(pdf_file):
         if end == len(full_text):
             break
         start = end - OVERLAP_CHARS
-    t1 = time.time()
-    debug_lines.append(f"✅ <b>Langkah 2</b> — Chunking teks: <b>{t1-t0:.4f}s</b> | {len(chunks)} chunk @ {MAX_CHUNK_CHARS} karakter")
 
-    # ── LANGKAH 3: NER model (ini yang biasanya paling lama) ──
-    t0 = time.time()
+    # NER
     all_ner_results = []
-    chunk_times = []
-    for i, chunk in enumerate(chunks):
-        tc = time.time()
-        chunk_results = ner(chunk.strip())
-        chunk_times.append(time.time() - tc)
-        all_ner_results.extend(chunk_results)
-    t1 = time.time()
-    avg_chunk = sum(chunk_times)/len(chunk_times) if chunk_times else 0
-    debug_lines.append(
-        f"✅ <b>Langkah 3</b> — NER model: <b>{t1-t0:.2f}s</b> total | "
-        f"{len(chunks)} chunk | avg/chunk: {avg_chunk:.2f}s | {len(all_ner_results)} entitas raw"
-    )
-    # Detail per chunk
-    for i, ct in enumerate(chunk_times):
-        debug_lines.append(f"&nbsp;&nbsp;&nbsp;↳ chunk {i+1}/{len(chunks)}: {ct:.2f}s ({len(chunks[i])} karakter)")
+    for chunk in chunks:
+        all_ner_results.extend(ner(chunk.strip()))
 
     if not all_ner_results:
-        dh = _build_debug_html(debug_lines)
-        return None, "", "<p style='color:#94a3b8;padding:16px;'>Tidak ada entitas ditemukan.</p>", dh
+        return "<p style='color:#94a3b8;padding:16px;'>Tidak ada entitas ditemukan.</p>", ""
 
-    # ── LANGKAH 4: Bangun entity_map ──
-    t0 = time.time()
+    # Bangun entity_map
     entity_map: dict[str, str] = {}
     for ent in all_ner_results:
         word = clean_word(ent["word"])
-        if len(word) < 2: continue
+        if len(word) < 2:
+            continue
         label_id = get_label_id(ent["entity_group"])
         w_lower = word.lower()
         if w_lower not in entity_map:
             entity_map[w_lower] = label_id
-    t1 = time.time()
-    debug_lines.append(f"✅ <b>Langkah 4</b> — Bangun entity_map: <b>{t1-t0:.4f}s</b> | {len(entity_map)} entitas unik")
 
     if not entity_map:
-        dh = _build_debug_html(debug_lines)
-        return None, "", "<p style='color:#94a3b8;padding:16px;'>Tidak ada entitas ditemukan.</p>", dh
-
-    # ── LANGKAH 5: Highlight → RTF content ──
-    t0 = time.time()
-    rtf_content = text_to_rtf_with_highlights(full_text, entity_map)
-    t1 = time.time()
-    debug_lines.append(f"✅ <b>Langkah 5</b> — Highlight & build RTF content: <b>{t1-t0:.4f}s</b>")
-
-    # ── LANGKAH 6: Build Download Link ──
-    t0 = time.time()
+        return "<p style='color:#94a3b8;padding:16px;'>Tidak ada entitas ditemukan.</p>", ""
 
-    rtf_full = wrap_rtf(rtf_content)
+    highlighted = highlight_html(full_text, entity_map)
 
-    rtf_bytes = rtf_full.encode(
-        "ascii",
-        errors="replace"
+    # Badge legend
+    found_labels = set(entity_map.values())
+    badges = "".join(
+        f'<span style="background:{LABEL_HEX.get(l,"#e2e8f0")};border-radius:20px;'
+        f'padding:3px 10px;font-size:12px;font-weight:600;color:#1e293b;margin:3px;">'
+        f'{LABEL_MAP.get(l, l)}</span>'
+        for l in sorted(found_labels)
     )
+    legend_html = f'<div style="margin-bottom:12px;line-height:2;">{badges}</div>'
 
-    b64 = base64.b64encode(
-        rtf_bytes
-    ).decode()
-
-    download_html = f"""
-    <div style="padding:12px 0;">
-        <a
-            href="data:application/rtf;base64,{b64}"
-            download="hasil_ner.rtf"
-            style="
-                display:inline-block;
-                background:linear-gradient(135deg,#1e40af,#6d28d9);
-                color:white;
-                text-decoration:none;
-                padding:12px 20px;
-                border-radius:10px;
-                font-weight:600;
-                font-family:'DM Sans',sans-serif;
-            ">
-            📥 Download Hasil RTF
-        </a>
-    </div>
-    """
-    t1 = time.time()
-
-    debug_lines.append(
-        f"✅ <b>Langkah 6</b> — Build Download Link: <b>{t1-t0:.4f}s</b>"
-    )
-
-    t_end = time.time()
-    debug_lines.append(f"⏱ <b>TOTAL WAKTU: {t_end-t_total:.2f}s</b>")
-
-    debug_html = _build_debug_html(debug_lines)
-    found_labels: set[str] = set(entity_map.values())
-
-    # Legend
-    legend_rows = ""
-    for label_id in sorted(found_labels):
-        label_idn = LABEL_MAP.get(label_id, label_id)
-        hex_color = LABEL_HEX.get(label_id, "#e2e8f0")
-        legend_rows += f"""
-        <tr>
-            <td style="padding:8px 14px;border-bottom:1px solid #e8edf5;">
-                <span style="display:inline-block;width:24px;height:24px;border-radius:6px;
-                             background:{hex_color};vertical-align:middle;border:1px solid rgba(0,0,0,0.1);"></span>
-            </td>
-            <td style="padding:8px 14px;border-bottom:1px solid #e8edf5;font-weight:600;color:#1e293b;font-size:13px;">{label_id}</td>
-            <td style="padding:8px 14px;border-bottom:1px solid #e8edf5;color:#475569;font-size:13px;">{label_idn}</td>
-        </tr>"""
-
-    legend_html = f"""
-    <div style="margin-top:4px;">
-        <p style="font-size:13px;color:#64748b;margin-bottom:10px;">
-            Ditemukan {len(found_labels)}jenis entitas
-            dari {page_count}halaman PDF.
-            Buka file rtf di Microsoft Word atau LibreOffice Writer.
-        </p>
-        <div style="overflow-x:auto;border-radius:12px;border:1px solid #e2e8f0;box-shadow:0 2px 12px rgba(0,0,0,0.06);">
-            <table style="width:100%;border-collapse:collapse;font-family:'Segoe UI',sans-serif;">
-                <thead>
-                    <tr style="background:linear-gradient(135deg,#1e40af,#6d28d9);">
-                        <th style="padding:11px 14px;color:#fff;font-size:12px;letter-spacing:0.05em;width:50px;">WARNA</th>
-                        <th style="padding:11px 14px;color:#fff;font-size:12px;letter-spacing:0.05em;text-align:left;">KODE</th>
-                        <th style="padding:11px 14px;color:#fff;font-size:12px;letter-spacing:0.05em;text-align:left;">NAMA ENTITAS</th>
-                    </tr>
-                </thead>
-                <tbody>{legend_rows}</tbody>
-            </table>
-        </div>
-    </div>"""
-
-    return download_html, legend_html, "", debug_html
+    return highlighted, legend_html     
 
 # CSS
 CUSTOM_CSS = """
@@ -462,58 +246,33 @@ with gr.Blocks(title="NER for Financial Statements", css=CUSTOM_CSS) as demo:
                 gr.HTML('<div class="section-heading" style="margin-top:20px;">Hasil Analisis Entitas</div>')
                 text_output = gr.HTML(value="<p style='color:#64748b;font-size:14px;font-family:DM Sans,sans-serif;padding:20px;text-align:center;'>Masukkan teks lalu klik Lakukan Analisis.</p>")
 
-        with gr.Tab("Analisis PDF"):
+        with gr.Tab("Analisis File"):
             with gr.Column(elem_id="center-col"):
-                gr.HTML(
-                    '''<div style="margin:0 0 16px;">
-                    <p style="color:#94a3b8;font-family:DM Sans,sans-serif;font-size:13px;margin:0 0 8px;">
-                        Unggah PDF laporan keuangan maks 5 halaman.
-                        Output berupa file.rtf dengan stabilo warna —
-                        bisa dibuka di Microsoft Word atau LibreOffice.
-                    </p>
-                    </div>'''
-                )
-                gr.HTML('<div class="section-heading">Unggah PDF</div>')
-                pdf_input = gr.File(label="", file_types=[".pdf"], type="filepath")
-                analyze_pdf_btn = gr.Button("Highlight Entitas → Download RTF", elem_id="analyze-pdf-btn")
-                pdf_warning = gr.HTML(value="")
-
-                with gr.Column(visible=False) as pdf_results_col:
-                    gr.HTML('<div class="section-heading" style="margin-top:20px;">Download Hasil RTF</div>')
-                    pdf_output = gr.HTML(value="")
-                    gr.HTML('<div class="section-heading" style="margin-top:20px;">Keterangan Warna Entitas</div>')
-                    pdf_legend = gr.HTML(value="")
-                    # Debug panel — selalu tampil setelah proses
-                    pdf_debug = gr.HTML(value="")
-
+                gr.HTML('<p style="color:#94a3b8;font-family:DM Sans,sans-serif;font-size:13px;margin:0 0 16px;">'
+                        'Unggah file .pdf (maks 5 halaman) atau .txt.</p>'
+                       )
+                gr.HTML('<div class="section-heading">Unggah File</div>')
+                pdf_input = gr.File(label="", file_types=[".pdf", ".txt"], type="filepath")
+                analyze_pdf_btn = gr.Button("Analisis & Highlight Entitas", elem_id="analyze-pdf-btn")
+                
+                gr.HTML('<div class="section-heading" style="margin-top:20px;">Entitas Ditemukan</div>')
+                pdf_legend = gr.HTML(value="")
+                gr.HTML('<div class="section-heading" style="margin-top:12px;">Teks Ter-highlight</div>')
+                pdf_output = gr.HTML(value="")
+                
     gr.HTML('<div id="footer">NER for Financial Statements 2026</div>')
-
+                
     # Wiring
     for btn, body in example_btns:
         btn.click(fn=lambda b=body: b, inputs=[], outputs=text_input)
-
+                    
     analyze_btn.click(fn=run_ner, inputs=text_input, outputs=text_output)
     text_input.submit(fn=run_ner, inputs=text_input, outputs=text_output)
 
-    def handle_pdf(pdf_file):
-        download_html, legend_html, warn_html, debug_html = run_ner_pdf(pdf_file)
-        
-        show = gr.Column(
-            visible=bool(download_html)
-        )
-        
-        return (
-            download_html,
-            legend_html,
-            warn_html,
-            debug_html,
-            show
-        )
-
     analyze_pdf_btn.click(
-        fn=handle_pdf,
+        fn=run_ner_file,
         inputs=pdf_input,
-        outputs=[pdf_output, pdf_legend, pdf_warning, pdf_debug, pdf_results_col],
+        outputs=[pdf_output, pdf_legend],
     )
 
 if __name__ == "__main__":