File size: 1,265 Bytes
83c875d d7bd362 83c875d |
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 |
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()
|