#!/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 import pandas as pd from pathlib import Path from tqdm import tqdm from transformers import EsmTokenizer, EsmModel, BertModel, BertTokenizer from transformers import T5Tokenizer, T5EncoderModel, AutoTokenizer, AutoModel, AutoModelForMaskedLM 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="Batch predict protein function for multiple sequences") # Model parameters parser.add_argument('--eval_method', type=str, default="freeze", choices=["full", "freeze", "plm-lora", "plm-qlora", "ses-adapter", 'plm-dora', 'plm-adalora', 'plm-ia3'], 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 and output parameters parser.add_argument('--input_file', type=str, required=True, help="Path to input CSV file with sequences") parser.add_argument('--output_dir', type=str, required=True, help="Path to output CSV file dir for predictions") parser.add_argument('--output_file', type=str, required=True, help="output CSV file name for predictions") 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") parser.add_argument('--batch_size', type=int, default=1, help="Batch size for prediction") parser.add_argument('--dataset', type=str, default="Protein-wise", help="Dataset name") args = parser.parse_args() 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).to(device).eval() 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).to(device).eval() 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).to(device).eval() 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).to(device).eval() 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).to(device).eval() 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}" if args.eval_method == "full": model_weights = torch.load(model_path) model.load_state_dict(model_weights['model_state_dict']) plm_model.load_state_dict(model_weights['plm_state_dict']) else: 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() return model, plm_model, tokenizer, device except Exception as e: print(f"Error: {str(e)}") raise def process_sequence(args, tokenizer, plm_model_name, aa_seq, foldseek_seq="", ss8_seq="", prosst_stru_token=None): """Process and prepare a single input sequence for prediction""" # Process amino acid sequence aa_seq = aa_seq.strip() if not aa_seq: raise ValueError("Amino acid sequence is empty") # Process structure sequences if needed foldseek_seq = foldseek_seq.strip() if foldseek_seq else "" ss8_seq = ss8_seq.strip() if ss8_seq else "" # Check if structure sequences are required but not provided if args.use_foldseek and not foldseek_seq: print(f"Warning: Foldseek sequence is required but not provided for sequence: {aa_seq[:20]}...") if args.use_ss8 and not ss8_seq: print(f"Warning: SS8 sequence is required but not provided for sequence: {aa_seq[:20]}...") # 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"], } if "ProSST" in plm_model_name and prosst_stru_token is not None: try: if isinstance(prosst_stru_token, str): seq_clean = prosst_stru_token.strip("[]").replace(" ","") tokens = list(map(int, seq_clean.split(','))) if seq_clean else [] elif isinstance(prosst_stru_token, (list, tuple)): tokens = [int(x) for x in 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"] return data_dict def predict_batch(model, plm_model, data_dict, device, args): """Run prediction on a batch of processed input data""" # 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) # Process outputs based on problem type if args.problem_type == "regression": predictions = outputs.squeeze().cpu().numpy() # 确保返回标量值 if np.isscalar(predictions): return {"predictions": predictions} else: # 如果是批处理，返回整个数组 return {"predictions": predictions.tolist() if isinstance(predictions, np.ndarray) else predictions} elif args.problem_type == "single_label_classification": probabilities = torch.nn.functional.softmax(outputs, dim=1) predicted_classes = torch.argmax(probabilities, dim=1).cpu().numpy() class_probs = probabilities.cpu().numpy() return { "predicted_classes": predicted_classes.tolist(), "probabilities": class_probs.tolist() } elif args.problem_type == "multi_label_classification": sigmoid_outputs = torch.sigmoid(outputs) predictions = (sigmoid_outputs > 0.5).int().cpu().numpy() probabilities = sigmoid_outputs.cpu().numpy() return { "predictions": predictions.tolist(), "probabilities": probabilities.tolist() } def main(): # Parse command line arguments args = parse_args() try: # Load model and tokenizer model, plm_model, tokenizer, device = load_model_and_tokenizer(args) # Read input CSV file print(f"---------- Reading input file: {args.input_file} ----------") try: df = pd.read_csv(args.input_file) print(f"Found {len(df)} sequences in input file") except Exception as e: print(f"Error reading input file: {str(e)}") sys.exit(1) # Check required columns required_columns = ["aa_seq"] if args.use_foldseek: required_columns.append("foldseek_seq") if args.use_ss8: required_columns.append("ss8_seq") missing_columns = [col for col in required_columns if col not in df.columns] if missing_columns: print(f"Error: Input file is missing required columns: {', '.join(missing_columns)}") sys.exit(1) # Initialize results dataframe results = [] # Process each sequence print("---------- Processing sequences ----------") for idx, row in tqdm(df.iterrows(), total=len(df), desc="Predicting"): try: # Get sequences from row aa_seq = row["aa_seq"] foldseek_seq = row["foldseek_seq"] if "foldseek_seq" in df.columns and args.use_foldseek else "" ss8_seq = row["ss8_seq"] if "ss8_seq" in df.columns and args.use_ss8 else "" # Process sequence data_dict = process_sequence(args, tokenizer, args.plm_model, aa_seq, foldseek_seq, ss8_seq) # Run prediction prediction_results = predict_batch(model, plm_model, data_dict, device, args) # Create result row result_row = {"aa_seq": aa_seq} # Add sequence ID if available if "id" in df.columns: result_row["id"] = row["id"] # Add prediction results based on problem type if args.problem_type == "regression": # result_row["prediction"] = prediction_results["predictions"][0] if isinstance(prediction_results["predictions"], (list, np.ndarray)): result_row["prediction"] = prediction_results["predictions"][0] else: result_row["prediction"] = prediction_results["predictions"] elif args.problem_type == "single_label_classification": result_row["predicted_class"] = prediction_results["predicted_classes"][0] # Add class probabilities for i, prob in enumerate(prediction_results["probabilities"][0]): result_row[f"class_{i}_prob"] = prob elif args.problem_type == "multi_label_classification": # Add binary predictions for i, pred in enumerate(prediction_results["predictions"][0]): result_row[f"label_{i}"] = pred # Add probabilities for i, prob in enumerate(prediction_results["probabilities"][0]): result_row[f"label_{i}_prob"] = prob results.append(result_row) except Exception as e: print(f"Error processing sequence at index {idx}: {str(e)}") # Add error row error_row = {"aa_seq": aa_seq, "error": str(e)} if "id" in df.columns: error_row["id"] = row["id"] results.append(error_row) # Create results dataframe results_df = pd.DataFrame(results) # Save results to output file if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) output_file = os.path.join(args.output_dir, args.output_file) print(f"---------- Saving results to {output_file} ----------") results_df.to_csv(output_file, index=False) print(f"Saved {len(results_df)} prediction results") print("---------- Batch prediction completed successfully ----------") except Exception as e: print(f"Error: {str(e)}") import traceback traceback.print_exc() sys.exit(1) if __name__ == "__main__": main()