Mhammad2023/Token-Classification

1import gradio as gr2import torch3from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline4 5# Load model and tokenizer from Hugging Face Hub6tokenizer = AutoTokenizer.from_pretrained("Mhammad2023/bert-finetuned-ner-torch")7model = AutoModelForTokenClassification.from_pretrained("Mhammad2023/bert-finetuned-ner-torch")8 9# Use aggregation_strategy="simple" to group B/I tokens10classifier = pipeline(11    "token-classification",12    model=model,13    tokenizer=tokenizer,14    aggregation_strategy="simple"15)16 17def predict(text):18    results = classifier(text)19    if not results:20        return "No entities found"21 22    output = []23    for entity in results:24        output.append(f"{entity['word']}: {entity['entity_group']} ({round(entity['score']*100, 2)}%)")25 26    return "\n".join(output)27 28gr.Interface(29    fn=predict,30    inputs="text",31    outputs="text",32    title="Named Entity Recognition"33).launch()34