FusionGDA

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst 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

.ipynb_checkpoints/gda_api-checkpoint.py

@@ -0,0 +1,60 @@
+# -*- 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()]
+        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,5 @@
+lightgbm
+pytorch-metric-learning
+torch
+transformers
+PyTDC

data/pretrain/disgenet_latest.csv

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

gda_api.py

@@ -0,0 +1,60 @@
+# -*- 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()]
+        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)

model.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

requirements.txt

@@ -0,0 +1,5 @@
+lightgbm
+pytorch-metric-learning
+torch
+transformers
+PyTDC

save_model_ckp/gda_infoNCE_2024_GPU3090/step_300_model.bin

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

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/finetune/finetune.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/downstream_disgenet.py

@@ -0,0 +1,148 @@
+import json
+import sys
+import os
+import torch
+from utils.data_loader import GDA_Dataset
+from sklearn.model_selection import train_test_split
+from sklearn.model_selection import KFold
+import numpy as np
+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')
+        test_data = pd.read_csv(input_csv_path)
+        
+        # test_data = pd.read_csv('/nfs/dpa_pretrain/data/downstream/GDA_Data/test.csv')
+        # valid_data, test_data = train_test_split(valid_data, test_size=1/3, random_state=42)
+        # train_data = pd.read_csv('/nfs/dpa_pretrain/data/downstream/test/train.csv')
+        # valid_data = pd.read_csv('/nfs/dpa_pretrain/data/downstream/test/valid.csv')
+
+        
+        # train_data = pd.read_csv('/nfs/dpa_pretrain/data/downstream/disgenet_finetune.csv')
+        # train_data, valid_data = train_test_split(train_data, test_size=0.2, random_state=42)
+        # valid_data, test_data = train_test_split(valid_data, test_size=1/3, random_state=42)
+        
+        # alzheimer and stomach dataset use [["proteinSeq", "diseaseDes", "Y"]].dropna()
+        
+        self.name = "DisGeNET"
+        self.train_dataset_df = train_data[["proteinSeq", "diseaseDes", "score"]].dropna()
+        self.val_dataset_df = valid_data[["proteinSeq", "diseaseDes", "score"]].dropna()
+        self.test_dataset_df = test_data[["proteinSeq", "diseaseDes", "score"]].dropna()
+        # self.test_dataset_df = test_data[["proteinSeq", "diseaseDes", "Y"]].dropna()
+
+        
+    def get_train_examples(self, test=False):
+        """get training examples
+
+        Args:
+            test (bool, optional): test can be int or bool. If test>1, will take test as the number of test examples. Defaults to False.
+
+        Returns:
+            _type_: _description_
+        """
+        if test == 1:  # Small testing set, to reduce the running time
+            return (
+                self.train_dataset_df["proteinSeq"].values[:4096],
+                self.train_dataset_df["diseaseDes"].values[:4096],
+                self.train_dataset_df["score"].values[:4096],
+            )
+        elif test > 1:
+            return (
+                self.train_dataset_df["proteinSeq"].values[:test],
+                self.train_dataset_df["diseaseDes"].values[:test],
+                self.train_dataset_df["score"].values[:test],
+            )
+        else:
+            return GDA_Dataset( (
+                self.train_dataset_df["proteinSeq"].values,
+                self.train_dataset_df["diseaseDes"].values,
+                self.train_dataset_df["score"].values,
+            ))
+
+    def get_val_examples(self, test=False):
+        """get validation examples
+
+        Args:
+            test (bool, optional): test can be int or bool. If test>1, will take test as the number of test examples. Defaults to False.
+
+        Returns:
+            _type_: _description_
+
+        """
+        if test == 1:  # Small testing set, to reduce the running time
+            return (
+                self.val_dataset_df["proteinSeq"].values[:1024],
+                self.val_dataset_df["diseaseDes"].values[:1024],
+                self.val_dataset_df["score"].values[:1024],
+            )
+        elif test > 1:
+            return (
+                self.val_dataset_df["proteinSeq"].values[:test],
+                self.val_dataset_df["diseaseDes"].values[:test],
+                self.val_dataset_df["score"].values[:test],
+            )
+        else:
+            return GDA_Dataset((
+                self.val_dataset_df["proteinSeq"].values,
+                self.val_dataset_df["diseaseDes"].values,
+                self.val_dataset_df["score"].values,
+            ))
+
+#     def get_test_examples(self, test=False):
+#         """get test examples
+
+#         Args:
+#             test (bool, optional): test can be int or bool. If test>1, will take test as the number of test examples. Defaults to False.
+
+#         Returns:
+#             _type_: _description_
+#         """
+#         if test == 1:  # Small testing set, to reduce the running time
+#             return (
+#                 self.test_dataset_df["proteinSeq"].values[:1024],
+#                 self.test_dataset_df["diseaseDes"].values[:1024],
+#                 self.test_dataset_df["Y"].values[:1024],
+#             )
+#         elif test > 1:
+#             return (
+#                 self.test_dataset_df["proteinSeq"].values[:test],
+#                 self.test_dataset_df["diseaseDes"].values[:test],
+#                 self.test_dataset_df["Y"].values[:test],
+#             )
+#         else:
+#             return GDA_Dataset( (
+#                 self.test_dataset_df["proteinSeq"].values,
+#                 self.test_dataset_df["diseaseDes"].values,
+#                 self.test_dataset_df["Y"].values,
+#             ))
+    def get_test_examples(self, test=False):
+        """get test examples
+
+        Args:
+            test (bool, optional): test can be int or bool. If test>1, will take test as the number of test examples. Defaults to False.
+
+        Returns:
+            _type_: _description_
+        """
+        if test == 1:  # Small testing set, to reduce the running time
+            return (
+                self.test_dataset_df["proteinSeq"].values[:1024],
+                self.test_dataset_df["diseaseDes"].values[:1024],
+                self.test_dataset_df["score"].values[:1024],
+            )
+        elif test > 1:
+            return (
+                self.test_dataset_df["proteinSeq"].values[:test],
+                self.test_dataset_df["diseaseDes"].values[:test],
+                self.test_dataset_df["score"].values[:test],
+            )
+        else:
+            return GDA_Dataset( (
+                self.test_dataset_df["proteinSeq"].values,
+                self.test_dataset_df["diseaseDes"].values,
+                self.test_dataset_df["score"].values,
+            ))

src/utils/metric_learning_models.py

@@ -0,0 +1,869 @@
+import logging
+import os
+import sys
+
+sys.path.append("../")
+from pytorch_metric_learning.distances import CosineSimilarity
+from pytorch_metric_learning.reducers import ThresholdReducer
+from pytorch_metric_learning.regularizers import LpRegularizer
+from pytorch_metric_learning import losses
+
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from pytorch_metric_learning import losses, miners
+from torch.cuda.amp import autocast
+from torch.nn import Module
+from tqdm import tqdm
+from utils.gd_model import GDANet
+from torch.nn import MultiheadAttention
+
+from transformers import BertModel
+from transformers import EsmModel, EsmConfig
+
+LOGGER = logging.getLogger(__name__)
+    
+class FusionModule(nn.Module):
+    def __init__(self, out_dim, num_head, dropout= 0.1):
+        super(FusionModule, self).__init__()
+        """FusionModule.
+
+        Args:
+            dropout= 0.1 is defaut
+            out_dim: model output dimension
+            num_head = 8: Multi-head Attention
+        """
+
+        self.out_dim = out_dim
+        self.num_head = num_head
+
+        self.WqS = nn.Linear(out_dim, out_dim)
+        self.WkS = nn.Linear(out_dim, out_dim)
+        self.WvS = nn.Linear(out_dim, out_dim)
+
+        self.WqT = nn.Linear(out_dim, out_dim)
+        self.WkT = nn.Linear(out_dim, out_dim)
+        self.WvT = nn.Linear(out_dim, out_dim)
+        self.multi_head_attention = nn.MultiheadAttention(out_dim, num_head, dropout=dropout)
+
+    def forward(self, zs, zt):
+        #  nn.MultiheadAttention The input representation is (token_length, batch_size, out_dim)
+        # zs = protein_representation.permute(1, 0, 2)
+        # zt = disease_representation.permute(1, 0, 2)   
+        
+        # Compute query, key and value representations
+        qs = self.WqS(zs)
+        ks = self.WkS(zs)
+        vs = self.WvS(zs)
+
+        qt = self.WqT(zt)
+        kt = self.WkT(zt)
+        vt = self.WvT(zt)
+        
+        #self.multi_head_attention() The function returns two values: the representation and the attention weight matrix, computed after multiple attentions. In this case, we only care about the computed representation and not the attention weight matrix, so "_" is used to indicate that we do not intend to use or store the second return value.
+        zs_attention1, _ = self.multi_head_attention(qs, ks, vs)
+        zs_attention2, _ = self.multi_head_attention(qs, kt, vt)
+        zt_attention1, _ = self.multi_head_attention(qt, kt, vt)
+        zt_attention2, _ = self.multi_head_attention(qt, ks, vs)
+
+        protein_fused = 0.5 * (zs_attention1 + zs_attention2)
+        dis_fused = 0.5 * (zt_attention1 + zt_attention2)
+        
+        return protein_fused, dis_fused
+
+class CrossAttentionBlock(nn.Module):
+
+    def __init__(self, hidden_dim, num_heads):
+        super(CrossAttentionBlock, self).__init__()
+        if hidden_dim % num_heads != 0:
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number of attention "
+                "heads (%d)" % (hidden_dim, num_heads))
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.head_size = hidden_dim // num_heads
+
+        self.query1 = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.key1 = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.value1 = nn.Linear(hidden_dim, hidden_dim, bias=False)
+
+        self.query2 = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.key2 = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.value2 = nn.Linear(hidden_dim, hidden_dim, bias=False)
+
+    def _alpha_from_logits(self, logits, mask_row, mask_col, inf=1e6):
+        N, L1, L2, H = logits.shape
+        mask_row = mask_row.view(N, L1, 1).repeat(1, 1, H)
+        mask_col = mask_col.view(N, L2, 1).repeat(1, 1, H)
+        mask_pair = torch.einsum('blh, bkh->blkh', mask_row, mask_col)
+
+        logits = torch.where(mask_pair, logits, logits - inf)
+        alpha = torch.softmax(logits, dim=2)
+        mask_row = mask_row.view(N, L1, 1, H).repeat(1, 1, L2, 1)
+        alpha = torch.where(mask_row, alpha, torch.zeros_like(alpha))
+        return alpha
+
+    def _heads(self, x, n_heads, n_ch):
+        s = list(x.size())[:-1] + [n_heads, n_ch]
+        return x.view(*s)
+
+    def forward(self, input1, input2, mask1, mask2):
+        query1 = self._heads(self.query1(input1), self.num_heads, self.head_size)
+        key1 = self._heads(self.key1(input1), self.num_heads, self.head_size)
+        query2 = self._heads(self.query2(input2), self.num_heads, self.head_size)
+        key2 = self._heads(self.key2(input2), self.num_heads, self.head_size)
+        logits11 = torch.einsum('blhd, bkhd->blkh', query1, key1)
+        logits12 = torch.einsum('blhd, bkhd->blkh', query1, key2)
+        logits21 = torch.einsum('blhd, bkhd->blkh', query2, key1)
+        logits22 = torch.einsum('blhd, bkhd->blkh', query2, key2)
+
+        alpha11 = self._alpha_from_logits(logits11, mask1, mask1)
+        alpha12 = self._alpha_from_logits(logits12, mask1, mask2)
+        alpha21 = self._alpha_from_logits(logits21, mask2, mask1)
+        alpha22 = self._alpha_from_logits(logits22, mask2, mask2)
+
+        value1 = self._heads(self.value1(input1), self.num_heads, self.head_size)
+        value2 = self._heads(self.value2(input2), self.num_heads, self.head_size)
+        output1 = (torch.einsum('blkh, bkhd->blhd', alpha11, value1).flatten(-2) +
+                   torch.einsum('blkh, bkhd->blhd', alpha12, value2).flatten(-2)) / 2
+        output2 = (torch.einsum('blkh, bkhd->blhd', alpha21, value1).flatten(-2) +
+                   torch.einsum('blkh, bkhd->blhd', alpha22, value2).flatten(-2)) / 2
+
+        return output1, output2
+
+class GDA_Metric_Learning(GDANet):
+    def __init__(
+            self, prot_encoder, disease_encoder, prot_out_dim, disease_out_dim, args
+    ):
+        """Constructor for the model.
+
+        Args:
+            prot_encoder (_type_): Protein encoder.
+            disease_encoder (_type_): Disease Textual encoder.
+            prot_out_dim (_type_): Dimension of the Protein encoder.
+            disease_out_dim (_type_): Dimension of the Disease encoder.
+            args (_type_): _description_
+        """
+        super(GDA_Metric_Learning, self).__init__(
+            prot_encoder,
+            disease_encoder,
+        )
+        self.prot_encoder = prot_encoder
+        self.disease_encoder = disease_encoder
+        self.loss = args.loss
+        self.use_miner = args.use_miner
+        self.miner_margin = args.miner_margin
+        self.agg_mode = args.agg_mode
+        self.prot_reg = nn.Linear(prot_out_dim, 1024)
+        # self.prot_reg = nn.Linear(prot_out_dim, disease_out_dim)
+        self.dis_reg = nn.Linear(disease_out_dim, 1024)
+        # self.prot_adapter_name = None
+        # self.disease_adapter_name = None
+        
+        self.fusion_layer = FusionModule(1024, num_head=8)
+        self.cross_attention_layer = CrossAttentionBlock(1024, 8)
+        
+        # # MMP Prediction Heads
+        # self.prot_pred_head = nn.Sequential(
+        #     nn.Linear(disease_out_dim, disease_out_dim),
+        #     nn.ReLU(),
+        #     nn.Linear(disease_out_dim, 1280)  #vocabulary size : prot model tokenize length 30   446
+        # )
+        # self.dise_pred_head = nn.Sequential(
+        #     nn.Linear(disease_out_dim, disease_out_dim),
+        #     nn.ReLU(),
+        #     nn.Linear(disease_out_dim, 768) #vocabulary size : disease model tokenize length 30522
+        # )
+        
+        if self.use_miner:
+            self.miner = miners.TripletMarginMiner(
+                margin=args.miner_margin, type_of_triplets="all"
+            )
+        else:
+            self.miner = None
+
+        if self.loss == "ms_loss":
+            self.loss = losses.MultiSimilarityLoss(
+                alpha=2, beta=50, base=0.5
+            )  # 1,2,3; 40,50,60
+            #1_40=1.5141 50=1.4988 60=1.4905 2_60=1.1786 50=1.1874 40=1.2008 3_40=1.1146 50=1.1012
+        elif self.loss == "circle_loss":
+            self.loss = losses.CircleLoss(
+                m=0.4, gamma=80
+            )
+        elif self.loss == "triplet_loss":
+            self.loss = losses.TripletMarginLoss( 
+                margin=0.05, swap=False, smooth_loss=False,
+                triplets_per_anchor="all")
+             # distance = CosineSimilarity(), 
+             # reducer = ThresholdReducer(high=0.3), 
+             # embedding_regularizer = LpRegularizer()  )
+            
+        elif self.loss == "infoNCE":
+            self.loss = losses.NTXentLoss(
+                temperature=0.07
+            )  # The MoCo paper uses 0.07, while SimCLR uses 0.5.
+        elif self.loss == "lifted_structure_loss":
+            self.loss = losses.LiftedStructureLoss(
+                neg_margin=1, pos_margin=0
+            )
+        elif self.loss == "nca_loss":
+            self.loss = losses.NCALoss(
+                softmax_scale=1
+            )
+        self.fusion = False
+        # self.stack = False
+        self.dropout = torch.nn.Dropout(args.dropout)
+        print("miner:", self.miner)
+        print("loss:", self.loss)
+
+    # def add_fusion(self):
+    #     adapter_setup = Fuse("prot_adapter", "disease_adapter")
+    #     self.prot_encoder.add_fusion(adapter_setup)
+    #     self.prot_encoder.set_active_adapters(adapter_setup)
+    #     self.prot_encoder.train_fusion(adapter_setup)
+    #     self.disease_encoder.add_fusion(adapter_setup)
+    #     self.disease_encoder.set_active_adapters(adapter_setup)
+    #     self.disease_encoder.train_fusion(adapter_setup)
+    #     self.fusion = True
+
+    # def add_stack_gda(self, reduction_factor):
+    #     self.add_gda_adapters(reduction_factor=reduction_factor)
+    #     # adapter_setup = Fuse("prot_adapter", "disease_adapter")
+    #     self.prot_encoder.active_adapters = Stack(
+    #         self.prot_adapter_name, self.gda_adapter_name
+    #     )
+    #     self.disease_encoder.active_adapters = Stack(
+    #         self.disease_adapter_name, self.gda_adapter_name
+    #     )
+    #     print("stacked adapters loaded.")
+    #     self.stack = True
+
+#     def load_adapters(
+#             self,
+#             prot_model_path,
+#             disease_model_path,
+#             prot_adapter_name="prot_adapter",
+#             disease_adapter_name="disease_adapter",
+#     ):
+#         if os.path.exists(prot_model_path):
+#             print(f"loading prot adapter from: {prot_model_path}")
+#             self.prot_adapter_name = prot_adapter_name
+#             self.prot_encoder.load_adapter(prot_model_path, load_as=prot_adapter_name)
+#             self.prot_encoder.set_active_adapters(prot_adapter_name)
+#             print(f"load protein adapters from: {prot_model_path} {prot_adapter_name}")
+#         else:
+#             print(f"{prot_model_path} not exits")
+
+#         if os.path.exists(disease_model_path):
+#             print(f"loading prot adapter from: {disease_model_path}")
+#             self.disease_adapter_name = disease_adapter_name
+#             self.disease_encoder.load_adapter(
+#                 disease_model_path, load_as=disease_adapter_name
+#             )
+#             self.disease_encoder.set_active_adapters(disease_adapter_name)
+#             print(
+#                 f"load disease adapters from: {disease_model_path} {disease_adapter_name}"
+#             )
+#         else:
+#             print(f"{disease_model_path} not exits")
+            
+    def non_adapters(
+            self,
+            prot_model_path,
+            disease_model_path,
+            
+    ):
+        if os.path.exists(prot_model_path):
+            # Load the entire model for prot_model
+            prot_model = torch.load(prot_model_path)
+            # Set the prot_encoder to the loaded model
+            self.prot_encoder = prot_model.prot_encoder
+            print(f"load protein from: {prot_model_path}")     
+        else:
+            print(f"{prot_model_path} not exits")
+
+        if os.path.exists(disease_model_path):
+            # Load the entire model for disease_model
+            disease_model = torch.load(disease_model_path)
+            # Set the disease_encoder to the loaded model
+            self.disease_encoder = disease_model.disease_encoder
+            print(f"load disease from: {disease_model_path}")
+
+        else:
+            print(f"{disease_model_path} not exits")
+
+
+#     def add_gda_adapters(
+#             self,
+#             gda_adapter_name="gda_adapter",
+#             reduction_factor=16,
+#     ):
+#         """Initialise adapters
+
+#         Args:
+#             prot_adapter_name (str, optional): _description_. Defaults to "prot_adapter".
+#             disease_adapter_name (str, optional): _description_. Defaults to "disease_adapter".
+#             reduction_factor (int, optional): _description_. Defaults to 16.
+#         """
+#         adapter_config = AdapterConfig.load(
+#             "pfeiffer", reduction_factor=reduction_factor
+#         )
+#         self.gda_adapter_name = gda_adapter_name
+#         self.prot_encoder.add_adapter(gda_adapter_name, config=adapter_config)
+#         self.prot_encoder.train_adapter([gda_adapter_name])
+#         self.disease_encoder.add_adapter(gda_adapter_name, config=adapter_config)
+#         self.disease_encoder.train_adapter([gda_adapter_name])
+
+#     def init_adapters(
+#             self,
+#             prot_adapter_name="gda_prot_adapter",
+#             disease_adapter_name="gda_disease_adapter",
+#             reduction_factor=16,
+#     ):
+#         """Initialise adapters
+
+#         Args:
+#             prot_adapter_name (str, optional): _description_. Defaults to "prot_adapter".
+#             disease_adapter_name (str, optional): _description_. Defaults to "disease_adapter".
+#             reduction_factor (int, optional): _description_. Defaults to 16.
+#         """
+#         adapter_config = AdapterConfig.load(
+#             "pfeiffer", reduction_factor=reduction_factor
+#         )
+
+#         self.prot_adapter_name = prot_adapter_name
+#         self.disease_adapter_name = disease_adapter_name
+#         self.prot_encoder.add_adapter(prot_adapter_name, config=adapter_config)
+#         self.prot_encoder.train_adapter([prot_adapter_name])
+#         self.disease_encoder.add_adapter(disease_adapter_name, config=adapter_config)
+#         self.disease_encoder.train_adapter([disease_adapter_name])
+#         print(f"adapter modules initialized")
+
+#     def save_adapters(self, save_path_prefix, total_step):
+#         """Save adapters into file.
+
+#         Args:
+#             save_path_prefix (string): saving path prefix.
+#             total_step (int): total step number.
+#         """
+#         prot_save_dir = os.path.join(
+#             save_path_prefix, f"prot_adapter_step_{total_step}"
+#         )# adapter
+#         disease_save_dir = os.path.join(
+#             save_path_prefix, f"disease_adapter_step_{total_step}"
+#         )
+#         os.makedirs(prot_save_dir, exist_ok=True)
+#         os.makedirs(disease_save_dir, exist_ok=True)
+#         self.prot_encoder.save_adapter(prot_save_dir, self.prot_adapter_name)
+#         prot_head_save_path = os.path.join(prot_save_dir, "prot_head.bin")
+#         torch.save(self.prot_reg, prot_head_save_path)
+#         self.disease_encoder.save_adapter(disease_save_dir, self.disease_adapter_name)
+#         disease_head_save_path = os.path.join(prot_save_dir, "disease_head.bin")
+#         torch.save(self.prot_reg, disease_head_save_path)
+#         if self.fusion:
+#             self.prot_encoder.save_all_adapters(prot_save_dir)
+#             self.disease_encoder.save_all_adapters(disease_save_dir)
+
+    def predict(self, query_toks1, query_toks2):
+        """
+        query : (N, h), candidates : (N, topk, h)
+        output : (N, topk)
+        """
+        # Extract input_ids and attention_mask for protein
+        prot_input_ids = query_toks1["input_ids"]
+        prot_attention_mask = query_toks1["attention_mask"]
+
+        # Extract input_ids and attention_mask for dis
+        dis_input_ids = query_toks2["input_ids"]
+        dis_attention_mask = query_toks2["attention_mask"]
+
+        # Process inputs through encoders
+        last_hidden_state1 = self.prot_encoder(
+            input_ids=prot_input_ids, attention_mask=prot_attention_mask, return_dict=True
+        ).last_hidden_state
+        last_hidden_state1 = self.prot_reg(last_hidden_state1)
+
+        last_hidden_state2 = self.disease_encoder(
+            input_ids=dis_input_ids, attention_mask=dis_attention_mask, return_dict=True
+        ).last_hidden_state
+        last_hidden_state2 = self.dis_reg(last_hidden_state2)
+       # Apply the cross-attention layer
+        prot_fused, dis_fused = self.cross_attention_layer(
+            last_hidden_state1, last_hidden_state2, prot_attention_mask, dis_attention_mask
+        )
+
+        # last_hidden_state1 = self.prot_encoder(
+        #     query_toks1, return_dict=True
+        # ).last_hidden_state
+        # last_hidden_state1 = self.prot_reg(
+        #     last_hidden_state1
+        # )  # transform the prot embedding into the same dimension as the disease embedding
+        # last_hidden_state2 = self.disease_encoder(
+        #     query_toks2, return_dict=True
+        # ).last_hidden_state
+        # last_hidden_state2 = self.dis_reg(
+        #     last_hidden_state2
+        # )  # transform the disease embedding into 1024
+        
+       # Apply the fusion layer and Recovery of representational shape
+       # prot_fused, dis_fused = self.fusion_layer(last_hidden_state1, last_hidden_state2)
+        
+        if self.agg_mode == "cls":
+            query_embed1 = prot_fused[:, 0]  # query : [batch_size, hidden]
+            query_embed2 = dis_fused[:, 0]  # query : [batch_size, hidden]
+        elif self.agg_mode == "mean_all_tok":
+            query_embed1 = prot_fused.mean(1)  # query : [batch_size, hidden]
+            query_embed2 = dis_fused.mean(1)  # query : [batch_size, hidden]
+        elif self.agg_mode == "mean":
+            query_embed1 = (
+                                   prot_fused * query_toks1["attention_mask"].unsqueeze(-1)
+                           ).sum(1) / query_toks1["attention_mask"].sum(-1).unsqueeze(-1)
+            query_embed2 = (
+                                   dis_fused * query_toks2["attention_mask"].unsqueeze(-1)
+                           ).sum(1) / query_toks2["attention_mask"].sum(-1).unsqueeze(-1)
+        else:
+            raise NotImplementedError()
+        
+        query_embed = torch.cat([query_embed1, query_embed2], dim=1)
+        return query_embed
+  
+    def forward(self, query_toks1, query_toks2, labels):
+        """
+        query : (N, h), candidates : (N, topk, h)
+        output : (N, topk)
+        """
+        # Extract input_ids and attention_mask for protein
+        prot_input_ids = query_toks1["input_ids"]
+        prot_attention_mask = query_toks1["attention_mask"]
+
+        # Extract input_ids and attention_mask for dis
+        dis_input_ids = query_toks2["input_ids"]
+        dis_attention_mask = query_toks2["attention_mask"]
+
+        # Process inputs through encoders
+        last_hidden_state1 = self.prot_encoder(
+            input_ids=prot_input_ids, attention_mask=prot_attention_mask, return_dict=True
+        ).last_hidden_state
+        last_hidden_state1 = self.prot_reg(last_hidden_state1)
+
+        last_hidden_state2 = self.disease_encoder(
+            input_ids=dis_input_ids, attention_mask=dis_attention_mask, return_dict=True
+        ).last_hidden_state
+        last_hidden_state2 = self.dis_reg(last_hidden_state2)
+       # Apply the cross-attention layer
+        prot_fused, dis_fused = self.cross_attention_layer(
+            last_hidden_state1, last_hidden_state2, prot_attention_mask, dis_attention_mask
+        )
+       #  last_hidden_state1 = self.prot_encoder(
+       #      query_toks1, return_dict=True
+       #  ).last_hidden_state
+        
+       #  last_hidden_state1 = self.prot_reg(
+       #      last_hidden_state1
+       #  )  # transform the prot embedding into the same dimension as the disease embedding
+       #  last_hidden_state2 = self.disease_encoder(
+       #      query_toks2, return_dict=True
+       #  ).last_hidden_state
+       #  last_hidden_state2 = self.dis_reg(
+       #      last_hidden_state2
+       #  )  # transform the disease embedding into 1024
+        
+       # # Apply the fusion layer and Recovery of representational shape
+       #  prot_fused, dis_fused = self.fusion_layer(last_hidden_state1, last_hidden_state2)
+        if self.agg_mode == "cls":
+            query_embed1 = prot_pred[:, 0]  # query : [batch_size, hidden]
+            query_embed2 = dise_pred[:, 0]  # query : [batch_size, hidden]
+        elif self.agg_mode == "mean_all_tok":
+            query_embed1 = prot_fused.mean(1)  # query : [batch_size, hidden]
+            query_embed2 = dis_fused.mean(1)  # query : [batch_size, hidden]
+        elif self.agg_mode == "mean":
+            query_embed1 = (
+                                   prot_pred * query_toks1["attention_mask"].unsqueeze(-1)
+                           ).sum(1) / query_toks1["attention_mask"].sum(-1).unsqueeze(-1)
+            query_embed2 = (
+                                   dis_fused * query_toks2["attention_mask"].unsqueeze(-1)
+                           ).sum(1) / query_toks2["attention_mask"].sum(-1).unsqueeze(-1)
+        else:
+            raise NotImplementedError()
+
+        # print("query_embed1 =", query_embed1.shape, "query_embed2 =", query_embed2.shape)
+        query_embed = torch.cat([query_embed1, query_embed2], dim=0)
+        # print("query_embed =", len(query_embed))
+        
+        labels = torch.cat([torch.arange(len(labels)), torch.arange(len(labels))], dim=0)
+        
+        if self.use_miner:
+            hard_pairs = self.miner(query_embed, labels)
+            return self.loss(query_embed, labels, hard_pairs)# + loss_mmp
+        else:
+            loss = self.loss(query_embed, labels)# + loss_mmp
+            # print('loss :', loss)
+            return loss
+
+    def get_embeddings(self, mentions, batch_size=1024):
+        """
+        Compute all embeddings from mention tokens.
+        """
+        embedding_table = []
+        with torch.no_grad():
+            for start in tqdm(range(0, len(mentions), batch_size)):
+                end = min(start + batch_size, len(mentions))
+                batch = mentions[start:end]
+                batch_embedding = self.vectorizer(batch)
+                batch_embedding = batch_embedding.cpu()
+                embedding_table.append(batch_embedding)
+        embedding_table = torch.cat(embedding_table, dim=0)
+        return embedding_table
+
+
+
+class DDA_Metric_Learning(Module):
+    def __init__(self, disease_encoder, args):
+        """Constructor for the model.
+
+        Args:
+            disease_encoder (_type_): disease encoder.
+            args (_type_): _description_
+        """
+        super(DDA_Metric_Learning, self).__init__()
+        self.disease_encoder = disease_encoder
+        self.loss = args.loss
+        self.use_miner = args.use_miner
+        self.miner_margin = args.miner_margin
+        self.agg_mode = args.agg_mode
+        self.disease_adapter_name = None
+        if self.use_miner:
+            self.miner = miners.TripletMarginMiner(
+                margin=args.miner_margin, type_of_triplets="all"
+            )
+        else:
+            self.miner = None
+
+        if self.loss == "ms_loss":
+            self.loss = losses.MultiSimilarityLoss(
+                alpha=1, beta=60, base=0.5
+            )  # 1,2,3; 40,50,60
+        elif self.loss == "circle_loss":
+            self.loss = losses.CircleLoss()
+        elif self.loss == "triplet_loss":
+            self.loss = losses.TripletMarginLoss()
+        elif self.loss == "infoNCE":
+            self.loss = losses.NTXentLoss(
+                temperature=0.07
+            )  # The MoCo paper uses 0.07, while SimCLR uses 0.5.
+        elif self.loss == "lifted_structure_loss":
+            self.loss = losses.LiftedStructureLoss()
+        elif self.loss == "nca_loss":
+            self.loss = losses.NCALoss()
+        self.reg = None
+        self.cls = None
+        self.dropout = torch.nn.Dropout(args.dropout)
+        print("miner:", self.miner)
+        print("loss:", self.loss)
+
+    def add_classification_head(self, disease_out_dim=768, out_dim=2):
+        """Add regression head.
+
+        Args:
+            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(disease_out_dim * 2, out_dim)
+
+    def load_disease_adapter(
+            self,
+            disease_model_path,
+            disease_adapter_name="disease_adapter",
+    ):
+        if os.path.exists(disease_model_path):
+            self.disease_adapter_name = disease_adapter_name
+            self.disease_encoder.load_adapter(
+                disease_model_path, load_as=disease_adapter_name
+            )
+            self.disease_encoder.set_active_adapters(disease_adapter_name)
+            print(
+                f"load disease adapters from: {disease_model_path} {disease_adapter_name}"
+            )
+        else:
+            print(f"{disease_adapter_name} not exits")
+
+    def init_adapters(
+            self,
+            disease_adapter_name="disease_adapter",
+            reduction_factor=16,
+    ):
+        """Initialise adapters
+
+        Args:
+            disease_adapter_name (str, optional): _description_. Defaults to "disease_adapter".
+            reduction_factor (int, optional): _description_. Defaults to 16.
+        """
+        adapter_config = AdapterConfig.load(
+            "pfeiffer", reduction_factor=reduction_factor
+        )
+        self.disease_adapter_name = disease_adapter_name
+        self.disease_encoder.add_adapter(disease_adapter_name, config=adapter_config)
+        self.disease_encoder.train_adapter([disease_adapter_name])
+
+    def save_adapters(self, save_path_prefix, total_step):
+        """Save adapters into file.
+
+        Args:
+            save_path_prefix (string): saving path prefix.
+            total_step (int): total step number.
+        """
+        disease_save_dir = os.path.join(
+            save_path_prefix, f"disease_adapter_step_{total_step}"
+        )
+        os.makedirs(disease_save_dir, exist_ok=True)
+        self.disease_encoder.save_adapter(disease_save_dir, self.disease_adapter_name)
+
+    def predict(self, x1, x2):
+        """
+        query : (N, h), candidates : (N, topk, h)
+        output : (N, topk)
+
+        """
+        if self.agg_mode == "cls":
+            x1 = self.disease_encoder(x1).last_hidden_state[:, 0]
+            x2 = self.disease_encoder(x2).last_hidden_state[:, 0]
+            x = torch.cat((x1, x2), 1)
+            return x
+        else:
+            x1 = self.disease_encoder(x1).last_hidden_state.mean(1)  # query : [batch_size, hidden]
+            x2 = self.disease_encoder(x2).last_hidden_state.mean(1)  # query : [batch_size, hidden]
+            x = torch.cat((x1, x2), 1)
+            return x
+
+    def module_predict(self, x1, x2):
+        """
+        query : (N, h), candidates : (N, topk, h)
+        output : (N, topk)
+
+        """
+        if self.agg_mode == "cls":
+            x1 = self.disease_encoder.module(x1).last_hidden_state[:, 0]
+            x2 = self.disease_encoder.module(x2).last_hidden_state[:, 0]
+            x = torch.cat((x1, x2), 1)
+            return x
+        else:
+            x1 = self.disease_encoder.module(x1).last_hidden_state.mean(1)  # query : [batch_size, hidden]
+            x2 = self.disease_encoder.module(x2).last_hidden_state.mean(1)  # query : [batch_size, hidden]
+            x = torch.cat((x1, x2), 1)
+            return x
+
+    @autocast()
+    def forward(self, query_toks1, query_toks2, labels):
+        """
+        query : (N, h), candidates : (N, topk, h)
+        output : (N, topk)
+        """
+        last_hidden_state1 = self.disease_encoder(
+            **query_toks1, return_dict=True
+        ).last_hidden_state
+        last_hidden_state2 = self.disease_encoder(
+            **query_toks2, return_dict=True
+        ).last_hidden_state
+        if self.agg_mode == "cls":
+            query_embed1 = last_hidden_state1[:, 0]  # query : [batch_size, hidden]
+            query_embed2 = last_hidden_state2[:, 0]  # query : [batch_size, hidden]
+        elif self.agg_mode == "mean_all_tok":
+            query_embed1 = last_hidden_state1.mean(1)  # query : [batch_size, hidden]
+            query_embed2 = last_hidden_state2.mean(1)  # query : [batch_size, hidden]
+        elif self.agg_mode == "mean":
+            query_embed1 = (
+                                   last_hidden_state1 * query_toks1["attention_mask"].unsqueeze(-1)
+                           ).sum(1) / query_toks1["attention_mask"].sum(-1).unsqueeze(-1)
+            query_embed2 = (
+                                   last_hidden_state2 * query_toks2["attention_mask"].unsqueeze(-1)
+                           ).sum(1) / query_toks2["attention_mask"].sum(-1).unsqueeze(-1)
+        else:
+            raise NotImplementedError()
+        query_embed = torch.cat([query_embed1, query_embed2], dim=0)
+
+        labels = torch.cat([labels, labels], dim=0)
+        if self.use_miner:
+            hard_pairs = self.miner(query_embed, labels)
+            print('miner used')
+            return self.loss(query_embed, labels, hard_pairs)
+        else:
+            print('no miner')
+            return self.loss(query_embed, labels)
+
+
+class PPI_Metric_Learning(Module):
+    def __init__(self, prot_encoder, args):
+        """Constructor for the model.
+
+        Args:
+            prot_encoder (_type_): Protein encoder.
+            prot_encoder (_type_): prot Textual encoder.
+            prot_out_dim (_type_): Dimension of the Protein encoder.
+            prot_out_dim (_type_): Dimension of the prot encoder.
+            args (_type_): _description_
+        """
+        super(PPI_Metric_Learning, self).__init__()
+        self.prot_encoder = prot_encoder
+        self.loss = args.loss
+        self.use_miner = args.use_miner
+        self.miner_margin = args.miner_margin
+        self.agg_mode = args.agg_mode
+        self.prot_adapter_name = None
+        if self.use_miner:
+            self.miner = miners.TripletMarginMiner(
+                margin=args.miner_margin, type_of_triplets="all"
+            )
+        else:
+            self.miner = None
+
+        if self.loss == "ms_loss":
+            self.loss = losses.MultiSimilarityLoss(
+                alpha=1, beta=60, base=0.5
+            )  # 1,2,3; 40,50,60
+        elif self.loss == "circle_loss":
+            self.loss = losses.CircleLoss()
+        elif self.loss == "triplet_loss":
+            self.loss = losses.TripletMarginLoss()
+        elif self.loss == "infoNCE":
+            self.loss = losses.NTXentLoss(
+                temperature=0.07
+            )  # The MoCo paper uses 0.07, while SimCLR uses 0.5.
+        elif self.loss == "lifted_structure_loss":
+            self.loss = losses.LiftedStructureLoss()
+        elif self.loss == "nca_loss":
+            self.loss = losses.NCALoss()
+        self.reg = None
+        self.cls = None
+        self.dropout = torch.nn.Dropout(args.dropout)
+        print("miner:", self.miner)
+        print("loss:", self.loss)
+
+    def add_classification_head(self, prot_out_dim=1024, out_dim=2):
+        """Add regression 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 + prot_out_dim, out_dim)
+
+    def load_prot_adapter(
+            self,
+            prot_model_path,
+            prot_adapter_name="prot_adapter",
+    ):
+        if os.path.exists(prot_model_path):
+            self.prot_adapter_name = prot_adapter_name
+            self.prot_encoder.load_adapter(prot_model_path, load_as=prot_adapter_name)
+            self.prot_encoder.set_active_adapters(prot_adapter_name)
+            print(f"load protein adapters from: {prot_model_path} {prot_adapter_name}")
+        else:
+            print(f"{prot_model_path} not exits")
+
+    def init_adapters(
+            self,
+            prot_adapter_name="prot_adapter",
+            reduction_factor=16,
+    ):
+        """Initialise adapters
+
+        Args:
+            prot_adapter_name (str, optional): _description_. Defaults to "prot_adapter".
+            reduction_factor (int, optional): _description_. Defaults to 16.
+        """
+        adapter_config = AdapterConfig.load(
+            "pfeiffer", reduction_factor=reduction_factor
+        )
+        self.prot_adapter_name = prot_adapter_name
+        self.prot_encoder.add_adapter(prot_adapter_name, config=adapter_config)
+        self.prot_encoder.train_adapter([prot_adapter_name])
+
+    def save_adapters(self, save_path_prefix, total_step):
+        """Save adapters into file.
+
+        Args:
+            save_path_prefix (string): saving path prefix.
+            total_step (int): total step number.
+        """
+        prot_save_dir = os.path.join(
+            save_path_prefix, f"prot_adapter_step_{total_step}"
+        )
+        os.makedirs(prot_save_dir, exist_ok=True)
+        self.prot_encoder.save_adapter(prot_save_dir, self.prot_adapter_name)
+
+    def predict(self, x1, x2):
+        """
+        query : (N, h), candidates : (N, topk, h)
+        output : (N, topk)
+
+        """
+
+        if self.agg_mode == "cls":
+            x1 = self.prot_encoder(x1).last_hidden_state[:, 0]
+            x2 = self.prot_encoder(x2).last_hidden_state[:, 0]
+            x = torch.cat((x1, x2), 1)
+            return x
+        else:
+            x1 = self.prot_encoder(x1).last_hidden_state.mean(1)  # query : [batch_size, hidden]
+            x2 = self.prot_encoder(x2).last_hidden_state.mean(1)  # query : [batch_size, hidden]
+            x = torch.cat((x1, x2), 1)
+            return x
+
+    def module_predict(self, x1, x2):
+        """
+        query : (N, h), candidates : (N, topk, h)
+        output : (N, topk)
+
+        """
+        if self.agg_mode == "cls":
+            x1 = self.prot_encoder.module(x1).last_hidden_state[:, 0]
+            x2 = self.prot_encoder.module(x2).last_hidden_state[:, 0]
+            x = torch.cat((x1, x2), 1)
+            return x
+        else:
+            x1 = self.prot_encoder.module(x1).last_hidden_state.mean(1)  # query : [batch_size, hidden]
+            x2 = self.prot_encoder.module(x2).last_hidden_state.mean(1)  # query : [batch_size, hidden]
+            x = torch.cat((x1, x2), 1)
+            return x
+
+    @autocast()
+    def forward(self, query_toks1, query_toks2, labels):
+        """
+        query : (N, h), candidates : (N, topk, h)
+        output : (N, topk)
+        """
+        last_hidden_state1 = self.prot_encoder(
+            **query_toks1, return_dict=True
+        ).last_hidden_state
+        last_hidden_state2 = self.prot_encoder(
+            **query_toks2, return_dict=True
+        ).last_hidden_state
+        if self.agg_mode == "cls":
+            query_embed1 = last_hidden_state1[:, 0]  # query : [batch_size, hidden]
+            query_embed2 = last_hidden_state2[:, 0]  # query : [batch_size, hidden]
+        elif self.agg_mode == "mean_all_tok":
+            query_embed1 = last_hidden_state1.mean(1)  # query : [batch_size, hidden]
+            query_embed2 = last_hidden_state2.mean(1)  # query : [batch_size, hidden]
+        elif self.agg_mode == "mean":
+            query_embed1 = (
+                                   last_hidden_state1 * query_toks1["attention_mask"].unsqueeze(-1)
+                           ).sum(1) / query_toks1["attention_mask"].sum(-1).unsqueeze(-1)
+            query_embed2 = (
+                                   last_hidden_state2 * query_toks2["attention_mask"].unsqueeze(-1)
+                           ).sum(1) / query_toks2["attention_mask"].sum(-1).unsqueeze(-1)
+        else:
+            raise NotImplementedError()
+        query_embed = torch.cat([query_embed1, query_embed2], dim=0)
+
+        labels = torch.cat([labels, labels], dim=0)
+        if self.use_miner:
+            hard_pairs = self.miner(query_embed, labels)
+            return self.loss(query_embed, labels, hard_pairs)
+        else:
+            return self.loss(query_embed, labels)
+