FusionGDA

.gitattributes

@@ -35,3 +35,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 step_300_model.bin filter=lfs diff=lfs merge=lfs -text
 disgenet_latest.csv filter=lfs diff=lfs merge=lfs -text
+train.csv filter=lfs diff=lfs merge=lfs -text
+valid.csv filter=lfs diff=lfs merge=lfs -text
+C0002395_disease.csv filter=lfs diff=lfs merge=lfs -text

.ipynb_checkpoints/app-checkpoint.py

@@ -0,0 +1,64 @@
+# -*- coding: utf-8 -*-
+import gradio as gr
+import pandas as pd
+import os
+import subprocess
+
+def predict_top_100_genes(disease_id):
+    # Initialize paths
+    input_csv_path = 'data/downstream/{}_disease.csv'.format(disease_id)
+    output_csv_path = 'data/downstream/{}_top100.csv'.format(disease_id)
+    
+    # Check if the output CSV already exists
+    if not os.path.exists(output_csv_path):
+        # Proceed with your existing code if the output file doesn't exist
+        df = pd.read_csv('data/pretrain/disgenet_latest.csv')
+        df = df[df['proteinSeq'].notna()]
+
+        # Check if the disease_id is present in the dataframe
+        if disease_id not in df['diseaseId'].values:
+            return f"Error: Disease ID '{disease_id}' not found in the database. Please check the ID and try again."
+
+        desired_diseaseDes = df[df['diseaseId'] == disease_id]['diseaseDes'].iloc[0]
+        related_proteins = df[df['diseaseDes'] == desired_diseaseDes]['proteinSeq'].unique()
+        df['score'] = df['proteinSeq'].isin(related_proteins).astype(int)
+        new_df = pd.DataFrame({
+            'diseaseId': disease_id,
+            'diseaseDes': desired_diseaseDes,
+            'geneSymbol': df['geneSymbol'],
+            'proteinSeq': df['proteinSeq'],
+            'score': df['score']
+        }).drop_duplicates().reset_index(drop=True)
+        
+        new_df.to_csv(input_csv_path, index=False)
+        
+        # Call the model script only if the output CSV does not exist
+        script_path = 'model.sh'
+        subprocess.run(['bash', script_path, input_csv_path, output_csv_path], check=True)
+
+    # Read the model output file or the existing file to get the top 100 genes
+    output_df = pd.read_csv(output_csv_path)
+    # Update here to select only the required columns and rename them
+    result_df = output_df[['geneSymbol', 'Prediction_score']].rename(columns={'geneSymbol': 'Gene', 'Prediction_score': 'Score'}).head(100)
+
+    return result_df
+
+iface = gr.Interface(
+    fn=predict_top_100_genes, 
+    inputs=gr.Textbox(lines=1, placeholder="Enter Disease ID Here...", label="Disease ID"), 
+    outputs=gr.Dataframe(label="Predicted Top 100 Related Genes"), 
+    title="Gene Disease Association Prediction",
+    description = (
+    "This AI model predicts the top 100 genes associated with a given disease based on 16,733 genes."
+    " To get started, you need a Disease ID (UMLS CUI), which can be obtained from the DisGeNET database. "
+    "\n\n**Steps to Obtain a Disease ID from DisGeNET:**\n"
+    "1. Visit the DisGeNET website: [https://www.disgenet.org/search](https://www.disgenet.org/search).\n"
+    "2. Use the search bar to enter your disease of interest. For instance, if you're interested in 'Alzheimer's Disease', type 'Alzheimer's Disease' into the search bar.\n"
+    "3. From the search results, identify the disease you're researching. The Disease ID (UMLS CUI) is listed alongside each disease name, e.g. C0002395.\n"
+    "4. Enter the Disease ID into the input box below and submit.\n\n"
+    "The DisGeNET database contains all known gene-disease associations and associated evidence. In addition, it is able to find the corresponding diseases based on a gene.\n"
+    "\n**The model will take about 18 minutes to inference a new disease.**\n"
+    )
+)
+
+iface.launch(share=True)

.ipynb_checkpoints/model-checkpoint.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+input_csv_path="$1"
+output_csv_path="$2"
+max_depth=6
+device='cuda:0'
+model_path_list=( 
+    "../../save_model_ckp/gda_infoNCE_2024_GPU3090" \
+)
+
+cd ../src/finetune/
+for save_model_path in ${model_path_list[@]}; do
+        num_leaves=$((2**($max_depth-1)))
+        python finetune.py \
+                --input_csv_path $input_csv_path \
+                --output_csv_path $output_csv_path \
+                --save_model_path $save_model_path \
+                --device $device \
+                --batch_size 128 \
+                --step "300" \
+                --use_pooled \
+                --num_leaves $num_leaves \
+                --max_depth $max_depth
+done

.ipynb_checkpoints/requirements-checkpoint.txt

@@ -0,0 +1,7 @@
+lightgbm
+pytorch-metric-learning
+torch
+transformers
+PyTDC
+gradio
+numpy

app.py

@@ -6,14 +6,19 @@ import subprocess
 
 def predict_top_100_genes(disease_id):
     # Initialize paths
-    input_csv_path = '/data/downstream/{}_disease.csv'.format(disease_id)
-    output_csv_path = '/data/downstream/{}_top100.csv'.format(disease_id)
+    input_csv_path = 'data/downstream/{}_disease.csv'.format(disease_id)
+    output_csv_path = 'data/downstream/{}_top100.csv'.format(disease_id)
     
     # Check if the output CSV already exists
     if not os.path.exists(output_csv_path):
         # Proceed with your existing code if the output file doesn't exist
-        df = pd.read_csv('/data/pretrain/disgenet_latest.csv')
+        df = pd.read_csv('data/pretrain/disgenet_latest.csv')
         df = df[df['proteinSeq'].notna()]
+
+        # Check if the disease_id is present in the dataframe
+        if disease_id not in df['diseaseId'].values:
+            return f"Error: Disease ID '{disease_id}' not found in the database. Please check the ID and try again."
+
         desired_diseaseDes = df[df['diseaseId'] == disease_id]['diseaseDes'].iloc[0]
         related_proteins = df[df['diseaseDes'] == desired_diseaseDes]['proteinSeq'].unique()
         df['score'] = df['proteinSeq'].isin(related_proteins).astype(int)
@@ -38,7 +43,6 @@ def predict_top_100_genes(disease_id):
 
     return result_df
 
-
 iface = gr.Interface(
     fn=predict_top_100_genes, 
     inputs=gr.Textbox(lines=1, placeholder="Enter Disease ID Here...", label="Disease ID"), 
@@ -54,7 +58,7 @@ iface = gr.Interface(
     "4. Enter the Disease ID into the input box below and submit.\n\n"
     "The DisGeNET database contains all known gene-disease associations and associated evidence. In addition, it is able to find the corresponding diseases based on a gene.\n"
     "\n**The model will take about 18 minutes to inference a new disease.**\n"
-)
+    )
 )
 
-iface.launch(share=True)
+iface.launch(share=True)

save_model_ckp/gda_infoNCE_2024_GPU3090/step_300_model.bin → data/downstream/C0002395_disease.csv

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:504129ccb1c717366e843df99e73d629b5c0bac0603deb8dbc6fb9b5479387b7
-size 3131981635
+oid sha256:2f7b5562dceae680af5fbe305e06e5ebacafb9bf8404ecebc04b8ecc60a3495d
+size 44085860

data/downstream/GDA_Data/train.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d0533480255f46747ca973110aaa031892bdfd5ca9b2f9bc989f91ce893385a2
+size 117023981

data/downstream/GDA_Data/valid.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03c90fd29ac0af5370c8dc66317cab3de004c1771f47f91301fcb6c11204815f
+size 29321915

src/finetune/.ipynb_checkpoints/finetune-checkpoint.py

@@ -0,0 +1,416 @@
+import argparse
+import os
+import random
+import string
+import sys
+import pandas as pd
+from datetime import datetime
+
+sys.path.append("../")
+import numpy as np
+import torch
+import lightgbm as lgb
+import sklearn.metrics as metrics
+from sklearn.utils import class_weight
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, precision_recall_curve, f1_score, precision_recall_fscore_support,roc_auc_score
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+from transformers import EsmTokenizer, EsmForMaskedLM, BertModel, BertTokenizer, AutoTokenizer, EsmModel
+from utils.downstream_disgenet import DisGeNETProcessor
+from utils.metric_learning_models import GDA_Metric_Learning
+
+def parse_config():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-f')
+    parser.add_argument("--step", type=int, default=0)
+    parser.add_argument(
+        "--save_model_path",
+        type=str,
+        default=None,
+        help="path of the pretrained disease model located",
+    )
+    parser.add_argument(
+        "--prot_encoder_path",
+        type=str,
+        default="facebook/esm2_t33_650M_UR50D",     
+        #"facebook/galactica-6.7b", "Rostlab/prot_bert"  “facebook/esm2_t33_650M_UR50D”
+        help="path/name of protein encoder model located",
+    )
+    parser.add_argument(
+        "--disease_encoder_path",
+        type=str,
+        default="microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext",
+        help="path/name of textual pre-trained language model",
+    )
+    parser.add_argument("--reduction_factor", type=int, default=8)
+    parser.add_argument(
+        "--loss",
+        help="{ms_loss|infoNCE|cosine_loss|circle_loss|triplet_loss}}",
+        default="infoNCE",
+    )
+    parser.add_argument(
+        "--input_feature_save_path",
+        type=str,
+        default="../../data/processed_disease",
+        help="path of tokenized training data",
+    )
+    parser.add_argument(
+        "--agg_mode", default="mean_all_tok", type=str, help="{cls|mean|mean_all_tok}"
+    )
+    parser.add_argument("--batch_size", type=int, default=256)
+    parser.add_argument("--patience", type=int, default=5)
+    parser.add_argument("--num_leaves", type=int, default=5)
+    parser.add_argument("--max_depth", type=int, default=5)
+    parser.add_argument("--lr", type=float, default=0.35)
+    parser.add_argument("--dropout", type=float, default=0.1)
+    parser.add_argument("--test", type=int, default=0)
+    parser.add_argument("--use_miner", action="store_true")
+    parser.add_argument("--miner_margin", default=0.2, type=float)
+    parser.add_argument("--freeze_prot_encoder", action="store_true")
+    parser.add_argument("--freeze_disease_encoder", action="store_true")
+    parser.add_argument("--use_adapter", action="store_true")
+    parser.add_argument("--use_pooled", action="store_true")
+    parser.add_argument("--device", type=str, default="cpu")
+    parser.add_argument(
+        "--use_both_feature",
+        help="use the both features of gnn_feature_v1_samples and pretrained models",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--use_v1_feature_only",
+        help="use the features of gnn_feature_v1_samples only",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--save_path_prefix",
+        type=str,
+        default="../../save_model_ckp/finetune/",
+        help="save the result in which directory",
+    )
+    parser.add_argument(
+        "--save_name", default="fine_tune", type=str, help="the name of the saved file"
+    )
+    # Add argument for input CSV file path
+    parser.add_argument("--input_csv_path", type=str, required=True, help="Path to the input CSV file.")
+
+    # Add argument for output CSV file path
+    parser.add_argument("--output_csv_path", type=str, required=True, help="Path to the output CSV file.")
+    return parser.parse_args()
+
+def get_feature(model, dataloader, args):
+    x = list()
+    y = list()
+    with torch.no_grad():
+        for step, batch in tqdm(enumerate(dataloader)):
+            prot_input_ids, prot_attention_mask, dis_input_ids, dis_attention_mask, y1 = batch
+            prot_input = {
+                'input_ids': prot_input_ids.to(args.device), 
+                'attention_mask': prot_attention_mask.to(args.device)
+            }
+            dis_input = {
+                'input_ids': dis_input_ids.to(args.device), 
+                'attention_mask': dis_attention_mask.to(args.device)
+            }
+            feature_output = model.predict(prot_input, dis_input)
+            x1 = feature_output.cpu().numpy()
+            x.append(x1)
+            y.append(y1.cpu().numpy())
+    x = np.concatenate(x, axis=0)
+    y = np.concatenate(y, axis=0)
+    return x, y
+
+
+def encode_pretrained_feature(args, disGeNET):
+    input_feat_file = os.path.join(
+        args.input_feature_save_path,
+        f"{args.model_short}_{args.step}_use_{'pooled' if args.use_pooled else 'cls'}_feat.npz",
+    )
+
+    if os.path.exists(input_feat_file):
+        print(f"load prior feature data from {input_feat_file}.")
+        loaded = np.load(input_feat_file)
+        x_train, y_train = loaded["x_train"], loaded["y_train"]
+        x_valid, y_valid = loaded["x_valid"], loaded["y_valid"]
+        # x_test, y_test = loaded["x_test"], loaded["y_test"]
+        
+        prot_tokenizer = EsmTokenizer.from_pretrained(args.prot_encoder_path, do_lower_case=False)
+        # prot_tokenizer = BertTokenizer.from_pretrained(args.prot_encoder_path, do_lower_case=False)
+        print("prot_tokenizer", len(prot_tokenizer))
+        disease_tokenizer = BertTokenizer.from_pretrained(args.disease_encoder_path)
+        print("disease_tokenizer", len(disease_tokenizer))
+
+        prot_model = EsmModel.from_pretrained(args.prot_encoder_path)
+        # prot_model = BertModel.from_pretrained(args.prot_encoder_path)
+        disease_model = BertModel.from_pretrained(args.disease_encoder_path)
+        
+        if args.save_model_path:
+            model = GDA_Metric_Learning(prot_model, disease_model, 1280, 768, args)
+
+            if args.use_adapter:
+                prot_model_path = os.path.join(
+                    args.save_model_path, f"prot_adapter_step_{args.step}"
+                )
+                disease_model_path = os.path.join(
+                    args.save_model_path, f"disease_adapter_step_{args.step}"
+                )
+                model.load_adapters(prot_model_path, disease_model_path)
+            else:
+                prot_model_path = os.path.join(
+                    args.save_model_path, f"step_{args.step}_model.bin"
+                )# , f"step_{args.step}_model.bin"
+                disease_model_path = os.path.join(
+                    args.save_model_path, f"step_{args.step}_model.bin"
+                )
+                model.non_adapters(prot_model_path, disease_model_path)
+                
+            model = model.to(args.device)
+            prot_model = model.prot_encoder
+            disease_model = model.disease_encoder
+            print(f"loaded prior model {args.save_model_path}.")
+            
+        def collate_fn_batch_encoding(batch):
+            query1, query2, scores = zip(*batch)
+            
+            query_encodings1 = prot_tokenizer.batch_encode_plus(
+                list(query1),
+                max_length=512,
+                padding="max_length",
+                truncation=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            query_encodings2 = disease_tokenizer.batch_encode_plus(
+                list(query2),
+                max_length=512,
+                padding="max_length",
+                truncation=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            scores = torch.tensor(list(scores))
+            attention_mask1 = query_encodings1["attention_mask"].bool()
+            attention_mask2 = query_encodings2["attention_mask"].bool()
+            
+            return query_encodings1["input_ids"], attention_mask1, query_encodings2["input_ids"], attention_mask2, scores
+
+        test_examples = disGeNET.get_test_examples(args.test)
+        print(f"get test examples: {len(test_examples)}")
+        
+        test_dataloader = DataLoader(
+            test_examples,
+            batch_size=args.batch_size,
+            shuffle=False,
+            collate_fn=collate_fn_batch_encoding,
+        )
+        print( f"dataset loaded: test-{len(test_examples)}")
+
+        x_test, y_test = get_feature(model, test_dataloader, args)
+
+    else:
+        prot_tokenizer = EsmTokenizer.from_pretrained(args.prot_encoder_path, do_lower_case=False)
+        # prot_tokenizer = BertTokenizer.from_pretrained(args.prot_encoder_path, do_lower_case=False)
+        print("prot_tokenizer", len(prot_tokenizer))
+        disease_tokenizer = BertTokenizer.from_pretrained(args.disease_encoder_path)
+        print("disease_tokenizer", len(disease_tokenizer))
+
+        prot_model = EsmModel.from_pretrained(args.prot_encoder_path)
+        # prot_model = BertModel.from_pretrained(args.prot_encoder_path)
+        disease_model = BertModel.from_pretrained(args.disease_encoder_path)
+        
+        if args.save_model_path:
+            model = GDA_Metric_Learning(prot_model, disease_model, 1280, 768, args)
+
+            if args.use_adapter:
+                prot_model_path = os.path.join(
+                    args.save_model_path, f"prot_adapter_step_{args.step}"
+                )
+                disease_model_path = os.path.join(
+                    args.save_model_path, f"disease_adapter_step_{args.step}"
+                )
+                model.load_adapters(prot_model_path, disease_model_path)
+            else:
+                prot_model_path = os.path.join(
+                    args.save_model_path, f"step_{args.step}_model.bin"
+                )# , f"step_{args.step}_model.bin"
+                disease_model_path = os.path.join(
+                    args.save_model_path, f"step_{args.step}_model.bin"
+                )
+                model.non_adapters(prot_model_path, disease_model_path)
+                
+            model = model.to(args.device)
+            prot_model = model.prot_encoder
+            disease_model = model.disease_encoder
+            print(f"loaded prior model {args.save_model_path}.")
+            
+        def collate_fn_batch_encoding(batch):
+            query1, query2, scores = zip(*batch)
+            
+            query_encodings1 = prot_tokenizer.batch_encode_plus(
+                list(query1),
+                max_length=512,
+                padding="max_length",
+                truncation=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            query_encodings2 = disease_tokenizer.batch_encode_plus(
+                list(query2),
+                max_length=512,
+                padding="max_length",
+                truncation=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            scores = torch.tensor(list(scores))
+            attention_mask1 = query_encodings1["attention_mask"].bool()
+            attention_mask2 = query_encodings2["attention_mask"].bool()
+            
+            return query_encodings1["input_ids"], attention_mask1, query_encodings2["input_ids"], attention_mask2, scores
+        
+        train_examples = disGeNET.get_train_examples(args.test)
+        print(f"get training examples: {len(train_examples)}")
+        valid_examples = disGeNET.get_val_examples(args.test)
+        print(f"get validation examples: {len(valid_examples)}")
+        test_examples = disGeNET.get_test_examples(args.test)
+        print(f"get test examples: {len(test_examples)}")
+
+        train_dataloader = DataLoader(
+            train_examples,
+            batch_size=args.batch_size,
+            shuffle=False,
+            collate_fn=collate_fn_batch_encoding,
+        )
+        valid_dataloader = DataLoader(
+            valid_examples,
+            batch_size=args.batch_size,
+            shuffle=False,
+            collate_fn=collate_fn_batch_encoding,
+        )
+        test_dataloader = DataLoader(
+            test_examples,
+            batch_size=args.batch_size,
+            shuffle=False,
+            collate_fn=collate_fn_batch_encoding,
+        )
+        print( f"dataset loaded: train-{len(train_examples)}; valid-{len(valid_examples)}; test-{len(test_examples)}")
+        
+        x_train, y_train = get_feature(model, train_dataloader, args)
+        x_valid, y_valid = get_feature(model, valid_dataloader, args)
+        x_test, y_test = get_feature(model, test_dataloader, args)
+
+        # Save input feature to reduce encoding time
+        np.savez_compressed(
+            input_feat_file,
+            x_train=x_train,
+            y_train=y_train,
+            x_valid=x_valid,
+            y_valid=y_valid,
+        )
+        print(f"save input feature into {input_feat_file}")
+        # Save input feature to reduce encoding time
+    return x_train, y_train, x_valid, y_valid, x_test, y_test
+
+
+def train(args):
+    # defining parameters
+    if args.save_model_path:
+        args.model_short = (
+            args.save_model_path.split("/")[-1]
+        )
+        print(f"model name {args.model_short}")
+
+    else:
+        args.model_short = (
+            args.disease_encoder_path.split("/")[-1] 
+        )
+        print(f"model name {args.model_short}")
+
+    # disGeNET = DisGeNETProcessor()
+    disGeNET = DisGeNETProcessor(input_csv_path=args.input_csv_path)
+
+         
+    x_train, y_train, x_valid, y_valid, x_test, y_test = encode_pretrained_feature(args, disGeNET)
+
+    print("train: ", x_train.shape, y_train.shape)
+    print("valid: ", x_valid.shape, y_valid.shape)
+    print("test: ", x_test.shape, y_test.shape)
+    
+    params = {
+        "task": "train",  # "predict"   train
+        "boosting": "gbdt", # "The options are "gbdt" (traditional Gradient Boosting Decision Tree), "rf" (Random Forest), "dart" (Dropouts meet Multiple Additive Regression Trees), or "goss" (Gradient-based One-Side Sampling). The default is "gbdt"."
+        "objective": "binary",
+        "num_leaves": args.num_leaves,
+        "early_stopping_round": 30,
+        "max_depth": args.max_depth,
+        "learning_rate": args.lr,
+        "metric": "binary_logloss", #"metric": "l2","binary_logloss"  "auc"
+        "verbose": 1,
+    }  
+    
+    lgb_train = lgb.Dataset(x_train, y_train) 
+    lgb_valid = lgb.Dataset(x_valid, y_valid)
+    lgb_eval = lgb.Dataset(x_test, y_test, reference=lgb_train)
+
+    # fitting the model
+    model = lgb.train(
+        params, train_set=lgb_train, valid_sets=lgb_valid)
+    
+    # prediction
+    valid_y_pred = model.predict(x_valid)
+    test_y_pred = model.predict(x_test)
+    
+    # predict liver fibrosis
+    predictions_df = pd.DataFrame(test_y_pred, columns=["Prediction_score"])
+    # data_test = pd.read_csv('/nfs/dpa_pretrain/data/downstream/GDA_Data/test_tdc.csv')
+    data_test = pd.read_csv(args.input_csv_path)
+    predictions = pd.concat([data_test, predictions_df], axis=1)
+    # filtered_dataset = test_dataset_with_predictions[test_dataset_with_predictions['diseaseId'] == 'C0009714']
+    predictions.sort_values(by='Prediction_score', ascending=False, inplace=True)
+    top_100_predictions = predictions.head(100)
+    top_100_predictions.to_csv(args.output_csv_path, index=False)
+    
+    # Accuracy
+    y_pred = model.predict(x_test, num_iteration=model.best_iteration)
+    y_pred[y_pred >= 0.5] = 1
+    y_pred[y_pred < 0.5] = 0
+    accuracy = accuracy_score(y_test, y_pred)
+    
+    # AUC
+    valid_roc_auc_score = metrics.roc_auc_score(y_valid, valid_y_pred)
+    valid_average_precision_score = metrics.average_precision_score(
+        y_valid, valid_y_pred
+    )
+    test_roc_auc_score = metrics.roc_auc_score(y_test, test_y_pred)
+    test_average_precision_score = metrics.average_precision_score(y_test, test_y_pred)
+    
+    # AUPR
+    valid_aupr = metrics.average_precision_score(y_valid, valid_y_pred)
+    test_aupr = metrics.average_precision_score(y_test, test_y_pred)
+
+    # Fmax
+    valid_precision, valid_recall, valid_thresholds = precision_recall_curve(y_valid, valid_y_pred)
+    valid_fmax = (2 * valid_precision * valid_recall / (valid_precision + valid_recall)).max()
+    test_precision, test_recall, test_thresholds = precision_recall_curve(y_test, test_y_pred)
+    test_fmax = (2 * test_precision * test_recall / (test_precision + test_recall)).max()
+
+    # F1
+    valid_f1 = f1_score(y_valid, valid_y_pred >= 0.5)
+    test_f1 = f1_score(y_test, test_y_pred >= 0.5)
+
+
+if __name__ == "__main__":
+    args = parse_config()
+    if torch.cuda.is_available():
+        print("cuda is available.")
+        print(f"current device {args}.")
+    else:
+        args.device = "cpu"
+    timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S")
+    random_str = "".join([random.choice(string.ascii_lowercase) for n in range(6)])
+    best_model_dir = (
+        f"{args.save_path_prefix}{args.save_name}_{timestamp_str}_{random_str}/"
+    )
+    os.makedirs(best_model_dir)
+    args.save_name = best_model_dir
+    train(args)

src/utils/data_loader.py

@@ -0,0 +1,328 @@
+import logging
+import sys
+import numpy as np
+sys.path.append("../")
+# from tdc.multi_pred import GDA
+import pandas as pd
+from torch.utils.data import Dataset
+
+LOGGER = logging.getLogger(__name__)
+
+class GDA_Dataset(Dataset):
+    """
+    Candidate Dataset for:
+        ALL gene-to-disease interactions
+    """ 
+    def __init__(self, data_examples):
+        self.protein_seqs = data_examples[0]
+        self.disease_dess = data_examples[1]
+        self.scores = data_examples[2]
+
+    def __getitem__(self, query_idx):
+
+        protein_seq = self.protein_seqs[query_idx]
+        disease_des = self.disease_dess[query_idx]
+        score = self.scores[query_idx]
+
+        return protein_seq, disease_des, score
+
+    def __len__(self):
+        return len(self.protein_seqs)
+
+
+class TDC_Pretrain_Dataset(Dataset):
+    """
+        Dataset of TDC:
+            ALL gene-disease associations
+    """
+    def __init__(self, data_dir="../../data/pretrain/", test=False):
+        LOGGER.info("Initializing TDC Pretraining Dataset ! ...")
+        
+        data = GDA(name="DisGeNET")  # , path=data_dir
+        data.neg_sample(frac = 1)
+        data.binarize(threshold = 0, order = 'ascending')
+        self.datasets = data.get_split()
+        self.name = "DisGeNET"
+        self.dataset_df = self.datasets['train']
+        # self.dataset_df = pd.read_csv(f"{data_dir}/disgenet_gda.csv")
+        self.dataset_df = self.dataset_df[
+            ["Gene", "Disease", "Y"]
+        ].dropna()  # Drop missing values.
+        # print(self.dataset_df.head())
+        print(
+            f"{data_dir}TDC training dataset loaded, found associations: {len(self.dataset_df.index)}"
+        )
+        self.protein_seqs = self.dataset_df["Gene"].values
+        self.disease_dess = self.dataset_df["Disease"].values
+        self.scores = len(self.dataset_df["Y"].values) * [1]
+
+    def __getitem__(self, query_idx):
+
+        protein_seq = self.protein_seqs[query_idx]
+        disease_des = self.disease_dess[query_idx]
+        score = self.scores[query_idx]
+
+        return protein_seq, disease_des, score
+
+    def __len__(self):
+        return len(self.protein_seqs)
+
+class GDA_Pretrain_Dataset(Dataset):
+    """
+    Candidate Dataset for:
+        ALL gene-disease associations
+    """
+
+    def __init__(self, data_dir="../../data/pretrain/", test=False, split="train", val_ratio=0.2):
+        LOGGER.info("Initializing GDA Pretraining Dataset ! ...")
+        self.dataset_df = pd.read_csv(f"{data_dir}/disgenet_gda.csv")
+        self.dataset_df = self.dataset_df[["proteinSeq", "diseaseDes", "score"]].dropna()
+        self.dataset_df = self.dataset_df.sample(frac=1, random_state=42).reset_index(drop=True)
+
+        num_val_samples = int(len(self.dataset_df) * val_ratio)
+        if split == "train":
+            self.dataset_df = self.dataset_df[:-num_val_samples]
+            print(f"{data_dir}disgenet_gda.csv loaded, found train associations: {len(self.dataset_df.index)}")
+        elif split == "val":
+            self.dataset_df = self.dataset_df[-num_val_samples:]
+            print(f"{data_dir}disgenet_gda.csv loaded, found valid associations: {len(self.dataset_df.index)}")
+            
+        if test:
+            self.protein_seqs = self.dataset_df["proteinSeq"].values[:128]
+            self.disease_dess = self.dataset_df["diseaseDes"].values[:128]
+            self.scores = 128 * [1]
+        else:
+            self.protein_seqs = self.dataset_df["proteinSeq"].values
+            self.disease_dess = self.dataset_df["diseaseDes"].values
+            self.scores = len(self.dataset_df["score"].values) * [1]
+
+    def __getitem__(self, query_idx):
+
+        protein_seq = self.protein_seqs[query_idx]
+        disease_des = self.disease_dess[query_idx]
+        score = self.scores[query_idx]
+
+        return protein_seq, disease_des, score
+
+    def __len__(self):
+        return len(self.protein_seqs)
+#         # 分离正负样本
+#         positive_samples = self.dataset_df[self.dataset_df["score"] == 1]
+#         negative_samples = self.dataset_df[self.dataset_df["score"] == 0]
+
+#         # 打乱并划分正样本
+#         positive_samples = positive_samples.sample(frac=1, random_state=42).reset_index(drop=True)
+#         num_pos_val_samples = int(len(positive_samples) * val_ratio)
+
+#         # 打乱并划分负样本
+#         negative_samples = negative_samples.sample(frac=1, random_state=42).reset_index(drop=True)
+#         num_neg_val_samples = int(len(negative_samples) * val_ratio)
+
+        # if split == "train":
+        #     self.dataset_df = pd.concat([positive_samples[:-num_pos_val_samples], negative_samples[:-num_neg_val_samples]])
+        #     print(f"{data_dir}disgenet_gda.csv loaded, found associations: {len(self.dataset_df.index)}")
+        # elif split == "val":
+        #     self.dataset_df = pd.concat([positive_samples[-num_pos_val_samples:], negative_samples[-num_neg_val_samples:]])
+        #     print(f"{data_dir}disgenet_gda.csv loaded, found associations: {len(self.dataset_df.index)}")
+        # Shuffle and split data
+
+# class GDA_Pretrain_Dataset(Dataset):
+#     """
+#     Candidate Dataset for:
+#         ALL gene-disease associations
+#     """
+
+#     def __init__(self, data_dir="../../data/pretrain/", test=False):
+#         LOGGER.info("Initializing GDA Pretraining Dataset ! ...")
+        # updated = pd.read_csv(f"{data_dir}/disgenet_updated.csv")     
+        
+        # data = GDA(name="DisGeNET")
+        # data = data.get_data()
+        # data = data[['Gene_ID','Disease_ID']].dropna()
+        # self.dataset_df = pd.read_csv(f"{data_dir}/disgenet_gda.csv")
+        
+        # num_unique_diseaseId = self.dataset_df['diseaseId'].nunique()
+        # num_unique_geneId = self.dataset_df['geneId'].nunique()
+
+        # print(f"Number of unique 'diseaseId': {num_unique_diseaseId}")
+        # print(f"Number of unique 'geneId': {num_unique_geneId}")
+        
+#         num_of_c0002395 = self.dataset_df[self.dataset_df['diseaseId'] == 'C0002395'].shape[0]
+        # print(f"Alzheimer Number in 2020：{num_of_c0002395}")
+        
+        # Convert 'Gene_ID' and 'Disease_ID' to str before merge
+        # data['Gene_ID'] = data['Gene_ID'].astype(str)
+        # data['Disease_ID'] = data['Disease_ID'].astype(str)
+
+        # Similarly for 'geneId' and 'diseaseId', if they're not already of type 'str'
+        # self.dataset_df['geneId'] = self.dataset_df['geneId'].astype(str)
+        # self.dataset_df['diseaseId'] = self.dataset_df['diseaseId'].astype(str)
+
+#         # 合并两个DataFrame并找出不同的行
+#         merged = df.merge(self.dataset_df, how='outer', indicator=True)
+#         differences = merged[merged['_merge'] != 'both']
+
+#         differences.to_csv('/nfs/dpa_pretrain/data/pretrain/differences.csv', index=False)
+
+        
+#         Check for overlap between TDC dataset and DisGeNET dataset
+#         merged_df = pd.merge(data, self.dataset_df, how='inner', left_on=['Gene_ID','Disease_ID'], right_on=['geneId','diseaseId'])
+        
+#         num_matched_pairs = merged_df.shape[0]
+
+#         print(f"Number of matched pairs TDC: {num_matched_pairs}")
+        
+#         merged_dis = pd.merge(data, updated, how='inner', left_on=['Gene','Disease'], right_on=['proteinSeq','diseaseDes'])
+        
+#         num_matched = merged_dis.shape[0]
+
+#         print(f"Number of matched pairs DisGeNET_test: {num_matched}")
+        
+        # self.dataset_df = self.dataset_df[
+        #     ["proteinSeq", "diseaseDes", "score"]
+        # ].dropna()  # Drop missing values.
+        # print(self.dataset_df.head())  "proteinSeq", "diseaseDes", "score"
+        
+        # print(
+        #     f"{data_dir}disgenet_gda.csv loaded, found associations: {len(self.dataset_df.index)}"
+        # )
+#         df1 = pd.read_csv(f"{data_dir}/disgenet_gda.csv")
+#         df1 = df1[
+#             ["proteinSeq", "diseaseDes", "score"]
+#         ].dropna()
+
+#         # 合并两个DataFrame并找出不同的行
+#         merged = df1.merge(self.dataset_df, how='outer', indicator=True)
+#         differences = merged[merged['_merge'] != 'both']
+
+#         # 将结果保存到新的文件中
+#         differences.to_csv('/nfs/dpa_pretrain/data/pretrain/differences.csv', index=False)
+
+#         if test:
+#             self.protein_seqs = self.dataset_df["proteinSeq"].values[:128]
+#             self.disease_dess = self.dataset_df["diseaseDes"].values[:128]
+#             self.scores = 128 * [1]
+#         else:
+#             self.protein_seqs = self.dataset_df["proteinSeq"].values
+#             self.disease_dess = self.dataset_df["diseaseDes"].values
+#             self.scores = len(self.dataset_df["score"].values) * [1]
+
+#     def __getitem__(self, query_idx):
+
+#         protein_seq = self.protein_seqs[query_idx]
+#         disease_des = self.disease_dess[query_idx]
+#         score = self.scores[query_idx]
+
+#         return protein_seq, disease_des, score
+
+#     def __len__(self):
+#         return len(self.protein_seqs)
+
+
+class PPI_Pretrain_Dataset(Dataset):
+    """
+    Candidate Dataset for:
+        ALL protein-to-protein interactions
+    """
+
+    def __init__(self, data_dir="../../data/pretrain/", test=False):
+        LOGGER.info("Initializing metric learning data set! ...")
+        self.dataset_df = pd.read_csv(f"{data_dir}/string_ppi_900_2m.csv")
+        self.dataset_df = self.dataset_df[["item_seq_a", "item_seq_b", "score"]]
+        self.dataset_df = self.dataset_df.dropna()
+        if test:
+            self.dataset_df = self.dataset_df.sample(100)
+        print(
+            f"{data_dir}/string_ppi_900_2m.csv loaded, found interactions: {len(self.dataset_df.index)}"
+        )
+        self.protein_seq1 = self.dataset_df["item_seq_a"].values
+        self.protein_seq2 = self.dataset_df["item_seq_b"].values
+        self.scores = len(self.dataset_df["score"].values) * [1]
+
+    def __getitem__(self, query_idx):
+
+        protein_seq1 = self.protein_seq1[query_idx]
+        protein_seq2 = self.protein_seq2[query_idx]
+        score = self.scores[query_idx]
+
+        return protein_seq1, protein_seq2, score
+
+    def __len__(self):
+        return len(self.protein_seq1)
+
+
+class PPI_Dataset(Dataset):
+    """
+    Candidate Dataset for:
+        ALL protein-to-protein interactions
+    """
+
+    def __init__(self, protein_seq1, protein_seq2, score):
+        self.protein_seq1 = protein_seq1
+        self.protein_seq2 = protein_seq2
+        self.scores = score
+
+    def __getitem__(self, query_idx):
+
+        protein_seq1 = self.protein_seq1[query_idx]
+        protein_seq2 = self.protein_seq2[query_idx]
+        score = self.scores[query_idx]
+
+        return protein_seq1, protein_seq2, score
+
+    def __len__(self):
+        return len(self.protein_seq1)
+
+
+class DDA_Dataset(Dataset):
+    """
+    Candidate Dataset for:
+        ALL disease-to-disease associations
+    """
+
+    def __init__(self, diseaseDes1, diseaseDes2, label):
+        self.diseaseDes1 = diseaseDes1
+        self.diseaseDes2 = diseaseDes2
+        self.label = label
+
+    def __getitem__(self, query_idx):
+
+        diseaseDes1 = self.diseaseDes1[query_idx]
+        diseaseDes2 = self.diseaseDes2[query_idx]
+        label = self.label[query_idx]
+
+        return diseaseDes1, diseaseDes2, label
+
+    def __len__(self):
+        return len(self.diseaseDes1)
+
+
+class DDA_Pretrain_Dataset(Dataset):
+    """
+    Candidate Dataset for:
+        ALL protein-to-protein interactions
+    """
+
+    def __init__(self, data_dir="../../data/pretrain/", test=False):
+        LOGGER.info("Initializing metric learning data set! ...")
+        self.dataset_df = pd.read_csv(f"{data_dir}disgenet_dda.csv")
+        self.dataset_df = self.dataset_df.dropna()  # Drop missing values.
+        if test:
+            self.dataset_df = self.dataset_df.sample(100)
+        print(
+            f"{data_dir}disgenet_dda.csv loaded, found associations: {len(self.dataset_df.index)}"
+        )
+        self.disease_des1 = self.dataset_df["diseaseDes1"].values
+        self.disease_des2 = self.dataset_df["diseaseDes2"].values
+        self.scores = len(self.dataset_df["jaccard_variant"].values) * [1]
+
+    def __getitem__(self, query_idx):
+
+        disease_des1 = self.disease_des1[query_idx]
+        disease_des2 = self.disease_des2[query_idx]
+        score = self.scores[query_idx]
+
+        return disease_des1, disease_des2, score
+
+    def __len__(self):
+        return len(self.disease_des1)

src/utils/downstream_disgenet.py

@@ -11,9 +11,9 @@ import pandas as pd
 sys.path.append("../")
 
 class DisGeNETProcessor:
-    def __init__(self,input_csv_path, data_dir="/nfs/dpa_pretrain/data/downstream/"):
-        train_data = pd.read_csv('/nfs/dpa_pretrain/data/downstream/GDA_Data/train.csv')
-        valid_data = pd.read_csv('/nfs/dpa_pretrain/data/downstream/GDA_Data/valid.csv')
+    def __init__(self,input_csv_path):
+        train_data = pd.read_csv('data/downstream/GDA_Data/train.csv')
+        valid_data = pd.read_csv('data/downstream/GDA_Data/valid.csv')
         test_data = pd.read_csv(input_csv_path)
         
         # test_data = pd.read_csv('/nfs/dpa_pretrain/data/downstream/GDA_Data/test.csv')

src/utils/gd_model.py

@@ -0,0 +1,77 @@
+import sys
+
+import torch
+import torch.nn as nn
+
+sys.path.append("../")
+
+
+class GDANet(torch.nn.Module):
+    def __init__(
+        self,
+        prot_encoder,
+        disease_encoder,
+    ):
+        """_summary_
+
+        Args:
+            prot_encoder (_type_): _description_
+            disease_encoder (_type_): _description_
+            prot_out_dim (int, optional): _description_. Defaults to 1024.
+            disease_out_dim (int, optional): _description_. Defaults to 768.
+            drop_out (int, optional): _description_. Defaults to 0.
+            freeze_prot_encoder (bool, optional): _description_. Defaults to True.
+            freeze_disease_encoder (bool, optional): _description_. Defaults to True.
+        """
+        super(GDANet, self).__init__()
+        self.prot_encoder = prot_encoder
+        self.disease_encoder = disease_encoder
+        self.cls = None
+        self.reg = None
+
+    def add_regression_head(self, prot_out_dim=1024, disease_out_dim=768):
+        """Add regression head.
+
+        Args:
+            prot_out_dim (_type_): protein encoder output dimension.
+            disease_out_dim (_type_): disease encoder output dimension.
+            drop_out (int, optional): dropout rate. Defaults to 0.
+        """
+        self.reg = nn.Linear(prot_out_dim + disease_out_dim, 1)
+    
+
+    def add_classification_head(
+        self, prot_out_dim=1024, disease_out_dim=768, out_dim=2
+    ):
+        """Add classification head.
+
+        Args:
+            prot_out_dim (_type_): protein encoder output dimension.
+            disease_out_dim (_type_): disease encoder output dimension.
+            out_dim (int, optional): output dimension. Defaults to 2.
+            drop_out (int, optional): dropout rate. Defaults to 0.
+        """
+        self.cls = nn.Linear(prot_out_dim + disease_out_dim, out_dim)
+        
+
+    def freeze_encoders(self, freeze_prot_encoder, freeze_disease_encoder):
+        """Freeze encoders.
+
+        Args:
+            freeze_prot_encoder (boolean): freeze protein encoder
+            freeze_disease_encoder (boolean): freeze disease textual encoder
+        """
+        if freeze_prot_encoder:
+            for param in self.prot_encoder.parameters():
+                param.requires_grad = False
+        else:
+            for param in self.disease_encoder.parameters():
+                param.requires_grad = True
+        if freeze_disease_encoder:
+            for param in self.disease_encoder.parameters():
+                param.requires_grad = False
+        else:
+            for param in self.disease_encoder.parameters():
+                param.requires_grad = True
+        print(f"freeze_prot_encoder:{freeze_prot_encoder}")
+        print(f"freeze_disease_encoder:{freeze_disease_encoder}")