import os import gradio as gr import pandas as pd import numpy as np import json import plotly.express as px from huggingface_hub import InferenceClient # Load or download spaCy English model dynamically import spacy try: nlp = spacy.load("en_core_web_sm") except OSError: import spacy.cli spacy.cli.download("en_core_web_sm") nlp = spacy.load("en_core_web_sm") def load_data(file_obj): """Safely loads CSV, Excel, or TXT file into a Pandas DataFrame.""" if file_obj is None: return None, gr.update(choices=[], visible=False), "Please upload a file." file_path = file_obj.name ext = os.path.splitext(file_path)[1].lower() try: if ext == '.csv': df = pd.read_csv(file_path) elif ext in ['.xls', '.xlsx']: df = pd.read_excel(file_path) elif ext == '.txt': with open(file_path, 'r', encoding='utf-8') as f: content = f.read() df = pd.DataFrame({'text': [content]}) else: return None, gr.update(choices=[], visible=False), "Unsupported file format. Please upload .csv, .xlsx, or .txt." # Find object/string columns for dropdown string_cols = [col for col in df.columns if df[col].dtype == 'object' or df[col].astype(str).str.len().mean() > 5] if not string_cols: string_cols = list(df.columns) return df, gr.update(choices=string_cols, value=string_cols[0], visible=True), f"Successfully loaded dataset with {len(df)} rows." except Exception as e: return None, gr.update(choices=[], visible=False), f"Error loading file: {str(e)}" def get_highlighted_text(text, entities): """Helper to convert entities into Gradio's HighlightedText list-of-tuples format.""" # entities list of dicts: {"start": int, "end": int, "label": str} # Sort entities by start index entities = sorted(entities, key=lambda x: x["start"]) highlighted = [] last_idx = 0 for ent in entities: start, end, label = ent["start"], ent["end"], ent["label"] if start < last_idx: continue # Avoid overlapping issues if start > last_idx: highlighted.append((text[last_idx:start], None)) highlighted.append((text[start:end], label)) last_idx = end if last_idx < len(text): highlighted.append((text[last_idx:], None)) return highlighted def run_spacy_ner(text): """Runs local SpaCy NER on a single string.""" doc = nlp(text) entities = [] for ent in doc.ents: entities.append({ "text": ent.text, "label": ent.label_, "start": ent.start_char, "end": ent.end_char }) return entities def run_transformer_ner_api(text, hf_token, model_name): """Runs state-of-the-art transformer NER using student's personal HF token.""" if not hf_token: raise ValueError("Hugging Face API Token is required for Transformer Mode.") client = InferenceClient(token=hf_token) # We use HF Token Classification API try: # returns list of dicts: [{'entity_group': 'PER', 'score': 0.99, 'word': '...', 'start': 0, 'end': 5}] response = client.token_classification(text, model=model_name) except Exception as e: raise RuntimeError(f"Hugging Face Inference API error: {str(e)}") entities = [] for item in response: # Standardize labels from CONLL/standard formats label = item.get("entity_group", item.get("entity", "ENTITY")) if label.startswith("B-") or label.startswith("I-"): label = label[2:] # Strip BIO prefixes for clean visualization entities.append({ "text": item.get("word", ""), "label": label, "start": item.get("start", 0), "end": item.get("end", 0) }) return entities def analyze_ner(text_input, file_obj, text_col, method, hf_token, hf_model): # Determine the input documents docs = [] if file_obj is not None: df, _, _ = load_data(file_obj) if df is not None and text_col in df.columns: docs = df[text_col].astype(str).fillna("").tolist() elif text_input and text_input.strip(): docs = [text_input] if not docs: return None, None, None, None, "Please enter text or upload a valid dataset first." all_extracted = [] # Process documents for doc_idx, doc_text in enumerate(docs): try: if method == "spaCy (Local & Fast)": ents = run_spacy_ner(doc_text) else: ents = run_transformer_ner_api(doc_text, hf_token, hf_model) for e in ents: all_extracted.append({ "Doc_Index": doc_idx + 1, "Entity_Text": e["text"], "Label": e["label"], "Start_Char": e["start"], "End_Char": e["end"], "Context": f"...{doc_text[max(0, e['start']-30):min(len(doc_text), e['end']+30)]}..." }) except Exception as e: return None, None, None, None, f"Error processing row {doc_idx + 1}: {str(e)}" if not all_extracted: return ( [("No entities found in the text.", None)], pd.DataFrame(), None, None, "Analysis finished: No named entities were detected." ) df_ents = pd.DataFrame(all_extracted) # 1. Visualization format for the first document (to show beautiful color-highlighted text in UI) first_doc_text = docs[0] first_doc_ents = [e for e in all_extracted if e["Doc_Index"] == 1] # Standardize keys highlight_ents = [{"start": e["Start_Char"], "end": e["End_Char"], "label": e["Label"]} for e in first_doc_ents] highlighted_output = get_highlighted_text(first_doc_text, highlight_ents) # 2. Statistics Bar Chart label_counts = df_ents["Label"].value_counts().reset_index() label_counts.columns = ["Entity Type", "Count"] fig = px.bar( label_counts, x="Entity Type", y="Count", color="Entity Type", title="Distribution of Extracted Entity Types", template="plotly_dark" ) fig.update_layout(height=350, margin=dict(l=20, r=20, t=40, b=20)) # 3. Save export files csv_path = "extracted_entities.csv" json_path = "extracted_entities.json" df_ents.to_csv(csv_path, index=False) # Save formatted JSON with open(json_path, 'w', encoding='utf-8') as f: json.dump(all_extracted, f, indent=4, ensure_ascii=False) # Clean table for UI display df_table = df_ents[["Doc_Index", "Entity_Text", "Label", "Context"]].copy() return highlighted_output, df_table, fig, csv_path, json_path custom_css = """ body { background-color: #0b0f19; color: #f3f4f6; } .gradio-container { font-family: 'Inter', sans-serif !important; } h1, h2 { color: #6366f1 !important; } """ with gr.Blocks(theme=gr.themes.Default(primary_hue="indigo", secondary_hue="slate"), css=custom_css) as demo: df_state = gr.State() gr.HTML("""

Interactive Named Entity Recognizer

Extract and analyze people, places, dates, and organizations from raw text or datasets. Runs locally on standard models, or unlocks state-of-the-art Transformer models using your personal Hugging Face Token.

""") with gr.Row(): # Left Panel: Input controls with gr.Column(scale=1): gr.Markdown("### 1. Choose Input Source") with gr.Tabs(): with gr.TabItem("Paste Raw Text"): text_input = gr.Textbox( label="Source Text", placeholder="Paste your text here (e.g., 'Apple Inc. was founded by Steve Jobs in Cupertino, California...').", lines=10 ) with gr.TabItem("Upload Dataset File"): file_input = gr.File(label="Upload (.csv, .xlsx, .txt)", file_types=[".csv", ".xlsx", ".txt"]) text_column_selector = gr.Dropdown( label="Target Text Column", choices=[], visible=False, interactive=True ) status_text = gr.Markdown("No file uploaded yet.") gr.Markdown("### 2. Configure Model") method_selector = gr.Radio( choices=["spaCy (Local & Fast)", "Transformers (API Mode)"], value="spaCy (Local & Fast)", label="Extraction Model" ) with gr.Group() as token_group: hf_token_input = gr.Textbox( label="Hugging Face API Token", placeholder="hf_...", type="password", visible=False, info="Required to call advanced transformer models. Get one free at huggingface.co." ) hf_model_input = gr.Dropdown( choices=[ "dbmdz/bert-large-cased-finetuned-conll03-english", "dslim/bert-base-NER", "Babelscape/wikineural-multilingual-ner" ], value="dbmdz/bert-large-cased-finetuned-conll03-english", label="Transformer Model (HF API)", visible=False ) run_btn = gr.Button("Extract Entities", variant="primary") # Right Panel: Results with gr.Column(scale=2): gr.Markdown("### 3. Extracted Named Entities") with gr.Tabs(): with gr.TabItem("Visual Color-Highlighting"): highlighted_output = gr.HighlightedText( label="First Document Entity Highlight", combine_adjacent=False ) with gr.TabItem("Full Analysis Table"): table_output = gr.Dataframe( headers=["Doc_Index", "Entity_Text", "Label", "Context"], datatype=["number", "str", "str", "str"], interactive=False, wrap=True ) with gr.TabItem("Statistics Chart"): chart_output = gr.Plot(label="Entity Frequency Plot") gr.Markdown("### 4. Export & Download") with gr.Row(): download_csv = gr.File(label="Download CSV Report") download_json = gr.File(label="Download JSON Report") # Show/hide token field depending on model def toggle_method_fields(method): if method == "Transformers (API Mode)": return gr.update(visible=True), gr.update(visible=True) else: return gr.update(visible=False), gr.update(visible=False) method_selector.change( fn=toggle_method_fields, inputs=method_selector, outputs=[hf_token_input, hf_model_input] ) file_input.change( fn=load_data, inputs=file_input, outputs=[df_state, text_column_selector, status_text] ) run_btn.click( fn=analyze_ner, inputs=[text_input, file_input, text_column_selector, method_selector, hf_token_input, hf_model_input], outputs=[highlighted_output, table_output, chart_output, download_csv, download_json] ) if __name__ == "__main__": demo.launch()