elnasharomar2 commited on
Commit
a138702
verified
1 Parent(s): 644ff91

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +38 -0
app.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+ import gradio as gr
3
+
4
+ get_completion = pipeline("ner", model="elnasharomar2/PUNCERT_multi_stages")
5
+
6
+ def merge_tokens(tokens):
7
+ merged_tokens = []
8
+ for token in tokens:
9
+ if merged_tokens and token['entity'].startswith('I-') and merged_tokens[-1]['entity'].endswith(token['entity'][2:]):
10
+ # If current token continues the entity of the last one, merge the two tokens
11
+ last_token = merged_tokens[-1]
12
+ last_token['word'] += token['word'].replace('##', '')
13
+ last_token['end'] = token['end']
14
+ last_token['score'] = (last_token['score'] + token['score']) / 2
15
+ else:
16
+ # Otherwise, add the token to the list
17
+ merged_tokens.append(token)
18
+
19
+ return merged_tokens
20
+
21
+ def ner(input):
22
+ output = get_completion(input)
23
+ merged_tokens = merge_tokens(output)
24
+ return {"text": input, "entities": merged_tokens}
25
+
26
+ gr.close_all()
27
+ demo = gr.Interface(fn=ner,
28
+ inputs=[gr.Textbox(label="Text to find entities", lines=2)],
29
+ outputs=[gr.HighlightedText(label="Text with entities")],
30
+ title="NER with dslim/bert-base-NER",
31
+ description="Find entities using the `BERT-base` model under the hood!",
32
+ allow_flagging="never",
33
+ examples=["My name is Raul and I live in Niter贸i, Rio de Janeiro, Brazil",
34
+ "Lionel Messi is the greatest footballer of the new century",
35
+ "Toronto is hockey capital of the world",
36
+ "S&P 500 has gained 400 points in last trailing 7 days",
37
+ "Paris is one of most visited cities in the world every year."])
38
+ demo.launch()