from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline
import gradio as gr

def ner_tagging(text):
  model_name = "hackathon-pln-es/jurisbert-finetuning-ner"
  tokenizer = AutoTokenizer.from_pretrained(model_name,       add_prefix_space=True)
  
  model = AutoModelForTokenClassification.from_pretrained(model_name)
  nlp = pipeline("ner", model=model, tokenizer=tokenizer)
  ner_results = nlp(text.lower())
  
  output = []
  
  text_2 = text.split(" ")
  
  for i in range(len(text_2)):
      ent = ner_results[i]["entity"]
      if ent != "O":
          output.extend([(text_2[i], ent), (" ", None)])
      else:
          output.extend([(text_2[i], None), (" ", None)])

  return output

def get_entities(example):
    model_name = "hackathon-pln-es/jurisbert-finetuning-ner"
    tokenizer = AutoTokenizer.from_pretrained(model_name,       add_prefix_space=True)
  
    model = AutoModelForTokenClassification.from_pretrained(model_name)
    token_classifier = pipeline("token-classification", aggregation_strategy="simple", model=model, tokenizer=tokenizer)
    results = token_classifier(example.lower())
    
    output = []

    i=0
    prev_item = None
    next_item = None
    while i < (len(results)):
        item = results[i]
        p=i-1
        n=i+1
        
        if p > 0:
            prev_item = results[p]
        
        
        if n<(len(results)):
            next_item = results[n]
        
        
        if (i==0):
            if item["start"]>0:
                output.extend([(example[0:item["start"]], None)])
        output.extend([(example[item["start"]:item["end"]], item["entity_group"])])
        if (next_item!=None):
            ##verificar el tramo entre actual y siguiente
            if(item["end"]!=next_item["start"]):
                output.extend([(example[item["end"]:next_item["start"]], None)])
        i=i+1

    if item["end"] < len(example):
        output.extend([(example[item["end"]:len(example)], None)])
    
    return output

def greet(name):
    return "Hello " + name + "!!"

iface = gr.Interface(fn=get_entities, inputs="text", outputs=['highlight'], examples=[['Esta Primera Sala de la Suprema Corte de Justicia de la Nación es competente para conocer de la presente Solicitud de Ejercicio de la Facultad de Atracción, en términos de lo dispuesto en los artículos 107, fracción VIII, penúltimo párrafo, de la Constitución Política de los Estados Unidos Mexicanos; 80 Bis de la Ley de Amparo; así como el precepto 21, fracción II, de la Ley Orgánica del Poder Judicial de la Federación, en relación con lo dispuesto en los puntos segundo, fracción IX, y tercero del Acuerdo General 5/2013, del Pleno de este Alto Tribunal, relativo a la determinación de los asuntos que el Tribunal Pleno conservará para su resolución y el envío de los de su competencia originaria a las Salas y a los tribunales colegiados de circuito.'], 
["Lo anterior es así, toda vez que, si bien es cierto, el artículo 1° de la Constitución Federal tiene como finalidad brindar la protección más amplia al gobernado, y que ello se logra garantizando el derecho a un recurso efectivo en términos del artículo 25 de la Convención Americana sobre Derechos Humanos, ello no significa que en cualquier caso el órgano jurisdiccional deba resolver el fondo del asunto sin verificar los requisitos de procedencia previstos en las leyes nacionales, ya que las formalidades procesales son la vía que hace posible arribar a una adecuada resolución."]], title="Test of jurisbert-finetuning-ner ",)
iface.launch()