File size: 1,936 Bytes
0a4290b 47ae213 daad8af 47ae213 19e2659 |
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 59 60 61 |
import gradio as gr
from transformers import pipeline
# Load the NER models
# Load the NER models
models = {
"dslim/bert-base-NER": pipeline(
"ner", model="dslim/bert-base-NER", grouped_entities=True
),
"dslim/bert-base-NER-uncased": pipeline(
"ner", model="dslim/bert-base-NER-uncased", grouped_entities=True
),
"dslim/bert-large-NER": pipeline(
"ner", model="dslim/bert-large-NER", grouped_entities=True
),
"dslim/distilbert-NER": pipeline(
"ner", model="dslim/distilbert-NER", grouped_entities=True
),
}
def process(text, model_name):
ner = models[model_name]
ner_results = ner(text)
highlighted_text = []
last_idx = 0
for entity in ner_results:
start = entity["start"]
end = entity["end"]
label = entity["entity_group"]
# Add non-entity text
if start > last_idx:
highlighted_text.append((text[last_idx:start], None))
# Add entity text
highlighted_text.append((text[start:end], label))
last_idx = end
# Add any remaining text after the last entity
if last_idx < len(text):
highlighted_text.append((text[last_idx:], None))
return highlighted_text
with gr.Blocks() as demo:
gr.Markdown("# Named Entity Recognition with BERT Models")
with gr.Row():
model_selector = gr.Dropdown(
choices=list(models.keys()),
value=list(models.keys())[0],
label="Select Model",
)
text_input = gr.Textbox(
label="Enter Text",
lines=5,
value="Hugging Face Inc. is a company based in New York City. Its headquarters are in DUMBO, therefore very close to the Manhattan Bridge.",
)
output = gr.HighlightedText(label="Named Entities")
analyze_button = gr.Button("Analyze")
analyze_button.click(process, inputs=[text_input, model_selector], outputs=output)
demo.launch()
|