eval_classifier.py

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2" 
import torch
from torch import cuda
from transformers import AutoTokenizer, AutoModelForSequenceClassification, default_data_collator
from datasets import Dataset
from sklearn.model_selection import train_test_split
import json
from tqdm import tqdm
from transformers import AutoTokenizer, EsmForProteinFolding
import torch
from PIL import Image
import torch
import argparse
from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OFProtein
from transformers.models.esm.openfold_utils.feats import atom14_to_atom37
import py3Dmol
from IPython.display import display
import subprocess
import tempfile
torch.backends.cuda.matmul.allow_tf32 = True


class Classifier:
    def __init__(self, model_checkpoint, device):
        self.model_checkpoint = model_checkpoint
        self.device = device
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_checkpoint)
        self.model = AutoModelForSequenceClassification.from_pretrained(self.model_checkpoint).to(self.device)
    
    def preprocess_data(self, sequences):
        batch_encoding = self.tokenizer(sequences, padding=True, truncation=True, return_tensors="pt")
        dataset = Dataset.from_dict(batch_encoding)
        return dataset

    def classify(self, sequences):
        dataset = self.preprocess_data(sequences)
        dataloader = torch.utils.data.DataLoader(
        dataset, shuffle=True, collate_fn=default_data_collator, batch_size=8, pin_memory=True)
        self.model.eval()
        results = []
        with torch.no_grad():
            for batch in dataloader:
                input_ids = batch["input_ids"].to(self.device)
                attention_mask = batch["attention_mask"].to(self.device)
                outputs = self.model(input_ids, attention_mask=attention_mask)
                logits = outputs.logits
                probabilities = torch.softmax(logits, dim=1)
                scores = probabilities[:, 1]
                results.extend(scores.cpu().numpy())
        return results


class Structure_Evaluate():
    def __init__(self):
        
        self.tokenizer = AutoTokenizer.from_pretrained("/data1/anonymity/Pre_Train_Model/esmfold_v1")
        self.model = EsmForProteinFolding.from_pretrained("/data1/anonymity/Pre_Train_Model/esmfold_v1", low_cpu_mem_usage=True)

        self.model = self.model.cuda()
        
    def get_plddt_scores(self, task_name, sequences,num = -1):
        results = []
        tmscores = []
        rmsds = []
        temp_sequences = sequences[ : num]
        
        db_path = os.path.join("/data1/anonymity/foldseek", task_name)
        for test_protein in temp_sequences:

            try:
                tokenized_input = self.tokenizer([test_protein], return_tensors="pt", add_special_tokens=False)['input_ids']
                tokenized_input = tokenized_input.cuda()
                with torch.no_grad():
                    outputs = self.model(tokenized_input)
                    pdb_content = self.convert_outputs_to_pdb(outputs)
                    pdb_content = pdb_content[0]
                    tmscore, rmsd = self.get_foldseek_result(pdb_content, db_path)
                tmscores.append(tmscore)
                rmsds.append(rmsd)
                results.append(self.compute_average_plddt(outputs))
            except:
                x=1
            
        
        return results, tmscores, rmsds
    
    
    def compute_average_plddt(self, outputs):
        x = outputs['plddt']
        return torch.mean(outputs['plddt']).item()
        
    def get_foldseek_result(self,pdb_content, db_path):

        output_content = self.run_foldseek(pdb_content, db_path)
        output_content = output_content.split("\n")[0]
        result = output_content.split("\t")
        name_list = ["query","target","fident","prob","lddt","alntmscore","qtmscore","ttmscore","rmsd"]
        
        result = dict(zip(name_list, result))
    
        return float(result["qtmscore"]), float(result["rmsd"])
    
        
    def run_foldseek(self, pdb_content, db_path):

        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.pdb') as tmp_query:
            tmp_query.write(pdb_content)
            tmp_query_path = tmp_query.name
        

        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.aln') as tmp_output:
            tmp_output_path = tmp_output.name

        try:
            
            command = [
                "foldseek", "easy-search",
                tmp_query_path, db_path, tmp_output_path, "../tmp",
                "--format-output", "query,target,fident,prob,lddt,alntmscore,qtmscore,ttmscore,rmsd"
            ]
            with open(os.devnull, 'w') as devnull:
                
                subprocess.run(command, check=True, stdout=devnull, stderr=devnull)
            
            with open(tmp_output_path, 'r') as f:
                output_content = f.read()
        finally:

            os.remove(tmp_query_path)
            os.remove(tmp_output_path)   
        return output_content  

    def convert_outputs_to_pdb(self, outputs):
        final_atom_positions = atom14_to_atom37(outputs["positions"][-1], outputs)
        outputs = {k: v.to("cpu").numpy() for k, v in outputs.items()}
        final_atom_positions = final_atom_positions.cpu().numpy()
        final_atom_mask = outputs["atom37_atom_exists"]
        pdbs = []
        for i in range(outputs["aatype"].shape[0]):
            aa = outputs["aatype"][i]
            pred_pos = final_atom_positions[i]
            mask = final_atom_mask[i]
            resid = outputs["residue_index"][i] + 1
            pred = OFProtein(
                aatype=aa,
                atom_positions=pred_pos,
                atom_mask=mask,
                residue_index=resid,
                b_factors=outputs["plddt"][i],
                chain_index=outputs["chain_index"][i] if "chain_index" in outputs else None,
            )
            pdbs.append(to_pdb(pred))
        return pdbs
    
    
def load_data(file_path):
       
    labels = []
    sequences = []
    with open(file_path, "r") as file:
        data = file.readlines()

        for line in data:
            line = line.rstrip()
            origin_data = json.loads(line)

            
            label = origin_data[0]
            sequence = origin_data[1]
            
            labels.append(label)
            sequences.append(sequence)
    return labels, sequences


def find_txt_files(folder_path):

    txt_files = []

    for root, dirs, files in os.walk(folder_path):

        for dirname in dirs:
            
            file_path = os.path.join(root, dirname, "result.txt")
            txt_files.append( ( dirname, file_path) )
        
    return txt_files



def get_subfolder_paths(dir_path):

    subfolder_paths = []

    for root, dirs, files in os.walk(dir_path):

        for subdir in dirs:
     
            sub_path = os.path.join(root, subdir)
            subfolder_paths.append(sub_path)

        break
    return subfolder_paths



def main():
    
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset_path", type=str, default="./single_function_result")
    args = parser.parse_args()
    folder_path = args.dataset_path

    args.dataset_path = "mlpo_model_result"

    task_name = ["function_0", "function_1"]


    for task  in task_name:
        folder_path = os.path.join(args.dataset_path,  task)
        subfolder_list = get_subfolder_paths(folder_path)
        subfolder_list = sorted(subfolder_list, key=lambda x: int(x.split('-')[-1]))
        
        process0_model_checkpoint = "/data1/anonymity/Pre_Train_Model/protein_evaluation/process_0"
        function0_model_checkpoint = "/data1/anonymity/Pre_Train_Model/protein_evaluation/function_0"
        component0_model_checkpoint = "/data1/anonymity/Pre_Train_Model/protein_evaluation/component_0"
        
        process1_model_checkpoint = "/data1/anonymity/Pre_Train_Model/protein_evaluation/process_1"
        function1_model_checkpoint = "/data1/anonymity/Pre_Train_Model/protein_evaluation/function_1"
        component1_model_checkpoint = "/data1/anonymity/Pre_Train_Model/protein_evaluation/component_1"
        
        
        device = "cuda" if cuda.is_available() else "cpu"
        process0_classifier = Classifier(process0_model_checkpoint, device)
        function0_classifier = Classifier(function0_model_checkpoint, device)
        component0_classifier = Classifier(component0_model_checkpoint, device)
        
        process1_classifier = Classifier(process1_model_checkpoint, device)
        function1_classifier = Classifier(function1_model_checkpoint, device)
        component1_classifier = Classifier(component1_model_checkpoint, device)
        
        structure_evaluate = Structure_Evaluate()
        result_path = os.path.join(folder_path, 'score.txt')
        with open(result_path,"w") as file:
            for subfolder in tqdm(subfolder_list):
                
                
                file_list = [os.path.join(subfolder,'result.txt')]
                
                for file_path in tqdm(file_list):   
                    file_name = os.path.basename(subfolder)
                    labels, sequences = load_data(file_path)
                    
                    function0_sequences = []
                    function0_scores = []
                    process0_sequences = []
                    process0_scores = []
                    component0_sequences = []
                    component0_scores = []
                    
                    function1_sequences = []
                    function1_scores = []
                    process1_sequences = []
                    process1_scores = []
                    component1_sequences = []
                    component1_scores = []
                    plddt_scores = []
                    tmscores = []
                    rmsds = []
                    
                   
                    plddt_scores, tmscores, rmsds = structure_evaluate.get_plddt_scores(task, sequences, num = 10)
                    
                    for i in range(len(sequences)):  
                        sequence = sequences[i]
                        sequence = sequence.replace('\n','')
                        label = labels[i]
                        
                        if label[0]==0:
                            function0_sequences.append(sequence)
                        elif label[0] == 1:
                            function1_sequences.append(sequence)
                        
                        if label[1]==0:
                            process0_sequences.append(sequence)
                        elif label[1] == 1:
                            process1_sequences.append(sequence)
                            
                        if label[2]==0:
                            component0_sequences.append(sequence)
                        elif label[2] == 1:
                            component1_sequences.append(sequence)
                            
                    
                    if function0_sequences:
                        function0_scores = function0_classifier.classify(function0_sequences)       
                    if process0_sequences:
                        process0_scores = process0_classifier.classify(process0_sequences)      
                    if component0_sequences:
                        component0_scores = component0_classifier.classify(component0_sequences)
                    
                    if function1_sequences:
                        function1_scores = function1_classifier.classify(function1_sequences)       
                    if process1_sequences:
                        process1_scores = process1_classifier.classify(process1_sequences)      
                    if component1_sequences:
                        component1_scores = component1_classifier.classify(component1_sequences)
                    
                    
                    
                    function_score = sum(function0_scores+function1_scores)/len(function0_scores+function1_scores) if len(function0_scores+function1_scores)!=0 else 0
                    process_score = sum(process0_scores + process1_scores)/len(process0_scores + process1_scores) if len(process0_scores + process1_scores)!=0 else 0
                    component_score = sum(component0_scores + component1_scores)/len(component0_scores + component1_scores) if len(component0_scores + component1_scores)!=0 else 0
                    
                    
                    avg_plddt_score = sum(plddt_scores)/len(plddt_scores) if len(plddt_scores)!=0 else 0
                    avg_tmscore = sum(tmscores)/len(tmscores) if len(tmscores)!=0 else 0
                    avg_rmsd = sum(rmsds)/len(rmsds) if len(rmsds)!=0 else 0
                    
                    avg_tmscore = round(avg_tmscore,4)
                    avg_rmsd = round(avg_rmsd,4)
                    
                    file.write(f"{file_name}, function:{function_score}, process:{process_score}, component:{component_score}, plddt:{avg_plddt_score}, tmscore:{avg_tmscore}, rmsd:{avg_rmsd}\n")
            
                    
        
    


if __name__ == "__main__":
    main()