|
import gradio as gr |
|
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline |
|
|
|
|
|
model_name = "guidobenb/DarkBERT-finetuned-ner" |
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
model = AutoModelForTokenClassification.from_pretrained(model_name) |
|
|
|
|
|
ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, grouped_entities=True) |
|
|
|
|
|
def extract_entities(text): |
|
entities = ner_pipeline(text) |
|
output = "" |
|
for entity in entities: |
|
word = entity.get("word", "") |
|
label = entity.get("entity_group", entity.get("entity", "")) |
|
score = entity.get("score", 0.0) |
|
output += f"{word} ({label}, {score:.2f})\n" |
|
return output if output else "No entities found." |
|
|
|
|
|
demo = gr.Interface( |
|
fn=extract_entities, |
|
inputs=gr.Textbox(lines=5, placeholder="Type your sentence here...", label="Input Text"), |
|
outputs=gr.Textbox(label="Named Entities"), |
|
title="🧠 DarkBERT NER", |
|
description="Named Entity Recognition using `guidobenb/DarkBERT-finetuned-ner`. Try something like: `The hacker from Raid Forums used malware in 2022.`" |
|
) |
|
|
|
if __name__ == "__main__": |
|
demo.launch() |
|
|