|
|
|
|
|
|
|
|
import sys |
|
|
import os |
|
|
import torch |
|
|
import pandas as pd |
|
|
from tqdm import tqdm |
|
|
from datetime import datetime |
|
|
from omegaconf import OmegaConf |
|
|
from transformers import AutoTokenizer, AutoModelForMaskedLM |
|
|
|
|
|
from src.lm.memdlm.diffusion_module import MembraneFlow |
|
|
from src.utils.model_utils import _print |
|
|
from src.sampling.guided_sampler import GuidedSampler |
|
|
from src.utils.generate_utils import ( |
|
|
mask_for_scaffold, |
|
|
calc_blosum_score, |
|
|
calc_ppl |
|
|
) |
|
|
|
|
|
config = OmegaConf.load("/home/a03-sgoel/MeMDLM_v2/src/configs/guidance.yaml") |
|
|
|
|
|
os.chdir(f'/home/a03-sgoel/MeMDLM_v2/results/infilling/guided/{config.lm.ft_evoflow}/test_set/') |
|
|
todays_date = datetime.today().strftime('%Y-%m-%d') |
|
|
csv_save_path = f'./{todays_date}_boltzmann-soft_new_clf_data_cleaned/' |
|
|
try: os.makedirs(csv_save_path, exist_ok=False) |
|
|
except FileExistsError: pass |
|
|
|
|
|
|
|
|
def main(): |
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(config.lm.pretrained_esm) |
|
|
esm_model = AutoModelForMaskedLM.from_pretrained(config.lm.pretrained_esm).eval().to(device) |
|
|
|
|
|
diffusion = MembraneFlow(config).to(device) |
|
|
state_dict = diffusion.get_state_dict(f"/home/a03-sgoel/MeMDLM_v2/checkpoints/{config.lm.ft_evoflow}/best_model.ckpt") |
|
|
diffusion.load_state_dict(state_dict) |
|
|
diffusion.eval().to(device) |
|
|
|
|
|
sampler = GuidedSampler(config, esm_model, tokenizer, diffusion, device) |
|
|
|
|
|
df = pd.read_csv('/home/a03-sgoel/MeMDLM_v2/data/classifier/test.csv') |
|
|
sequences = df['Sequence'].tolist() |
|
|
|
|
|
gen_seqs, ppls, blosums = [], [], [] |
|
|
|
|
|
|
|
|
for seq in tqdm(sequences, desc='Infilling Sequences'): |
|
|
masked_seq = mask_for_scaffold(seq, generate_type='uppercase', mask_token='<mask>') |
|
|
tokens = tokenizer(masked_seq, return_tensors='pt') |
|
|
input_ids, attn_masks = tokens['input_ids'].to(device), tokens['attention_mask'].to(device) |
|
|
|
|
|
soluble_idxs = [i for i in range(len(seq)) if seq[i].isupper()] |
|
|
infilled_tokens = sampler.optimize_sequence( |
|
|
input_ids=input_ids, |
|
|
attn_masks=attn_masks, |
|
|
soluble_indices=soluble_idxs, |
|
|
) |
|
|
infilled_seq = tokenizer.decode(infilled_tokens).replace(" ", "")[5:-5] |
|
|
|
|
|
bl = calc_blosum_score(seq.upper(), infilled_seq, soluble_idxs) |
|
|
try: |
|
|
ppl = calc_ppl(esm_model, tokenizer, infilled_seq, [i for i in range(len(seq))], model_type='esm') |
|
|
except: |
|
|
ppl = float('inf') |
|
|
|
|
|
gen_seqs.append(infilled_seq) |
|
|
ppls.append(ppl) |
|
|
blosums.append(bl) |
|
|
|
|
|
_print(seq) |
|
|
_print(infilled_seq) |
|
|
_print(ppl) |
|
|
_print(bl) |
|
|
_print('\n') |
|
|
|
|
|
|
|
|
df['MeMDLM Sequence'] = gen_seqs |
|
|
df['MeMDLM PPL'] = ppls |
|
|
df['MeMDLM BLOSUM'] = blosums |
|
|
|
|
|
_print(df) |
|
|
df.to_csv(f'./{csv_save_path}/t=0.7_new-data-cleaned_infilled_seqs.csv', index=False) |
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
main() |
|
|
|
|
|
|
