from transformers import pipeline
import gradio as gr

# 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:]):
            # 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:
            # De lo contrario, agrega el token a la lista
            merged_tokens.append(token)
    return merged_tokens

# Función de NER
def ner(input):
    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)