File size: 3,182 Bytes
0a4290b 47ae213 daad8af 43ed631 47ae213 43ed631 47ae213 43ed631 47ae213 43ed631 47ae213 43ed631 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
import gradio as gr
from transformers import pipeline
# Model names (keeping it programmatic)
model_names = [
"dslim/bert-base-NER",
"dslim/bert-base-NER-uncased",
"dslim/bert-large-NER",
"dslim/distilbert-NER",
]
example_sent = (
"Nim Chimpsky was a chimpanzee at Columbia University named after Noam Chomsky."
)
# Programmatically build the model info dict
model_info = {
model_name: {
"link": f"https://huggingface.co/{model_name}",
"usage": f"""from transformers import pipeline
ner = pipeline("ner", model="{model_name}", grouped_entities=True)
result = ner("{example_sent}")
print(result)""",
}
for model_name in model_names
}
# Load models into a dictionary programmatically for the analyze function
models = {
model_name: pipeline("ner", model=model_name, grouped_entities=True)
for model_name in model_names
}
# Function to display model info (link and usage code)
def display_model_info(model_name):
info = model_info[model_name]
usage_code = info["usage"]
link_button = f'[Open model page for {model_name} ]({info["link"]})'
return usage_code, link_button
# Function to run NER on input text
def analyze_text(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 (NER) with BERT Models")
# Dropdown for model selection
model_selector = gr.Dropdown(
choices=list(model_info.keys()),
value=list(model_info.keys())[0],
label="Select Model",
)
# Textbox for input text
text_input = gr.Textbox(
label="Enter Text",
lines=5,
value=example_sent,
)
analyze_button = gr.Button("Run NER Model")
output = gr.HighlightedText(label="NER Result", combine_adjacent=True)
# Outputs: usage code, model page link, and analyze button
code_output = gr.Code(label="Use this model", visible=True)
link_output = gr.Markdown(
f"[Open model page for {model_selector} ]({model_selector})"
)
# Button for analyzing the input text
analyze_button.click(
analyze_text, inputs=[text_input, model_selector], outputs=output
)
# Trigger the code output and model link when model is changed
model_selector.change(
display_model_info, inputs=[model_selector], outputs=[code_output, link_output]
)
# Call the display_model_info function on load to set initial values
demo.load(
fn=display_model_info,
inputs=[model_selector],
outputs=[code_output, link_output],
)
demo.launch()
|