Spaces:
Running
Running
| import gradio as gr | |
| import transformers | |
| from transformers import AutoModelForTokenClassification, AutoTokenizer | |
| import torch | |
| # model large | |
| model_name = "pucpr/clinicalnerpt-chemical" | |
| model_large = AutoModelForTokenClassification.from_pretrained(model_name) | |
| tokenizer_large = AutoTokenizer.from_pretrained(model_name) | |
| # model base | |
| model_name = "pucpr/clinicalnerpt-chemical" | |
| model_base = AutoModelForTokenClassification.from_pretrained(model_name) | |
| tokenizer_base = AutoTokenizer.from_pretrained(model_name) | |
| # css | |
| background_colors_entity_word = { | |
| 'ChemicalDrugs': "#fae8ff", | |
| } | |
| background_colors_entity_tag = { | |
| 'ChemicalDrugs': "#d946ef", | |
| } | |
| 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['ChemicalDrugs'])};padding:0.5em;" | |
| list_EN += "'>ChemicalDrugs</span>" | |
| # infos | |
| title = "BioBERTpt - Chemical entities" | |
| description = "BioBERTpt - Chemical entities" | |
| allow_screenshot = False | |
| allow_flagging = False | |
| examples = [ | |
| ["Dispneia venoso central em subclavia D duplolumen recebendo solução salina e glicosada em BI."], | |
| ["Paciente com Sepse pulmonar em D8 tazocin (paciente não recebeu por 2 dias Atb)."], | |
| ["FOI REALIZADO CURSO DE ATB COM LEVOFLOXACINA POR 7 DIAS."], | |
| ] | |
| 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="NER1"),gr.outputs.HTML(label="NER2")], | |
| examples=examples | |
| ) | |
| iface.launch() |