Update app.py

app.py

@@ -1,3 +1,60 @@
 import gradio as gr
+import pandas as pd
+from transformers import pipeline
 
-gr.load("models/facebook/bart-large-mnli").launch()
+# Initialize the zero-shot classifier
+classifier = pipeline('zero-shot-classification', model='distilbert-base-uncased')
+
+# Define the candidate labels for injury classification
+candidate_labels = ["ACL Tear", "Meniscus Tear", "Achilles Tear", "Fracture", "Hamstring", "Foot", "Shoulder", "Hip", "Calf", "Hand", "Wrist"]
+
+def process_injury_notes(file):
+    # Read CSV file
+    df = pd.read_csv(file.name)
+    
+    # Limit to a sample (for performance in demonstration)
+    new_df = df.head(100).copy()
+    
+    # Classify each note and save results
+    classifications = classifier(new_df['Notes'].tolist(), candidate_labels)
+    new_df['Classifications'] = classifications
+    
+    # Extract top classification and score
+    new_df['Top Classification'] = new_df['Classifications'].apply(lambda x: x['labels'][0] if isinstance(x, dict) else None)
+    new_df['Top Score'] = new_df['Classifications'].apply(lambda x: x['scores'][0] if isinstance(x, dict) else None)
+    
+    # Initialize the 'Specific Injury' column with default value
+    new_df['Specific Injury'] = None
+    
+    # Function to extract specific injury classification based on keywords
+    def extract_specific_injury(note, injury):
+        note = note.lower()
+        if "left" in note:
+            return f"left {injury.lower()} injury"
+        elif "right" in note:
+            return f"right {injury.lower()} injury"
+        else:
+            return f"{injury.lower()} injury"
+    
+    # Apply specific injury classification based on keywords
+    for injury in candidate_labels:
+        new_df.loc[new_df['Top Classification'].str.contains(injury, case=False, na=False), 'Specific Injury'] = \
+            new_df['Notes'].apply(lambda x: extract_specific_injury(x, injury) if injury.lower() in x.lower() else None)
+    
+    # Sort by 'Top Score' in descending order
+    new_df_sorted = new_df.sort_values(by='Top Score', ascending=False)
+    
+    # Return sorted DataFrame
+    return new_df_sorted[['Notes', 'Top Classification', 'Top Score', 'Specific Injury']]
+
+# Set up Gradio interface
+iface = gr.Interface(
+    fn=process_injury_notes,
+    inputs=gr.File(label="Upload CSV File (must have 'Notes' column)"),
+    outputs="dataframe",
+    title="Injury Classification App",
+    description="Upload a CSV file with injury notes to classify injuries based on given categories. Displays top classification and specificity."
+)
+
+# Launch the Gradio app
+iface.launch()