scripts/inference.py

"""
Run EpHod to predict pHopt for enzyme sequences
"""


import numpy as np
import pandas as pd
from sklearn.svm import SVR
import torch
from torch.nn.parallel import DataParallel
import torch.nn as nn
import random
import tqdm
import argparse
import joblib
import os
import sys
import subprocess
import warnings
from pathlib import Path
warnings.filterwarnings('ignore')

import esm

PROJECT_ROOT = Path(__file__).resolve().parents[1]
WEIGHT_DIR = PROJECT_ROOT / 'weight'
sys.path.insert(0, str(PROJECT_ROOT))
from model.ephod.training import nn_models

MIN_ASSET_BYTES = {
    'esm1v_t33_650M_UR90S_1.pt': 1_000_000_000,
    'ESM1v-RLATtr.pt': 150_000_000,
    'ESM1v-SVR.pkl': 50_000_000,
}


def require_complete_asset(filename):
    """Return a local model asset and reject obvious partial transfers."""
    path = WEIGHT_DIR / filename
    if not path.is_file():
        raise FileNotFoundError(f'Missing model asset: {path}')
    minimum_size = MIN_ASSET_BYTES[filename]
    if path.stat().st_size < minimum_size:
        raise RuntimeError(
            f'Model asset appears incomplete: {path} '
            f'({path.stat().st_size} bytes; expected at least {minimum_size} bytes). '
            'Replace it with the complete file before running inference.'
        )
    return path




def parse_arguments():
    '''Parse command-line training arguments'''
    
    parser = argparse.ArgumentParser(description="Predict pHopt of enzymes with EpHod")
    parser.add_argument('--fasta_path', type=str,  
                        help='Path to fasta file of enzyme sequences')
    parser.add_argument('--save_dir', type=str, default='./',
                        help='Directory to which prediction results will be written')
    parser.add_argument('--csv_name', type=str, default='prediction.csv', 
                        help='Name of csv file to which prediction results will be written')
    parser.add_argument('--output_path', type=str, default=None,
                        help='Full path of the prediction CSV; overrides --save_dir and --csv_name')
    parser.add_argument('--verbose', default=1, type=int,
                        help='Whether to print out prediction progress to terminal')
    parser.add_argument('--save_attention_weights', default=0, type=int,
                        help="Whether to write RLAT attention weights for each sequence")
    parser.add_argument('--save_embeddings', default=0, type=int,
                        help="Whether to save 2560-dim EpHod embeddings for each sequence")
    args = parser.parse_args()

    return args




def write_attention_weights(accs, seqs, attention_weights, attention_dir, attention_mode='average'):
    '''Write RLAT attention weights for each sequence'''
    
    for i, (acc,seq) in enumerate(zip(accs, seqs)):
        seqlen = len(seq)
        weights = attention_weights[i,:,:seqlen]
        if attention_mode == 'average':
            weights = weights.mean(axis=0).transpose()
        elif attention_mode == 'max':
            weights = weights.max(axis=0).transpose()
        else:
            raise ValueError("attention_mode must be either 'average' or 'max'")
        weights = pd.DataFrame(weights.transpose(), index=list(seq), columns=['weights'])
        weights.to_csv(f'{attention_dir}/{acc}.csv')
    
        

               
def read_fasta(fasta, return_as_dict=False):
    '''Read the protein sequences in a fasta file. If return_as_dict, return a dictionary
    with headers as keys and sequences as values, else return a tuple, 
    (list_of_headers, list_of_sequences)'''
    
    headers, sequences = [], []
    with open(fasta, 'r') as fast:
        for line in fast:
            if line.startswith('>'):
                head = line.replace('>','').strip()
                headers.append(head)
                sequences.append('')
            else :
                seq = line.strip()
                if len(seq) > 0:
                    sequences[-1] += seq
    if return_as_dict:
        return dict(zip(headers, sequences))
    else:
        return (headers, sequences) 




def replace_noncanonical(seq, replace_char='X'):
    '''Replace all non-canonical amino acids with a specific character'''

    for char in ['B', 'J', 'O', 'U', 'Z']:
        seq = seq.replace(char, replace_char)
    return seq




class EpHodModel():
    
    def __init__(self, seed=0):

        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        if self.device != 'cuda':
            print('WARNING: You are not using a GPU. Inference will be slow')
        self.set_seed(seed=seed)
        self.esm1v_model, self.esm1v_batch_converter = self.load_ESM1v_model()
        self.svr_model, self.svr_stats = self.load_SVR_model()
        self.rlat_model = self.load_RLAT_model()
        self.esm1v_model.eval()
        self.rlat_model.eval()

    
    def set_seed(self, seed):
    
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        if self.device == 'cuda':
            torch.cuda.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        
        
        
    def load_ESM1v_model(self):
        '''Return pretrained ESM1v model weights and batch converter'''

        model_path = require_complete_asset('esm1v_t33_650M_UR90S_1.pt')
        model, alphabet = esm.pretrained.load_model_and_alphabet_local(str(model_path))
        model = model.to(self.device)
        batch_converter = alphabet.get_batch_converter()
        
        return model, batch_converter
    
    
    def get_ESM1v_embeddings(self, accs, seqs):
        '''Return per-residue embeddings (padded) for protein sequences from ESM1v model'''

        seqs = [replace_noncanonical(seq, 'X') for seq in seqs]
        data = [(accs[i], seqs[i]) for i in range(len(accs))]
        batch_labels, batch_strs, batch_tokens = self.esm1v_batch_converter(data)
        batch_tokens = batch_tokens.to(device=self.device, non_blocking=True)
        emb = self.esm1v_model(batch_tokens, repr_layers=[33], return_contacts=False)
        emb = emb["representations"][33]
        emb = emb.transpose(2,1) # From (batch, seqlen, features) to (batch, features, seqlen)

        return emb
    
    
    def load_RLAT_model(self):
        '''Return residual light attention top model'''
        
        model = nn_models.ResidualLightAttention(dim=1280, kernel_size=7, dropout=0.0, res_blocks=4, activation='elu')
        model = model.to(self.device)
        model_path = require_complete_asset('ESM1v-RLATtr.pt')
        model_dict = torch.load(model_path, map_location=self.device, weights_only=False)
        model_dict = {key[len('module.'):]: value for key, value in model_dict.items()} # Remove DataParallel suffix
        model.load_state_dict(model_dict)

        return model

    
    def load_SVR_model(self):
        '''Return SVR top model'''
        
        path = require_complete_asset('ESM1v-SVR.pkl')
        svr_model, svr_stats = joblib.load(path)

        return svr_model, svr_stats
        
    
    def predict(self, accs, seqs):
        '''Predict pHopt of sequences with EpHod'''
        
        # Get ESM1v embeddings and run RLATtr model
        emb_esm1v = self.get_ESM1v_embeddings(accs, seqs)
        maxlen = emb_esm1v.shape[-1]
        masks = [[1] * len(seqs[i]) + [0] * (maxlen - len(seqs[i])) \
                 for i in range(len(seqs))]
        masks = torch.tensor(masks, dtype=torch.int32)
        masks = masks.to(self.device)
        out = self.rlat_model(emb_esm1v, masks)
        rlat_pred, rlat_emb, rlat_attn = [item.cpu().numpy() for item in out]
    
        # Run SVR
        emb_pool = emb_esm1v.cpu().numpy().mean(axis=-1) # (batch, features, seqlen)
        emb_pool = (emb_pool - self.svr_stats[:,0]) / (self.svr_stats[:,1] + 1e-8) # Normalize with means/std.dev
        svr_pred = self.svr_model.predict(emb_pool) # Note that batch size > 1 affects this pooling
        ensemble_pred = (rlat_pred + svr_pred) / 2
        outdict = dict(rlat_pred=rlat_pred, rlat_emb=rlat_emb, rlat_attn=rlat_attn, 
                       svr_pred=svr_pred, ensemble_pred=ensemble_pred)

        return outdict
            


    
def main():
    '''Run inference with EpHod model'''
    
    args = parse_arguments()
    
    # Read enzyme sequence data
    assert os.path.exists(args.fasta_path), f"File not found in {args.fasta_path}"
    headers, sequences = read_fasta(args.fasta_path)
    accessions = [head.split()[0] for head in headers]
    headers, sequences, accessions = [np.array(item) for item in (headers, sequences, accessions)]
    assert len(accessions) == len(headers) == len(sequences), 'Fasta file has unequal headers and sequences'
    numseqs = len(sequences)
    if args.verbose:
        print(f'Reading {numseqs} sequences from {args.fasta_path}')
        
    # Check sequence lengths
    lengths = np.array([len(seq) for seq in sequences])
    if max(lengths) > 1022:
        long_count = np.sum(lengths > 1022)
        warning = f"{long_count} sequences are longer than 1022 residues and will be truncated"
        print(warning)
        sequences = np.asarray([item[:1022] for item in sequences])    
    
    # Directory and CSV path to which predictions will be written.
    if args.output_path:
        phout_file = Path(args.output_path).expanduser()
        output_dir = phout_file.parent
    else:
        output_dir = Path(args.save_dir).expanduser()
        phout_file = output_dir / args.csv_name
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Directory to write RLATtr attention weights
    if args.save_attention_weights:
        attention_dir = output_dir / 'attention_weights'
        attention_dir.mkdir(parents=True, exist_ok=True)

    # CSV file to write EpHod embeddings
    embed_file = output_dir / 'embeddings.csv'

    # Initialize EpHod model
    ephod_model = EpHodModel()
    if args.verbose:
        print('Initializing EpHod model')
        print(f'Device is {ephod_model.device}')

    # Batch prediction
    batch_size = 1 # Use batch_size of 1, since >1 will lead to wrong results in pooling (line 194)
    num_batches = int(np.ceil(numseqs / batch_size))
    all_ypred, all_emb_ephod = np.empty((0,3)), np.empty((0, 2560))
    
    with torch.no_grad():
        batches = range(num_batches)
        if args.verbose:
            batches = tqdm.tqdm(batches, desc="Predicting pHopt")
        
        for batch_step in batches:
            
            # Batch sequences
            start_idx = batch_step * batch_size
            stop_idx = (batch_step + 1) * batch_size
            accs = accessions[start_idx : stop_idx] 
            seqs = sequences[start_idx : stop_idx]
            
            # Predict with EpHod model
            out = ephod_model.predict(accs, seqs) # dict_keys(['rlat_pred', 'rlat_emb', 'rlat_attn', 'svr_pred', 'ensemble_pred'])
            all_ypred = np.vstack((all_ypred, np.array([out['rlat_pred'], out['svr_pred'], out['ensemble_pred']]).transpose()))
            all_emb_ephod = np.vstack((all_emb_ephod, out['rlat_emb']))
            if args.save_attention_weights:
                _ = write_attention_weights(accs, seqs, out['rlat_attn'], attention_dir)
            
    if args.save_embeddings:
        all_emb_ephod = pd.DataFrame(np.array(all_emb_ephod), index=accessions)
        all_emb_ephod.to_csv(embed_file)

    if args.verbose:
        print('Prediction completed.')
        print(f'Prediction CSV: {phout_file}')
        
    # Save predictions
    all_ypred = pd.DataFrame(all_ypred, index=accessions, columns=['RLATtr', 'SVR', 'Ensemble'])
    all_ypred.to_csv(phout_file)
    
    
    

if __name__ == '__main__':
    main()