raul-padua commited on
Commit
0d8634d
1 Parent(s): 9f96e9d

Update app.py

Browse files

added merge_tokens function.

Files changed (1) hide show
  1. app.py +17 -1
app.py CHANGED
@@ -3,9 +3,25 @@ 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
  gr.close_all()
11
  demo = gr.Interface(fn=ner,
 
3
 
4
  get_completion = pipeline("ner", model="dslim/bert-base-NER")
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 them
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,