Update app.py

app.py

@@ -19,7 +19,6 @@ medical_specialties = [
 ]
 
 # Initialize the zero-shot classification pipeline
-# A better-performing, fine-tuned model could be used here.
 classifier = pipeline(
     "zero-shot-classification",
     model="facebook/bart-large-mnli",
@@ -28,7 +27,7 @@ classifier = pipeline(
 
 def classify_medical_text(text):
     """
-    Classifies a medical text into one of the predefined medical specialties.
+    Classifies a medical text into one of the predefined medical specialties and returns the top 3 predictions.
     """
     if not text:
         return {"Error": "Please provide some text to classify."}
@@ -36,12 +35,24 @@ def classify_medical_text(text):
     # Perform zero-shot classification
     result = classifier(text, medical_specialties)
 
-    # Format the output for better display
-    labels = result['labels']
-    scores = result['scores']
+    # Combine labels and scores
+    combined_results = zip(result['labels'], result['scores'])
     
-    # Return the results as a dictionary for Gradio to display
-    return {label: score for label, score in zip(labels, scores)}
+    # Sort the results by score in descending order and get the top 3
+    top_3_predictions = sorted(combined_results, key=lambda x: x[1], reverse=True)[:3]
+    
+    # Format the output as a dictionary for Gradio
+    top_3_dict = {label: score for label, score in top_3_predictions}
+    
+    return top_3_dict
+
+# Define example medical texts
+examples = [
+    "Patient presenting with chest pain, shortness of breath, and palpitations. ECG shows atrial fibrillation.",
+    "Aspiration of the knee joint was performed due to swelling and suspected septic arthritis.",
+    "Post-operative report for a patient who underwent a hysterectomy due to uterine fibroids.",
+    "Neurological examination revealed a positive Babinski sign and nystagmus, suggesting a central nervous system disorder."
+]
 
 # Create the Gradio interface
 iface = gr.Interface(
@@ -51,9 +62,10 @@ iface = gr.Interface(
         placeholder="Paste a medical document or text here...", 
         label="Medical Text"
     ),
-    outputs=gr.Label(num_top_classes=len(medical_specialties)),
+    outputs=gr.Label(num_top_classes=3),
     title="Medical Document Classifier",
-    description="This application uses a zero-shot classification model to predict the medical specialty of a given text."
+    description="This application uses a zero-shot classification model to predict the medical specialty of a given text. It will display the top 3 most likely specialties. Click on an example below to get started!",
+    examples=examples
 )
 
 # Launch the interface