import os import re import gradio as gr from transformers import pipeline import torch # 1. Thread Optimization & GPU/CPU Configuration # Force PyTorch to use 1 thread for CPU inference to avoid excessive context switching in constrained envs torch.set_num_threads(1) # Global lazy-loaded pipeline to minimize startup latency _ner_pipeline = None def get_ner_pipeline(): global _ner_pipeline if _ner_pipeline is None: # Standard NER model fine-tuned on CoNLL-2003 # Extremely accurate and lightweight _ner_pipeline = pipeline( "ner", model="dslim/bert-base-NER", aggregation_strategy="simple" ) return _ner_pipeline # 2. Date Extraction Logic def extract_dates(text): """ Extracts dates using regex patterns to supplement standard NER models which typically do not extract dates under CoNLL-2003 standard. """ date_patterns = [ # YYYY-MM-DD or YYYY/MM/DD r'\b\d{4}[-/]\d{1,2}[-/]\d{1,2}\b', # DD-MM-YYYY or MM/DD/YYYY or DD/MM/YY r'\b\d{1,2}[-/]\d{1,2}[-/]\d{2,4}\b', # Month Day, Year (e.g., January 15, 2021 or Jan 15 2021 or January 15th, 2021) r'\b(?:Jan(?:uary)?|Feb(?:ruary)?|Mar(?:ch)?|Apr(?:il)?|May|Jun(?:e)?|Jul(?:y)?|Aug(?:ust)?|Sep(?:tember)?|Oct(?:ober)?|Nov(?:ember)?|Dec(?:ember)?)\s+\d{1,2}(?:st|nd|rd|th)?(?:,\s+\d{4})?\b', # Day of Month (e.g. 15th of January, 2021) r'\b\d{1,2}(?:st|nd|rd|th)?\s+of\s+(?:Jan(?:uary)?|Feb(?:ruary)?|Mar(?:ch)?|Apr(?:il)?|May|Jun(?:e)?|Jul(?:y)?|Aug(?:ust)?|Sep(?:tember)?|Oct(?:ober)?|Nov(?:ember)?|Dec(?:ember)?)(?:\s+\d{4})?\b', # Month Year (e.g., January 2021) r'\b(?:Jan(?:uary)?|Feb(?:ruary)?|Mar(?:ch)?|Apr(?:il)?|May|Jun(?:e)?|Jul(?:y)?|Aug(?:ust)?|Sep(?:tember)?|Oct(?:ober)?|Nov(?:ember)?|Dec(?:ember)?)\s+\d{4}\b', ] dates = [] for pattern in date_patterns: for match in re.finditer(pattern, text, re.IGNORECASE): dates.append({ 'entity_group': 'DATE', 'word': match.group(), 'start': match.start(), 'end': match.end(), 'score': 1.0 }) return dates # 3. Label Mapping & Resolution LABEL_MAPPING = { 'PER': 'Person', 'ORG': 'Organization', 'LOC': 'Location', 'MISC': 'Miscellaneous', 'DATE': 'Date' } def merge_and_resolve_entities(ner_entities, date_entities): """ Merges transformer-detected entities and regex-extracted dates, resolving overlapping spans by selecting the longer span. """ # Filter to only keep entities we care about valid_ner = [e for e in ner_entities if e['entity_group'] in LABEL_MAPPING] all_entities = sorted(valid_ner + date_entities, key=lambda x: x['start']) resolved = [] last_end = -1 for ent in all_entities: if ent['start'] >= last_end: resolved.append(ent) last_end = ent['end'] else: # Overlap detected, keep the longer span if resolved: prev = resolved[-1] prev_len = prev['end'] - prev['start'] curr_len = ent['end'] - ent['start'] if curr_len > prev_len: resolved[-1] = ent last_end = ent['end'] return resolved def format_for_highlighted_text(text, entities): """ Formats entities into the standard list of tuples format required by gr.HighlightedText. """ entities = sorted(entities, key=lambda x: x['start']) result = [] last_idx = 0 for ent in entities: start = ent['start'] end = ent['end'] label = ent['entity_group'] if start > last_idx: result.append((text[last_idx:start], None)) friendly_label = LABEL_MAPPING.get(label, label) result.append((text[start:end], friendly_label)) last_idx = end if last_idx < len(text): result.append((text[last_idx:], None)) return result def analyze_entities(text): if not text or not text.strip(): return [], [["No entities", "N/A", "0.0%"]], {} try: # Get NER pipeline and predict ner_pipe = get_ner_pipeline() ner_results = ner_pipe(text) # Extract dates date_results = extract_dates(text) # Merge and resolve overlaps resolved_entities = merge_and_resolve_entities(ner_results, date_results) # Format highlight text representation highlight_data = format_for_highlighted_text(text, resolved_entities) # Format table data table_data = [] stats = {} for ent in resolved_entities: friendly_label = LABEL_MAPPING.get(ent['entity_group'], ent['entity_group']) confidence = f"{ent['score']:.1%}" table_data.append([ent['word'], friendly_label, confidence]) stats[friendly_label] = stats.get(friendly_label, 0) + 1 if not table_data: table_data = [["None detected", "N/A", "0.0%"]] return highlight_data, table_data, stats except Exception as e: return [(f"Error: {str(e)}", None)], [["Error", "N/A", "0.0%"]], {} # 5. Theme Configuration # Let Gradio handle the style natively with a standard clean Soft theme theme = gr.themes.Soft(primary_hue="indigo") with gr.Blocks(theme=theme, title="Minimalist NER Engine") as demo: gr.Markdown( """ # Entity Recognition Engine A minimalist pipeline trained to identify Persons, Organizations, Locations, Dates, and Miscellaneous entities instantly. """ ) with gr.Row(): with gr.Column(scale=5): input_text = gr.Textbox( label="Input Text", placeholder="Enter text to analyze here (e.g., 'On September 4, 1998, Google was founded by Larry Page and Sergey Brin in Menlo Park, California.')...", lines=8 ) submit_btn = gr.Button("Analyze Text", variant="primary") gr.Examples( examples=[ ["Microsoft was founded by Bill Gates and Paul Allen in Albuquerque, New Mexico on April 4, 1975. Today, Satya Nadella leads the company from Redmond, Washington."], ["Neil Armstrong and Buzz Aldrin landed the Apollo 11 Lunar Module on July 20, 1969. The entire mission was coordinated by NASA in Houston, Texas."], ["Alice visited the British Museum in London yesterday. She spoke with Dr. Vance about the new Greek history exhibition scheduled for October 12, 2026."] ], inputs=input_text, label="Sample Texts" ) with gr.Column(scale=5): gr.Markdown("### Extracted Entities") highlighted_output = gr.HighlightedText( label="Inline Highlights", combine_adjacent=False, show_legend=True, color_map={ "Person": "#dbeafe", # Light blue "Organization": "#fef3c7", # Light amber "Location": "#dcfce7", # Light green "Date": "#f3e8ff", # Light purple "Miscellaneous": "#f1f5f9" # Light slate } ) with gr.Tabs(): with gr.Tab("Structured Table"): entities_table = gr.Dataframe( headers=["Entity", "Type", "Confidence"], datatype=["str", "str", "str"], interactive=False, wrap=True ) with gr.Tab("Distribution"): stats_output = gr.Label( label="Entity Type Counts" ) submit_btn.click( fn=analyze_entities, inputs=input_text, outputs=[highlighted_output, entities_table, stats_output] ) gr.Markdown( """ *Powered by Hugging Face & Gradio* """ ) if __name__ == "__main__": port = int(os.environ.get("PORT", 7860)) demo.queue().launch( server_name="0.0.0.0", server_port=port, share=False )