Spaces:
Sleeping
Sleeping
File size: 1,657 Bytes
1f8bdf0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | 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()
|