app.py

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

# model large
model_name = "pierreguillou/ner-bert-large-cased-pt-lenerbr"
model_large = AutoModelForTokenClassification.from_pretrained(model_name)
tokenizer_large = AutoTokenizer.from_pretrained(model_name)

# model base
model_name = "pierreguillou/ner-bert-base-cased-pt-lenerbr"
model_base = AutoModelForTokenClassification.from_pretrained(model_name)
tokenizer_base = AutoTokenizer.from_pretrained(model_name)

# css
background_colors_entity_word = {
    'JURISPRUDENCIA': "#ccfbf1",
    'LEGISLACAO': "#ede9fe",
    'LOCAL': "#ffedd5", 
    'ORGANIZACAO': "#fae8ff",
    'PESSOA': "#e0f2fe", 
    'TEMPO': "#fefde0",  
}

background_colors_entity_tag = {
    'JURISPRUDENCIA': "#14b8a6", 
    'LEGISLACAO': "#8b5cf6", 
    'LOCAL': "#f97316", 
    'ORGANIZACAO': "#d946ef", 
    'PESSOA': "#0ea5e9", 
    'TEMPO': "#e9c70e", 
}

css = {
'entity_word': 'color:#000000;background: #xxxxxx; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 2.5; border-radius: 0.35em;',
'entity_tag': 'color:#fff;background: #xxxxxx; font-size: 0.8em; font-weight: bold; line-height: 2.5; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5em;'
}

list_EN = "<span style='"
list_EN += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['JURISPRUDENCIA'])};padding:0.5em;"
list_EN += "'>JURISPRUDENCIA</span>"
list_EN += "<span style='"
list_EN += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['LEGISLACAO'])};padding:0.5em;"
list_EN += "'>LEGISLACAO</span>"
list_EN += "<span style='"
list_EN += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['LOCAL'])};padding:0.5em;"
list_EN += "'>LOCAL</span>"
list_EN += "<span style='"
list_EN += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['ORGANIZACAO'])};padding:0.5em;"
list_EN += "'>ORGANIZACAO</span>"
list_EN += "<span style='"
list_EN += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['PESSOA'])};padding:0.5em;"
list_EN += "'>PESSOA</span>"
list_EN += "<span style='"
list_EN += f"{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['TEMPO'])};padding:0.5em;"
list_EN += "'>TEMPO</span>"

# list_EN += | <span style='{css['entity_tag'].replace('#xxxxxx',background_colors_entity_tag['LEGISLACAO'])}">LEGISLACAO</span>"

# infos
title = "LeNER-Br | NER App"
description = "(24/12/2021) Este aplicativo NER permite comparar 2 modelos de BERT ajustados (finetuned) na tarefa NER com o dataset jurídico LeNER-Br (NER BERT base com f1 de 0.893 e NER BERT large com f1 de 0.908). O download desses 2 modelos pode demorar alguns segundos."
article = f"<div style='text-align: center; font-size: 90%;'><p>Lista das Entidades Nomeadas (EN): {list_EN}</p><p>Links para os modelos BERT NER: <a style='color:blue;' href='https://huggingface.co/pierreguillou/ner-bert-base-cased-pt-lenerbr' target='_blank'>base</a> e <a style='color:blue;' href='https://huggingface.co/pierreguillou/ner-bert-large-cased-pt-lenerbr' target='_blank'>large</a></p><p>Blog post: <a style='color:blue;' href='https://medium.com/@pierre_guillou/nlp-modelos-e-web-app-para-reconhecimento-de-entidade-nomeada-ner-no-dom%C3%ADnio-jur%C3%ADdico-b658db55edfb' target='_blank'>NLP | Modelos e Web App para Reconhecimento de Entidade Nomeada (NER) no domínio jurídico brasileiro</a> (29/12/2021)</p><p>LeNER-Br NER App by <a style='color:blue;' href='https://www.linkedin.com/in/pierreguillou/' target='_blank'>Pierre GUILLOU</a> | <a style='color:blue;' href='https://github.com/piegu/language-models#language-models' target='_blank'>Github Repo</a></p></div>"
allow_screenshot = False
allow_flagging = False
examples = [
["Ao Instituto Médico Legal da jurisdição do acidente ou da residência cumpre fornecer, no prazo de 90 dias, laudo à vítima (art. 5, § 5, Lei n. 6.194/74  de 19 de dezembro de 1974), função técnica que pode ser suprida por prova pericial realizada por ordem do juízo da causa, ou por prova técnica realizada no âmbito administrativo que se mostre coerente com os demais elementos de prova constante dos autos."],
["Acrescento que não há de se falar em violação do artigo 114, § 3º, da Constituição Federal, posto que referido dispositivo revela-se impertinente, tratando da possibilidade de ajuizamento de dissídio coletivo pelo Ministério Público do Trabalho nos casos de greve em atividade essencial."],
["Todavia, entendo que extrair da aludida norma o sentido expresso na redação acima implica desconstruir o significado do texto constitucional, o que é absolutamente vedado ao intérprete. Nesse sentido, cito Dimitri Dimoulis: ‘(...) ao intérprete não é dado escolher significados que não estejam abarcados pela moldura da norma. Interpretar não pode significar violentar a norma.’ (Positivismo Jurídico. São Paulo: Método, 2006, p. 220).59. Dessa forma, deve-se tomar o sentido etimológico como limite da atividade interpretativa, a qual não pode superado, a ponto de destruir a própria norma a ser interpretada. Ou, como diz Konrad Hesse, ‘o texto da norma é o limite insuperável da atividade interpretativa.’ (Elementos de Direito Constitucional da República Federal da Alemanha, Porto Alegre: Sergio Antonio Fabris, 2003, p. 71)."],
]

def ner(input_text):

  num = 0
  for tokenizer,model in zip([tokenizer_large,tokenizer_base],[model_large,model_base]):

    # tokenization
    inputs = tokenizer(input_text, max_length=512, truncation=True, return_tensors="pt")
    tokens = inputs.tokens()

    # get predictions
    outputs = model(**inputs).logits
    predictions = torch.argmax(outputs, dim=2)
    preds = [model_base.config.id2label[prediction] for prediction in predictions[0].numpy()]

    # variables
    groups_pred = dict()
    group_indices = list()
    group_label = ''
    pred_prec = ''
    group_start =  ''
    count = 0

    # group the NEs
    for i,en in enumerate(preds):

      if en == 'O': 

        if len(group_indices) > 0:
          groups_pred[count] = {'indices':group_indices,'en':group_label}
          group_indices = list()
          group_label = ''
          count += 1

      if en.startswith('B'):

        if len(group_indices) > 0:
          groups_pred[count] = {'indices':group_indices,'en':group_label}
          group_indices = list()
          group_label = '' 
          count += 1 

        group_indices.append(i)
        group_label = en.replace('B-','')
        pred_prec = en

      elif en.startswith('I'):
        
        if len(group_indices) > 0:
          if en.replace('I-','') == group_label:
            group_indices.append(i)
          else:
            groups_pred[count] = {'indices':group_indices,'en':group_label}
            group_indices = [i]
            group_label = en.replace('I-','')
            count += 1 
        else:
          group_indices = [i]
          group_label = en.replace('I-','')

      if i == len(preds) - 1 and len(group_indices) > 0:
        groups_pred[count] = {'indices':group_indices,'en':group_label}
        group_indices = list()
        group_label = ''
        count += 1

    # there is at least one NE
    len_groups_pred = len(groups_pred)
    inputs = inputs['input_ids'][0].numpy()#[1:-1]
    if len_groups_pred > 0:
      for pred_num in range(len_groups_pred):
        en = groups_pred[pred_num]['en']
        indices = groups_pred[pred_num]['indices']
        if pred_num == 0:
          if indices[0] > 0:
            output = tokenizer.decode(inputs[:indices[0]]) + f'<span style="{css["entity_word"].replace("#xxxxxx",background_colors_entity_word[en])}">' + tokenizer.decode(inputs[indices[0]:indices[-1]+1]) + f'<span style="{css["entity_tag"].replace("#xxxxxx",background_colors_entity_tag[en])}">' + en + '</span></span> '
          else:
            output = f'<span style="{css["entity_word"].replace("#xxxxxx",background_colors_entity_word[en])}">' + tokenizer.decode(inputs[indices[0]:indices[-1]+1]) + f'<span style="{css["entity_tag"].replace("#xxxxxx",background_colors_entity_tag[en])}">' + en + '</span></span> '
        else:
          output += tokenizer.decode(inputs[indices_prev[-1]+1:indices[0]]) + f'<span style="{css["entity_word"].replace("#xxxxxx",background_colors_entity_word[en])}">' + tokenizer.decode(inputs[indices[0]:indices[-1]+1]) + f'<span style="{css["entity_tag"].replace("#xxxxxx",background_colors_entity_tag[en])}">' + en + '</span></span> '
        indices_prev = indices
      output += tokenizer.decode(inputs[indices_prev[-1]+1:])
    else:
      output = input_text

    # output
    output = output.replace('[CLS]','').replace(' [SEP]','').replace('##','')
    output = "<div style='max-width:100%; max-height:360px; overflow:auto'>" + output + "</div>"

    if num == 0: 
      output_large = output
      num += 1
    else: output_base = output

  return output_large, output_base

# interface gradio
iface = gr.Interface(
    title=title,
    description=description,
    article=article,
    allow_screenshot=allow_screenshot,
    allow_flagging=allow_flagging,
    fn=ner,
    inputs=gr.inputs.Textbox(placeholder="Digite uma frase aqui ou clique em um exemplo:", lines=5),
    outputs=[gr.outputs.HTML(label="NER com BERT large"),gr.outputs.HTML(label="NER com BERT base")],
    examples=examples
)

iface.launch()