Upload 2 files

VisionBERT.py

@@ -0,0 +1,533 @@
+import os
+from typing import Dict
+from datasets import Dataset
+import torch
+from sklearn.metrics import confusion_matrix, precision_recall_fscore_support, accuracy_score
+from transformers import AutoTokenizer, TrainingArguments, Trainer, AutoModelForSequenceClassification, DataCollatorWithPadding
+import pandas as pd
+import numpy as np
+from sklearn.preprocessing import LabelEncoder, StandardScaler
+from sklearn.model_selection import train_test_split
+from sklearn.cluster import KMeans
+from torch.nn import CrossEntropyLoss
+import pickle
+
+os.environ['OMP_NUM_THREADS'] = '7'
+
+
+class WeightedTrainer(Trainer):
+    def compute_loss(self, model, inputs, return_outputs: bool = False, num_items_in_batch: int = None):
+        """
+        Custom loss computation with sample weights
+        """
+        labels = inputs.get("labels")
+        weights = inputs.get("weight")
+
+        # Forward pass
+        outputs = model(**{k: v for k, v in inputs.items()
+                           if k not in ["weight", "labels"]})
+        logits = outputs.get("logits")
+
+        # Add labels back to outputs
+        outputs["labels"] = labels
+
+        # Compute weighted loss
+        if weights is not None:
+            weights = weights.to(logits.device)
+            loss_fct = CrossEntropyLoss(reduction='none')
+            loss = loss_fct(logits.view(-1, self.model.config.num_labels),
+                            labels.view(-1))
+
+            # Adjust weights if num_items_in_batch is provided
+            if num_items_in_batch:
+                weights = weights[:num_items_in_batch]
+
+            loss = (loss * weights.view(-1)).mean()
+        else:
+            loss_fct = CrossEntropyLoss(label_smoothing=0.1)
+            loss = loss_fct(logits.view(-1, self.model.config.num_labels),
+                            labels.view(-1))
+
+        outputs["loss"] = loss
+        return (loss, outputs) if return_outputs else loss
+
+
+def create_feature_vector(df):
+    """Create numerical feature vector for clustering with sample size weighting, handling missing/unseen labels."""
+
+    # Initialize LabelEncoders
+    le_gender = LabelEncoder()
+    le_race = LabelEncoder()
+    le_risk = LabelEncoder()
+
+    # Fit and transform while handling missing values
+    gender_encoded = le_gender.fit(df['Gender'].unique()).transform(df['Gender'].fillna('Unknown'))
+    race_encoded = le_race.fit(df['RaceEthnicity'].unique()).transform(df['RaceEthnicity'].fillna('Unknown'))
+    risk_encoded = le_risk.fit(df['RiskFactor'].unique()).transform(df['RiskFactor'].fillna('Unknown'))
+
+    # Create age groups numerical representation with a default for missing values
+    age_map = {
+        '12-17 years': 0,
+        '18-39 years': 1,
+        '40-64 years': 2,
+        '65-79 years': 3,
+        '80 years and older': 4  # Include all possible labels, even if missing
+    }
+
+    # Use `.get()` with a default value for missing/unseen age groups
+    age_encoded = df['Age'].map(lambda x: age_map.get(x, -1))
+
+    # Combine features
+    features = np.column_stack([
+        age_encoded,
+        gender_encoded,
+        race_encoded,
+        risk_encoded,
+        df['Sample_Size'].values  # Add sample size as a feature
+    ])
+
+    # Scale features
+    scaler = StandardScaler()
+    features_scaled = scaler.fit_transform(features)
+
+    return features_scaled, scaler
+
+
+def weighted_kmeans(X, sample_weights, n_clusters, max_iter=300, random_state=42):
+    """Custom K-means implementation that considers sample weights"""
+    n_samples = X.shape[0]
+
+    # Initialize centroids randomly from the weighted distribution
+    rng = np.random.RandomState(random_state)
+    weighted_indices = rng.choice(n_samples, size=n_clusters, p=sample_weights / sample_weights.sum())
+    centroids = X[weighted_indices]
+
+    for _ in range(max_iter):
+        # Assign points to nearest centroid
+        distances = np.sqrt(((X[:, np.newaxis] - centroids) ** 2).sum(axis=2))
+        labels = np.argmin(distances, axis=1)
+
+        # Update centroids using weighted means
+        new_centroids = np.zeros_like(centroids)
+        for k in range(n_clusters):
+            mask = labels == k
+            if mask.any():
+                weights_k = sample_weights[mask]
+                new_centroids[k] = np.average(X[mask], axis=0, weights=weights_k)
+
+        # Check for convergence
+        if np.allclose(centroids, new_centroids):
+            break
+
+        centroids = new_centroids
+
+    return labels, centroids
+
+
+def prepare_data(file_path='data/Vision_Survey_Cleaned.csv'):
+    """Load and prepare the vision health dataset with sample-size-aware clustering."""
+    print("\nLoading and preparing data...")
+    df = pd.read_csv(file_path)
+
+    # Filter data
+    vision_cat = ['Best-corrected visual acuity']
+    df = df[df['Question'].isin(vision_cat)].copy()
+    df = df[df["RiskFactor"] != "All participants"]
+    df = df[df["RiskFactorResponse"] != "Total"]
+
+    # Reset index after filtering
+    df = df.reset_index(drop=True)
+
+    # Create feature vectors for clustering
+    features_scaled, scaler = create_feature_vector(df)
+
+    # Normalize sample sizes for weights
+    sample_weights = df['Sample_Size'].values
+    sample_weights = sample_weights / sample_weights.sum()
+
+    # Apply weighted clustering
+    n_clusters = min(5, len(df))
+    clusters, centroids = weighted_kmeans(
+        features_scaled,
+        sample_weights,
+        n_clusters=n_clusters
+    )
+
+    # Add clusters as a column
+    df['cluster'] = clusters
+
+    # Calculate cluster importance based on total sample size in each cluster
+    cluster_total_samples = df.groupby('cluster')['Sample_Size'].sum()
+    cluster_weights = cluster_total_samples / cluster_total_samples.sum()
+
+    # Enhanced feature engineering with clustering information
+    df['doc'] = df.apply(
+        lambda x: f"""
+        Patient Demographics:
+        - Age Category: {x['Age']}
+        - Gender: {x['Gender']}
+        - Race/Ethnicity: {x['RaceEthnicity']}
+
+        Risk Factors:
+        - {x['RiskFactor']}: {x['RiskFactorResponse']}
+
+        Additional Information:
+        - Sample Size: {x['Sample_Size']}
+        - Cluster Profile: {x['cluster']} (Weight: {cluster_weights.get(x['cluster'], 0):.3f})
+        """.strip(),
+        axis=1
+    )
+
+    # Encode labels
+    le = LabelEncoder()
+    df['labels'] = le.fit_transform(df['Response'].astype(str))
+
+    # Combine sample size weights with cluster importance
+    df['weight'] = df.apply(
+        lambda x: (x['Sample_Size'] / df['Sample_Size'].sum()) *
+                  cluster_weights.get(x['cluster'], 0),
+        axis=1
+    )
+
+    # Create train and test splits with stratification
+    train_df, test_df = train_test_split(
+        df,
+        test_size=0.2,
+        stratify=df['labels'],
+        random_state=42
+    )
+
+    # Convert to dict format
+    train_data = {
+        'doc': train_df['doc'].tolist(),
+        'labels': train_df['labels'].tolist(),
+        'weight': train_df['weight'].tolist()
+    }
+
+    test_data = {
+        'doc': test_df['doc'].tolist(),
+        'labels': test_df['labels'].tolist(),
+        'weight': test_df['weight'].tolist()
+    }
+
+    # Convert to datasets
+    train_dataset = Dataset.from_dict(train_data)
+    test_dataset = Dataset.from_dict(test_data)
+
+    dataset_dict = {
+        'train': train_dataset,
+        'test': test_dataset
+    }
+
+    # Print detailed dataset statistics
+    print("\nDataset Summary:")
+    print(f"Training samples: {len(train_dataset)}")
+    print(f"Test samples: {len(test_dataset)}")
+
+    print("\nCluster Distribution:")
+    for i in range(n_clusters):
+        cluster_mask = df['cluster'] == i
+        cluster_samples = df[cluster_mask]['Sample_Size'].sum()
+        print(f"\nCluster {i} (Total samples: {cluster_samples:,}, Weight: {cluster_weights.get(i, 0):.3f}):")
+        print("Most common characteristics:")
+        for col in ['Age', 'Gender', 'RaceEthnicity', 'RiskFactor']:
+            values = df[col][cluster_mask].value_counts().head(3)
+            samples = df[cluster_mask].groupby(col)['Sample_Size'].sum().sort_values(ascending=False).head(3)
+            print(f"{col}:")
+            for val, count in values.items():
+                sample_count = samples.get(val, 0)  # Use .get() for safety
+                print(f"  - {val}: {count} groups ({sample_count:,} individuals)")
+
+    print("\nLabel Distribution:")
+    for label, idx in zip(le.classes_, range(len(le.classes_))):
+        count = (df['labels'] == idx).sum()
+        total_size = df[df['labels'] == idx]['Sample_Size'].sum()
+        print(f"{label}: {count} groups, {total_size:,} individuals")
+
+    return dataset_dict, le
+
+
+
+def main():
+    # Setup
+    output_dir = "models/vision-classifier"
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Load the dataset
+    dataset_dict, label_encoder = prepare_data()
+
+    # Initialize the tokenizer
+    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+
+    # Define tokenization function within main to have access to tokenizer
+    def tokenize_function(examples):
+        """Tokenize the input texts and maintain the correct column names"""
+        tokenized = tokenizer(
+            examples["doc"],
+            truncation=True,
+            padding='max_length',
+            max_length=128,
+            return_tensors=None
+        )
+        # Keep the additional columns
+        tokenized['labels'] = examples['labels']
+        tokenized['weight'] = examples['weight']
+        return tokenized
+
+    # Tokenize the datasets
+    tokenized_datasets = {}
+    for split, dataset in dataset_dict.items():
+        tokenized_datasets[split] = dataset.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=['doc']
+        )
+
+    # Print sample to verify
+    print("\nSample tokenized data:", tokenized_datasets["train"][0])
+
+    # Initialize the model
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "distilbert-base-uncased",
+        num_labels=len(label_encoder.classes_),
+        id2label={i: label for i, label in enumerate(label_encoder.classes_)},
+        label2id={label: i for i, label in enumerate(label_encoder.classes_)},
+    )
+
+    # Data collator
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+
+    # Check device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"\nTraining on device: {device}")
+
+    # Move model to device
+    model.to(device)
+
+    # Set up training arguments
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        learning_rate=3e-5,
+        per_device_train_batch_size=8,
+        per_device_eval_batch_size=8,
+        num_train_epochs=7,
+        weight_decay=0.01,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+        remove_unused_columns=False,
+        push_to_hub=True,
+    )
+
+    # Create the Trainer
+    trainer = WeightedTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_datasets["train"],
+        eval_dataset=tokenized_datasets["test"],
+        data_collator=data_collator,
+    )
+
+    # Train the model
+    print("\nStarting training...")
+    trainer.train()
+
+    # Save the model
+    print("\nSaving model...")
+    trainer.save_model(output_dir=os.path.join(output_dir, "model"))
+
+    # Save the tokenizer
+    tokenizer.save_pretrained(os.path.join(output_dir, "tokenizer"))
+
+    # Save the label encoder
+    label_encoder_path = os.path.join(output_dir, "label_encoder.pkl")
+    with open(label_encoder_path, 'wb') as f:
+        pickle.dump(label_encoder, f)
+
+    return trainer, model, tokenizer, label_encoder
+
+
+def evaluate_model(model, eval_dataset, tokenizer, label_encoder, device) -> Dict:
+    """
+    Evaluate model performance using multiple metrics
+    """
+    model.eval()
+    all_predictions = []
+    all_labels = []
+
+    # Process each example in evaluation dataset
+    for item in eval_dataset:
+        # Tokenize input
+        inputs = tokenizer(
+            item['doc'],
+            truncation=True,
+            padding=True,
+            return_tensors="pt"
+        )
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+
+        # Get predictions
+        with torch.no_grad():
+            outputs = model(**inputs)
+            predictions = torch.argmax(outputs.logits, dim=1)
+
+        all_predictions.extend(predictions.cpu().numpy())
+        all_labels.append(item['labels'])
+
+    # Calculate metrics
+    accuracy = accuracy_score(all_labels, all_predictions)
+    precision, recall, f1, support = precision_recall_fscore_support(
+        all_labels,
+        all_predictions,
+        average='weighted'
+    )
+
+    # Calculate per-class metrics
+    per_class_precision, per_class_recall, per_class_f1, _ = precision_recall_fscore_support(
+        all_labels,
+        all_predictions,
+        average=None
+    )
+
+    # Create confusion matrix
+    conf_matrix = confusion_matrix(all_labels, all_predictions)
+
+    # Combine metrics
+    metrics = {
+        'accuracy': accuracy,
+        'weighted_precision': precision,
+        'weighted_recall': recall,
+        'weighted_f1': f1,
+        'confusion_matrix': conf_matrix,
+        'per_class_metrics': {
+            label: {
+                'precision': p,
+                'recall': r,
+                'f1': f
+            } for label, p, r, f in zip(
+                label_encoder.classes_,
+                per_class_precision,
+                per_class_recall,
+                per_class_f1
+            )
+        }
+    }
+
+    return metrics
+
+
+def print_evaluation_report(metrics: Dict, label_encoder):
+    """
+    Print formatted evaluation report
+    """
+    print("\n" + "=" * 50)
+    print("MODEL EVALUATION REPORT")
+    print("=" * 50)
+
+    print("\nOverall Metrics:")
+    print(f"Accuracy: {metrics['accuracy']:.4f}")
+    print(f"Weighted Precision: {metrics['weighted_precision']:.4f}")
+    print(f"Weighted Recall: {metrics['weighted_recall']:.4f}")
+    print(f"Weighted F1-Score: {metrics['weighted_f1']:.4f}")
+
+    print("\nPer-Class Metrics:")
+    print("-" * 50)
+    print(f"{'Class':<30} {'Precision':>10} {'Recall':>10} {'F1-Score':>10}")
+    print("-" * 50)
+
+    for label, class_metrics in metrics['per_class_metrics'].items():
+        print(
+            f"{label:<30} {class_metrics['precision']:>10.4f} {class_metrics['recall']:>10.4f} {class_metrics['f1']:>10.4f}")
+
+    print("\nConfusion Matrix:")
+    print("-" * 50)
+    conf_matrix = metrics['confusion_matrix']
+    print(conf_matrix)
+
+
+if __name__ == "__main__":
+    output_dir = "models/vision-classifier"
+    model_path = os.path.join(output_dir, "model")
+    tokenizer_path = os.path.join(output_dir, "tokenizer")
+
+    if os.path.exists(model_path):
+        print("\nLoading pre-trained model...")
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+            model = AutoModelForSequenceClassification.from_pretrained(model_path)
+            label_encoder_path = os.path.join(output_dir, "label_encoder.pkl")
+            if os.path.exists(label_encoder_path):
+                with open(label_encoder_path, 'rb') as f:
+                    label_encoder = pickle.load(f)
+            else:
+                print("Warning: Label encoder not found. Running full training...")
+                trainer, model, tokenizer, label_encoder = main()
+
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            model.to(device)
+            print(f"Model loaded successfully and moved to {device}")
+
+            # Load test dataset for evaluation
+            dataset_dict, _ = prepare_data()
+
+            # Run evaluation
+            print("\nEvaluating model performance...")
+            eval_metrics = evaluate_model(
+                model,
+                dataset_dict['test'],
+                tokenizer,
+                label_encoder,
+                device
+            )
+
+            # Print evaluation report
+            print_evaluation_report(eval_metrics, label_encoder)
+
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            print("Running full training instead...")
+            trainer, model, tokenizer, label_encoder = main()
+    else:
+        print("\nNo pre-trained model found. Running training...")
+        trainer, model, tokenizer, label_encoder = main()
+
+
+    def predict_vision_status(text, model, tokenizer, label_encoder):
+        """Make prediction using the loaded/trained model"""
+        inputs = tokenizer(
+            text,
+            truncation=True,
+            padding=True,
+            return_tensors="pt"
+        )
+
+        device = next(model.parameters()).device
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = model(**inputs)
+            # Apply softmax to get probabilities
+            probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
+
+            # Convert to numpy array
+            probabilities = probabilities.cpu().numpy()[0]
+
+            # Create list of (label, probability) tuples
+            predictions = []
+            for idx, prob in enumerate(probabilities):
+                label = label_encoder.inverse_transform([idx])[0]
+                predictions.append((label, float(prob)))
+
+            # Sort by probability in descending order
+            predictions.sort(key=lambda x: x[1], reverse=True)
+
+        return predictions
+
+    example_text = "Age: 40-64 years, Gender: Female, Race: White, non-Hispanic, Diabetes: No"
+    predictions = predict_vision_status(example_text, model, tokenizer, label_encoder)
+
+    print(f"\nPredictions for: {example_text}")
+    print("\nLabel Confidence Scores:")
+    print("-" * 50)
+    for label, confidence in predictions:
+        print(f"{label:<30} {confidence:.2%}")