from flair.data import Sentence
from flair.models import SequenceTagger
import gradio as gr

title = "Spanish Named Entity Recognition (NER) - Demo App"

# Load tagger
tagger1 = SequenceTagger.load("aymurai/anonymizer-beto-cased-flair")
tagger2 = SequenceTagger.load("flair/ner-spanish-large")

# Examples
mytext1="1. DECLARAR EXTINGUIDA LA ACCIÓN PENAL en este caso por cumplimiento de la suspensión del proceso a prueba, y SOBRESEER a EZEQUIEL CAMILO MARCONNI, DNI 11.222.333, en orden a los delitos de lesiones leves agravadas, amenazas simples y agravadas por el uso de armas, en contra de Chuchita Perez de 50 años."
mytext2="El sombrío Prudhon, imbuído, sin duda, en las ideas de los Santos Padres de la Iglesia que predicaban el desden por los bienes terrenales, decía que la pobreza es una ley de nuestra naturaleza, ley bajo la cual hemos sido constituídos, de donde se deduce que el pauperismo es mal que no tiene remedio ni cura."

# Function with NER Models
def ner_builder(model_, text_):
  if model_=="NER (BETO->Judicial)":
    # make example sentence
    sentence = Sentence(text_)
    # predict NER tags
    tagger1.predict(sentence)
    tags_tokens = []
    for entity in sentence.get_spans('ner'):
      tags_tokens.append(entity)
  elif model_=="NER (CoNLL-03 Spanish)":
    # make example sentence
    sentence = Sentence(text_)
    # predict NER tags
    tagger2.predict(sentence)
    tags_tokens = []    
    # iterate over entities and print
    for entity in sentence.get_spans('ner'):
      tags_tokens.append(entity)
  # return predicted NER spans
  return f"""Con {model_} se encontraron las etiquetas: {tags_tokens} """

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.HTML("<h1>"+title+"</h1>"+"""
            <p>Aplicación de Procesamiento de Lenguaje Natural (PNL), para el reconocimiento de entidades (nombres, organizaciones, ubicaciones, entre otras), dentro de un texto.</p>
      """)
    gr.Interface(
        ner_builder,
        [
            #gr.Slider(1, 10, value=2, label="Número", info="Choose between 1 and 10"),
            gr.Dropdown(
                ["NER (BETO->Judicial)", "NER (CoNLL-03 Spanish)"], label="Modelo", info="Elige un modelo (NER pipeline)"
            ),
            gr.Textbox(placeholder="Escribe tu texto aquí...", label="Texto", info="Pega un texto o da clic en un ejemplo"),
        ],
        "text",
        #theme="soft",
        #title=title,
        examples=[
            ["NER (BETO->Judicial)", str(mytext1)],
            ["NER (CoNLL-03 Spanish)", str(mytext2)],
        ]
    )
    gr.HTML("""
      <h2>¿Por qué usar Flair AI?<h2>
      <a href="https://towardsdatascience.com/benchmark-ner-algorithm-d4ab01b2d4c3">Por sus resultados superiores o similares a otros modelos pero sin límite de 512 tokens</a>
    """)


demo.launch()