bert / app.py
qpoire's picture
Update app.py
d7bd362 verified
import gradio as gr
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
# Load the DarkBERT NER model
model_name = "guidobenb/DarkBERT-finetuned-ner"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(model_name)
# Create the NER pipeline
ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, grouped_entities=True)
# Function to extract entities from text
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."
# Gradio interface
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()