Spaces:
Sleeping
Sleeping
| import pandas as pd | |
| import polars as pl | |
| import torch | |
| import pickle | |
| import warnings | |
| from typing import Dict, Optional, Tuple, List | |
| import numpy as np | |
| from utils import MemoryOptimizer, UncertaintyEstimator, load_and_clean_data, preprocess_ultra_fast | |
| from models import EnhancedTabTransformerWithImprovements | |
| from config import BEST_PARAMS | |
| from torch.utils.data import DataLoader, TensorDataset | |
| class _PredictorUnpickler(pickle.Unpickler): | |
| """Custom unpickler to map legacy __main__ references to actual modules.""" | |
| def find_class(self, module: str, name: str): | |
| if name == 'EnhancedTabTransformerWithImprovements': | |
| return EnhancedTabTransformerWithImprovements | |
| if name == 'UncertaintyEstimator': | |
| return UncertaintyEstimator | |
| return super().find_class(module, name) | |
| warnings.filterwarnings("ignore") | |
| class DrugInteractionPredictor: | |
| """ | |
| A reusable class for predicting drug-drug interactions. | |
| This class encapsulates all the functionality needed to: | |
| - Load and configure a trained model | |
| - Preprocess drug interaction data | |
| - Make predictions with uncertainty estimates | |
| - Format results into human-readable text | |
| """ | |
| def __init__(self, | |
| model_path: str = "/mnt/data/drugbank/model_weights/da_train_enhanced_model_params_drop_2220.pkl", | |
| labels_csv: str = "/mnt/data/drugbank/converted_ddi_types.csv", | |
| features_file: str = "list_of_all_features_ascending_order.txt", | |
| test_data_path: str = "/mnt/data/drugbank/data_splits/test_extracted_lowercase.csv", | |
| device: Optional[torch.device] = None, | |
| confidence_threshold: float = 0.8): | |
| """ | |
| Initialize the Drug Interaction Predictor. | |
| Args: | |
| model_path: Path to the trained model weights | |
| labels_csv: Path to the label mapping CSV file | |
| features_file: Path to the feature list file | |
| test_data_path: Path to the test data CSV | |
| device: PyTorch device (defaults to CUDA if available) | |
| confidence_threshold: Minimum confidence for high-confidence predictions | |
| """ | |
| self.model_path = model_path | |
| self.test_data_path = test_data_path | |
| self.confidence_threshold = confidence_threshold | |
| # Set device | |
| self.device = device if device is not None else torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| # Load features to drop | |
| self.columns_to_drop = self._load_feature_list(features_file) | |
| # Load label mappings | |
| self.converted_labels_dict = self._load_label_mappings(labels_csv) | |
| # Initialize model (will be loaded later) | |
| self.model = None | |
| self.label_encoder = None | |
| self.preprocessors = None | |
| self.categorical_cols = [] | |
| self.numerical_cols = [] | |
| self.feature_names = [] | |
| self.num_classes = None | |
| print(f"DrugInteractionPredictor initialized on device: {self.device}") | |
| def _load_feature_list(self, features_file: str) -> List[str]: | |
| """Load the list of features to drop from file.""" | |
| try: | |
| with open(features_file, "r", encoding="utf-8") as f: | |
| columns_to_drop = [line.strip() for line in f if line.strip()] | |
| print(f"Loaded {len(columns_to_drop)} features to drop") | |
| return columns_to_drop | |
| except FileNotFoundError: | |
| print(f"Warning: Feature file {features_file} not found. Using empty list.") | |
| return [] | |
| def _load_label_mappings(self, labels_csv: str) -> Dict[int, str]: | |
| """Load the label mappings from CSV.""" | |
| converted_labels = pd.read_csv(labels_csv) | |
| labels_dict = dict(zip(converted_labels['class'], converted_labels['converted_interaction'])) | |
| print(f"Loaded {len(labels_dict)} label mappings") | |
| return labels_dict | |
| def load_model(self) -> None: | |
| """Load the trained model from disk.""" | |
| print(f"Loading model from {self.model_path}...") | |
| def _safe_load_pickle(path: str): | |
| with open(path, "rb") as fh: | |
| return _PredictorUnpickler(fh).load() | |
| loaded_model = None | |
| try: | |
| loaded_model = _safe_load_pickle(self.model_path) | |
| except Exception as exc: | |
| print(f"Direct load failed ({exc}); falling back to manual loader.") | |
| if isinstance(loaded_model, EnhancedTabTransformerWithImprovements): | |
| self.model = loaded_model | |
| print("Loaded EnhancedTabTransformerWithImprovements directly from pickle.") | |
| else: | |
| self.model = EnhancedTabTransformerWithImprovements( | |
| categories=[], | |
| num_continuous=0, | |
| num_classes=178, | |
| device=self.device, | |
| best_params=BEST_PARAMS | |
| ) | |
| self.model.load_models(self.model_path) | |
| # Ensure underlying models are on the desired device | |
| if hasattr(self.model, "models"): | |
| for mdl in self.model.models: | |
| if hasattr(mdl, "to"): | |
| mdl.to(self.device) | |
| if hasattr(mdl, "eval"): | |
| mdl.eval() | |
| self.model.device = self.device | |
| if not getattr(self.model, "models", None): | |
| raise RuntimeError( | |
| "Model file did not contain any trained sub-models. " | |
| "Please verify the pickle at " + self.model_path | |
| ) | |
| self.num_classes = getattr(self.model, "num_classes", None) | |
| model_count = len(self.model.models) | |
| print(f"Model loaded: ensemble of {model_count} model(s) with " | |
| f"{len(getattr(self.model, 'categories', []))} categorical features, " | |
| f"{getattr(self.model, 'num_continuous', 0)} continuous features, " | |
| f"{getattr(self.model, 'num_classes', 'unknown')} classes") | |
| def preprocess_data(self, drug_a: str, drug_b: str) -> Tuple[pd.DataFrame, Optional[np.ndarray], Dict]: | |
| """ | |
| Load and preprocess data for a specific drug pair. | |
| Args: | |
| drug_a: Name of the first drug | |
| drug_b: Name of the second drug | |
| Returns: | |
| Tuple of (preprocessed_features, labels, preprocessors) | |
| """ | |
| # Load and filter data for the specific drug pair | |
| test_df = load_and_clean_data(self.test_data_path, [], drug_a, drug_b) | |
| # Drop unnecessary columns | |
| columns_to_drop_actual = [ | |
| "drugsmiles-drug_a", | |
| 'drugsmiles-drug_b', | |
| 'drugid-drug_a', | |
| 'drugid-drug_b', | |
| 'drugname-drug_a', | |
| 'drugname-drug_b' | |
| ] | |
| test_df = test_df.drop([col for col in columns_to_drop_actual if col in test_df.columns], axis=1) | |
| # Preprocess the data | |
| X_test, y_test, categorical_cols, numerical_cols, preprocessors = preprocess_ultra_fast(test_df) | |
| # Reorder columns so categorical features come first (matches TabTransformer expectations) | |
| ordered_cols = categorical_cols + [col for col in X_test.columns if col not in categorical_cols] | |
| X_test = X_test[ordered_cols] | |
| self.label_encoder = preprocessors['label_encoder'] | |
| self.preprocessors = preprocessors | |
| self.categorical_cols = categorical_cols | |
| self.numerical_cols = numerical_cols | |
| self.feature_names = list(X_test.columns) | |
| print(f"Data preprocessed: {X_test.shape}") | |
| # Validate feature count | |
| if self.model is not None: | |
| if not getattr(self.model, "categories", None) and categorical_cols: | |
| # Populate categories cardinalities if missing | |
| cat_encoders = preprocessors.get('categorical_encoders', {}) or {} | |
| self.model.categories = [ | |
| len(cat_encoders.get(col, {}).get('mapping', {}) or {0: None}) | |
| for col in categorical_cols | |
| ] | |
| if getattr(self.model, "num_continuous", None) in (None, 0) and numerical_cols: | |
| self.model.num_continuous = len(numerical_cols) | |
| if getattr(self.model, "num_classes", None) is None and self.label_encoder is not None: | |
| self.model.num_classes = len(self.label_encoder.classes_) | |
| self.num_classes = self.model.num_classes | |
| expected_features = len(getattr(self.model, "categories", []) or []) + (getattr(self.model, "num_continuous", 0) or 0) | |
| actual_features = X_test.shape[1] | |
| if actual_features != expected_features: | |
| if expected_features == 0: | |
| # Fallback: assume all features are continuous if metadata missing | |
| self.model.num_continuous = actual_features | |
| expected_features = actual_features | |
| print("Model metadata missing; assuming all features are continuous.") | |
| else: | |
| raise ValueError( | |
| f"Feature count mismatch! Model expects {expected_features} features " | |
| f"({len(self.model.categories)} categorical + {self.model.num_continuous} continuous), " | |
| f"but data has {actual_features} features.\n" | |
| f"The model was trained with a different number of dropped features." | |
| ) | |
| print(f"Feature validation passed: {actual_features} features match model expectation") | |
| return X_test, y_test, preprocessors | |
| def predict(self, drug_a: str, drug_b: str, return_all_predictions: bool = False) -> Dict: | |
| """ | |
| Predict drug-drug interaction for a given drug pair. | |
| Args: | |
| drug_a: Name of the first drug | |
| drug_b: Name of the second drug | |
| return_all_predictions: If True, return all predictions; if False, only high-confidence ones | |
| Returns: | |
| Dictionary containing: | |
| - predictions: Array of predicted classes | |
| - probabilities: Array of prediction probabilities | |
| - uncertainties: Dictionary of uncertainty metrics | |
| - confidence_scores: Array of confidence scores | |
| - high_confidence_mask: Boolean mask for high-confidence predictions | |
| - formatted_predictions: List of human-readable prediction texts | |
| - drug_names: Tuple of (drug_a, drug_b) | |
| """ | |
| # Ensure model is loaded | |
| if self.model is None: | |
| self.load_model() | |
| # Preprocess data | |
| X_test, y_test, preprocessors = self.preprocess_data(drug_a, drug_b) | |
| # Make predictions | |
| test_df_for_prediction = X_test.copy() | |
| test_results = self.model.batch_predict_with_uncertainty(test_df_for_prediction) | |
| predictions = test_results['predictions'] | |
| probabilities = test_results['probabilities'] | |
| uncertainties = test_results['uncertainties'] | |
| confidence_scores = uncertainties['confidence'] | |
| high_conf_mask = confidence_scores > self.confidence_threshold | |
| # Format predictions | |
| formatted_predictions = [] | |
| indices_to_format = np.where(high_conf_mask)[0] if not return_all_predictions else range(len(predictions)) | |
| for idx in indices_to_format: | |
| encoded_pred = predictions[idx] | |
| if encoded_pred in self.converted_labels_dict: | |
| text_pred = self.converted_labels_dict[encoded_pred] | |
| text_pred = text_pred.replace("drug_a", drug_a).replace("drug_b", drug_b) | |
| formatted_predictions.append({ | |
| 'text': text_pred, | |
| 'encoded_class': encoded_pred, | |
| 'confidence': float(confidence_scores[idx]), | |
| 'probability': float(probabilities[idx].max()) | |
| }) | |
| return { | |
| 'predictions': predictions, | |
| 'probabilities': probabilities, | |
| 'uncertainties': uncertainties, | |
| 'confidence_scores': confidence_scores, | |
| 'high_confidence_mask': high_conf_mask, | |
| 'formatted_predictions': formatted_predictions, | |
| 'drug_names': (drug_a, drug_b), | |
| 'ground_truth': y_test | |
| } | |
| # def predict_single_hien(self, drug_a: str, drug_b: str) -> Optional[Dict]: | |
| def predict_single(self, drug_a: str, drug_b: str) -> Optional[Dict]: | |
| """ | |
| Get the single best high-confidence prediction for a drug pair. | |
| Args: | |
| drug_a: Name of the first drug | |
| drug_b: Name of the second drug | |
| Returns: | |
| Dictionary with prediction, confidence, and features for LIME/SHAP: | |
| { | |
| 'text': str, # Formatted prediction text | |
| 'encoded_class': int, # Encoded class label | |
| 'confidence': float, # Confidence score | |
| 'probability': float, # Max probability | |
| 'features': pd.Series, # Feature vector for this sample | |
| 'feature_names': List[str], # List of feature names | |
| 'X_test': pd.DataFrame, # Full preprocessed dataframe (for batch LIME) | |
| 'y_test': np.ndarray, # Ground truth labels (if available) | |
| } | |
| """ | |
| # Preprocess data to get features | |
| X_test, y_test, preprocessors = self.preprocess_data(drug_a, drug_b) | |
| # Make predictions | |
| results = self.predict(drug_a, drug_b, return_all_predictions=False) | |
| if results['formatted_predictions']: | |
| best_prediction = results['formatted_predictions'][0] | |
| # Add feature information for LIME/SHAP | |
| best_prediction['features'] = X_test.iloc[0] # First row as pandas Series | |
| best_prediction['feature_names'] = list(X_test.columns) | |
| best_prediction['X_test'] = X_test # Full dataframe (in case multiple rows) | |
| best_prediction['y_test'] = y_test # Ground truth (if available) | |
| return best_prediction | |
| else: | |
| return None | |
| # ==================== SCRIPT EXECUTION (for backward compatibility) ==================== | |
| if __name__ == "__main__": | |
| # Example usage: Predict interaction between two drugs | |
| drugname_a, drugname_b = "trimethoprim", "fenbufen" | |
| # Initialize predictor | |
| predictor = DrugInteractionPredictor() | |
| # Make prediction | |
| results = predictor.predict(drugname_a, drugname_b) | |
| # Display results | |
| print("\n" + "="*80) | |
| print("PREDICTION RESULTS") | |
| print("="*80) | |
| print(f"Drug pair: {drugname_a} + {drugname_b}") | |
| print(f"Total predictions: {len(results['predictions'])}") | |
| print(f"High-confidence predictions: {results['high_confidence_mask'].sum()}") | |
| if results['formatted_predictions']: | |
| print("\nFormatted Predictions:") | |
| for i, pred in enumerate(results['formatted_predictions'], 1): | |
| print(f"\n{i}. {pred['text']}") | |
| print(f" Confidence: {pred['confidence']:.4f}") | |
| print(f" Probability: {pred['probability']:.4f}") | |
| print(f" Encoded class: {pred['encoded_class']}") | |
| else: | |
| print("\nNo high-confidence predictions found.") | |
| # Show ground truth if available | |
| if results['ground_truth'] is not None: | |
| high_conf_indices = np.where(results['high_confidence_mask'])[0] | |
| if len(high_conf_indices) > 0: | |
| print(f"\nGround truth label: {results['ground_truth'][high_conf_indices[0]]}") | |
| print("="*80) | |