Update app.py

app.py

@@ -1,35 +1,60 @@
-import requests
+import gradio as gr
 import pandas as pd
+from transformers import pipeline
 
-# Read your data
-data = pd.read_csv('/path/to/torn.csv')
-notes = data['Notes'].tolist()  # Ensure 'Notes' is a column in your CSV
+# Initialize the zero-shot classifier
+classifier = pipeline('zero-shot-classification', model='distilbert-base-uncased')
 
-# Define your candidate labels for classification
+# Define the possible categories (more granular categories)
 candidate_labels = ["ACL Tear", "Meniscus Tear", "Achilles Tear", "Fracture", "Hamstring", "Foot", "Shoulder", "Hip", "Calf", "Hand", "Wrist"]
 
-# Hugging Face API details
-API_URL = "https://api-inference.huggingface.co/models/facebook/bart-large-mnli"
-headers = {"Authorization": f"Bearer YOUR_HUGGING_FACE_API_KEY"}  # Replace with your actual API key
+def classify_injuries(file):
+    # Load the uploaded CSV file
+    df = pd.read_csv(file.name)
+    
+    # Limit to a sample (e.g., first 100 rows) if necessary for performance
+    new_df = df.head(100).copy()
 
-def classify_with_api(note):
-    # Prepare payload
-    payload = {
-        "inputs": note,
-        "parameters": {"candidate_labels": candidate_labels},
-    }
-    # Make request to API
-    response = requests.post(API_URL, headers=headers, json=payload)
-    if response.status_code == 200:
-        return response.json()
-    else:
-        print("Error:", response.status_code, response.text)
-        return None
+    # Apply zero-shot classification to each note in the 'Notes' column
+    classifications = classifier(new_df['Notes'].tolist(), candidate_labels)
+    
+    # Add the classification results to the DataFrame
+    new_df['Classifications'] = classifications
+    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
+    
+    # Define a function to determine the specific injury 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 a subset of columns for clarity
+    return new_df_sorted[['Notes', 'Top Classification', 'Top Score', 'Specific Injury']]
 
-# Classify each note and store results
-data['Classifications'] = data['Notes'].apply(lambda x: classify_with_api(x))
-data['Top Classification'] = data['Classifications'].apply(lambda x: x['labels'][0] if x else None)
-data['Top Score'] = data['Classifications'].apply(lambda x: x['scores'][0] if x else None)
+# Set up the Gradio interface
+iface = gr.Interface(
+    fn=classify_injuries,
+    inputs=gr.File(label="Upload CSV File (must have a 'Notes' column)"),
+    outputs="dataframe",
+    title="Injury Classification App",
+    description="Upload a CSV file with injury notes. The app classifies each note based on specified injury types and provides specific classifications where possible."
+)
 
-# Display results
-print(data[['Notes', 'Top Classification', 'Top Score']])
+# Launch the Gradio app
+iface.launch()