File size: 931 Bytes
d394f6d
311f263
7e000d1
5667b16
7e000d1
f4ce35e
f77bff8
 
 
 
 
 
 
 
311f263
d394f6d
5667b16
7aa02fc
 
 
f77bff8
7aa02fc
 
f77bff8
 
7aa02fc
d394f6d
f77bff8
 
 
 
 
 
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
import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline

# Load model and tokenizer from Hugging Face Hub
tokenizer = AutoTokenizer.from_pretrained("Mhammad2023/bert-finetuned-ner-torch")
model = AutoModelForTokenClassification.from_pretrained("Mhammad2023/bert-finetuned-ner-torch")

# Use aggregation_strategy="simple" to group B/I tokens
classifier = pipeline(
    "token-classification",
    model=model,
    tokenizer=tokenizer,
    aggregation_strategy="simple"
)

def predict(text):
    results = classifier(text)
    if not results:
        return "No entities found"

    output = []
    for entity in results:
        output.append(f"{entity['word']}: {entity['entity_group']} ({round(entity['score']*100, 2)}%)")

    return "\n".join(output)

gr.Interface(
    fn=predict,
    inputs="text",
    outputs="text",
    title="Named Entity Recognition"
).launch()