Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import pipeline | |
| # Model names (keeping it programmatic) | |
| model_names = [ | |
| "cahya/NusaBert-ner-v1.3", | |
| "cahya/bert-base-indonesian-NER", | |
| ] | |
| example_sent = "Desember 1618, laksamana Inggris Thomas Dale mengusir Jan Pieterszoon Coen dari pelabuhan Jayakarta. Coen lari ke Maluku, saat itu pangkalan utama VOC. Kemudian Dale, dibantu Wijayakrama, mengepung benteng VOC." | |
| # 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 NusaBERT") | |
| # 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() | |