update

app.py

@@ -1,31 +1,33 @@
-# import gradio as gr
-
-# def greet(name):
-#     return "Hello " + name + "!!"
-
-# demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-# demo.launch()
-
 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"
+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"
 )
-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']} ({round(entity['score']*100, 2)}%)")
-    
+        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()
+gr.Interface(
+    fn=predict,
+    inputs="text",
+    outputs="text",
+    title="Named Entity Recognition"
+).launch()