yunkexiang commited on
Commit
50dcc1c
1 Parent(s): bf2ccea

create the app.py

Browse files
Files changed (2) hide show
  1. app.py +45 -0
  2. requirements.txt +4 -0
app.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+ import gradio as gr
3
+
4
+ get_completion = pipeline("ner", model="dslim/bert-base-NER")
5
+
6
+ def ner(input):
7
+ output = get_completion(input)
8
+ return {"text": input, "entities": output}
9
+
10
+
11
+ def merge_tokens(tokens):
12
+ merged_tokens = []
13
+ for token in tokens:
14
+ if merged_tokens and token['entity'].startswith('I-') and merged_tokens[-1]['entity'].endswith(token['entity'][2:]):
15
+ # If current token continues the entity of the last one, merge them
16
+ last_token = merged_tokens[-1]
17
+ last_token['word'] += token['word'].replace('##', '')
18
+ last_token['end'] = token['end']
19
+ last_token['score'] = (last_token['score'] + token['score']) / 2
20
+ else:
21
+ # Otherwise, add the token to the list
22
+ merged_tokens.append(token)
23
+
24
+ return merged_tokens
25
+
26
+ def ner(input):
27
+ output = get_completion(input)#, parameters=None, ENDPOINT_URL=API_URL
28
+ merged_tokens = merge_tokens(output)
29
+ return {"text": input, "entities": merged_tokens}
30
+
31
+
32
+
33
+ if __name__ == "__main__":
34
+ # Create the Gradio interface
35
+ demo = gr.Interface(fn=ner,
36
+ inputs=[gr.Textbox(label="Text to find entities", lines=2)],
37
+ outputs=[gr.HighlightedText(label="Text with entities")],
38
+ title="Sanofi Digital M&S Demo: Name Entity Recognition",
39
+ description="Find entities using the `dslim/bert-base-NER` model under the hood!",
40
+ allow_flagging="never",
41
+ examples=["My name is Yunke, I lead a data science team at Sanofi. I live in Cambridge, Massachusattes", "My address is: 450 water street, Cambridge, MA"])
42
+ # Launch the interface
43
+ demo.launch() #share=True
44
+
45
+
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ transformers
2
+ gradio
3
+ torch
4
+ # dotenv