import os, sys, warnings, argparse, math, tqdm, datetime import pytorch_lightning as pl import torch from pytorch_lightning.trainer import Trainer import pytorch_lightning.callbacks as plc import pytorch_lightning.loggers as plog from model_interface import MInterface from data_interface import DInterface from src.tools.logger import SetupCallback, BackupCodeCallback from shutil import ignore_patterns from transformers import AutoTokenizer import numpy as np import yaml import wandb warnings.filterwarnings("ignore") def create_parser(): parser = argparse.ArgumentParser() parser.add_argument('--infer_path', type=str, help='Path where to read the data to be predicted and where to save the predictions.') # Set-up parameters parser.add_argument('--res_dir', default='./train/results', type=str) parser.add_argument('--ex_name', default='debug', type=str) parser.add_argument('--check_val_every_n_epoch', default=1, type=int) parser.add_argument('--stage', default='predict', type=str) #'fit', 'test' or 'predict' parser.add_argument('--val_check_interval', default=0.5, type=float, help='Validation check interval') parser.add_argument('--dataset', default='PDBInference') # AF2DB_dataset, CATH_dataset parser.add_argument('--model_name', default='ProteinMPNN', choices=['StructGNN', 'GraphTrans', 'GVP', 'GCA', 'AlphaDesign', 'ESMIF', 'PiFold', 'ProteinMPNN', 'KWDesign', 'E3PiFold']) # parser.add_argument('--lr', default=4e-4, type=float, help='Learning rate') # parser.add_argument('--lr_scheduler', default='onecycle') # parser.add_argument('--offline', default=1, type=int) parser.add_argument('--seed', default=111, type=int) parser.add_argument('--num_workers', default=12, type=int) parser.add_argument('--pad', default=1024, type=int) parser.add_argument('--min_length', default=40, type=int) parser.add_argument('--data_root', default='./data/') # Training parameters # parser.add_argument('--epoch', default=10, type=int, help='end epoch') parser.add_argument('--augment_eps', default=0.0, type=float, help='noise level') # parser.add_argument('--gpus', default=1, type=int, help='how many GPUs to train on') # parser.add_argument('--weight_decay', default=0.0, type=float, help='Weight decay for optimizer') # # Eval parameters # parser.add_argument('--eval_sequences_sampled', default=1, type=int, help='How many sequences to sample in evaluation.') # parser.add_argument('--eval_sequences_temperature', default=0, type=float, help='What temperature to use for the sampling in evaluation.') # parser.add_argument('--eval_output_dir', default=None, type=str, help='Where to save the evaluation output.') # Model parameters parser.add_argument('--use_dist', default=1, type=int) parser.add_argument('--use_product', default=0, type=int) parser.add_argument('--use_pmpnn_checkpoint', default=1, type=int, help='By 1 or 0 decide whether to start with pretrained ProteinMPNN.') parser.add_argument('--checkpoint_path', type=str, default=None, help='Path to the model checkpoint to load weights from') # Dynamics aware parameters parser.add_argument('--use_dynamics', default=0, type=int) # parser.add_argument('--flex_loss_coeff', default=0.5, type=float) # parser.add_argument('--get_gt_flex_onthefly', default=0, type=int, help='Flag to get ground truth flexibility on-the-fly (with subsequent caching)') parser.add_argument('--init_flex_features', default=1, type=int, help="Set to 0 if no flexibility information should be passed on input to the node features h_V") # parser.add_argument('--loss_fn', default='MSE', type=str, help= 'Define what loss to use. Choose MSE, L1 or DPO.') # parser.add_argument('--grad_normalization', default=1, type=int, help="Set to 0 if the gradients of the seq and flex losses should not be normalized.") # parser.add_argument('--test_engineering', default=0, type=int, help="In this main.py should be set to 0 to not overwrite the training dataset.") args = parser.parse_args() return args if __name__ == "__main__": args = create_parser() args.batch_size = 1 print('In the predict stage, defaulting batch size to 1.') assert args.use_dynamics == 0, "In the inference script this should be set to 0." if not os.path.exists(args.infer_path): os.makedirs(args.infer_path) if (len(args.infer_path) > 0 or args.dataset=='PDBInference') and (len(args.infer_path) == 0 or args.dataset!='PDBInference'): raise ValueError("You should only use --infer_path with --dataset 'PDBInference' and vice versa.") # Load model weights from checkpoint if provided if args.checkpoint_path is not None: trained_model_path = args.checkpoint_path print(f"Loading model weights from checkpoint passed by argument: {trained_model_path}") else: with open('configs/Flexpert-Design-inference.yaml', 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) trained_model_path = config['pmpnn_model_path'] print(f"Loading model weights from checkpoint specified in Flexpert-Design-inference.yaml: {trained_model_path}") if os.path.exists(trained_model_path): print(f"Rewriting the path to the Flexpert-Design trained ProteinMPNN weights in the model interface.") args.starting_checkpoint_path = trained_model_path else: raise FileNotFoundError(f"Checkpoint file not found at {trained_model_path}") pl.seed_everything(args.seed) data_module = DInterface(**vars(args)) data_module.setup(stage='predict') model = MInterface(**vars(args)) trainer_config = { 'devices': 1, 'max_epochs': 1, 'num_nodes': 1, "strategy": 'ddp', "precision": '32', 'accelerator': 'gpu', 'val_check_interval': args.val_check_interval, 'check_val_every_n_epoch': args.check_val_every_n_epoch } trainer = Trainer(**trainer_config) predictions = trainer.predict(model, data_module) tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D", cache_dir='./cache_dir/') # mask token: 32 serializable_predictions = [] for pred_idx, pred in enumerate(predictions): logprobs = pred['log_probs'].cpu().numpy()[0] # [L, 21] pmpnn_alphabet_tokens_argmax = logprobs.argmax(axis=-1) # [L] aa_sequence = ''.join(tokenizer.decode(pmpnn_alphabet_tokens_argmax, skip_special_tokens=True).split()) # Get probability of the predicted sequence seq_probs = np.exp(logprobs.max(axis=-1)) # [L] avg_prob = float(np.mean(seq_probs)) serializable_predictions.append({ 'prediction_id': pred['batch']['title'][0], 'amino_acid_sequence': aa_sequence }) with open(f'{args.infer_path}/predictions.txt', 'w') as f: for pred in serializable_predictions: f.write(f'>{pred["prediction_id"]}\n{pred["amino_acid_sequence"]}\n')