Update app.py

app.py

@@ -31,10 +31,10 @@ class BertClassifier(nn.Module):
 # Load the tokenizer and model
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 bert_model = BertModel.from_pretrained('bert-base-uncased')
-model = BertClassifier(bert_model, num_classes=7)
+model = BertClassifier(bert_model, num_classes=10)
 
 # Load the model weights from the .pkl file
-model.load_state_dict(torch.load('bert_classifier_mltc.pkl', map_location=torch.device('cpu')))
+model.load_state_dict(torch.load('mimic4_icd_classifier.pkl', map_location=torch.device('cpu')))
 model.eval()
 
 # Define prediction function
@@ -50,38 +50,33 @@ def predict(text):
     confidence_scores = outputs[0].tolist()
     
     # Convert to a dictionary mapping labels to confidence scores
-    labels = ['toxic', 'severe_toxic', 'obscene', 'threat', 'insult', 'identity_hate', 'clean']
+    labels = ['Cardiovascular', 'Respiratory', 'Neurological', 'Infectious', 'Endocrine', 'Musculoskeletal', 'Gastrointestinal', 'Renal', 'Psychiatric', 'Other']
     prediction = {label: score for label, score in zip(labels, confidence_scores)}
     
     return prediction
 
 # Add example texts
 examples = [
-    ["This is a positive and friendly comment."],
-    ["I'm not happy with your service."],
-    ["You are a great person!"],
-    ["How can one defame someone who thinks the Fort Hood Shooting is justified by God's law?"],
-    ["A Bisexual, like a homosexual or a heterosexual, is not defined by sexual activity. (Much like a 15 year old boy who is attracted to a girl sexually but has never had sex is still straight). A person who is actually sexually attracted/aroused by the same sex as well as the opposite sex is bisexual."],
-    ["I'LL SAY NIGGER IF I WANT TO BRO"],
-    ["Do you need a verifiable source for that statement?  Check your own user page you fat, miserable balding fuck."],
-    ["Fuck you bitch! I will go to your house and kill you with my gun as well! Prepare to die!"],
-    ["you're hot, i will rape you ;)"],
-    ["WOULDN'T BE THE FIRST TIME BITCH. FUCK YOU I'LL FIND OUT WHERE YOU LIVE, SODOMIZE YOUR WIFE AND THEN BURN YOUR HOUSE DOWN. FUCK YOU YOU FUCKING QUEER."],
-    ["What a motherfucking piece of crap those fuckheads for blocking us!"],
-    ["Get a life you animal fucker. Shut up you fucking nigger. Fuck off and shit your pants full of all the shit you can fill them with. 144.131.176.126"],
-    ["HOPE YOUR HEAD GETS CUT OFF AND SOMEONE WIPS THERE ASS WITH IT AND THEN STABS YOU IN YOUR HEART"],
-    ["you people are pretty overzealous with this whole free thing. get a fucking life, you fucking niggers !!!23 16!!!"],
-    ["Stupid peace of shit stop deleting my stuff asshole go die and fall in a hole go to hell!"],
-    ["Bye! Don't look, come or think of comming back! Tosser."]
+    ["Patient admitted with chest pain, shortness of breath, and abnormal ECG findings."],
+    ["Elderly patient presented with symptoms of confusion, fever, and elevated white blood cell count."],
+    ["Patient complains of persistent cough, wheezing, and history of asthma."],
+    ["Admitted with severe abdominal pain, nausea, and vomiting. Suspected appendicitis."],
+    ["Patient has a history of diabetes mellitus and presented with high blood glucose levels and dehydration."],
+    ["Patient admitted following a fall, showing signs of fracture in the left femur."],
+    ["Patient experiencing severe headaches, dizziness, and a history of epilepsy."],
+    ["Acute kidney injury suspected due to elevated creatinine and reduced urine output."],
+    ["Patient diagnosed with major depressive disorder, experiencing prolonged sadness and loss of interest in activities."],
+    ["Presented with a bacterial skin infection requiring intravenous antibiotics."]
 ]
+
 # Create Gradio interface
 iface = gr.Interface(
     fn=predict,
-    inputs=gr.Textbox(lines=10, placeholder="Enter a comment here..."),
-    outputs=gr.Label(num_top_classes=7),
+    inputs=gr.Textbox(lines=10, placeholder="Enter clinical notes here..."),
+    outputs=gr.Label(num_top_classes=10),
     examples=examples,
-    title="Toxic Comment Classification",
-    description="Classify comments into toxic and non-toxic categories using BERT and GNN model.",
+    title="MIMIC-IV ICD Code Prediction",
+    description="Predict ICD code categories based on clinical text using a BERT-based model. The model outputs confidence scores for ten common ICD code categories.",
 )
 
 iface.launch()