File size: 3,881 Bytes
d061944 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import torch
import torch.nn.functional as F
import math
import random
import sys
import pandas as pd
from mlm_generate_utils import mask_for_scaffold, calculate_cosine_sim, calculate_hamming_dist
from diffusion import Diffusion
import hydra
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModel, pipeline
def masking_test(sequence: str, generate_case: str, tokenizer, mask_prob: float = 0.50):
"""
Masks 50% of the tokens in the sequence.
"""
tokens = list(sequence.upper())
num_tokens_to_mask = int(mask_prob * len(tokens)) # Select some fraction of the tokens
print(num_tokens_to_mask,len(tokens))
# Get random indices to mask
mask_indices = random.sample(range(len(tokens)), num_tokens_to_mask)
for idx in mask_indices:
tokens[idx] = tokenizer.mask_token # Replace with mask token
return ''.join(tokens)
@torch.no_grad()
def generate_scaffold_mdlm(sequence: str, generate_case: str, tokenizer, mdlm: Diffusion):
# # Mask soluble or transmembrane domains
# masked_sequence = mask_for_scaffold(sequence, generate_case)
# # Test out different masking rates
# masked_sequence = masking_test(sequence, generate_case, tokenizer)
# 100% masking rate, de novo generation
masked_sequence = len(sequence) * "<mask>"
print(masked_sequence)
inputs = tokenizer(masked_sequence, return_tensors="pt").to(mdlm.device)
logits = mdlm._sample(x_input=inputs) # using sample, change config.sampling.steps to determine robustness
# logits = mdlm.forward(inputs)
# print(tokenizer.decode(logits.squeeze(), skip_special_tokens=True))
return tokenizer.decode(logits.squeeze()), masked_sequence
@hydra.main(version_base=None, config_path='configs', config_name='config')
def mdlm_motif_benchmark(config):
path = "/workspace/sg666/MDpLM"
test_sequences = pd.read_csv(path + "/data/membrane/test.csv")['Sequence'].tolist()
tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
mdlm_model = Diffusion.load_from_checkpoint(config.eval.checkpoint_path, config=config, tokenizer=tokenizer)
esm_model = AutoModel.from_pretrained("facebook/esm2_t6_8M_UR50D") # model used for functionality testing
mdlm_model.eval()
esm_model.eval()
print("loaded models...")
device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
mdlm_model.to(device)
esm_model.to(device)
for generate_case in ['uppercase', 'lowercase']:
case_results = []
for original_sequence in tqdm(test_sequences, desc=f"scaffolding ({generate_case}): "):
generated_sequence, masked_input = generate_scaffold_mdlm(original_sequence, generate_case, tokenizer, mdlm_model)
generated_sequence = generated_sequence[5:-5].replace(" ", "") # Remove bos/eos tokens
perplexity = mdlm_model.compute_masked_perplexity([original_sequence], masked_input)
cos_sim = calculate_cosine_sim(original_sequence, generated_sequence, tokenizer, esm_model, device)
hamming_distance = calculate_hamming_dist(original_sequence, generated_sequence)
case_results.append([original_sequence, generated_sequence, perplexity, cos_sim, hamming_distance])
print("perplexity: ", perplexity, "cos sim: ", cos_sim, "hamming: ", hamming_distance)
print(f"generated sequence: {generated_sequence}")
print(f"original sequence: {original_sequence.upper()}")
sys.stdout.flush()
df = pd.DataFrame(case_results, columns=['Original Sequence', 'Generated Sequence', 'Perplexity', 'Cosine Similarity', 'Hamming Distance'])
df.to_csv(path + f'/benchmarks/MLM/mlm_{generate_case}_results.csv', index=False)
if __name__ == "__main__":
mdlm_motif_benchmark() |