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()