Dengue_Drug_Repourposing_FCNN_Model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d910198-3521-400a-8891-138cef64abe6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Set up the working directory and import necessary libraries\n",
    "import os\n",
    "import re\n",
    "import math\n",
    "import torch\n",
    "import umap\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from rdkit import Chem\n",
    "import torch.nn as nn\n",
    "import seaborn as sns\n",
    "import torch.optim as optim\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.nn import GATConv\n",
    "from torch_geometric.loader import NeighborLoader\n",
    "import matplotlib.pyplot as plt\n",
    "from rdkit.Chem import AllChem\n",
    "from tape import ProteinBertModel, TAPETokenizer\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.manifold import TSNE\n",
    "from sklearn.model_selection import train_test_split, KFold\n",
    "from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset\n",
    "from transformers import AutoTokenizer, AutoModel, T5Tokenizer, T5Model, T5EncoderModel,AutoModelForSeq2SeqLM\n",
    "\n",
    "# Define the working directory\n",
    "working_dir = r\"/home/Desktop/re/DB2\"\n",
    "\n",
    "# Check if the directory exists\n",
    "if not os.path.exists(working_dir):\n",
    "    raise FileNotFoundError(f\"Directory {working_dir} does not exist\")\n",
    "\n",
    "# Load the SMILES and proteins data into pandas dataframes\n",
    "smiles_df = pd.read_csv(os.path.join(working_dir, 'unique_smiles.csv'))\n",
    "proteins_df = pd.read_csv(os.path.join(working_dir, 'unique_proteins.csv'))\n",
    "\n",
    "# Print the first few rows to ensure data is loaded correctly\n",
    "print(\"First few rows of smiles.csv:\")\n",
    "print(smiles_df.head())\n",
    "\n",
    "print(\"\\nFirst few rows of proteins.csv:\")\n",
    "print(proteins_df.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fdfd8fa-9f22-4db5-b735-f00fee53a2eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#########################################################Step1#########################################################################\n",
    "# Set up device\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "# Initialize ChemBERTa for drugs\n",
    "tokenizer_chemberta = AutoTokenizer.from_pretrained(\"seyonec/ChemBERTa-zinc-base-v1\")\n",
    "model_chemberta = AutoModel.from_pretrained(\"seyonec/ChemBERTa-zinc-base-v1\")\n",
    "\n",
    "# Function to extract ChemBERTa embeddings for a SMILES string\n",
    "def get_chemberta_embedding(smiles):\n",
    "    inputs = tokenizer_chemberta(smiles, return_tensors=\"pt\", padding=True, truncation=True)\n",
    "    with torch.no_grad():\n",
    "        outputs = model_chemberta(**inputs)\n",
    "    # Taking the mean of the hidden states as the embedding\n",
    "    return outputs.last_hidden_state.mean(dim=1).detach().numpy().flatten()\n",
    "\n",
    "# Process drugs\n",
    "drug_embeddings = []\n",
    "for _, row in tqdm(smiles_df.iterrows(), total=len(smiles_df), desc=\"Processing drugs\"):\n",
    "    smiles = row['Ligand SMILES']\n",
    "    drug_id = row['DrugBank ID']\n",
    "    \n",
    "    # Get ChemBERTa embeddings\n",
    "    chemberta_emb = get_chemberta_embedding(smiles)\n",
    "    \n",
    "    # Append the DrugBank ID and the embeddings as a row\n",
    "    drug_embeddings.append([drug_id] + chemberta_emb.tolist())\n",
    "\n",
    "# Convert to a DataFrame with proper column names\n",
    "embedding_columns = [f\"embedding_dim_{i+1}\" for i in range(chemberta_emb.shape[0])]\n",
    "drug_embeddings_df = pd.DataFrame(drug_embeddings, columns=['DrugBank ID'] + embedding_columns)\n",
    "\n",
    "# Save drug embeddings to CSV\n",
    "drug_embeddings_df.to_csv(os.path.join(working_dir, 'ChemBERT_drug_embeddings.csv'), index=False)\n",
    "\n",
    "# Initialize the tokenizer and model for TAPE (Protein embeddings)\n",
    "tokenizer_tape = TAPETokenizer(vocab=\"iupac\")\n",
    "model_tape = ProteinBertModel.from_pretrained(\"bert-base\")\n",
    "\n",
    "# Function to compute Morgan fingerprints (for drugs)\n",
    "def get_morgan_fingerprint(smiles, radius=2, n_bits=2048):\n",
    "    mol = Chem.MolFromSmiles(smiles)\n",
    "    if mol is None:\n",
    "        print(f\"Invalid SMILES: {smiles}\")\n",
    "        return np.zeros(n_bits)  # Return a zero vector for invalid SMILES\n",
    "    fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=n_bits)\n",
    "    return np.array(fp)\n",
    "\n",
    "# Function to clean and validate protein sequences\n",
    "def clean_sequence(sequence):\n",
    "    valid_amino_acids = set(\"ACDEFGHIKLMNPQRSTVWY\")\n",
    "    cleaned_sequence = ''.join([aa for aa in sequence if aa.upper() in valid_amino_acids])\n",
    "    return cleaned_sequence.upper()\n",
    "\n",
    "# Function to compute TAPE embeddings (for proteins)\n",
    "def get_tape_embedding(sequence):\n",
    "    # Clean and validate the sequence\n",
    "    cleaned_sequence = clean_sequence(sequence)\n",
    "    if len(cleaned_sequence) == 0:\n",
    "        print(f\"Invalid sequence: {sequence}\")\n",
    "        return np.zeros(model_tape.config.hidden_size)  # Return a zero vector for invalid sequences\n",
    "\n",
    "    # Tokenize the sequence\n",
    "    tokens = tokenizer_tape.tokenize(cleaned_sequence)\n",
    "    input_ids = torch.tensor([tokenizer_tape.convert_tokens_to_ids(tokens)]).to(torch.long)\n",
    "    \n",
    "    # Forward pass through TAPE model\n",
    "    with torch.no_grad():\n",
    "        outputs = model_tape(input_ids)\n",
    "    \n",
    "    # Access the first element in the tuple (final hidden states)\n",
    "    hidden_states = outputs[0]\n",
    "    \n",
    "    # Use mean pooling over all tokens to get a fixed-size embedding\n",
    "    embedding = hidden_states.mean(dim=1).squeeze().cpu().numpy()\n",
    "    return embedding\n",
    "\n",
    "# Initialize lists to store the results\n",
    "drug_morgan_fps = []\n",
    "drug_ids = []\n",
    "protein_tape_embeddings = []\n",
    "protein_ids = []\n",
    "\n",
    "# Extract Morgan fingerprints for drugs\n",
    "print(\"Processing Morgan fingerprints for drugs...\")\n",
    "for _, row in tqdm(smiles_df.iterrows(), total=len(smiles_df)):\n",
    "    smiles = row['Ligand SMILES']\n",
    "    drug_id = row['DrugBank ID']\n",
    "    \n",
    "    # Compute Morgan fingerprint\n",
    "    morgan_fp = get_morgan_fingerprint(smiles)\n",
    "    \n",
    "    # Store the results\n",
    "    drug_morgan_fps.append(morgan_fp)\n",
    "    drug_ids.append(drug_id)\n",
    "\n",
    "# Convert drug fingerprints to a numpy array and save to a CSV file\n",
    "drug_morgan_fps = np.array(drug_morgan_fps)\n",
    "drug_fps_df = pd.DataFrame(drug_morgan_fps, columns=[f'fingerprint_{i+1}' for i in range(drug_morgan_fps.shape[1])])\n",
    "drug_fps_df.insert(0, 'DrugBank ID', drug_ids)  # Insert DrugBank ID as the first column\n",
    "drug_fps_df.to_csv(os.path.join(working_dir, 'drug_morgan_fingerprints.csv'), index=False)\n",
    "print(\"Morgan fingerprints for drugs saved to 'drug_morgan_fingerprints.csv'.\")\n",
    "\n",
    "# Extract TAPE embeddings for proteins\n",
    "print(\"Processing TAPE embeddings for proteins...\")\n",
    "for _, row in tqdm(proteins_df.iterrows(), total=len(proteins_df)):\n",
    "    sequence = row['Sequence']\n",
    "    protein_id = row['UniProt ID']\n",
    "    \n",
    "    # Compute TAPE embedding\n",
    "    try:\n",
    "        tape_emb = get_tape_embedding(sequence)\n",
    "        protein_tape_embeddings.append(tape_emb)\n",
    "        protein_ids.append(protein_id)\n",
    "    except Exception as e:\n",
    "        print(f\"Error processing protein {protein_id}: {e}\")\n",
    "        continue  # Skip this protein and move on to the next\n",
    "\n",
    "# Convert protein embeddings to a numpy array and save to a CSV file\n",
    "protein_tape_embeddings = np.array(protein_tape_embeddings)\n",
    "protein_emb_df = pd.DataFrame(protein_tape_embeddings, columns=[f'embedding_dim_{i+1}' for i in range(protein_tape_embeddings.shape[1])])\n",
    "protein_emb_df.insert(0, 'UniProt ID', protein_ids)  # Insert UniProt ID as the first column\n",
    "protein_emb_df.to_csv(os.path.join(working_dir, 'protein_tape_embeddings.csv'), index=False)\n",
    "print(\"TAPE embeddings for proteins saved to 'protein_tape_embeddings.csv'.\")\n",
    "\n",
    "# Load the protein data from the CSV file\n",
    "proteins_df = pd.read_csv(os.path.join(working_dir, 'unique_proteins.csv'))\n",
    "\n",
    "# Initialize the tokenizer and model for ProstT5 (from Hugging Face)\n",
    "model_name = \"Rostlab/ProstT5\"\n",
    "tokenizer = T5Tokenizer.from_pretrained(model_name, do_lower_case=False)\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(model_name)\n",
    "\n",
    "# Set up the device to use (GPU if available, otherwise CPU)\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "model = model.to(device)\n",
    "\n",
    "# Use half-precision on GPUs, full-precision on CPU\n",
    "if device == 'cpu':\n",
    "    model.float()\n",
    "else:\n",
    "    model.half()\n",
    "\n",
    "# Function to clean and validate protein sequences\n",
    "def clean_sequence(sequence):\n",
    "    # Keep only valid amino acids (A-Z)\n",
    "    valid_amino_acids = set(\"ACDEFGHIKLMNPQRSTVWY\")\n",
    "    cleaned_sequence = ''.join([aa for aa in sequence if aa.upper() in valid_amino_acids])\n",
    "    return cleaned_sequence.upper()\n",
    "\n",
    "# Function to get ProstT5 embeddings for a protein sequence\n",
    "def get_prostt5_embedding(sequence):\n",
    "    # Clean and validate the sequence\n",
    "    cleaned_sequence = clean_sequence(sequence)\n",
    "    \n",
    "    if len(cleaned_sequence) == 0:\n",
    "        print(f\"Invalid sequence: {sequence}\")\n",
    "        return np.zeros(model.config.d_model)  # Return a zero vector if the sequence is invalid\n",
    "\n",
    "    # Add pre-fix for amino acid sequence\n",
    "    sequence_with_prefix = \"<AA2fold> \" + cleaned_sequence\n",
    "\n",
    "    # Tokenize the sequence and convert to tensor\n",
    "    inputs = tokenizer(sequence_with_prefix, return_tensors=\"pt\", padding=True, truncation=True).to(device)\n",
    "    \n",
    "    # Forward pass through the ProstT5 model to get the hidden states\n",
    "    with torch.no_grad():\n",
    "        outputs = model.encoder(**inputs)\n",
    "    \n",
    "    # Extract the embeddings (mean of the last hidden states)\n",
    "    hidden_states = outputs.last_hidden_state\n",
    "    embedding = hidden_states.mean(dim=1).squeeze().cpu().numpy()  # Mean pooling over all tokens\n",
    "    \n",
    "    return embedding\n",
    "\n",
    "# Initialize lists to store protein embeddings and protein IDs\n",
    "protein_embeddings = []\n",
    "protein_ids = []\n",
    "\n",
    "# Loop over the protein sequences and extract embeddings\n",
    "for _, row in tqdm(proteins_df.iterrows(), total=len(proteins_df), desc=\"Processing proteins\"):\n",
    "    sequence = row['Sequence']\n",
    "    protein_id = row['UniProt ID']\n",
    "    \n",
    "    # Get the ProstT5 embedding for each protein sequence\n",
    "    try:\n",
    "        embedding = get_prostt5_embedding(sequence)\n",
    "        protein_embeddings.append(embedding)\n",
    "        protein_ids.append(protein_id)\n",
    "    except Exception as e:\n",
    "        print(f\"Error processing protein {protein_id}: {e}\")\n",
    "        continue  # Skip this protein and move on to the next\n",
    "\n",
    "# Convert the list of embeddings into a numpy array and then into a DataFrame\n",
    "protein_embeddings = np.array(protein_embeddings)\n",
    "embedding_columns = [f'embedding_dim_{i+1}' for i in range(protein_embeddings.shape[1])]\n",
    "protein_emb_df = pd.DataFrame(protein_embeddings, columns=embedding_columns)\n",
    "\n",
    "# Insert the UniProt IDs as the first column\n",
    "protein_emb_df.insert(0, 'UniProt ID', protein_ids)\n",
    "\n",
    "# Save the embeddings to a CSV file\n",
    "protein_emb_df.to_csv(os.path.join(working_dir, 'protein_prostt5_embeddings.csv'), index=False)\n",
    "\n",
    "print(\"ProstT5 embeddings for proteins saved to 'protein_prostt5_embeddings.csv'.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66ac9e9d-9cff-4e67-aec1-36f37a0a6e85",
   "metadata": {},
   "outputs": [],
   "source": [
    "##########################################Step3#############################################################################################\n",
    "#70-15-15 split\n",
    "#Load the embeddings and Kd values from CSV files\n",
    "chemberta_drug_embeddings = pd.read_csv(os.path.join(working_dir, 'ChemBERT_drug_embeddings.csv'))\n",
    "morgan_drug_fingerprints = pd.read_csv(os.path.join(working_dir, 'drug_morgan_fingerprints.csv'))\n",
    "prostt5_protein_embeddings = pd.read_csv(os.path.join(working_dir, 'protein_prostt5_embeddings.csv'))\n",
    "tape_protein_embeddings = pd.read_csv(os.path.join(working_dir, 'protein_tape_embeddings.csv'))\n",
    "kd_values = pd.read_csv(os.path.join(working_dir, 'Links_Kd_Scores.csv'))  # with 'DrugBank ID', 'UniProt ID', 'Kd'\n",
    "\n",
    "#Log-normalize the Kd values\n",
    "kd_values['Log_Kd'] = kd_values['Kd(nM)'].apply(lambda x: math.log10(x + 1))  # Adding 1 to avoid log(0)\n",
    "\n",
    "#Merge embeddings\n",
    "# Merge drug embeddings (ChemBERTa + Morgan fingerprints) using 'DrugBank ID'\n",
    "drug_embeddings = pd.merge(chemberta_drug_embeddings, morgan_drug_fingerprints, on='DrugBank ID')\n",
    "\n",
    "# Merge protein embeddings (ProstT5 + TAPE embeddings) using 'UniProt ID'\n",
    "protein_embeddings = pd.merge(prostt5_protein_embeddings, tape_protein_embeddings, on='UniProt ID')\n",
    "\n",
    "# Merge the drug-protein pairs with Kd values\n",
    "merged_df = pd.merge(kd_values, drug_embeddings, on='DrugBank ID')\n",
    "merged_df = pd.merge(merged_df, protein_embeddings, on='UniProt ID')\n",
    "\n",
    "#Z-score normalization for ChemBERTa, ProstT5, and TAPE embeddings\n",
    "chemberta_cols = [col for col in chemberta_drug_embeddings.columns if 'embedding_dim_' in col]\n",
    "prostt5_cols = [col for col in prostt5_protein_embeddings.columns if 'embedding_dim_' in col]\n",
    "tape_cols = [col for col in tape_protein_embeddings.columns if 'embedding_dim_' in col]\n",
    "morgan_cols = [col for col in morgan_drug_fingerprints.columns if 'fingerprint_' in col]  # No normalization for Morgan\n",
    "\n",
    "# Normalizing ChemBERTa, ProstT5, and TAPE embeddings\n",
    "scaler_chemberta = StandardScaler()\n",
    "scaler_prostt5 = StandardScaler()\n",
    "scaler_tape = StandardScaler()\n",
    "\n",
    "# Normalize embeddings\n",
    "merged_df[chemberta_cols] = scaler_chemberta.fit_transform(merged_df[chemberta_cols])\n",
    "merged_df[prostt5_cols] = scaler_prostt5.fit_transform(merged_df[prostt5_cols])\n",
    "merged_df[tape_cols] = scaler_tape.fit_transform(merged_df[tape_cols])\n",
    "\n",
    "#Attention mechanism to combine embeddings\n",
    "class AttentionCombiner(nn.Module):\n",
    "    def __init__(self, chemberta_dim, morgan_dim, prostt5_dim, tape_dim):\n",
    "        super(AttentionCombiner, self).__init__()\n",
    "        self.chemberta_fc = nn.Linear(chemberta_dim, chemberta_dim)\n",
    "        self.morgan_fc = nn.Linear(morgan_dim, morgan_dim)\n",
    "        self.prostt5_fc = nn.Linear(prostt5_dim, prostt5_dim)\n",
    "        self.tape_fc = nn.Linear(tape_dim, tape_dim)\n",
    "        \n",
    "        self.attention_layer = nn.Linear(chemberta_dim + morgan_dim + prostt5_dim + tape_dim, 1)\n",
    "\n",
    "    def forward(self, chemberta, morgan, prostt5, tape):\n",
    "        # Apply linear layers\n",
    "        chemberta = self.chemberta_fc(chemberta)\n",
    "        morgan = self.morgan_fc(morgan)\n",
    "        prostt5 = self.prostt5_fc(prostt5)\n",
    "        tape = self.tape_fc(tape)\n",
    "        \n",
    "        # Concatenate embeddings\n",
    "        combined = torch.cat([chemberta, morgan, prostt5, tape], dim=1)\n",
    "        \n",
    "        # Apply attention\n",
    "        attention_weights = torch.softmax(self.attention_layer(combined), dim=1)\n",
    "        \n",
    "        # Compute weighted combination\n",
    "        combined_embedding = attention_weights * combined\n",
    "        return combined_embedding\n",
    "\n",
    "# Convert embeddings to tensors and apply attention mechanism\n",
    "chemberta_tensor = torch.tensor(merged_df[chemberta_cols].values, dtype=torch.float32)\n",
    "morgan_tensor = torch.tensor(merged_df[morgan_cols].values, dtype=torch.float32)\n",
    "prostt5_tensor = torch.tensor(merged_df[prostt5_cols].values, dtype=torch.float32)\n",
    "tape_tensor = torch.tensor(merged_df[tape_cols].values, dtype=torch.float32)\n",
    "\n",
    "# Initialize attention combiner\n",
    "combiner = AttentionCombiner(\n",
    "    chemberta_dim=len(chemberta_cols),\n",
    "    morgan_dim=len(morgan_cols),\n",
    "    prostt5_dim=len(prostt5_cols),\n",
    "    tape_dim=len(tape_cols)\n",
    ")\n",
    "\n",
    "# Apply attention mechanism to get the combined embedding for each drug-protein pair\n",
    "with torch.no_grad():\n",
    "    combined_embeddings = combiner(chemberta_tensor, morgan_tensor, prostt5_tensor, tape_tensor)\n",
    "\n",
    "# Convert combined embeddings back to DataFrame\n",
    "combined_embeddings_df = pd.DataFrame(combined_embeddings.numpy())\n",
    "\n",
    "# Add combined embeddings and log-normalized Kd values to final dataset\n",
    "final_dataset = pd.concat([merged_df[['DrugBank ID', 'UniProt ID', 'Log_Kd']], combined_embeddings_df], axis=1)\n",
    "\n",
    "#Split the dataset into training (70%) and remaining (30%)\n",
    "train_df, remaining_df = train_test_split(final_dataset, test_size=0.3, random_state=42)\n",
    "\n",
    "#Split the remaining dataset into validation (15%) and test (15%)\n",
    "val_df, test_df = train_test_split(remaining_df, test_size=0.5, random_state=42)\n",
    "\n",
    "# Save the splits to CSV files in the working directory\n",
    "train_df.to_csv(os.path.join(working_dir, 'train_dataset.csv'), index=False)\n",
    "val_df.to_csv(os.path.join(working_dir, 'val_dataset.csv'), index=False)\n",
    "test_df.to_csv(os.path.join(working_dir, 'test_dataset.csv'), index=False)\n",
    "\n",
    "# Print a message to confirm successful completion\n",
    "print(\"Data processing complete! Training, validation, and test datasets saved in the working directory.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8c63694-2658-49c9-afaa-428b1919061f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#################################################################Analyssis#################################################\n",
    "# Load the datasets\n",
    "train_df = pd.read_csv(os.path.join(working_dir, 'train_dataset.csv'))\n",
    "val_df = pd.read_csv(os.path.join(working_dir, 'val_dataset.csv'))\n",
    "test_df = pd.read_csv(os.path.join(working_dir, 'test_dataset.csv'))\n",
    "\n",
    "# Function to assess dataset\n",
    "def assess_dataset(dataset, name):\n",
    "    print(f\"\\nAssessing {name} dataset:\")\n",
    "    \n",
    "    # Check the shape of the dataset\n",
    "    print(f\"Shape: {dataset.shape}\")\n",
    "    \n",
    "    # Check for NaN values\n",
    "    nan_values = dataset.isna().sum()\n",
    "    if nan_values.any():\n",
    "        print(f\"Missing values (NaNs) detected in the following columns:\\n{nan_values[nan_values > 0]}\")\n",
    "    else:\n",
    "        print(\"No missing values (NaNs) detected.\")\n",
    "    \n",
    "    # Check for columns with constant values (variance = 0)\n",
    "    constant_columns = dataset.columns[dataset.nunique() == 1]\n",
    "    if len(constant_columns) > 0:\n",
    "        print(f\"Columns with constant values: {constant_columns.tolist()}\")\n",
    "    else:\n",
    "        print(\"No columns with constant values detected.\")\n",
    "    \n",
    "    # Summary statistics\n",
    "    print(\"\\nSummary statistics:\")\n",
    "    print(dataset.describe())\n",
    "    \n",
    "    # Check data types of each column\n",
    "    print(\"\\nData types of each column:\")\n",
    "    print(dataset.dtypes)\n",
    "    \n",
    "    # Check for any infinite values\n",
    "    inf_values = (dataset == float('inf')).sum()\n",
    "    if inf_values.any():\n",
    "        print(f\"Infinite values detected in the following columns:\\n{inf_values[inf_values > 0]}\")\n",
    "    else:\n",
    "        print(\"No infinite values detected.\")\n",
    "    \n",
    "    # Check if the number of features (columns) is consistent across all datasets\n",
    "    return dataset.shape[1]\n",
    "\n",
    "# Assess each dataset\n",
    "train_columns = assess_dataset(train_df, 'Training')\n",
    "val_columns = assess_dataset(val_df, 'Validation')\n",
    "test_columns = assess_dataset(test_df, 'Test')\n",
    "\n",
    "# Check if the number of columns is consistent across train, val, and test sets\n",
    "if train_columns == val_columns == test_columns:\n",
    "    print(\"\\nAll datasets have the same number of features.\")\n",
    "else:\n",
    "    print(\"\\nWarning: The number of features (columns) is inconsistent across the datasets.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02478b48-8505-4dc7-8344-9dc6c80d1782",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the device to GPU if available\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "train_file = os.path.join(working_dir, 'train_dataset.csv')\n",
    "val_file = os.path.join(working_dir, 'val_dataset.csv')\n",
    "test_file = os.path.join(working_dir, 'test_dataset.csv')\n",
    "\n",
    "# Dataset class\n",
    "class KDDataset(Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.drug_protein_embedding = data.iloc[:, 3:].values  # All embeddings (drug + protein)\n",
    "        self.log_kd = data['Log_Kd'].values\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.log_kd)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return torch.tensor(self.drug_protein_embedding[idx], dtype=torch.float32), torch.tensor(self.log_kd[idx], dtype=torch.float32)\n",
    "\n",
    "# Load datasets\n",
    "def load_data(file_path):\n",
    "    data = pd.read_csv(file_path)\n",
    "    return KDDataset(data)\n",
    "\n",
    "# Load the datasets\n",
    "train_dataset = load_data(train_file)\n",
    "val_dataset = load_data(val_file)\n",
    "test_dataset = load_data(test_file)\n",
    "\n",
    "# DataLoader\n",
    "batch_size = 64\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "# FCNN Model\n",
    "class FCNN(nn.Module):\n",
    "    def __init__(self, input_dim):\n",
    "        super(FCNN, self).__init__()\n",
    "        self.fc1 = nn.Linear(input_dim, 1024)\n",
    "        self.fc2 = nn.Linear(1024, 512)\n",
    "        self.fc3 = nn.Linear(512, 256)\n",
    "        self.fc4 = nn.Linear(256, 1)  # Output layer\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = torch.relu(self.fc1(x))\n",
    "        x = torch.relu(self.fc2(x))\n",
    "        x = torch.relu(self.fc3(x))\n",
    "        x = self.fc4(x)\n",
    "        return x\n",
    "\n",
    "# Get the input dimension\n",
    "input_dim = train_dataset[0][0].shape[0]\n",
    "\n",
    "# Instantiate the model\n",
    "model = FCNN(input_dim).to(device)\n",
    "\n",
    "# Define optimizer and loss function\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "criterion = nn.MSELoss()\n",
    "\n",
    "# Training loop (without early stopping)\n",
    "def train_model(model, train_loader, val_loader, num_epochs=50):\n",
    "    train_losses, val_losses = [], []\n",
    "    \n",
    "    for epoch in range(num_epochs):\n",
    "        model.train()\n",
    "        train_loss = 0.0\n",
    "        \n",
    "        for batch in train_loader:\n",
    "            drug_embedding, log_kd = batch\n",
    "            drug_embedding, log_kd = drug_embedding.to(device), log_kd.to(device)\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(drug_embedding)\n",
    "            loss = criterion(outputs.squeeze(), log_kd)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            train_loss += loss.item()\n",
    "        \n",
    "        train_losses.append(train_loss / len(train_loader))\n",
    "\n",
    "        # Validation loop\n",
    "        model.eval()\n",
    "        val_loss = 0.0\n",
    "        with torch.no_grad():\n",
    "            for batch in val_loader:\n",
    "                drug_embedding, log_kd = batch\n",
    "                drug_embedding, log_kd = drug_embedding.to(device), log_kd.to(device)\n",
    "                outputs = model(drug_embedding)\n",
    "                loss = criterion(outputs.squeeze(), log_kd)\n",
    "                val_loss += loss.item()\n",
    "\n",
    "        val_losses.append(val_loss / len(val_loader))\n",
    "\n",
    "        print(f\"Epoch [{epoch + 1}/{num_epochs}] - Train Loss: {train_losses[-1]:.4f}, Val Loss: {val_losses[-1]:.4f}\")\n",
    "\n",
    "    return model, train_losses, val_losses\n",
    "\n",
    "# Save the metrics plot and evaluation results\n",
    "def plot_metrics_and_save(loss, mae, mse, r2, working_dir):\n",
    "    metrics = {\n",
    "        'Loss': loss,\n",
    "        'MAE': mae,\n",
    "        'MSE': mse,\n",
    "        'R²': r2\n",
    "    }\n",
    "    \n",
    "    # Plot the metrics\n",
    "    plt.figure()\n",
    "    plt.bar(metrics.keys(), metrics.values())\n",
    "    plt.xlabel('Metrics')\n",
    "    plt.ylabel('Value')\n",
    "    plt.title('Evaluation Metrics on Test Set')\n",
    "    \n",
    "    # Save the plot as a high-quality image\n",
    "    image_path = os.path.join(working_dir, 'evaluation_metrics.png')\n",
    "    plt.savefig(image_path, dpi=600, bbox_inches='tight')\n",
    "    plt.show()\n",
    "\n",
    "# Updated evaluation function\n",
    "def evaluate_model(model, test_loader, working_dir):\n",
    "    model.eval()\n",
    "    test_loss = 0.0\n",
    "    criterion = nn.MSELoss()\n",
    "\n",
    "    predictions = []\n",
    "    targets = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch in test_loader:\n",
    "            drug_embedding, log_kd = batch\n",
    "            drug_embedding, log_kd = drug_embedding.to(device), log_kd.to(device)\n",
    "\n",
    "            outputs = model(drug_embedding)\n",
    "\n",
    "            loss = criterion(outputs.squeeze(), log_kd)\n",
    "            test_loss += loss.item()\n",
    "\n",
    "            predictions.append(outputs.cpu().numpy())\n",
    "            targets.append(log_kd.cpu().numpy())\n",
    "\n",
    "    predictions = np.concatenate(predictions)\n",
    "    targets = np.concatenate(targets)\n",
    "\n",
    "    # Calculate metrics\n",
    "    loss = test_loss / len(test_loader)\n",
    "    mse = mean_squared_error(targets, predictions)\n",
    "    mae = mean_absolute_error(targets, predictions)\n",
    "    r2 = r2_score(targets, predictions)\n",
    "\n",
    "    print(f\"Test Loss: {loss:.4f}\")\n",
    "    print(f\"Test MAE: {mae:.4f}, Test MSE: {mse:.4f}, Test R²: {r2:.4f}\")\n",
    "\n",
    "    # Plot and save metrics\n",
    "    plot_metrics_and_save(loss, mae, mse, r2, working_dir)\n",
    "\n",
    "    return loss, mae, mse, r2\n",
    "\n",
    "# Function to plot training and validation loss\n",
    "def plot_training_history(train_losses, val_losses):\n",
    "    plt.figure()\n",
    "    plt.plot(train_losses, label='Training Loss')\n",
    "    plt.plot(val_losses, label='Validation Loss')\n",
    "    plt.xlabel('Epochs')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.title('Training and Validation Loss over Epochs')\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "\n",
    "# Train the model\n",
    "trained_model, train_losses, val_losses = train_model(model, train_loader, val_loader)\n",
    "\n",
    "# Plot training and validation loss\n",
    "plot_training_history(train_losses, val_losses)\n",
    "\n",
    "# Evaluate the model and save results\n",
    "loss, mae, mse, r2 = evaluate_model(trained_model, test_loader, working_dir)\n",
    "\n",
    "# Save the model\n",
    "torch.save(trained_model.state_dict(), os.path.join(working_dir, 'fcnn_model.pth'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deeb28e2-7aea-4df1-ac03-aeea09ba90ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set device to GPU if available\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# Define the working directory and file paths\n",
    "new_drug_file = os.path.join(working_dir, 'New_drugs.csv')\n",
    "\n",
    "# Load new drugs\n",
    "new_drugs_df = pd.read_csv(new_drug_file)\n",
    "\n",
    "################################### ChemBERTa Embeddings ###################################\n",
    "\n",
    "# Load ChemBERTa model and tokenizer\n",
    "tokenizer_chemberta = AutoTokenizer.from_pretrained(\"seyonec/ChemBERTa-zinc-base-v1\")\n",
    "model_chemberta = AutoModel.from_pretrained(\"seyonec/ChemBERTa-zinc-base-v1\").to(device)\n",
    "\n",
    "# Function to extract ChemBERTa embeddings\n",
    "def get_chemberta_embedding(smiles):\n",
    "    inputs = tokenizer_chemberta(smiles, return_tensors=\"pt\", padding=True, truncation=True).to(device)\n",
    "    with torch.no_grad():\n",
    "        outputs = model_chemberta(**inputs)\n",
    "    return outputs.last_hidden_state.mean(dim=1).cpu().numpy().flatten()\n",
    "\n",
    "# Initialize lists to store embeddings\n",
    "new_drug_embeddings = []\n",
    "\n",
    "# Loop through each drug and extract ChemBERTa embeddings\n",
    "for _, row in new_drugs_df.iterrows():\n",
    "    smiles = row['Ligand SMILES']\n",
    "    drug_id = row['DrugBank ID']\n",
    "    \n",
    "    # Get ChemBERTa embedding\n",
    "    chemberta_emb = get_chemberta_embedding(smiles)\n",
    "    \n",
    "    # Store embedding with DrugBank ID\n",
    "    new_drug_embeddings.append([drug_id] + chemberta_emb.tolist())\n",
    "\n",
    "# Convert embeddings to DataFrame\n",
    "chemberta_df = pd.DataFrame(new_drug_embeddings, columns=['DrugBank ID'] + [f'embedding_dim_{i+1}' for i in range(len(new_drug_embeddings[0]) - 1)])\n",
    "\n",
    "# Save ChemBERTa embeddings to CSV\n",
    "chemberta_output_file = os.path.join(working_dir, 'New_ChemBERT_drug_embeddings.csv')\n",
    "chemberta_df.to_csv(chemberta_output_file, index=False)\n",
    "\n",
    "print(f\"ChemBERTa embeddings saved to {chemberta_output_file}\")\n",
    "\n",
    "################################### Morgan Fingerprints ###################################\n",
    "\n",
    "# Function to compute Morgan fingerprints\n",
    "def get_morgan_fingerprint(smiles, radius=2, n_bits=2048):\n",
    "    mol = Chem.MolFromSmiles(smiles)\n",
    "    if mol is None:\n",
    "        return np.zeros(n_bits)  # Return a zero vector for invalid SMILES\n",
    "    fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=n_bits)\n",
    "    return np.array(fp)\n",
    "\n",
    "# Initialize list to store fingerprints\n",
    "new_morgan_fingerprints = []\n",
    "\n",
    "# Loop through each drug and compute Morgan fingerprints\n",
    "for _, row in new_drugs_df.iterrows():\n",
    "    smiles = row['Ligand SMILES']\n",
    "    drug_id = row['DrugBank ID']\n",
    "    \n",
    "    # Get Morgan fingerprint\n",
    "    morgan_fp = get_morgan_fingerprint(smiles)\n",
    "    \n",
    "    # Store fingerprint with DrugBank ID\n",
    "    new_morgan_fingerprints.append([drug_id] + morgan_fp.tolist())\n",
    "\n",
    "# Convert fingerprints to DataFrame\n",
    "morgan_df = pd.DataFrame(new_morgan_fingerprints, columns=['DrugBank ID'] + [f'fingerprint_{i+1}' for i in range(len(new_morgan_fingerprints[0]) - 1)])\n",
    "\n",
    "# Save Morgan fingerprints to CSV\n",
    "morgan_output_file = os.path.join(working_dir, 'New_drug_morgan_fingerprints.csv')\n",
    "morgan_df.to_csv(morgan_output_file, index=False)\n",
    "\n",
    "print(f\"Morgan fingerprints saved to {morgan_output_file}\")\n",
    "\n",
    "import os\n",
    "import torch\n",
    "import pandas as pd\n",
    "from transformers import T5Tokenizer, T5EncoderModel\n",
    "from tape import ProteinBertModel, TAPETokenizer\n",
    "\n",
    "# Set device to GPU if available\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "new_protein_file = os.path.join(working_dir, 'Book.csv')\n",
    "\n",
    "# Load new proteins\n",
    "new_proteins_df = pd.read_csv(new_protein_file)\n",
    "\n",
    "################################### ProstT5 Embeddings ###################################\n",
    "\n",
    "# Load ProstT5 tokenizer and model\n",
    "tokenizer_prostt5 = T5Tokenizer.from_pretrained('Rostlab/ProstT5', do_lower_case=False)\n",
    "model_prostt5 = T5EncoderModel.from_pretrained('Rostlab/ProstT5').to(device)\n",
    "\n",
    "# Function to preprocess and tokenize protein sequences for ProstT5\n",
    "def preprocess_and_tokenize_prostt5(sequences):\n",
    "    sequences = [\" \".join(list(seq)) for seq in sequences]  # Add space between each amino acid\n",
    "    sequences = [\"<AA2fold> \" + seq for seq in sequences]   # Add special token \"<AA2fold>\"\n",
    "    return tokenizer_prostt5.batch_encode_plus(sequences, add_special_tokens=True, padding=\"longest\", return_tensors='pt').to(device)\n",
    "\n",
    "# Function to extract ProstT5 embeddings\n",
    "def extract_prostt5_embeddings(sequences):\n",
    "    tokenized_seqs = preprocess_and_tokenize_prostt5(sequences)\n",
    "    with torch.no_grad():\n",
    "        embeddings = model_prostt5(input_ids=tokenized_seqs.input_ids, attention_mask=tokenized_seqs.attention_mask)\n",
    "    per_protein_embeddings = []\n",
    "    for i in range(len(sequences)):\n",
    "        seq_len = (tokenized_seqs.attention_mask[i] == 1).sum() - 1  # Exclude padding\n",
    "        per_protein_emb = embeddings.last_hidden_state[i, 1:seq_len+1].mean(dim=0).cpu().numpy()  # Mean pooling\n",
    "        per_protein_embeddings.append(per_protein_emb)\n",
    "    return per_protein_embeddings\n",
    "\n",
    "# Extract ProstT5 embeddings for all proteins\n",
    "new_protein_embeddings_prostt5 = []\n",
    "for _, row in new_proteins_df.iterrows():\n",
    "    sequence = row['Sequence']\n",
    "    protein_id = row['UniProt ID']\n",
    "    \n",
    "    # Extract ProstT5 embedding\n",
    "    prostt5_emb = extract_prostt5_embeddings([sequence])[0]\n",
    "    \n",
    "    # Store embedding with UniProt ID\n",
    "    new_protein_embeddings_prostt5.append([protein_id] + prostt5_emb.tolist())\n",
    "\n",
    "# Convert ProstT5 embeddings to DataFrame\n",
    "prostt5_df = pd.DataFrame(new_protein_embeddings_prostt5, columns=['UniProt ID'] + [f'embedding_dim_{i+1}' for i in range(len(new_protein_embeddings_prostt5[0]) - 1)])\n",
    "\n",
    "# Save ProstT5 embeddings to CSV\n",
    "prostt5_output_file = os.path.join(working_dir, 'New_protein_prostt5_embeddings.csv')\n",
    "prostt5_df.to_csv(prostt5_output_file, index=False)\n",
    "\n",
    "print(f\"ProstT5 embeddings saved to {prostt5_output_file}\")\n",
    "\n",
    "################################### TAPE Embeddings ###################################\n",
    "\n",
    "# Load TAPE tokenizer and model\n",
    "tokenizer_tape = TAPETokenizer(vocab=\"iupac\")\n",
    "model_tape = ProteinBertModel.from_pretrained(\"bert-base\")\n",
    "\n",
    "# Function to extract TAPE embeddings\n",
    "def extract_tape_embeddings(sequence):\n",
    "    tokens = tokenizer_tape.tokenize(sequence)\n",
    "    input_ids = torch.tensor([tokenizer_tape.convert_tokens_to_ids(tokens)]).to(torch.long)\n",
    "    with torch.no_grad():\n",
    "        outputs = model_tape(input_ids)\n",
    "    return outputs[0].mean(dim=1).squeeze().cpu().numpy()  # Mean pooling over the tokens\n",
    "\n",
    "# Extract TAPE embeddings for all proteins\n",
    "new_protein_embeddings_tape = []\n",
    "for _, row in new_proteins_df.iterrows():\n",
    "    sequence = row['Sequence']\n",
    "    protein_id = row['UniProt ID']\n",
    "    \n",
    "    # Extract TAPE embedding\n",
    "    tape_emb = extract_tape_embeddings(sequence)\n",
    "    \n",
    "    # Store embedding with UniProt ID\n",
    "    new_protein_embeddings_tape.append([protein_id] + tape_emb.tolist())\n",
    "\n",
    "# Convert TAPE embeddings to DataFrame\n",
    "tape_df = pd.DataFrame(new_protein_embeddings_tape, columns=['UniProt ID'] + [f'embedding_dim_{i+1}' for i in range(len(new_protein_embeddings_tape[0]) - 1)])\n",
    "\n",
    "# Save TAPE embeddings to CSV\n",
    "tape_output_file = os.path.join(working_dir, 'New_protein_tape_embeddings.csv')\n",
    "tape_df.to_csv(tape_output_file, index=False)\n",
    "\n",
    "print(f\"TAPE embeddings saved to {tape_output_file}\")\n",
    "\n",
    "import pandas as pd\n",
    "import os\n",
    "from itertools import product\n",
    "\n",
    "# Load the drugs and proteins data into pandas dataframes\n",
    "drugs_df = pd.read_csv(os.path.join(working_dir, 'New_drugs.csv'))  # New set of drugs\n",
    "proteins_df = pd.read_csv(os.path.join(working_dir, 'Book.csv'))  # New set of proteins\n",
    "\n",
    "# Generate all possible combinations of drug and protein pairs using Cartesian product\n",
    "all_pairs = list(product(drugs_df['DrugBank ID'], proteins_df['UniProt ID']))\n",
    "\n",
    "# Create a dataframe to store these pairs\n",
    "pairs_df = pd.DataFrame(all_pairs, columns=['DrugBank ID', 'UniProt ID'])\n",
    "\n",
    "# Save the pairs to a CSV file\n",
    "pairs_file = os.path.join(working_dir, 'all_drug_protein_pairs.csv')\n",
    "pairs_df.to_csv(pairs_file, index=False)\n",
    "\n",
    "print(f\"All possible drug-protein pairs saved to: {pairs_file}\")\n",
    "\n",
    "\n",
    "# Define AttentionCombiner class\n",
    "class AttentionCombiner(nn.Module):\n",
    "    def __init__(self, chemberta_dim, morgan_dim, prostt5_dim, tape_dim):\n",
    "        super(AttentionCombiner, self).__init__()\n",
    "        self.chemberta_fc = nn.Linear(chemberta_dim, chemberta_dim)\n",
    "        self.morgan_fc = nn.Linear(morgan_dim, morgan_dim)\n",
    "        self.prostt5_fc = nn.Linear(prostt5_dim, prostt5_dim)\n",
    "        self.tape_fc = nn.Linear(tape_dim, tape_dim)\n",
    "        \n",
    "        # Attention layer to learn the importance of different embeddings\n",
    "        self.attention_layer = nn.Linear(chemberta_dim + morgan_dim + prostt5_dim + tape_dim, 1)\n",
    "\n",
    "    def forward(self, chemberta, morgan, prostt5, tape):\n",
    "        # Apply linear layers to each type of embedding\n",
    "        chemberta = self.chemberta_fc(chemberta)\n",
    "        morgan = self.morgan_fc(morgan)\n",
    "        prostt5 = self.prostt5_fc(prostt5)\n",
    "        tape = self.tape_fc(tape)\n",
    "        \n",
    "        # Concatenate embeddings\n",
    "        combined = torch.cat([chemberta, morgan, prostt5, tape], dim=1)\n",
    "        \n",
    "        # Apply attention weights\n",
    "        attention_weights = torch.softmax(self.attention_layer(combined), dim=1)\n",
    "        \n",
    "        # Compute the final combined embedding\n",
    "        combined_embedding = attention_weights * combined\n",
    "        return combined_embedding\n",
    "\n",
    "# File paths for separate embeddings and drug-protein pairs\n",
    "pair_file = os.path.join(working_dir, 'all_drug_protein_pairs.csv')  # Generated drug-protein pair file\n",
    "chemberta_file = os.path.join(working_dir, 'New_ChemBERT_drug_embeddings.csv')\n",
    "morgan_file = os.path.join(working_dir, 'New_drug_morgan_fingerprints.csv')\n",
    "prostt5_file = os.path.join(working_dir, 'New_protein_prostt5_embeddings.csv')\n",
    "tape_file = os.path.join(working_dir, 'New_protein_tape_embeddings.csv')\n",
    "\n",
    "# Load the separate embeddings and the pairs file\n",
    "pairs_df = pd.read_csv(pair_file)\n",
    "chemberta_df = pd.read_csv(chemberta_file)\n",
    "morgan_df = pd.read_csv(morgan_file)\n",
    "prostt5_df = pd.read_csv(prostt5_file)\n",
    "tape_df = pd.read_csv(tape_file)\n",
    "\n",
    "# Extract embedding columns for each type\n",
    "chemberta_cols = [col for col in chemberta_df.columns if 'embedding_dim_' in col]\n",
    "morgan_cols = [col for col in morgan_df.columns if 'fingerprint_' in col]\n",
    "prostt5_cols = [col for col in prostt5_df.columns if 'embedding_dim_' in col]\n",
    "tape_cols = [col for col in tape_df.columns if 'embedding_dim_' in col]\n",
    "\n",
    "#Z-score normalization for ChemBERTa, ProstT5, and TAPE embeddings (Morgan remains unchanged)\n",
    "scaler_chemberta = StandardScaler()\n",
    "scaler_prostt5 = StandardScaler()\n",
    "scaler_tape = StandardScaler()\n",
    "\n",
    "chemberta_df[chemberta_cols] = scaler_chemberta.fit_transform(chemberta_df[chemberta_cols])\n",
    "prostt5_df[prostt5_cols] = scaler_prostt5.fit_transform(prostt5_df[prostt5_cols])\n",
    "tape_df[tape_cols] = scaler_tape.fit_transform(tape_df[tape_cols])\n",
    "\n",
    "#Initialize lists to store mapped embeddings\n",
    "mapped_chemberta = []\n",
    "mapped_morgan = []\n",
    "mapped_prostt5 = []\n",
    "mapped_tape = []\n",
    "drug_ids = []\n",
    "protein_ids = []\n",
    "\n",
    "#Map embeddings to drug-protein pairs\n",
    "for _, pair in pairs_df.iterrows():\n",
    "    drug_id = pair['DrugBank ID']\n",
    "    protein_id = pair['UniProt ID']\n",
    "    \n",
    "    # Find the row in each embedding DataFrame based on DrugBank ID and UniProt ID\n",
    "    drug_row_chemberta = chemberta_df[chemberta_df['DrugBank ID'] == drug_id]\n",
    "    drug_row_morgan = morgan_df[morgan_df['DrugBank ID'] == drug_id]\n",
    "    protein_row_prostt5 = prostt5_df[prostt5_df['UniProt ID'] == protein_id]\n",
    "    protein_row_tape = tape_df[tape_df['UniProt ID'] == protein_id]\n",
    "    \n",
    "    if drug_row_chemberta.empty or drug_row_morgan.empty or protein_row_prostt5.empty or protein_row_tape.empty:\n",
    "        continue  # Skip if no matching row is found\n",
    "    \n",
    "    # Extract embeddings for the drug and protein\n",
    "    chemberta_embedding = drug_row_chemberta[chemberta_cols].values.flatten()\n",
    "    morgan_embedding = drug_row_morgan[morgan_cols].values.flatten()\n",
    "    prostt5_embedding = protein_row_prostt5[prostt5_cols].values.flatten()\n",
    "    tape_embedding = protein_row_tape[tape_cols].values.flatten()\n",
    "    \n",
    "    # Append embeddings to lists\n",
    "    mapped_chemberta.append(chemberta_embedding)\n",
    "    mapped_morgan.append(morgan_embedding)\n",
    "    mapped_prostt5.append(prostt5_embedding)\n",
    "    mapped_tape.append(tape_embedding)\n",
    "    drug_ids.append(drug_id)\n",
    "    protein_ids.append(protein_id)\n",
    "\n",
    "#Convert mapped embeddings to PyTorch tensors\n",
    "chemberta_tensor = torch.tensor(np.array(mapped_chemberta), dtype=torch.float32)\n",
    "morgan_tensor = torch.tensor(np.array(mapped_morgan), dtype=torch.float32)\n",
    "prostt5_tensor = torch.tensor(np.array(mapped_prostt5), dtype=torch.float32)\n",
    "tape_tensor = torch.tensor(np.array(mapped_tape), dtype=torch.float32)\n",
    "\n",
    "#Initialize the AttentionCombiner and combine embeddings\n",
    "combiner = AttentionCombiner(\n",
    "    chemberta_dim=len(chemberta_cols),\n",
    "    morgan_dim=len(morgan_cols),\n",
    "    prostt5_dim=len(prostt5_cols),\n",
    "    tape_dim=len(tape_cols)\n",
    ")\n",
    "\n",
    "# Apply attention mechanism to combine embeddings\n",
    "combined_embeddings = []\n",
    "with torch.no_grad():\n",
    "    for i in range(chemberta_tensor.size(0)):\n",
    "        combined_emb = combiner(\n",
    "            chemberta_tensor[i].unsqueeze(0),\n",
    "            morgan_tensor[i].unsqueeze(0),\n",
    "            prostt5_tensor[i].unsqueeze(0),\n",
    "            tape_tensor[i].unsqueeze(0)\n",
    "        ).squeeze(0).cpu().numpy().flatten()\n",
    "        combined_embeddings.append(combined_emb)\n",
    "\n",
    "#Save combined embeddings with drug-protein pairs\n",
    "combined_embeddings_df = pd.DataFrame(combined_embeddings)\n",
    "combined_embeddings_df.insert(0, 'DrugBank ID', drug_ids)\n",
    "combined_embeddings_df.insert(1, 'UniProt ID', protein_ids)\n",
    "\n",
    "# Save the combined embeddings to a CSV file\n",
    "combined_embeddings_file = os.path.join(working_dir, 'combined_drug_protein_embeddings.csv')\n",
    "combined_embeddings_df.to_csv(combined_embeddings_file, index=False)\n",
    "\n",
    "print(f\"Combined embeddings saved to: {combined_embeddings_file}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}