Update app.py

app.py

@@ -3,7 +3,8 @@ from transformers import pipeline
 
 playground = gr.Blocks()
 image_pipe = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-get_completion = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
+summary_pipe = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
+ner_pipe = pipeline("ner", model="dslim/bert-base-NER")
 
 examples_summarization = [
     ["The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct."], 
@@ -24,11 +25,31 @@ def translate(input_text, source, target):
       return "", f"Error: Translation direction {source_readable} to {target} is not supported by Helsinki Translation Models"
 
 def summarize(input):
-    output = get_completion(input)
+    output = summary_pipe(input)
     summary_origin = output[0]['summary_text']
     summary_translated = translate(summary_origin,'en','fr')
     return summary_origin, summary_translated[0]
 
+def merge_tokens(tokens):
+    merged_tokens = []
+    for token in tokens:
+        if merged_tokens and token['entity'].startswith('I-') and merged_tokens[-1]['entity'].endswith(token['entity'][2:]):
+            # If current token continues the entity of the last one, merge them
+            last_token = merged_tokens[-1]
+            last_token['word'] += token['word'].replace('##', '')
+            last_token['end'] = token['end']
+            last_token['score'] = (last_token['score'] + token['score']) / 2
+        else:
+            # Otherwise, add the token to the list
+            merged_tokens.append(token)
+
+    return merged_tokens
+
+def ner(input):
+    output = ner_pipe(input)
+    merged_tokens = merge_tokens(output)
+    return {"text": input, "entities": merged_tokens}
+    
 def create_playground_header():
     gr.Markdown("""
                 # 🤗 Hugging Face Labs
@@ -111,6 +132,15 @@ with playground:
             gr.Markdown("""
                         > Name Entity Recognition
                         """)
+
+            iface = gr.Interface(fn=ner,
+                    inputs=[gr.Textbox(label="Text to find entities", lines=2)],
+                    outputs=[gr.HighlightedText(label="Text with entities")],
+                    title="NER with dslim/bert-base-NER",
+                    description="Find entities using the `dslim/bert-base-NER` model!",
+                    allow_flagging="never",
+                    examples=["My name is Ray, I'm learning through Hugging Face and DeepLearning.AI and I live in Caversham, Reading", "My name is Raymond, I work at A&O IT Group"])
+
             
     create_playground_footer()