Spaces:
Running on Zero
Running on Zero
| from machine_learning.Dataset import BindingResidueDataset | |
| from machine_learning.ModelManager import load_classifier_torch | |
| from setup.configProcessor import get_weight_dir, get_structure_cutoff | |
| from setup.generalSetup import select_device | |
| from torch_geometric.loader import DataLoader | |
| import os | |
| import torch | |
| import numpy as np | |
| def run_prediction(model_parameters): | |
| structure_cutoff = get_structure_cutoff(only_first_value=True) | |
| dataset = BindingResidueDataset(structure_cutoff) | |
| predictions = {} | |
| sigmoid = torch.nn.Sigmoid() | |
| models = os.listdir(get_weight_dir()) | |
| for model in models: | |
| model_path = os.path.join(get_weight_dir(), model) | |
| pretrained_model = load_classifier_torch(model_path, model_parameters) | |
| pretrained_model.to(select_device()) | |
| data_loader = DataLoader(dataset, batch_size=1, shuffle=True) | |
| pretrained_model.eval() | |
| for protein_graph, protein_id in data_loader: | |
| with torch.no_grad(): | |
| protein_id = protein_id[0] | |
| protein_graph = protein_graph.to(select_device()) | |
| prediction = pretrained_model.forward( | |
| protein_graph.x, | |
| protein_graph.edge_index, | |
| protein_graph.edge_index_cutoff, | |
| protein_graph.edge_features, | |
| protein_graph.dssp_features | |
| ) | |
| prediction = sigmoid(prediction) | |
| prediction = prediction.detach().cpu() | |
| prediction = np.array(prediction) | |
| if protein_id not in predictions: | |
| predictions[protein_id] = [prediction] | |
| else: | |
| predictions[protein_id] += [prediction] | |
| return predictions | |