|
|
import gradio as gr
|
|
|
import torch
|
|
|
from transformers import AutoTokenizer, BertForSequenceClassification, AutoModel
|
|
|
from torch import nn
|
|
|
import re
|
|
|
|
|
|
|
|
|
def paragraph_leveling(text):
|
|
|
model_name = "./trained_model/fine_tunning_encoder_v2"
|
|
|
model = AutoModel.from_pretrained(model_name)
|
|
|
model.to('cuda')
|
|
|
tokenizer = AutoTokenizer.from_pretrained('zzxslp/RadBERT-RoBERTa-4m')
|
|
|
|
|
|
class MLP(nn.Module):
|
|
|
def __init__(self, target_size=3, input_size=768):
|
|
|
super(MLP, self).__init__()
|
|
|
self.num_classes = target_size
|
|
|
self.input_size = input_size
|
|
|
self.fc1 = nn.Linear(input_size, target_size)
|
|
|
|
|
|
def forward(self, x):
|
|
|
out = self.fc1(x)
|
|
|
return out
|
|
|
|
|
|
classifier = MLP(target_size=3, input_size=768)
|
|
|
classifier.load_state_dict(torch.load('./trained_model/fine_tunning_classifier'))
|
|
|
classifier.cuda()
|
|
|
classifier.eval()
|
|
|
|
|
|
output_list = []
|
|
|
text_list = text.split(".")
|
|
|
result = []
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
output_list.append(("\n", None))
|
|
|
|
|
|
for idx_sentence in text_list:
|
|
|
train_encoding = tokenizer(
|
|
|
idx_sentence,
|
|
|
return_tensors='pt',
|
|
|
padding='max_length',
|
|
|
truncation=True,
|
|
|
max_length=120)
|
|
|
output = model(**train_encoding.to('cuda'))
|
|
|
output = classifier(output[1])
|
|
|
output = output[0]
|
|
|
|
|
|
if output.argmax(-1) == 0:
|
|
|
output_list.append((idx_sentence, 'abnormal'))
|
|
|
result.append(0)
|
|
|
elif output.argmax(-1) == 1:
|
|
|
output_list.append((idx_sentence, 'normal'))
|
|
|
result.append(1)
|
|
|
else:
|
|
|
output_list.append((idx_sentence, 'not much information'))
|
|
|
result.append(2)
|
|
|
|
|
|
output_list.append(('\n', None))
|
|
|
if 0 in result:
|
|
|
output_list.append(('FINAL LABEL: ', None))
|
|
|
output_list.append(('ABNORMAL', 'abnormal'))
|
|
|
|
|
|
else:
|
|
|
output_list.append(('FINAL LABEL: ', None))
|
|
|
output_list.append(('NORMAL', 'normal'))
|
|
|
|
|
|
return output_list
|
|
|
|
|
|
|
|
|
demo = gr.Interface(
|
|
|
paragraph_leveling,
|
|
|
[
|
|
|
gr.Textbox(
|
|
|
label="Medical Report",
|
|
|
info="You can put any types of medical report",
|
|
|
lines=20,
|
|
|
value=" ",
|
|
|
),
|
|
|
],
|
|
|
gr.HighlightedText(
|
|
|
label="labeling",
|
|
|
show_legend = True,
|
|
|
show_label = True,
|
|
|
color_map={"abnormal": "violet", "normal": "lightgreen", "not much information": "lightgray"}),
|
|
|
theme=gr.themes.Base()
|
|
|
)
|
|
|
if __name__ == "__main__":
|
|
|
demo.launch(share=True)
|
|
|
|
|
|
|
