sd99 commited on
Commit
f9ddd28
1 Parent(s): 9a37f4c

Added Highlight to the output

Browse files
Files changed (1) hide show
  1. app.py +9 -2
app.py CHANGED
@@ -1,6 +1,7 @@
1
  import gradio as gr
2
  from transformers import AutoTokenizer, AutoModelForTokenClassification
3
  from transformers import pipeline
 
4
 
5
  description = "Named Entity Recognition Using BERT"
6
  title = "NERBERT"
@@ -10,7 +11,13 @@ def findNER(example):
10
  tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
11
  model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")
12
  ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
13
- return ner_pipeline(example)
 
 
 
 
 
 
14
 
15
- interface = gr.Interface(fn=findNER, inputs='text', outputs='text', examples=examples, description=description, title=title)
16
  interface.launch()
 
1
  import gradio as gr
2
  from transformers import AutoTokenizer, AutoModelForTokenClassification
3
  from transformers import pipeline
4
+ import re
5
 
6
  description = "Named Entity Recognition Using BERT"
7
  title = "NERBERT"
 
11
  tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
12
  model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")
13
  ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
14
+ pipeline_output = ner_pipeline(example)
15
+ final_output = []
16
+ # all_words = re.split(r'[^a-zA-Z0-9\s]', example)
17
+ for _ in pipeline_output:
18
+ final_output.extend([(_['word'], _['entity'])])
19
+
20
+ return final_output
21
 
22
+ interface = gr.Interface(fn=findNER, inputs='text', outputs=['highlight'], examples=examples, description=description, title=title, interpretation='default')
23
  interface.launch()