app.py

#!/usr/bin/env python
# coding: utf-8



from datasets import load_dataset
import pandas as pd
import re
import json

# Load dataset
ds = load_dataset("AGBonnet/augmented-clinical-notes")
df = ds["train"].to_pandas()  # Convert to pandas DataFrame for easier manipulation




from snorkel.labeling import labeling_function

# Define pulmonary keywords (ICD-10 inspired)
PULMONARY_TERMS = {
    "asthma", "copd", "pneumonia", "pulmonary fibrosis", "bronchitis",
    "tuberculosis", "lung cancer", "emphysema", "pneumothorax",
    "cystic fibrosis", "ARDS", "pulmonary embolism", "chronic bronchitis"
}

# Define regex patterns for variations (e.g., "COPD exacerbation")
PULMONARY_REGEX = re.compile(
    r'('
    r'\b(asthma|asthmatic|bronchial asthma)\b|'
    r'\b(COPD|chronic obstructive pulmonary disease|chronic obstructive lung disease)\b|'
    r'\b(pneumonia|CAP|HAP|VAP|community-acquired pneumonia|hospital-acquired pneumonia|ventilator-associated pneumonia)\b|'
    r'\b(pulmonary embolism|PE|pulmonary thromboembolism)\b|'
    r'\b(tuberculosis|TB|mycobacterium tuberculosis|pulmonary TB)\b|'
    r'\b(lung cancer|lung carcinoma|bronchogenic carcinoma|NSCLC|SCLC|non-small cell lung cancer|small cell lung cancer)\b|'
    r'\b(bronchitis|acute bronchitis|chronic bronchitis)\b|'
    r'\b(pulmonary fibrosis|idiopathic pulmonary fibrosis|IPF)\b|'
    r'\b(cystic fibrosis|CF)\b|'
    r'\b(pneumothorax|collapsed lung)\b|'
    r'\b(ARDS|acute respiratory distress syndrome)\b|'
    r'\b(emphysema|pulmonary emphysema)\b|'
    r'\b(interstitial lung disease|ILD)\b|'
    r'\b(pulmonary hypertension|PH)\b|'
    r'\b(pleural effusion|hydrothorax)\b|'
    r'\b(silicosis|occupational lung disease)\b|'
    r'\b(COVID-19|SARS-CoV-2|coronavirus)\b'
    r')', 
    flags=re.IGNORECASE  # Match case-insensitively
)

# Labeling Function 1: Check structured JSON summary for diagnoses
@labeling_function()
def lf_summary_diagnosis(row):
    try:
        summary = json.loads(row["summary"])
        diagnoses = summary.get("diagnosis", [])
        # Ensure diagnoses is a list
        if not isinstance(diagnoses, list):
            diagnoses = [diagnoses]
        for d in diagnoses:
            if any(term in d.lower() for term in PULMONARY_TERMS):
                return 1
    except Exception as e:
        pass
    return 0  # non-pulmonary

# Labeling Function 2: Keyword search in notes
@labeling_function()
def lf_note_keywords(row):
    note_text = ((row.get("note") or "") + " " + (row.get("full_note") or "")).lower()
    if any(term in note_text for term in PULMONARY_TERMS):
        return 1
    return 0

# Improved negation-aware regex (checks for negation near pulmonary terms)
NEGATION_REGEX = re.compile(
    r'\b(no history of|ruled out|denies|negative for|no|without)\b\s*'  # Negation trigger
    r'(?:\w+\s+){0,5}'  # Allow up to 5 words between negation and pulmonary term
    r'(' + PULMONARY_REGEX.pattern + r')',  # Pulmonary terms from your regex
    flags=re.IGNORECASE
)

@labeling_function()
def lf_note_regex(row):
    note_text = row["note"] + " " + row["full_note"]
    # Check for pulmonary terms
    pulmonary_match = PULMONARY_REGEX.search(note_text)
    if not pulmonary_match:
        return 0  # No pulmonary term found
    
    # Check if the pulmonary term is negated
    if NEGATION_REGEX.search(note_text):
        return 0  # Pulmonary term is negated
    return 1  # Pulmonary term is affirmed




from snorkel.labeling import PandasLFApplier, LFAnalysis
from snorkel.labeling.model import LabelModel

# Combine labeling functions
lfs = [lf_summary_diagnosis, lf_note_keywords, lf_note_regex]

# Apply labeling functions to the DataFrame using Snorkel's PandasLFApplier
applier = PandasLFApplier(lfs)
L_train = applier.apply(df)

# Analyze LF performance (coverage, conflicts)
analysis = LFAnalysis(L_train, lfs)
analysis.lf_summary()  # This prints a summary of your labeling functions

# Train a LabelModel to combine LF outputs
label_model = LabelModel(cardinality=2, verbose=True)
label_model.fit(L_train, n_epochs=500, log_freq=100)

# Predict probabilistic labels; here, tie_break_policy="abstain" will mark ties as abstentions (-1)
df["label_pulmonary"] = label_model.predict(L_train, tie_break_policy="abstain")

# Filter for pulmonary cases (label == 1)
pulmonary_df = df[df["label_pulmonary"] == 1].reset_index(drop=True)

# Optionally, inspect the results
print("Pulmonary cases:", len(pulmonary_df))




# Display a random sample of rows
print(df[['note', 'summary', 'label_pulmonary']].sample(10))




# Define regex patterns for target conditions
CONDITION_REGEX = {
    "Asthma": re.compile(
        r'\b(asthma|asthmatic|bronchial asthma)\b', 
        flags=re.IGNORECASE
    ),
    "COPD": re.compile(
        r'\b(COPD|chronic obstructive pulmonary disease|chronic obstructive lung disease|emphysema|chronic bronchitis)\b', 
        flags=re.IGNORECASE
    ),
    "Pneumonia": re.compile(
        r'\b(pneumonia|CAP|HAP|VAP|community-acquired pneumonia|hospital-acquired pneumonia|ventilator-associated pneumonia)\b', 
        flags=re.IGNORECASE
    ),
    "Lung Cancer": re.compile(
        r'\b(lung cancer|lung carcinoma|bronchogenic carcinoma|NSCLC|SCLC|non-small cell lung cancer|small cell lung cancer)\b', 
        flags=re.IGNORECASE
    ),
    "Tuberculosis": re.compile(
        r'\b(tuberculosis|TB|mycobacterium tuberculosis|pulmonary TB)\b', 
        flags=re.IGNORECASE
    ),
    "Pleural Effusion": re.compile(
        r'\b(pleural effusion|hydrothorax)\b', 
        flags=re.IGNORECASE
    )
}

# Negation regex (improved to check proximity to condition terms)
NEGATION_REGEX = re.compile(
    r'\b(no history of|ruled out|denies|negative for|no|without)\b\s*'  # Negation trigger
    r'(?:\w+\s+){0,5}'  # Allow up to 5 words between negation and condition
    r'(' + '|'.join([pattern.pattern for pattern in CONDITION_REGEX.values()]) + r')',  # Combined condition terms
    flags=re.IGNORECASE
)

def get_condition_labels(row):
    note_text = row["note"] + " " + row["full_note"]
    labels = []
    
    # Check for negations first
    negation_match = NEGATION_REGEX.search(note_text)
    
    for condition, pattern in CONDITION_REGEX.items():
        # Skip if the condition term is negated
        if negation_match and pattern.search(negation_match.group(0)):
            continue
        # Check if condition is mentioned
        if pattern.search(note_text):
            labels.append(condition)
    
    return labels

# Apply labeling to pulmonary cases
pulmonary_df["conditions"] = pulmonary_df.apply(get_condition_labels, axis=1)

# Classify remaining cases as "Other Pulmonary"
pulmonary_df["conditions"] = pulmonary_df["conditions"].apply(
    lambda x: x if x else ["Other Pulmonary"]
)




from collections import defaultdict

label_counts = defaultdict(int)
for labels in pulmonary_df["conditions"]:
    for label in labels:
        label_counts[label] += 1

print("Label distribution:")
for k, v in label_counts.items():
    print(f"{k}: {v}")




import pandas as pd

# Label distribution data
label_counts = {
    "Asthma": 509,
    "Pneumonia": 1294,
    "Other Pulmonary": 1907,
    "Tuberculosis": 851,
    "Pleural Effusion": 743,
    "COPD": 697,
    "Lung Cancer": 415
}

# Convert to DataFrame for easier plotting
df_counts = pd.DataFrame(list(label_counts.items()), columns=["Condition", "Count"])

import matplotlib.pyplot as plt
import seaborn as sns

# Set style
sns.set(style="whitegrid")

# Create bar plot
plt.figure(figsize=(10, 6))
sns.barplot(x="Condition", y="Count", data=df_counts, palette="viridis")

# Add labels and title
plt.title("Distribution of Pulmonary Conditions", fontsize=16)
plt.xlabel("Condition", fontsize=14)
plt.ylabel("Count", fontsize=14)
plt.xticks(rotation=45, ha="right")  # Rotate x-axis labels for readability

# Show plot
plt.tight_layout()
plt.show()




import nltk
from nltk.corpus import stopwords
from wordcloud import WordCloud
import matplotlib.pyplot as plt

# Download stop words from nltk (do this once)
nltk.download('stopwords')

# Get the list of stop words
stop_words = set(stopwords.words('english'))

# Combine all notes into one large string
text = " ".join(pulmonary_df['note'].dropna())  # Combine all notes into a single string

# Tokenize the text and remove stop words
filtered_words = [word for word in text.split() if word.lower() not in stop_words]

# Join the filtered words back into a single string
cleaned_text = " ".join(filtered_words)

# Create a WordCloud object
wordcloud = WordCloud(width=800, height=400, background_color='white').generate(cleaned_text)

# Plot the WordCloud image
plt.figure(figsize=(10, 5))
plt.imshow(wordcloud, interpolation='bilinear')
plt.axis('off')
plt.show()




from sklearn.feature_extraction.text import TfidfVectorizer

vectorizer = TfidfVectorizer(max_features=5000, stop_words="english",ngram_range=(1, 2))  
X = vectorizer.fit_transform(pulmonary_df['note'])  # Note column




#Transform Object data type to string
pulmonary_df["conditions"] = pulmonary_df["conditions"].apply(lambda x: x[0])

from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report

X_train, X_test, y_train, y_test = train_test_split(X, pulmonary_df['conditions'], test_size=0.1, random_state=42)

#Logistic Regression Classification Report
model = LogisticRegression(max_iter=1000)
model.fit(X_train, y_train)
y_pred_before_smote = model.predict(X_test)
print(classification_report(y_test, y_pred_before_smote))
report_before_smote = classification_report(y_test, y_pred_before_smote, output_dict=True)




from imblearn.over_sampling import SMOTE

# Logistic Regression after SMOTE
smote = SMOTE(random_state=42)
X_train_res, y_train_res = smote.fit_resample(X_train, y_train)
model.fit(X_train_res, y_train_res)
y_pred_after_smote = model.predict(X_test)
print(classification_report(y_test, y_pred_after_smote))
report_after_smote = classification_report(y_test, y_pred_after_smote, output_dict=True)




# Convert y_resampled to a pandas Series to get the distribution
y_resampled = pd.Series(y_train_res)

# Get class distribution after SMOTE
label_counts_smote = y_resampled.value_counts()

print("Label distribution after SMOTE:")
print(label_counts_smote)




from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report
from sklearn.model_selection import train_test_split
from prettytable import PrettyTable
import pandas as pd

# Split the data before applying SMOTE
X_train, X_test, y_train, y_test = train_test_split(X, pulmonary_df['conditions'], test_size=0.3, random_state=42)

# Train Random Forest without SMOTE
rf_model_before_smote = RandomForestClassifier(n_estimators=100, random_state=42)
rf_model_before_smote.fit(X_train, y_train)

# Make Predictions
y_pred_before_smote = rf_model_before_smote.predict(X_test)

# Generate classification report
report_before_smote = classification_report(y_test, y_pred_before_smote, output_dict=True)

# Convert to DataFrame
df_report_before_smote = pd.DataFrame(report_before_smote).transpose()

# Use PrettyTable for a more structured look
table_before_smote = PrettyTable()
table_before_smote.field_names = ["Class", "Precision", "Recall", "F1-Score", "Support"]

for index, row in df_report_before_smote.iterrows():
    table_before_smote.add_row([index, round(row['precision'], 2), round(row['recall'], 2), round(row['f1-score'], 2), int(row['support'])])

print(table_before_smote)




# %pip install prettytable
from prettytable import PrettyTable

#Random Forest Now Model with SMOTE

from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report
from imblearn.over_sampling import SMOTE
from sklearn.model_selection import train_test_split

# Apply SMOTE
smote = SMOTE(random_state=42)
X_resampled, y_resampled = smote.fit_resample(X, pulmonary_df['conditions'])

# Split the data
X_train, X_test, y_train, y_test = train_test_split(X_resampled, y_resampled, test_size=0.3, random_state=42)

# Train Random Forest
rf_model = RandomForestClassifier(n_estimators=100, random_state=42)
rf_model.fit(X_train, y_train)

# Make Predictions
y_pred = rf_model.predict(X_test)

# Generate classification report
report = classification_report(y_test, y_pred, output_dict=True)

# Convert to DataFrame
df_report = pd.DataFrame(report).transpose()

# Use PrettyTable for a more structured look
table = PrettyTable()
table.field_names = ["Class", "Precision", "Recall", "F1-Score", "Support"]

for index, row in df_report.iterrows():
    table.add_row([index, round(row['precision'], 2), round(row['recall'], 2), round(row['f1-score'], 2), int(row['support'])])

print(table)




import matplotlib.pyplot as plt
import pandas as pd

# Assuming report_after_smote and df_report are already generated as DataFrames

# Convert the necessary columns to DataFrame for easy plotting
df_lr = pd.DataFrame(report_after_smote).transpose()  # Logistic Regression after SMOTE
df_rf = pd.DataFrame(report).transpose()  # Random Forest

# Extract relevant columns (precision, recall, and f1-score)
metrics = ['precision', 'recall', 'f1-score']

# Set up the figure for the plot
plt.figure(figsize=(10, 6))

# Plot for each metric
for metric in metrics:
    plt.plot(df_lr.index, df_lr[metric], label=f'LR After SMOTE - {metric.capitalize()}', marker='o')
    plt.plot(df_rf.index, df_rf[metric], label=f'RF - {metric.capitalize()}', marker='x')

# Add labels and title
plt.title('Comparison of Logistic Regression and Random Forest Performance')
plt.xlabel('Class Labels')
plt.ylabel('Score')
plt.legend(title="Model and Metric")

# Show the plot
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()




from sklearn.utils import resample

#Separate majority and minority classes
grouped = pulmonary_df.groupby("Conditions")
max_size = grouped.size().max()

#Oversample each class to the same count as the majority class
oversampled_df = grouped.apply(
    lambda x: resample(x, replace=True, n_samples=max_size, random_state=42)
    
).reset_index(drop=True)

print(oversampled_df["conditions"].value_counts())





from datasets import Dataset
from transformers import AutoTokenizer

dataset = Dataset.from_pandas(oversampled_df)

# Tokenize using ClinicalBERT
model_checkpoint = "emilyalsentzer/Bio_ClinicalBERT"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

def tokenize_function(example):
    return tokenizer(example["note"], truncation=True, padding="max_length", max_length=512)

tokenized_dataset = dataset.map(tokenize_function, batched=True)

# Label encoding 
from sklearn.preprocessing import LabelEncoder
label_encoder = LabelEncoder()
tokenized_dataset = tokenized_dataset.add_column("label", label_encoder.fit_transform(tokenized_dataset["conditions"]))

# Final formatting
tokenized_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])




import torch

# print("Number of GPU: ", torch.cuda.device_count())
# print("GPU Name: ", torch.cuda.get_device_name())




# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# print('Using device:', device)




from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from transformers import Trainer, AutoModelForSequenceClassification, AutoTokenizer


split_dataset = tokenized_dataset.train_test_split(test_size=0.1, seed=42)
train_dataset = split_dataset["train"]
eval_dataset = split_dataset["test"]


# Load saved model + tokenizer
model = AutoModelForSequenceClassification.from_pretrained("./trained_clinicalbert")
tokenizer = AutoTokenizer.from_pretrained("./trained_clinicalbert")

# Load model with the correct number of classes
# num_classes = len(label_encoder.classes_)
# model = AutoModelForSequenceClassification.from_pretrained(
#     model_checkpoint,
#     num_labels=num_classes
# )

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./results",
    eval_strategy="epoch",
    save_strategy="epoch",
    logging_strategy="epoch",
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=4,
    learning_rate=2e-5,
    weight_decay=0.01,
    load_best_model_at_end=True,
    metric_for_best_model="f1",
)

def compute_metrics(p):
    preds = p.predictions.argmax(axis=1)
    labels = p.label_ids
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average="weighted")
    acc = accuracy_score(labels, preds)
    return {"accuracy": acc, "f1": f1, "precision": precision, "recall": recall}


trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics
)

# trainer.train()

# trainer.save_model("./trained_clinicalbert")
# tokenizer.save_pretrained("./trained_clinicalbert")

trainer.evaluate()




predictions_output = trainer.predict(eval_dataset)
y_pred = predictions_output.predictions.argmax(axis=1)
y_prob = predictions_output.predictions  # softmax scores (for ROC/AUC)
y_true = predictions_output.label_ids

from sklearn.metrics import confusion_matrix, classification_report, roc_auc_score, roc_curve, precision_recall_curve, average_precision_score
import matplotlib.pyplot as plt
import seaborn as sns

# Confusion Matrix
cm = confusion_matrix(y_true, y_pred)
plt.figure(figsize=(8,6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
            xticklabels=label_encoder.classes_,
            yticklabels=label_encoder.classes_)
plt.xlabel("Predicted")
plt.ylabel("Actual")
plt.title("Confusion Matrix")
plt.show()

# Classification Report (includes F1, Precision, Recall per class)
print(classification_report(y_true, y_pred, target_names=label_encoder.classes_))




from sklearn.preprocessing import label_binarize
from scipy.special import softmax


# Apply softmax to get probabilities
y_probs = softmax(y_prob, axis=1)

# Binarize the true labels for multi-class ROC (One-vs-Rest)
y_true_bin = label_binarize(y_true, classes=list(range(len(label_encoder.classes_))))

# Plot ROC curve per class
plt.figure(figsize=(10, 6))

for i, class_name in enumerate(label_encoder.classes_):
    fpr, tpr, _ = roc_curve(y_true_bin[:, i], y_probs[:, i])
    auc_score = roc_auc_score(y_true_bin[:, i], y_probs[:, i])
    plt.plot(fpr, tpr, label=f"{class_name} (AUC = {auc_score:.2f})")

# Plot random classifier line
plt.plot([0, 1], [0, 1], 'k--', label="Random (AUC = 0.50)")

# Plot formatting
plt.xlabel("False Positive Rate")
plt.ylabel("True Positive Rate")
plt.title("ROC Curves (One-vs-Rest for 7 Classes)")
plt.legend(loc="lower right")
plt.grid(True)
plt.tight_layout()
plt.show()




#PR Curves

#Plot PR Curve per class
plt.figure(figsize=(8, 6))

for i, class_name in enumerate(label_encoder.classes_):
    precision, recall, _ = precision_recall_curve(y_true_bin[:, i], y_probs[:, i])
    pr_auc = average_precision_score(y_true_bin[:, i], y_probs[:, i])
    plt.plot(recall, precision, label=f"{class_name} (AP = {pr_auc:.2f})")

# Plot formatting
plt.xlabel("Recall")
plt.ylabel("Precision")
plt.title("Precision-Recall Curves (One-vs-Rest for 7 Classes)")
plt.legend(loc="lower left")
plt.grid(True)
plt.tight_layout()
plt.show()




import pandas as pd
from scipy.special import softmax

# Sample clinical notes
demo_notes = [
    "Patient presents with high fever, chills, shortness of breath, and crackles heard on auscultation. Chest X-ray shows consolidation in the right lower lobe.",
    "Patient complains of chest tightness and wheezing that worsens at night and after physical activity. Symptoms relieved by use of albuterol inhaler.",
    "The patient is a 68-year-old male with a 40-pack-year smoking history who presents with worsening shortness of breath over the past 6 months. He reports a chronic productive cough that is worse in the mornings, occasional wheezing, and fatigue with mild exertion. On physical examination, breath sounds are diminished bilaterally with prolonged expiratory phase. Pulmonary function tests show reduced FEV1/FVC ratio consistent with obstructive lung disease. There are no signs of active infection. He denies fever or chills. Chest X-ray reveals hyperinflated lungs and flattened diaphragms.",
    "Patient has persistent cough, night sweats, weight loss, and hemoptysis. Sputum test positive for acid-fast bacilli.",
    "Patient presents with shortness of breath and pleuritic chest pain. Physical exam shows decreased breath sounds and dullness to percussion on the left side. Ultrasound confirms fluid accumulation."
]

# Predict function
def batch_predict(notes, model, tokenizer, label_encoder):
    predictions = []
    for note in notes:
        inputs = tokenizer(note, return_tensors="pt", truncation=True, padding="max_length", max_length=512)
        inputs = {key: val.to(model.device) for key, val in inputs.items()}
        outputs = model(**inputs)
        probs = softmax(outputs.logits.detach().cpu().numpy(), axis=1)
        pred_idx = probs.argmax(axis=1)[0]
        pred_class = label_encoder.inverse_transform([pred_idx])[0]
        confidence = probs[0][pred_idx]
        predictions.append((pred_class, round(float(confidence), 4)))
    return predictions

# Create DataFrame
demo_df = pd.DataFrame({"Clinical Note": demo_notes})
demo_df[["Predicted Label", "Confidence"]] = batch_predict(demo_notes, model, tokenizer, label_encoder)

# View table
demo_df




import gradio as gr

# Extract class names dynamically from the DataFrame
classes = sorted(oversampled_df["conditions"].unique().tolist())

# Load model and tokenizer
model_path = "trained_clinicalbert"
model = AutoModelForSequenceClassification.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)

model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Prediction function
def predict_clinical_note(note):
    inputs = tokenizer(note, return_tensors="pt", truncation=True, padding="max_length", max_length=512)
    inputs = {key: val.to(device) for key, val in inputs.items()}
    with torch.no_grad():
        outputs = model(**inputs)
    probs = softmax(outputs.logits.cpu().numpy(), axis=1)
    pred_idx = probs.argmax(axis=1)[0]
    pred_class = classes[pred_idx]
    confidence = float(probs[0][pred_idx])
    return f"{pred_class} (Confidence: {confidence:.2f})"

# Gradio interface
iface = gr.Interface(
    fn=predict_clinical_note,
    inputs=gr.Textbox(lines=6, placeholder="Paste clinical note here..."),
    outputs="text",
    title="Pulmonary Disease Classifier",
    description="Enter a clinical note to predict pulmonary condition (e.g., COPD, Pneumonia, Tuberculosis...)"
)

if __name__ == "__main__":
    iface.launch()