Spaces:
Sleeping
Sleeping
| """ | |
| Named Entity Recognition (NER) Module | |
| Extracts and categorizes entities: persons, organizations, dates, | |
| monetary values, and locations. | |
| Uses spaCy (en_core_web_trf or en_core_web_lg) for high accuracy NER. | |
| Falls back to en_core_web_sm if transformer model is unavailable. | |
| Combines regex patterns for money/date extraction to improve recall. | |
| """ | |
| import logging | |
| import re | |
| from collections import defaultdict | |
| from typing import Dict, List | |
| logger = logging.getLogger(__name__) | |
| class EntityExtractor: | |
| """ | |
| High-precision NER using spaCy with regex augmentation. | |
| Entity categories: | |
| - persons → PERSON | |
| - organizations → ORG | |
| - dates → DATE | |
| - amounts → MONEY | |
| - locations → GPE, LOC, FAC | |
| """ | |
| # spaCy label → our category | |
| LABEL_MAP = { | |
| "PERSON": "persons", | |
| "ORG": "organizations", | |
| "DATE": "dates", | |
| "MONEY": "amounts", | |
| "GPE": "locations", | |
| "LOC": "locations", | |
| "FAC": "locations", | |
| } | |
| # Regex patterns for augmenting spaCy (higher recall for structured patterns) | |
| MONEY_PATTERNS = [ | |
| r"\$\s*[\d,]+(?:\.\d{1,2})?(?:\s*(?:million|billion|trillion|thousand|M|B|K))?\b", | |
| r"(?:USD|EUR|GBP|INR|JPY)\s*[\d,]+(?:\.\d{1,2})?(?:\s*(?:million|billion|trillion))?\b", | |
| r"[\d,]+(?:\.\d{1,2})?\s*(?:dollars?|euros?|pounds?|rupees?|crore|lakh)\b", | |
| ] | |
| DATE_PATTERNS = [ | |
| r"\b(?:January|February|March|April|May|June|July|August|September|October|November|December)" | |
| r"\s+\d{1,2},?\s+\d{4}\b", | |
| r"\b\d{1,2}[/-]\d{1,2}[/-]\d{2,4}\b", | |
| r"\b\d{4}[/-]\d{1,2}[/-]\d{1,2}\b", | |
| r"\b(?:Q[1-4]|FY|FQ)\s*\d{4}\b", | |
| r"\b\d{4}\b(?=\s*(?:to|through|until|-)\s*\d{4})", # Year ranges | |
| r"\b(?:Jan|Feb|Mar|Apr|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\.?\s+\d{1,2},?\s+\d{4}\b", | |
| ] | |
| _nlp = None # Class-level singleton | |
| def __init__(self): | |
| self._load_nlp() | |
| self._compile_patterns() | |
| def _load_nlp(cls): | |
| """Load spaCy NLP model (tries best available model).""" | |
| if cls._nlp is not None: | |
| return | |
| model_preference = [ | |
| "en_core_web_trf", # Transformer-based (best accuracy) | |
| "en_core_web_lg", # Large (good accuracy) | |
| "en_core_web_md", # Medium | |
| "en_core_web_sm", # Small fallback | |
| ] | |
| import spacy | |
| for model_name in model_preference: | |
| try: | |
| cls._nlp = spacy.load(model_name) | |
| logger.info(f"Loaded spaCy model: {model_name}") | |
| return | |
| except OSError: | |
| continue | |
| # Last resort: download small model | |
| logger.warning("No spaCy model found. Downloading en_core_web_sm...") | |
| try: | |
| import subprocess, sys | |
| subprocess.run( | |
| [sys.executable, "-m", "spacy", "download", "en_core_web_sm"], | |
| check=True, capture_output=True | |
| ) | |
| cls._nlp = spacy.load("en_core_web_sm") | |
| except Exception as e: | |
| logger.error(f"Failed to load any spaCy model: {e}") | |
| cls._nlp = None | |
| def _compile_patterns(self): | |
| """Pre-compile regex patterns for performance.""" | |
| self._money_re = [re.compile(p, re.IGNORECASE) for p in self.MONEY_PATTERNS] | |
| self._date_re = [re.compile(p, re.IGNORECASE) for p in self.DATE_PATTERNS] | |
| def extract(self, text: str) -> Dict[str, List[str]]: | |
| """ | |
| Extract named entities from text. | |
| Returns a dict with keys: persons, organizations, dates, amounts, locations. | |
| """ | |
| result = defaultdict(set) | |
| # ── spaCy NER ────────────────────────────────────────────────────────── | |
| if self._nlp and text.strip(): | |
| spacy_entities = self._extract_spacy(text) | |
| for category, entities in spacy_entities.items(): | |
| result[category].update(entities) | |
| # ── Regex augmentation ───────────────────────────────────────────────── | |
| regex_money = self._extract_regex(text, self._money_re) | |
| regex_dates = self._extract_regex(text, self._date_re) | |
| result["amounts"].update(regex_money) | |
| result["dates"].update(regex_dates) | |
| # ── Deduplicate and sort ─────────────────────────────────────────────── | |
| return { | |
| "persons": sorted(self._dedupe(result["persons"])), | |
| "organizations": sorted(self._dedupe(result["organizations"])), | |
| "dates": sorted(self._dedupe(result["dates"])), | |
| "amounts": sorted(self._dedupe(result["amounts"])), | |
| "locations": sorted(self._dedupe(result["locations"])), | |
| } | |
| def _extract_spacy(self, text: str) -> Dict[str, List[str]]: | |
| """Run spaCy NER on text, handling long documents by chunking.""" | |
| result = defaultdict(list) | |
| max_chars = 1_000_000 # spaCy default limit | |
| # Process in chunks for very long documents | |
| for chunk_start in range(0, len(text), max_chars - 100): | |
| chunk = text[chunk_start: chunk_start + max_chars] | |
| try: | |
| doc = self._nlp(chunk) | |
| for ent in doc.ents: | |
| category = self.LABEL_MAP.get(ent.label_) | |
| if category: | |
| cleaned = self._clean_entity(ent.text) | |
| if cleaned: | |
| result[category].append(cleaned) | |
| except Exception as e: | |
| logger.warning(f"spaCy processing chunk failed: {e}") | |
| return result | |
| def _extract_regex(self, text: str, patterns: List[re.Pattern]) -> List[str]: | |
| """Extract entities using regex patterns.""" | |
| found = [] | |
| for pattern in patterns: | |
| matches = pattern.findall(text) | |
| found.extend(m.strip() for m in matches if m.strip()) | |
| return found | |
| def _clean_entity(self, text: str) -> str: | |
| """ | |
| Clean an entity string: | |
| - Strip whitespace and punctuation artifacts | |
| - Normalize internal whitespace | |
| - Skip single-character or pure-number entities for persons/orgs | |
| """ | |
| text = text.strip() | |
| # Normalize internal whitespace | |
| text = re.sub(r"\s+", " ", text) | |
| # Remove leading/trailing punctuation | |
| text = text.strip(".,;:\"'()-") | |
| return text.strip() | |
| def _dedupe(self, items) -> List[str]: | |
| """ | |
| Deduplicate entity list. | |
| Also removes sub-strings that are already covered by longer matches. | |
| """ | |
| seen = list(dict.fromkeys(item for item in items if item)) | |
| # Remove items that are substrings of other items | |
| result = [] | |
| for item in seen: | |
| if not any( | |
| item.lower() != other.lower() and item.lower() in other.lower() | |
| for other in seen | |
| ): | |
| result.append(item) | |
| return result | |