Spaces:
Sleeping
Sleeping
| from transformers import pipeline | |
| import gradio as gr | |
| import pandas as pd | |
| # Load NER model | |
| ner_model = pipeline( | |
| "ner", | |
| model="dslim/bert-base-NER", | |
| aggregation_strategy="simple" | |
| ) | |
| def create_df(): | |
| return pd.DataFrame( | |
| columns=["Entity", "Confidence (%)"] | |
| ) | |
| def ner_inference(text): | |
| entities = ner_model(text) | |
| per_rows, org_rows, loc_rows = [], [], [] | |
| for ent in entities: | |
| row = [ent["word"], round(ent["score"] * 100, 2)] | |
| if ent["entity_group"] == "PER": | |
| per_rows.append(row) | |
| elif ent["entity_group"] == "ORG": | |
| org_rows.append(row) | |
| elif ent["entity_group"] == "LOC": | |
| loc_rows.append(row) | |
| df_per = pd.DataFrame(per_rows, columns=["Person", "Confidence (%)"]) if per_rows else create_df() | |
| df_org = pd.DataFrame(org_rows, columns=["Organization", "Confidence (%)"]) if org_rows else create_df() | |
| df_loc = pd.DataFrame(loc_rows, columns=["Location", "Confidence (%)"]) if loc_rows else create_df() | |
| return df_per, df_org, df_loc | |
| # Gradio UI | |
| interface = gr.Interface( | |
| fn=ner_inference, | |
| inputs=gr.Textbox( | |
| lines=5, | |
| placeholder="Enter text here...", | |
| label="Input Text" | |
| ), | |
| outputs=[ | |
| gr.Dataframe(label="👤 Persons", interactive=False), | |
| gr.Dataframe(label="🏢 Organizations", interactive=False), | |
| gr.Dataframe(label="📍 Locations", interactive=False), | |
| ], | |
| title="Named Entity Recognition (NER)", | |
| description="NER results grouped by entity type for better readability and usability.", | |
| theme="dark" | |
| ) | |
| if __name__ == "__main__": | |
| interface.launch() | |