| from transformers import BertTokenizer, BertForSequenceClassification | |
| import torch | |
| import gradio as gr | |
| model_path = 'elliealbertson/identifying_pregnancy_clinical_notes' | |
| tokenizer = BertTokenizer.from_pretrained(model_path) | |
| model = BertForSequenceClassification.from_pretrained(model_path) | |
| def predict(text): | |
| inputs = tokenizer(text, return_tensors="pt") | |
| num_tokens = inputs['input_ids'].size(1) | |
| if num_tokens <= 512: | |
| outputs = model(**inputs) | |
| predicted_class_id = torch.argmax(outputs.logits).item() | |
| probability_of_predicted_class = round(torch.nn.functional.softmax(outputs.logits, dim=1)[0, predicted_class_id].item(),2) | |
| if (predicted_class_id == 0) & (probability_of_predicted_class >= 0.5): | |
| predicted_class_label = "No, the note does not discuss the patient's pregnancy based on the model's assessment." | |
| elif (predicted_class_id == 1) & (probability_of_predicted_class >= 0.5): | |
| predicted_class_label = "Yes, the note discusses the patient's pregnancy based on the model's assessment." | |
| else: | |
| predicted_class_label = "The model was unable to determine with high certainty whether or not the note discusses the patient's pregnancy. Please provide additional text or a different note." | |
| return predicted_class_label | |
| else: | |
| error_message = 'Unfortunately the model is limited in how much text it can process at once. Please enter a shorter note.' | |
| return error_message | |
| with gr.Blocks() as interface: | |
| gr.HTML("<div style='text-align:center;'><div><h1>Identifying Pregnancy in Clinical Notes</h1></div>") | |
| gr.HTML("<p align='center'>Use this app to classify a clinical note as discussing or not discussing the patient's pregnancy.</p>") | |
| gr.HTML("<p align='center'>The model was fine-tuned on a small number of clinical notes agumented by limited synthetic data. As a result, it may give inaccurate results at times.</p>") | |
| with gr.Row(): | |
| with gr.Column(): | |
| inputs = gr.Textbox(label='Input a clinical note here:', lines=4) | |
| button = gr.Button('Assess Note') | |
| gr.Examples(['The patient is pregnant.', 'She has high cholesterol and hypertension.', 'Normal vaginal delivery.', 'Fetus development normal.', 'Presented with nausea.', 'Broken arm and leg.'], inputs) | |
| with gr.Column(): | |
| outputs=gr.Textbox(label="Does the note discuss the patient's pregnancy?", lines=4) | |
| button.click(fn=predict, inputs=inputs, outputs=outputs, queue=False) | |
| gr.HTML("<p align='center'>Model fine-tuned from <a href='https://huggingface.co./emilyalsentzer/Bio_ClinicalBERT' target='_blank'> Bio+ClinicalBERT </a>.</p>") | |
| gr.HTML("<p align='center'>Repository available on <a href='https://github.com/elliealbertson/identifying_pregnancy_clinical_notes' target='_blank'> GitHub </a>.</p>") | |
| interface.launch() | |