Delete siamese

siamese/siamese_ppi_decoy.py

@@ -1,187 +0,0 @@
-import os
-import pdb
-import torch
-import torch.nn as nn
-import torch.optim as optim
-from torch.utils.data import Dataset, DataLoader
-from transformers import EsmModel, EsmTokenizer
-from sklearn.model_selection import train_test_split
-import pandas as pd
-from peft import BOFTConfig, get_peft_model
-from datasets import load_from_disk
-import time
-
-os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
-
-# Hyperparameters
-HYPERPARAMS = {
-    'learning_rate': 0.001,
-    'batch_size': 32,
-    'num_epochs': 10,
-    # 'boft_block_size': 8,
-    # 'boft_n_butterfly_factor': 1,
-    # 'boft_dropout': 0.1,
-    # 'boft_bias': 'boft_only',
-    # 'boft_modules_to_save': [],  # List any specific modules to save if needed
-    # 'boft_target_modules': ["query", "value", "key", "output.dense", "mlp.fc1", "mlp.fc2"],
-    'margin': 1.0
-}
-
-
-# Siamese NN
-class SiameseNetwork(nn.Module):
-    def __init__(self, encoder):
-        super(SiameseNetwork, self).__init__()
-        self.encoder = encoder
-        self.embedding_dim = encoder.config.hidden_size
-        self.projection = nn.Linear(self.embedding_dim * 2, self.embedding_dim)
-
-    def forward(self, target_tokens, binder_tokens, decoy_tokens):
-        target_embedding = self.encoder(**target_tokens).last_hidden_state[:, 0, :]
-        binder_embedding = self.encoder(**binder_tokens).last_hidden_state[:, 0, :]
-        decoy_embedding = self.encoder(**decoy_tokens).last_hidden_state[:, 0, :]
-
-        # Compute joint embeddings
-        anchor_embedding = torch.cat((target_embedding, binder_embedding), dim=-1)
-        positive_embedding = torch.cat((binder_embedding, target_embedding), dim=-1)
-        negative_embedding = torch.cat((decoy_embedding, binder_embedding), dim=-1)
-
-        # Project joint embeddings back to original dimensions
-        anchor_embedding = self.projection(anchor_embedding)
-        positive_embedding = self.projection(positive_embedding)
-        negative_embedding = self.projection(negative_embedding)
-
-        return anchor_embedding, positive_embedding, negative_embedding
-
-
-# Generate scores for candidate binders
-def generate_scores(siamese_net, tokenizer, target_seq, candidate_binders, decoy_seq):
-    siamese_net.eval()
-    scores = []
-
-    with torch.no_grad():
-        target_tokens = tokenizer(target_seq, return_tensors="pt", padding=True, truncation=True).to(device)
-        decoy_tokens = tokenizer(decoy_seq, return_tensors="pt", padding=True, truncation=True).to(device)
-
-        for binder_seq in candidate_binders:
-            binder_tokens = tokenizer(binder_seq, return_tensors="pt", padding=True, truncation=True).to(device)
-            target_embedding, binder_embedding, decoy_embedding = siamese_net(target_tokens, binder_tokens, decoy_tokens)
-            target_binder_similarity = torch.cosine_similarity(target_embedding, binder_embedding)
-            target_decoy_similarity = torch.cosine_similarity(target_embedding, decoy_embedding)
-            score = target_binder_similarity - target_decoy_similarity
-            scores.append(score.item())
-
-    return scores
-
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-distributed = torch.cuda.device_count() > 1
-
-
-# Load the pre-trained ESM-2-650M model and tokenizer
-model_name = "facebook/esm2_t33_650M_UR50D"
-tokenizer = EsmTokenizer.from_pretrained(model_name)
-model = EsmModel.from_pretrained(model_name)
-
-siamese_ppi_net = SiameseNetwork(model).to(device)
-if distributed:
-    siamese_ppi_net = torch.nn.DataParallel(siamese_ppi_net)
-
-
-# Define the triplet loss function
-criterion = nn.TripletMarginLoss(margin=HYPERPARAMS['margin']).to(device)
-
-# Define the optimizer
-optimizer = optim.Adam(siamese_ppi_net.parameters(), lr=HYPERPARAMS['learning_rate'])
-
-# Load dataset
-train_dataset = load_from_disk('/home/tc415/muPPIt/dataset/train_mut')
-val_dataset = load_from_disk('/home/tc415/muPPIt/dataset/val_mut')
-test_dataset = load_from_disk('/home/tc415/muPPIt/dataset/test_mut')
-
-# Training loop
-for epoch in range(HYPERPARAMS['num_epochs']):
-    # Training
-    siamese_ppi_net.train()
-    train_loss = 0.0
-    # for target_tokens, binder_tokens, decoy_tokens in train_dataloader:
-    for batch in train_dataset:
-        # pdb.set_trace()
-        start = time.time()
-        target_tokens = {'input_ids': torch.tensor(batch['anchor_input_ids']).to(device),
-                         'attention_mask': torch.tensor(batch['anchor_attention_mask']).to(device)}
-        binder_tokens = {'input_ids': torch.tensor(batch['positive_input_ids']).to(device),
-                         'attention_mask': torch.tensor(batch['positive_attention_mask']).to(device)}
-        decoy_tokens  = {'input_ids': torch.tensor(batch['negative_input_ids']).to(device),
-                         'attention_mask': torch.tensor(batch['negative_attention_mask']).to(device)}
-
-
-        # pdb.set_trace()
-        # Forward pass
-        target_embedding, binder_embedding, decoy_embedding = siamese_ppi_net(target_tokens, binder_tokens, decoy_tokens)
-
-        # Compute the triplet loss
-        loss = criterion(target_embedding, binder_embedding, decoy_embedding)
-
-        # Backward pass and optimization
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-        train_loss += loss.item()
-
-        print(f"loss = {loss.item()}, time = {time.time()-start}s")
-
-    train_loss /= len(train_dataset)
-
-    # Validation
-    siamese_ppi_net.eval()
-    val_loss = 0.0
-    with torch.no_grad():
-        for batch in val_dataset:
-            target_tokens = {'input_ids': torch.tensor(batch['anchor_input_ids']).to(device),
-                             'attention_mask': torch.tensor(batch['anchor_attention_mask']).to(device)}
-            binder_tokens = {'input_ids': torch.tensor(batch['positive_input_ids']).to(device),
-                             'attention_mask': torch.tensor(batch['positive_attention_mask']).to(device)}
-            decoy_tokens = {'input_ids': torch.tensor(batch['negative_input_ids']).to(device),
-                            'attention_mask': torch.tensor(batch['negative_attention_mask']).to(device)}
-            target_embedding, binder_embedding, decoy_embedding = siamese_ppi_net(target_tokens, binder_tokens, decoy_tokens)
-            loss = criterion(target_embedding, binder_embedding, decoy_embedding)
-            val_loss += loss.item()
-    val_loss /= len(val_dataset)
-
-    print(f"Epoch [{epoch+1}/{HYPERPARAMS['num_epochs']}], Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")
-
-# Testing
-siamese_ppi_net.eval()
-test_loss = 0.0
-with torch.no_grad():
-    for batch in test_dataset:
-        target_tokens = {'input_ids': torch.tensor(batch['anchor_input_ids']).to(device),
-                         'attention_mask': torch.tensor(batch['anchor_attention_mask']).to(device)}
-        binder_tokens = {'input_ids': torch.tensor(batch['positive_input_ids']).to(device),
-                         'attention_mask': torch.tensor(batch['positive_attention_mask']).to(device)}
-        decoy_tokens = {'input_ids': torch.tensor(batch['negative_input_ids']).to(device),
-                        'attention_mask': torch.tensor(batch['negative_attention_mask']).to(device)}
-        target_embedding, binder_embedding, decoy_embedding = siamese_ppi_net(target_tokens, binder_tokens, decoy_tokens)
-        loss = criterion(target_embedding, binder_embedding, decoy_embedding)
-        test_loss += loss.item()
-test_loss /= len(test_dataset)
-
-print(f"Test Loss: {test_loss:.4f}")
-
-# Save the trained model
-torch.save(siamese_ppi_net.state_dict(), "siamese_ppi_model.pth")
-
-# # Example: Scoring for candidate binders
-# target_seq = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
-# candidate_binders = [
-#     "KTVNELEKVIKKQGKRAKLIIAIIMIIIIIIVV",
-#     "ATVRELEKQIKKQRKRAKLIIAIVMIFIIVVVV",
-#     "KTVNELEKQIKKQGKRAKLIIAIVMIIIIVVVV"
-# ]
-# decoy_seq = "MHIKPLLSRLAQAAANASATPPPPPPPPPGPAVAEEPLHRPTNPGASSGCHKQPLKQSDCPKRPR"
-
-# scores = generate_scores(siamese_ppi_net, tokenizer, target_seq, candidate_binders, decoy_seq)
-# print("Candidate Binder Scores:")
-# for binder, score in zip(candidate_binders, scores):
-#     print(f"Binder: {binder}, Score: {score:.4f}")