Mhammad Ibrahim commited on
Commit
7aa02fc
·
1 Parent(s): 1176a10

Fix pipeline type to token-classification for NER model

Browse files
Files changed (1) hide show
  1. app.py +11 -4
app.py CHANGED
@@ -10,10 +10,17 @@ import gradio as gr
10
  from transformers import pipeline
11
 
12
  # Load model from Hugging Face Hub
13
- classifier = pipeline("text-classification", model="Mhammad2023/bert-finetuned-ner", from_tf=True)
14
 
15
  def predict(text):
16
- result = classifier(text)[0]
17
- return f"{result['label']} ({round(result['score']*100, 2)}%)"
 
 
 
 
 
 
 
18
 
19
- gr.Interface(fn=predict, inputs="text", outputs="text", title="Text Classification").launch()
 
10
  from transformers import pipeline
11
 
12
  # Load model from Hugging Face Hub
13
+ classifier = pipeline("token-classification", model="Mhammad2023/bert-finetuned-ner", from_tf=True)
14
 
15
  def predict(text):
16
+ results = classifier(text)
17
+ if not results:
18
+ return "No entities found"
19
+
20
+ output = []
21
+ for entity in results:
22
+ output.append(f"{entity['word']}: {entity['entity']} ({round(entity['score']*100, 2)}%)")
23
+
24
+ return "\n".join(output)
25
 
26
+ gr.Interface(fn=predict, inputs="text", outputs="text", title="Named Entity Recognition").launch()