Spaces:
Runtime error
Runtime error
File size: 8,485 Bytes
608709a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 |
from flask import Flask, request, jsonify
import torch
import torch.nn as nn
import pandas as pd
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from rdkit import Chem
from rdkit.Chem.rdFingerprintGenerator import GetMorganGenerator
import joblib
import pickle
import pubchempy as pcp
import logging
import os
app = Flask(__name__)
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Device setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info(f"Using device: {device}")
# Model name
model_name = "Fredaaaaaa/smiles"
# Load dataset
dataset_path = "/kaggle/input/labeled-data/labeled_severity.csv"
try:
df = pd.read_csv(dataset_path, encoding='latin1')
df.rename(columns={"Interaction Description": "interaction_description"}, inplace=True)
df['Drug 1_normalized'] = df['Drug 1_normalized'].str.lower()
df['Drug 2_normalized'] = df['Drug 2_normalized'].str.lower()
logger.info("Dataset loaded successfully")
except Exception as e:
logger.error(f"Failed to load dataset: {e}")
df = pd.DataFrame()
# Load model components
model_dir = "/kaggle/working/drug_interaction_model"
try:
# Load tokenizer and BioBERT model
if os.path.exists(model_dir):
tokenizer = AutoTokenizer.from_pretrained(model_dir)
text_model = AutoModelForSequenceClassification.from_pretrained(model_dir).to(device)
else:
logger.warning(f"Local model directory {model_dir} not found, falling back to {model_name}")
tokenizer = AutoTokenizer.from_pretrained(model_name)
text_model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=3).to(device)
text_model.eval()
logger.info("BioBERT and tokenizer loaded")
# Load custom model components
checkpoint = torch.load(os.path.join(model_dir, 'custom_model.pt'), map_location=device)
input_size = checkpoint['input_size']
dropout_rate = checkpoint['dropout_rate']
# Define HybridModel
class HybridModel(nn.Module):
def __init__(self, text_model, input_size, dropout_rate=0.4):
super(HybridModel, self).__init__()
self.text_model = text_model
self.drug_branch = nn.Sequential(
nn.Linear(input_size, 2048),
nn.ReLU(),
nn.BatchNorm1d(2048),
nn.Dropout(dropout_rate),
nn.Linear(2048, 1024),
nn.ReLU(),
nn.BatchNorm1d(1024),
nn.Dropout(dropout_rate),
nn.Linear(1024, 512),
nn.ReLU(),
nn.BatchNorm1d(512),
nn.Dropout(dropout_rate),
nn.Linear(512, 256),
nn.ReLU(),
nn.BatchNorm1d(256),
nn.Dropout(dropout_rate),
nn.Linear(256, 128),
nn.ReLU()
)
self.fusion = nn.Sequential(
nn.Linear(128 + 3, 1024),
nn.ReLU(),
nn.BatchNorm1d(1024),
nn.Dropout(dropout_rate),
nn.Linear(1024, 512),
nn.ReLU(),
nn.BatchNorm1d(512),
nn.Dropout(dropout_rate),
nn.Linear(512, 256),
nn.ReLU(),
nn.BatchNorm1d(256),
nn.Dropout(dropout_rate),
nn.Linear(256, 128),
nn.ReLU(),
nn.BatchNorm1d(128),
nn.Dropout(dropout_rate),
nn.Linear(128, 3)
)
def forward(self, input_ids, attention_mask, drug_features):
text_outputs = self.text_model(input_ids=input_ids, attention_mask=attention_mask)
text_features = text_outputs.logits
drug_features = self.drug_branch(drug_features)
combined = torch.cat((text_features, drug_features), dim=1)
output = self.fusion(combined)
return output
model = HybridModel(text_model, input_size, dropout_rate).to(device)
model.drug_branch.load_state_dict(checkpoint['drug_branch_state_dict'])
model.fusion.load_state_dict(checkpoint['fusion_state_dict'])
model.eval()
logger.info("HybridModel loaded")
# Load Random Forest
rf_model = joblib.load(os.path.join(model_dir, 'rf_model.joblib'))
logger.info("Random Forest model loaded")
# Load label encoder
with open(os.path.join(model_dir, 'label_encoder.pkl'), 'rb') as f:
label_encoder = pickle.load(f)
logger.info("Label encoder loaded")
except Exception as e:
logger.error(f"Failed to load model components: {e}")
raise
# Function to fetch SMILES from PubChem
def get_smiles(drug_name):
try:
compounds = pcp.get_compounds(drug_name, 'name')
if compounds:
return compounds[0].canonical_smiles
logger.warning(f"No SMILES found for {drug_name}")
return None
except Exception as e:
logger.error(f"PubChem API error for {drug_name}: {e}")
return None
# Function to compute Morgan fingerprints
def preprocess_smiles(smiles):
try:
mol = Chem.MolFromSmiles(smiles)
if mol is None:
return np.zeros(1024)
morgan_gen = GetMorganGenerator(radius=2, fpSize=1024)
fingerprint = morgan_gen.GetFingerprint(mol)
return np.array(fingerprint)
except:
return np.zeros(1024)
# Prediction function
def predict_interaction(drug1, drug2, interaction_description):
drug1 = drug1.lower()
drug2 = drug2.lower()
# Check dataset for SMILES
smiles1 = None
smiles2 = None
if not df.empty:
drug1_matches = df[df['Drug 1_normalized'] == drug1]
drug2_matches = df[df['Drug 2_normalized'] == drug2]
if not drug1_matches.empty:
smiles1 = drug1_matches['SMILES'].iloc[0]
if not drug2_matches.empty:
smiles2 = drug2_matches['SMILES_2'].iloc[0]
# Fetch SMILES from PubChem if not in dataset
if smiles1 is None:
smiles1 = get_smiles(drug1)
if smiles2 is None:
smiles2 = get_smiles(drug2)
# Validate SMILES
if not smiles1 or not smiles2:
return {"error": "Unable to retrieve SMILES for one or both drugs"}
# Preprocess SMILES
drug1_features = preprocess_smiles(smiles1)
drug2_features = preprocess_smiles(smiles2)
drug_features = np.hstack([drug1_features, drug2_features])
drug_features_tensor = torch.tensor(drug_features, dtype=torch.float32).unsqueeze(0).to(device)
# Tokenize interaction description
encodings = tokenizer(interaction_description, truncation=True, padding=True, max_length=128, return_tensors='pt')
input_ids = encodings['input_ids'].to(device)
attention_mask = encodings['attention_mask'].to(device)
# Model prediction
with torch.no_grad():
outputs = model(input_ids, attention_mask, drug_features_tensor)
nn_pred = torch.argmax(outputs, dim=1).cpu().numpy()[0]
# Random Forest prediction
rf_pred = rf_model.predict(drug_features.reshape(1, -1))[0]
# Ensemble prediction
votes = [nn_pred] * 9 + [rf_pred] * 1
ensemble_pred = max(set(votes), key=votes.count)
# Decode prediction
severity = label_encoder.inverse_transform([ensemble_pred])[0]
return {"severity": severity}
# Flask routes
@app.route('/')
def index():
return """
<h1>Drug Interaction Severity Prediction</h1>
<form method="POST" action="/predict">
<label>Drug 1:</label><br>
<input type="text" name="drug1" required><br>
<label>Drug 2:</label><br>
<input type="text" name="drug2" required><br>
<label>Interaction Description:</label><br>
<textarea name="interaction_description" required></textarea><br>
<input type="submit" value="Predict">
</form>
"""
@app.route('/predict', methods=['POST'])
def predict():
try:
drug1 = request.form['drug1']
drug2 = request.form['drug2']
interaction_description = request.form['interaction_description']
result = predict_interaction(drug1, drug2, interaction_description)
return jsonify(result)
except Exception as e:
logger.error(f"Prediction error: {e}")
return jsonify({"error": str(e)}), 500
if __name__ == '__main__':
app.run(debug=True, host='0.0.0.0', port=5000) |