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)