Spaces:
Sleeping
Sleeping
File size: 3,698 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 | """
NER Engine — uses dslim/bert-large-NER (F1 ~92.8 on CoNLL-2003)
instead of a basic pretrained pipeline, giving significantly higher accuracy.
"""
from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
import re
# ── Model selection ──────────────────────────────────────────────────────────
# dslim/bert-large-NER: F1 92.8% on CoNLL-2003 (vs ~85% for basic spacy en_core_web_sm)
# Falls back to bert-base-NER if large model unavailable (slower machines)
PRIMARY_MODEL = "dslim/bert-large-NER"
FALLBACK_MODEL = "dslim/bert-base-NER"
_nlp_pipeline = None
def load_model(model_name=PRIMARY_MODEL):
global _nlp_pipeline
if _nlp_pipeline is None:
try:
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(model_name)
_nlp_pipeline = pipeline(
"ner",
model=model,
tokenizer=tokenizer,
aggregation_strategy="max", # merges subword tokens → cleaner entities
device=-1 # CPU; change to 0 for GPU
)
except Exception:
# Fallback to smaller model
_nlp_pipeline = pipeline(
"ner",
model=FALLBACK_MODEL,
aggregation_strategy="max",
device=-1
)
return _nlp_pipeline
# ── Label mapping — human-readable ──────────────────────────────────────────
LABEL_MAP = {
"PER": "Person",
"ORG": "Organization",
"LOC": "Location",
"MISC": "Miscellaneous"
}
LABEL_COLORS = {
"Person": "#4FC3F7", # blue
"Organization": "#81C784", # green
"Location": "#FFB74D", # orange
"Miscellaneous": "#CE93D8", # purple
}
def run_ner(text: str) -> list[dict]:
"""
Run NER on text. Returns list of entity dicts:
{word, label, score, start, end}
"""
if not text or not text.strip():
return []
nlp = load_model()
raw = nlp(text)
entities = []
for ent in raw:
label = LABEL_MAP.get(ent["entity_group"], ent["entity_group"])
entities.append({
"word": ent["word"],
"label": label,
"score": round(float(ent["score"]) * 100, 1),
"start": ent["start"],
"end": ent["end"],
})
return entities
def get_highlighted_html(text: str, entities: list[dict]) -> str:
"""
Returns HTML string with entities highlighted inline.
Handles overlapping spans safely by processing right-to-left.
"""
if not entities:
return f"<p style='font-family:sans-serif;line-height:1.8'>{text}</p>"
# Sort by start position descending to insert HTML without shifting indices
sorted_ents = sorted(entities, key=lambda e: e["start"], reverse=True)
result = text
for ent in sorted_ents:
color = LABEL_COLORS.get(ent["label"], "#E0E0E0")
span = (
f'<mark style="background:{color};padding:2px 6px;border-radius:4px;'
f'margin:0 1px;font-family:sans-serif">'
f'{result[ent["start"]:ent["end"]]}'
f'<sup style="font-size:10px;margin-left:3px;font-weight:bold">'
f'{ent["label"]}</sup></mark>'
)
result = result[:ent["start"]] + span + result[ent["end"]:]
return f"<div style='font-family:sans-serif;line-height:2;font-size:15px'>{result}</div>"
|