Spaces:
Sleeping
Sleeping
| """ | |
| SmartDoc NER Pro — Upgraded Named Entity Recognition System | |
| Author: Nafees Ahmad | PAF-IAST | |
| Modules: | |
| 1. Text NER — direct text input with inline highlighting | |
| 2. Document NER — upload PDF / DOCX / TXT | |
| 3. Batch Processing — multiple documents at once | |
| 4. Entity Analytics — frequency, co-occurrence, confidence stats | |
| 5. Search & Filter — search within extracted entities | |
| 6. Export — CSV, JSON, annotated DOCX | |
| Model: dslim/bert-large-NER (F1 ~92.8% on CoNLL-2003) | |
| vs spacy en_core_web_sm (F1 ~85%) — significant accuracy gain | |
| """ | |
| import gradio as gr | |
| import pandas as pd | |
| import sys | |
| import os | |
| sys.path.insert(0, os.path.dirname(__file__)) | |
| from modules.ner_engine import run_ner, get_highlighted_html, LABEL_COLORS | |
| from modules.doc_reader import read_document, chunk_text | |
| from modules.analytics import (build_summary_table, label_counts, | |
| top_entities, co_occurrence, | |
| confidence_stats, deduplicate_entities) | |
| from modules.exporter import export_csv, export_json, export_annotated_docx | |
| from modules.entity_search import search_entities, filter_by_label, filter_by_confidence | |
| from modules.batch_processor import process_single, process_batch, aggregate_entities | |
| # ── Global state (simple in-memory store for current session) ──────────────── | |
| _state = { | |
| "entities": [], | |
| "text": "", | |
| "meta": {}, | |
| "batch_results": [], | |
| } | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| # TAB 1 — Text NER | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| def process_text(text: str, min_confidence: float): | |
| if not text.strip(): | |
| return "<p style='color:gray'>Please enter some text.</p>", [], "", "" | |
| entities = run_ner(text) | |
| entities = filter_by_confidence(entities, min_confidence) | |
| _state["entities"] = entities | |
| _state["text"] = text | |
| _state["meta"] = {"type": "Text input", "word_count": len(text.split()), "filename": "text_input"} | |
| html = get_highlighted_html(text, entities) | |
| rows = build_summary_table(entities) | |
| stats = confidence_stats(entities) | |
| counts = label_counts(entities) | |
| stats_text = ( | |
| f"Total entities: {stats['total']} | " | |
| f"Avg confidence: {stats['avg']}% | " | |
| f"Min: {stats['min']}% | Max: {stats['max']}%\n" | |
| f"Breakdown: {counts}" | |
| ) | |
| headers = ["Entity", "Type", "Confidence (%)"] | |
| df = pd.DataFrame(rows, columns=headers) if rows else pd.DataFrame(columns=headers) | |
| return html, df, stats_text, f"{len(entities)} entities extracted" | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| # TAB 2 — Document NER | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| def process_document(file, min_confidence: float): | |
| if file is None: | |
| return "<p style='color:gray'>Please upload a file.</p>", [], "" | |
| try: | |
| text, meta = read_document(file.name) | |
| except ValueError as e: | |
| return f"<p style='color:red'>{e}</p>", [], "" | |
| chunks = chunk_text(text, max_chars=400) | |
| entities = [] | |
| offset = 0 | |
| for chunk in chunks: | |
| ents = run_ner(chunk) | |
| for e in ents: | |
| e["start"] += offset | |
| e["end"] += offset | |
| entities.extend(ents) | |
| offset += len(chunk) + 1 | |
| entities = filter_by_confidence(entities, min_confidence) | |
| _state["entities"] = entities | |
| _state["text"] = text | |
| _state["meta"] = meta | |
| html = get_highlighted_html(text[:3000], entities) # display first 3000 chars | |
| rows = build_summary_table(entities) | |
| counts = label_counts(entities) | |
| headers = ["Entity", "Type", "Confidence (%)"] | |
| df = pd.DataFrame(rows, columns=headers) if rows else pd.DataFrame(columns=headers) | |
| info = ( | |
| f"File: {meta['filename']} | Type: {meta['type']} | " | |
| f"Pages: {meta['pages']} | Words: {meta['word_count']} | " | |
| f"Entities found: {len(entities)} | Breakdown: {counts}" | |
| ) | |
| return html, df, info | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| # TAB 3 — Export | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| def do_export(export_type: str): | |
| entities = _state["entities"] | |
| text = _state["text"] | |
| meta = _state["meta"] | |
| name = meta.get("filename", "document").replace(".", "_") | |
| if not entities: | |
| return None, "No entities to export. Run NER first." | |
| if export_type == "CSV": | |
| path = export_csv(entities, name) | |
| elif export_type == "JSON": | |
| path = export_json(entities, meta, name) | |
| elif export_type == "Annotated DOCX": | |
| path = export_annotated_docx(text, entities, name) | |
| else: | |
| return None, "Unknown export type." | |
| return path, f"Exported {len(entities)} entities to {os.path.basename(path)}" | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| # TAB 4 — Analytics | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| def show_analytics(): | |
| entities = _state["entities"] | |
| if not entities: | |
| return "Run NER on a document first.", "", "" | |
| deduped = deduplicate_entities(entities) | |
| top = top_entities(entities, top_n=5) | |
| pairs = co_occurrence(entities) | |
| stats = confidence_stats(entities) | |
| counts = label_counts(entities) | |
| top_text = "TOP ENTITIES PER TYPE:\n" + "-"*40 + "\n" | |
| for label, items in top.items(): | |
| top_text += f"\n{label}:\n" | |
| for word, freq in items: | |
| top_text += f" • {word} (×{freq})\n" | |
| pairs_text = "" | |
| if pairs: | |
| pairs_text = "\nPERSON — ORGANIZATION CO-OCCURRENCES:\n" + "-"*40 + "\n" | |
| for person, org in pairs[:10]: | |
| pairs_text += f" {person} ↔ {org}\n" | |
| stats_text = ( | |
| f"\nCONFIDENCE STATISTICS:\n" + "-"*40 + "\n" | |
| f" Total entities: {stats['total']}\n" | |
| f" Unique entities: {len(deduped)}\n" | |
| f" Avg confidence: {stats['avg']}%\n" | |
| f" Max confidence: {stats['max']}%\n" | |
| f" Min confidence: {stats['min']}%\n\n" | |
| f"LABEL DISTRIBUTION:\n" + "-"*40 + "\n" | |
| ) | |
| for label, count in counts.items(): | |
| bar = "█" * int(count / max(counts.values()) * 20) | |
| stats_text += f" {label:<20} {bar} {count}\n" | |
| return top_text, pairs_text, stats_text | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| # TAB 5 — Search & Filter | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| def search_and_filter(query: str, label_filter: list, min_conf: float): | |
| entities = _state["entities"] | |
| if not entities: | |
| return [], "Run NER first." | |
| filtered = search_entities(entities, query) | |
| filtered = filter_by_label(filtered, label_filter) if label_filter else filtered | |
| filtered = filter_by_confidence(filtered, min_conf) | |
| rows = build_summary_table(filtered) | |
| headers = ["Entity", "Type", "Confidence (%)"] | |
| df = pd.DataFrame(rows, columns=headers) if rows else pd.DataFrame(columns=headers) | |
| info = f"{len(filtered)} entities match your filters (from {len(entities)} total)" | |
| return df, info | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| # TAB 6 — Batch Processing | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| def process_batch_files(files, min_confidence: float): | |
| if not files: | |
| return [], "Upload files first." | |
| paths = [f.name for f in files] | |
| results = process_batch(paths) | |
| summary_rows = [] | |
| for res in results: | |
| if "error" in res: | |
| summary_rows.append([res["path"], "ERROR", res["error"], 0, 0, 0, 0]) | |
| else: | |
| c = res["counts"] | |
| summary_rows.append([ | |
| res["meta"]["filename"], | |
| res["meta"]["type"], | |
| res["meta"]["word_count"], | |
| c.get("Person", 0), | |
| c.get("Organization", 0), | |
| c.get("Location", 0), | |
| sum(c.values()), | |
| ]) | |
| _state["batch_results"] = results | |
| all_ents = aggregate_entities(results) | |
| _state["entities"] = all_ents | |
| _state["text"] = " ".join(r.get("text", "") for r in results) | |
| _state["meta"] = {"filename": "batch", "type": "Batch", "pages": "N/A", | |
| "word_count": sum(r.get("meta", {}).get("word_count", 0) for r in results)} | |
| headers = ["File", "Type", "Words", "Persons", "Orgs", "Locations", "Total"] | |
| df = pd.DataFrame(summary_rows, columns=headers) | |
| info = f"Processed {len(results)} files. Total entities: {len(all_ents)}. Results available in Analytics & Export tabs." | |
| return df, info | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| # BUILD GRADIO UI | |
| # ════════════════════════════════════════════════════════════════════════════ | |
| with gr.Blocks( | |
| title="SmartDoc NER Pro", | |
| theme=gr.themes.Soft(primary_hue="teal", secondary_hue="blue"), | |
| ) as app: | |
| gr.Markdown(""" | |
| # 🔍 SmartDoc NER Pro | |
| **Advanced Named Entity Recognition** — Persons · Organizations · Locations · Miscellaneous | |
| **Model:** `dslim/bert-large-NER` (F1 ≈ 92.8% on CoNLL-2003) | |
| **Supports:** Text · PDF · DOCX · TXT · Batch Processing · Export (CSV / JSON / DOCX) | |
| *Built by Nafees Ahmad | PAF-IAST Pakistan* | |
| """) | |
| with gr.Tabs(): | |
| # ── Tab 1: Text Input ────────────────────────────────────────────── | |
| with gr.Tab("📝 Text NER"): | |
| gr.Markdown("### Enter text directly to extract named entities") | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| txt_input = gr.Textbox( | |
| label="Input Text", | |
| placeholder="Paste any text here... e.g. 'Nafees Ahmad studied at PAF-IAST in Haripur, Pakistan.'", | |
| lines=8 | |
| ) | |
| txt_min_conf = gr.Slider(50, 100, value=70, step=5, label="Min Confidence (%)") | |
| txt_btn = gr.Button("Extract Entities", variant="primary") | |
| with gr.Column(scale=1): | |
| txt_status = gr.Textbox(label="Status", interactive=False) | |
| txt_stats = gr.Textbox(label="Statistics", lines=3, interactive=False) | |
| txt_html = gr.HTML(label="Highlighted Text") | |
| txt_table = gr.DataFrame(label="Extracted Entities", interactive=False) | |
| txt_btn.click( | |
| process_text, | |
| inputs=[txt_input, txt_min_conf], | |
| outputs=[txt_html, txt_table, txt_stats, txt_status] | |
| ) | |
| # Quick example texts | |
| gr.Examples( | |
| examples=[ | |
| ["Nafees Ahmad is a Software Engineering student at PAF-IAST in Haripur Hazara, Pakistan. He won the Pak Angels Generative AI Hackathon in 2024 and completed internships at TIERS Limited and Advanced Telecom Services.", 70], | |
| ["Elon Musk, CEO of Tesla and SpaceX, announced a new factory in Austin, Texas. Amazon's Jeff Bezos also unveiled plans for a facility in Berlin, Germany.", 70], | |
| ["The World Health Organization (WHO) and UNICEF signed a new agreement in Geneva, Switzerland, to support healthcare initiatives in South Asia and sub-Saharan Africa.", 70], | |
| ], | |
| inputs=[txt_input, txt_min_conf] | |
| ) | |
| # ── Tab 2: Document Upload ───────────────────────────────────────── | |
| with gr.Tab("📄 Document NER"): | |
| gr.Markdown("### Upload a PDF, DOCX, or TXT file") | |
| with gr.Row(): | |
| doc_file = gr.File(label="Upload Document", file_types=[".pdf", ".docx", ".txt"]) | |
| doc_min_conf = gr.Slider(50, 100, value=70, step=5, label="Min Confidence (%)") | |
| doc_btn = gr.Button("Process Document", variant="primary") | |
| doc_info = gr.Textbox(label="Document Info", interactive=False) | |
| doc_html = gr.HTML(label="Highlighted Text (first 3000 chars)") | |
| doc_table = gr.DataFrame(label="Extracted Entities", interactive=False) | |
| doc_btn.click( | |
| process_document, | |
| inputs=[doc_file, doc_min_conf], | |
| outputs=[doc_html, doc_table, doc_info] | |
| ) | |
| # ── Tab 3: Batch Processing ──────────────────────────────────────── | |
| with gr.Tab("📦 Batch Processing"): | |
| gr.Markdown("### Process multiple documents at once") | |
| batch_files = gr.File(label="Upload Multiple Files", file_count="multiple", | |
| file_types=[".pdf", ".docx", ".txt"]) | |
| batch_min_conf = gr.Slider(50, 100, value=70, step=5, label="Min Confidence (%)") | |
| batch_btn = gr.Button("Process All", variant="primary") | |
| batch_info = gr.Textbox(label="Batch Status", interactive=False) | |
| batch_table = gr.DataFrame(label="Per-Document Summary", interactive=False) | |
| batch_btn.click( | |
| process_batch_files, | |
| inputs=[batch_files, batch_min_conf], | |
| outputs=[batch_table, batch_info] | |
| ) | |
| # ── Tab 4: Search & Filter ───────────────────────────────────────── | |
| with gr.Tab("🔎 Search & Filter"): | |
| gr.Markdown("### Search and filter entities from the last processed document") | |
| with gr.Row(): | |
| sf_query = gr.Textbox(label="Search keyword", placeholder="e.g. Ahmad, Microsoft, Pakistan") | |
| sf_labels = gr.CheckboxGroup( | |
| ["Person", "Organization", "Location", "Miscellaneous"], | |
| label="Filter by type", value=[] | |
| ) | |
| sf_conf = gr.Slider(0, 100, value=0, step=5, label="Min Confidence (%)") | |
| sf_btn = gr.Button("Search", variant="primary") | |
| sf_info = gr.Textbox(label="Results", interactive=False) | |
| sf_table = gr.DataFrame(label="Filtered Entities", interactive=False) | |
| sf_btn.click( | |
| search_and_filter, | |
| inputs=[sf_query, sf_labels, sf_conf], | |
| outputs=[sf_table, sf_info] | |
| ) | |
| # ── Tab 5: Analytics ─────────────────────────────────────────────── | |
| with gr.Tab("📊 Analytics"): | |
| gr.Markdown("### Entity analytics from the last processed document or batch") | |
| an_btn = gr.Button("Run Analytics", variant="primary") | |
| with gr.Row(): | |
| an_top = gr.Textbox(label="Top Entities per Type", lines=15, interactive=False) | |
| an_pairs = gr.Textbox(label="Person–Organization Co-occurrences", lines=15, interactive=False) | |
| an_stats = gr.Textbox(label="Confidence & Distribution Statistics", lines=12, interactive=False) | |
| an_btn.click( | |
| show_analytics, | |
| inputs=[], | |
| outputs=[an_top, an_pairs, an_stats] | |
| ) | |
| # ── Tab 6: Export ───────────────────────────────────────────────── | |
| with gr.Tab("💾 Export"): | |
| gr.Markdown("### Export extracted entities (from last processed document)") | |
| with gr.Row(): | |
| ex_type = gr.Radio(["CSV", "JSON", "Annotated DOCX"], label="Export format", value="CSV") | |
| ex_btn = gr.Button("Export", variant="primary") | |
| ex_info = gr.Textbox(label="Export status", interactive=False) | |
| ex_file = gr.File(label="Download exported file") | |
| ex_btn.click( | |
| do_export, | |
| inputs=[ex_type], | |
| outputs=[ex_file, ex_info] | |
| ) | |
| gr.Markdown(""" | |
| --- | |
| **Legend:** | |
| <span style='background:#4FC3F7;padding:2px 8px;border-radius:4px'>Person</span> | |
| <span style='background:#81C784;padding:2px 8px;border-radius:4px;margin-left:6px'>Organization</span> | |
| <span style='background:#FFB74D;padding:2px 8px;border-radius:4px;margin-left:6px'>Location</span> | |
| <span style='background:#CE93D8;padding:2px 8px;border-radius:4px;margin-left:6px'>Miscellaneous</span> | |
| """) | |
| if __name__ == "__main__": | |
| app.launch() | |