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