Update app.py

app.py

@@ -1,37 +1,67 @@
 from transformers import pipeline
 import gradio as gr
 
-get_completion = pipeline("ner", model="gyr66/Ernie-3.0-base-chinese-finetuned-ner")
+# Cargar modelos
+model1 = "gyr66/RoBERTa-ext-large-crf-chinese-finetuned-ner-v2"
+model2 = "gyr66/Ernie-3.0-large-chinese-finetuned-ner"
+model3 = "gyr66/Ernie-3.0-base-chinese-finetuned-ner"
 
+get_completion1 = pipeline("ner", model1)
+get_completion2 = pipeline("ner", model2)
+get_completion3 = pipeline("ner", model3)
+
+# Función para fusionar tokens
 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
+            # Si el token continúa la entidad del anterior, fusiónalos
             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
+            # De lo contrario, agrega el token a la lista
             merged_tokens.append(token)
-
     return merged_tokens
 
+# Función de NER
 def ner(input):
-    output = get_completion(input)
-    merged_tokens = merge_tokens(output)
-    print(output)
-    print(merged_tokens)
-    return {"text": input, "entities": merged_tokens}
-
-gr.close_all()
-demo = 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 under the hood!",
-                    allow_flagging="never",
-                    examples=["My name is Andrew, I'm building DeeplearningAI and I live in California", "My name is Poli, I live in Vienna and work at HuggingFace"])
+    output1 = get_completion1(input)
+    output2 = get_completion2(input)
+    output3 = get_completion3(input)
+
+    merged_tokens1 = merge_tokens(output1)
+    merged_tokens2 = merge_tokens(output2)
+    merged_tokens3 = merge_tokens(output3)
+
+    # Formatear la salida para Gradio
+    entities1 = [{"entity": t['entity'], "start": t['start'], "end": t['end']} for t in merged_tokens1]
+    entities2 = [{"entity": t['entity'], "start": t['start'], "end": t['end']} for t in merged_tokens2]
+    entities3 = [{"entity": t['entity'], "start": t['start'], "end": t['end']} for t in merged_tokens3]
+
+    return (
+        {"text": input, "entities": entities1},
+        {"text": input, "entities": entities2},
+        {"text": input, "entities": entities3}
+    )
+
+# Crear interfaz Gradio
+demo = gr.Interface(
+    fn=ner,
+    inputs=gr.Textbox(label="Text to find entities", lines=2),
+    outputs=[
+        gr.HighlightedText(label=f"NER Output - Model 1"),
+        gr.HighlightedText(label=f"NER Output - Model 2"),
+        gr.HighlightedText(label=f"NER Output - Model 3")
+    ],
+    title="NER with Multiple Models",
+    description="Extract entities using three different models.",
+    allow_flagging="never",
+    examples=[
+        "My name is Andrew, I'm building DeeplearningAI and I live in California",
+        "My name is Poli, I live in Vienna and work at HuggingFace"
+    ]
+)
 
 demo.launch(inline=False)