src/predict.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src')))
import argparse
import torch
import re
import json
import os
import warnings
import numpy as np
from pathlib import Path
from transformers import EsmTokenizer, EsmModel, BertModel, BertTokenizer, AutoModelForMaskedLM
from transformers import T5Tokenizer, T5EncoderModel, AutoTokenizer, AutoModel
from transformers import logging
from peft import PeftModel
# Import project modules
from models.adapter_model import AdapterModel
from models.lora_model import LoraModel
from models.pooling import MeanPooling, Attention1dPoolingHead, LightAttentionPoolingHead

# Ignore warning information
logging.set_verbosity_error()
warnings.filterwarnings("ignore")

def parse_args():
    parser = argparse.ArgumentParser(description="Predict protein function for a single sequence")
    
    # Model parameters
    parser.add_argument('--eval_method', type=str, default="freeze", choices=["freeze", "plm-lora", "plm-qlora", "ses-adapter"], help="Evaluation method")
    parser.add_argument('--model_path', type=str, required=True, help="Path to the trained model")
    parser.add_argument('--plm_model', type=str, required=True, help="Pretrained language model name or path")
    parser.add_argument('--pooling_method', type=str, default="mean", choices=["mean", "attention1d", "light_attention"], help="Pooling method")
    parser.add_argument('--problem_type', type=str, default="single_label_classification", 
                        choices=["single_label_classification", "multi_label_classification", "regression"], 
                        help="Problem type")
    parser.add_argument('--num_labels', type=int, default=2, help="Number of labels")
    parser.add_argument('--hidden_size', type=int, default=None, help="Embedding hidden size of the model")
    parser.add_argument('--num_attention_head', type=int, default=8, help="Number of attention heads")
    parser.add_argument('--attention_probs_dropout', type=float, default=0, help="Attention probs dropout prob")
    parser.add_argument('--pooling_dropout', type=float, default=0.25, help="Pooling dropout")
    
    # Input sequence parameters
    parser.add_argument('--aa_seq', type=str, required=True, help="Amino acid sequence")
    parser.add_argument('--foldseek_seq', type=str, default="", help="Foldseek sequence (optional)")
    parser.add_argument('--ss8_seq', type=str, default="", help="Secondary structure sequence (optional)")
    parser.add_argument('--dataset', type=str, default="single", help="Dataset name (optional)")
    parser.add_argument('--use_foldseek', action='store_true', help="Use foldseek sequence")
    parser.add_argument('--use_ss8', action='store_true', help="Use secondary structure sequence")
    parser.add_argument('--structure_seq', type=str, default=None, help="Structure sequence types to use (comma-separated)")
    
    # Other parameters
    parser.add_argument('--max_seq_len', type=int, default=1024, help="Maximum sequence length")
    
    args = parser.parse_args()
    
    # Automatically determine whether to use structure sequences based on input
    args.use_foldseek = bool(args.foldseek_seq)
    args.use_ss8 = bool(args.ss8_seq)
    
    return args

def load_model_and_tokenizer(args):
    print("---------- Loading Model and Tokenizer ----------")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Check if model file exists
    if not os.path.exists(args.model_path):
        raise FileNotFoundError(f"Model file not found: {args.model_path}")
    
    # Load model configuration if available
    config_path = os.path.join(os.path.dirname(args.model_path), "config.json")
    try:
        with open(config_path, "r") as f:
            config = json.load(f)
            print(f"Loaded configuration from {config_path}")
            
            # Update args with config values if they exist
            if "pooling_method" in config:
                args.pooling_method = config["pooling_method"]
            if "problem_type" in config:
                args.problem_type = config["problem_type"]
            if "num_labels" in config:
                args.num_labels = config["num_labels"]
            if "num_attention_head" in config:
                args.num_attention_head = config["num_attention_head"]
            if "attention_probs_dropout" in config:
                args.attention_probs_dropout = config["attention_probs_dropout"]
            if "pooling_dropout" in config:
                args.pooling_dropout = config["pooling_dropout"]
    except FileNotFoundError:
        print(f"Model config not found at {config_path}. Using command line arguments.")
    
    # Build tokenizer and protein language model
    if "esm" in args.plm_model:
        tokenizer = EsmTokenizer.from_pretrained(args.plm_model)
        plm_model = EsmModel.from_pretrained(args.plm_model)
        args.hidden_size = plm_model.config.hidden_size
    elif "bert" in args.plm_model:
        tokenizer = BertTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
        plm_model = BertModel.from_pretrained(args.plm_model)
        args.hidden_size = plm_model.config.hidden_size
    elif "prot_t5" in args.plm_model:
        tokenizer = T5Tokenizer.from_pretrained(args.plm_model, do_lower_case=False)
        plm_model = T5EncoderModel.from_pretrained(args.plm_model)
        args.hidden_size = plm_model.config.d_model
    elif "ankh" in args.plm_model:
        tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
        plm_model = T5EncoderModel.from_pretrained(args.plm_model)
        args.hidden_size = plm_model.config.d_model
    elif "ProSST" in args.plm_model:
        tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
        plm_model = AutoModelForMaskedLM.from_pretrained(args.plm_model).to(device).eval()
        args.hidden_size = plm_model.config.hidden_size
    elif "Prime" in args.plm_model:
        tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
        plm_model = AutoModelForMaskedLM.from_pretrained(args.plm_model).to(device).eval()
        args.hidden_size = plm_model.config.hidden_size
    else:
        tokenizer = AutoTokenizer.from_pretrained(args.plm_model)
        plm_model = AutoModel.from_pretrained(args.plm_model)
        args.hidden_size = plm_model.config.hidden_size
    
    args.vocab_size = plm_model.config.vocab_size
    
    # Determine structure sequence types
    if args.structure_seq is None:
        args.structure_seq = ""
        print("Warning: structure_seq was None, setting to empty string")
    
    # Auto-set structure sequence flags based on structure_seq parameter
    if 'foldseek_seq' in args.structure_seq:
        args.use_foldseek = True
        print("Enabled foldseek_seq based on structure_seq parameter")
    if 'ss8_seq' in args.structure_seq:
        args.use_ss8 = True
        print("Enabled ss8_seq based on structure_seq parameter")
    
    # If flags are set but structure_seq is not, update structure_seq
    structure_seq_list = []
    if args.use_foldseek and 'foldseek_seq' not in args.structure_seq:
        structure_seq_list.append("foldseek_seq")
    if args.use_ss8 and 'ss8_seq' not in args.structure_seq:
        structure_seq_list.append("ss8_seq")
    
    if structure_seq_list and not args.structure_seq:
        args.structure_seq = ",".join(structure_seq_list)
    
    print(f"Training method: {args.eval_method}")  # Default for prediction
    print(f"Structure sequence: {args.structure_seq}")
    print(f"Use foldseek: {args.use_foldseek}")
    print(f"Use ss8: {args.use_ss8}")
    print(f"Problem type: {args.problem_type}")
    print(f"Number of labels: {args.num_labels}")
    print(f"Number of attention heads: {args.num_attention_head}")
    
    # Create and load model
    try:
        if args.eval_method in ["full", "ses-adapter", "freeze"]:
            model = AdapterModel(args)
        # ! lora/ qlora
        elif args.eval_method in ['plm-lora', 'plm-qlora', 'plm-dora', 'plm-adalora', 'plm-ia3']:
            model = LoraModel(args)
        if args.model_path is not None:
            model_path = args.model_path
        else:
            model_path = f"{args.output_root}/{args.output_dir}/{args.output_model_name}"
        model.load_state_dict(torch.load(model_path))
        model.to(device).eval()
        # ! lora/ qlora
        if args.eval_method == 'plm-lora':
            lora_path = model_path.replace(".pt", "_lora")
            plm_model = PeftModel.from_pretrained(plm_model,lora_path)
            plm_model = plm_model.merge_and_unload()
        elif args.eval_method == 'plm-qlora':
            lora_path = model_path.replace(".pt", "_qlora")
            plm_model = PeftModel.from_pretrained(plm_model,lora_path)
            plm_model = plm_model.merge_and_unload()
        elif args.eval_method == "plm-dora":
            dora_path = model_path.replace(".pt", "_dora")
            plm_model = PeftModel.from_pretrained(plm_model, dora_path)
            plm_model = plm_model.merge_and_unload()
        elif args.eval_method == "plm-adalora":
            adalora_path = model_path.replace(".pt", "_adalora")
            plm_model = PeftModel.from_pretrained(plm_model, adalora_path)
            plm_model = plm_model.merge_and_unload()
        elif args.eval_method == "plm-ia3":
            ia3_path = model_path.replace(".pt", "_ia3")
            plm_model = PeftModel.from_pretrained(plm_model, ia3_path)
            plm_model = plm_model.merge_and_unload()
            plm_model.to(device).eval()  
        plm_model.to(device).eval()  
        return model, plm_model, tokenizer, device
    
    except Exception as e:
        print(f"Error: {str(e)}")
        raise

def process_sequences(args, tokenizer, plm_model_name):
    """Process and prepare input sequences for prediction"""
    print("---------- Processing Input Sequences ----------")
    
    # Process amino acid sequence
    aa_seq = args.aa_seq.strip()
    if not aa_seq:
        raise ValueError("Amino acid sequence is empty")
    
    # Process structure sequences if needed
    foldseek_seq = args.foldseek_seq.strip() if args.foldseek_seq else ""
    ss8_seq = args.ss8_seq.strip() if args.ss8_seq else ""
    
    # Check if structure sequences are required but not provided
    if args.use_foldseek and not foldseek_seq:
        print("Warning: Foldseek sequence is required but not provided.")
    if args.use_ss8 and not ss8_seq:
        print("Warning: SS8 sequence is required but not provided.")
    
    # Format sequences based on model type
    if 'prot_bert' in plm_model_name or "prot_t5" in plm_model_name:
        aa_seq = " ".join(list(aa_seq))
        aa_seq = re.sub(r"[UZOB]", "X", aa_seq)
        if args.use_foldseek and foldseek_seq:
            foldseek_seq = " ".join(list(foldseek_seq))
        if args.use_ss8 and ss8_seq:
            ss8_seq = " ".join(list(ss8_seq))
    elif 'ankh' in plm_model_name:
        aa_seq = list(aa_seq)
        if args.use_foldseek and foldseek_seq:
            foldseek_seq = list(foldseek_seq)
        if args.use_ss8 and ss8_seq:
            ss8_seq = list(ss8_seq)
    
    # Truncate sequences if needed
    if args.max_seq_len:
        aa_seq = aa_seq[:args.max_seq_len]
        if args.use_foldseek and foldseek_seq:
            foldseek_seq = foldseek_seq[:args.max_seq_len]
        if args.use_ss8 and ss8_seq:
            ss8_seq = ss8_seq[:args.max_seq_len]
    
    # Tokenize sequences
    if 'ankh' in plm_model_name:
        aa_inputs = tokenizer.batch_encode_plus([aa_seq], add_special_tokens=True, padding=True, is_split_into_words=True, return_tensors="pt")
        if args.use_foldseek and foldseek_seq:
            foldseek_inputs = tokenizer.batch_encode_plus([foldseek_seq], add_special_tokens=True, padding=True, is_split_into_words=True, return_tensors="pt")
        if args.use_ss8 and ss8_seq:
            ss8_inputs = tokenizer.batch_encode_plus([ss8_seq], add_special_tokens=True, padding=True, is_split_into_words=True, return_tensors="pt")
    else:
        aa_inputs = tokenizer([aa_seq], return_tensors="pt", padding=True, truncation=True)
        if args.use_foldseek and foldseek_seq:
            foldseek_inputs = tokenizer([foldseek_seq], return_tensors="pt", padding=True, truncation=True)
        if args.use_ss8 and ss8_seq:
            ss8_inputs = tokenizer([ss8_seq], return_tensors="pt", padding=True, truncation=True)
    
    # Prepare data dictionary
    data_dict = {
        "aa_seq_input_ids": aa_inputs["input_ids"],
        "aa_seq_attention_mask": aa_inputs["attention_mask"],
    }
    
    # 只有 ProSST 模型需要结构标记
    if "ProSST" in plm_model_name and hasattr(args, 'prosst_stru_token') and args.prosst_stru_token:
        try:
            # 处理 ProSST 结构标记
            if isinstance(args.prosst_stru_token, str):
                seq_clean = args.prosst_stru_token.strip("[]").replace(" ","")
                tokens = list(map(int, seq_clean.split(','))) if seq_clean else []
            elif isinstance(args.prosst_stru_token, (list, tuple)):
                tokens = [int(x) for x in args.prosst_stru_token]
            else:
                tokens = []
                
            # 添加到数据字典
            if tokens:
                stru_tokens = torch.tensor([tokens], dtype=torch.long)
                data_dict["aa_seq_stru_tokens"] = stru_tokens
            else:
                # 如果没有标记，则使用零填充
                data_dict["aa_seq_stru_tokens"] = torch.zeros_like(aa_inputs["input_ids"], dtype=torch.long)
        except Exception as e:
            print(f"Warning: Failed to process ProSST structure tokens: {e}")
            # 使用零填充
            data_dict["aa_seq_stru_tokens"] = torch.zeros_like(aa_inputs["input_ids"], dtype=torch.long)
    
    if args.use_foldseek and foldseek_seq:
        data_dict["foldseek_seq_input_ids"] = foldseek_inputs["input_ids"]
    
    if args.use_ss8 and ss8_seq:
        data_dict["ss8_seq_input_ids"] = ss8_inputs["input_ids"]
    
    print("Processed input sequences with keys:", data_dict.keys())
    return data_dict

def predict(model, data_dict, device, args, plm_model):
    """Run prediction on the processed input data"""
    print("---------- Running Prediction ----------")
    
    # Move data to device
    for k, v in data_dict.items():
        data_dict[k] = v.to(device)
    
    # Run model inference
    with torch.no_grad():
        outputs = model(plm_model, data_dict)  # Pass the actual plm_model instead of None
        
        # Process outputs based on problem type
        if args.problem_type == "regression":
            predictions = outputs.squeeze().item()
            print(f"Prediction result: {predictions}")
            return {"prediction": predictions}
        
        elif args.problem_type == "single_label_classification":
            probabilities = torch.nn.functional.softmax(outputs, dim=1)
            predicted_class = torch.argmax(probabilities, dim=1).item()
            class_probs = probabilities.squeeze().tolist()
            
            # Ensure class_probs is a list
            if not isinstance(class_probs, list):
                class_probs = [class_probs]
            
            print(f"Predicted class: {predicted_class}")
            print(f"Class probabilities: {class_probs}")
            
            return {
                "predicted_class": predicted_class,
                "probabilities": class_probs
            }
        
        elif args.problem_type == "multi_label_classification":
            sigmoid_outputs = torch.sigmoid(outputs)
            predictions = (sigmoid_outputs > 0.5).int().squeeze().tolist()
            probabilities = sigmoid_outputs.squeeze().tolist()
            
            # Ensure predictions and probabilities are lists
            if not isinstance(predictions, list):
                predictions = [predictions]
            if not isinstance(probabilities, list):
                probabilities = [probabilities]
            
            print(f"Predicted labels: {predictions}")
            print(f"Label probabilities: {probabilities}")
            
            return {
                "predictions": predictions,
                "probabilities": probabilities
            }

def main():
    try:
        # Parse arguments
        args = parse_args()
        
        # Load model, tokenizer and get device
        model, plm_model, tokenizer, device = load_model_and_tokenizer(args)
        
        # Process input sequences
        data_dict = process_sequences(args, tokenizer, args.plm_model)
        
        # Run prediction
        results = predict(model, data_dict, device, args, plm_model)
        
        # Output results
        print("\n---------- Prediction Results ----------")
        print(json.dumps(results, indent=2))
        
        return 0
    except Exception as e:
        print(f"Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == "__main__":
    sys.exit(main())