Spaces:
Sleeping
Sleeping
File size: 4,007 Bytes
fbd78fc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """
Export Module — NEW feature
Exports extracted entities to:
1. CSV — for spreadsheet analysis
2. JSON — for downstream NLP pipelines
3. Annotated DOCX — entities highlighted in a Word document
"""
import csv
import json
import os
from datetime import datetime
from docx import Document
from docx.shared import RGBColor, Pt
from docx.oxml.ns import qn
from docx.oxml import OxmlElement
EXPORT_DIR = os.path.join(os.path.dirname(__file__), "..", "exports")
os.makedirs(EXPORT_DIR, exist_ok=True)
def _timestamp() -> str:
return datetime.now().strftime("%Y%m%d_%H%M%S")
# ── DOCX highlight colors (RGB) ─────────────────────────────────────────────
DOCX_COLORS = {
"Person": RGBColor(0x4F, 0xC3, 0xF7), # blue
"Organization": RGBColor(0x81, 0xC7, 0x84), # green
"Location": RGBColor(0xFF, 0xB7, 0x4D), # orange
"Miscellaneous": RGBColor(0xCE, 0x93, 0xD8), # purple
}
def export_csv(entities: list[dict], source_name: str = "document") -> str:
"""Export entities to CSV. Returns file path."""
fname = os.path.join(EXPORT_DIR, f"ner_{source_name}_{_timestamp()}.csv")
with open(fname, "w", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(f, fieldnames=["word", "label", "score", "start", "end"])
writer.writeheader()
writer.writerows(entities)
return fname
def export_json(entities: list[dict], meta: dict, source_name: str = "document") -> str:
"""Export entities + metadata to JSON. Returns file path."""
fname = os.path.join(EXPORT_DIR, f"ner_{source_name}_{_timestamp()}.json")
payload = {
"source": meta,
"entities": entities,
"count": len(entities),
}
with open(fname, "w", encoding="utf-8") as f:
json.dump(payload, f, indent=2, ensure_ascii=False)
return fname
def export_annotated_docx(
original_text: str,
entities: list[dict],
source_name: str = "document"
) -> str:
"""
Create a Word document with entities highlighted in context.
Returns file path.
"""
fname = os.path.join(EXPORT_DIR, f"annotated_{source_name}_{_timestamp()}.docx")
doc = Document()
# Title
title = doc.add_heading("SmartDoc NER — Annotated Output", level=1)
title.runs[0].font.color.rgb = RGBColor(0x1D, 0x6E, 0x6E)
doc.add_paragraph(f"Source: {source_name} | Entities found: {len(entities)}")
doc.add_paragraph("")
# Build annotated paragraph by inserting runs
para = doc.add_paragraph()
sorted_ents = sorted(entities, key=lambda e: e["start"])
cursor = 0
for ent in sorted_ents:
# Normal text before entity
if ent["start"] > cursor:
run = para.add_run(original_text[cursor:ent["start"]])
run.font.size = Pt(11)
# Highlighted entity run
run = para.add_run(original_text[ent["start"]:ent["end"]])
run.font.size = Pt(11)
run.font.bold = True
color = DOCX_COLORS.get(ent["label"], RGBColor(0xE0, 0xE0, 0xE0))
run.font.color.rgb = color
# Label superscript
run2 = para.add_run(f"[{ent['label']}]")
run2.font.size = Pt(8)
run2.font.color.rgb = color
cursor = ent["end"]
# Remaining text
if cursor < len(original_text):
run = para.add_run(original_text[cursor:])
run.font.size = Pt(11)
# Entity summary table
doc.add_heading("Entity Summary", level=2)
table = doc.add_table(rows=1, cols=3)
table.style = "Table Grid"
hdr = table.rows[0].cells
hdr[0].text, hdr[1].text, hdr[2].text = "Entity", "Type", "Confidence"
for ent in sorted(entities, key=lambda e: e["score"], reverse=True):
row = table.add_row().cells
row[0].text = ent["word"]
row[1].text = ent["label"]
row[2].text = f'{ent["score"]}%'
doc.save(fname)
return fname
|