Spaces:
Sleeping
Sleeping
| """ | |
| 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>" | |